from flask import Flask, render_template, Response
from flask_sock import Sock
import numpy as np
import numpy.typing as npt
from simple_websocket import Server
import time

import matplotlib.pyplot as plt
import io
from PIL import Image
import logging

from ..aoa import AoA

import matplotlib

matplotlib.use("agg")

app = Flask(__name__)
sock = Sock(app)

logger = logging.getLogger(__name__)


data: npt.NDArray[np.complex128] = np.array([], dtype=complex)
aoa: AoA = AoA()


@app.route("/preprocessed")
def preprocessed():
    return render_template("preprocessed.html")


type Subscriber = Server
subscriber_settings: dict[Subscriber, tuple[int, int, int]] = {}


@sock.route("/data")
def get_data(sock: Subscriber):
    while True:
        msg = sock.receive()
        if len(msg.split()) != 3:
            break
        subcarrier, rx, tx = map(int, msg.split())
        subscriber_settings[sock] = (subcarrier, rx, tx)


def add_data(
    raw_data: npt.NDArray[np.complex128], new_data: npt.NDArray[np.complex128]
):
    # magn = np.abs(raw_data)
    # phase = np.angle(raw_data)

    # new_mag = np.abs(new_data)
    # new_phase = np.angle(new_data)

    # fig, axs = plt.subplots(2, 2)
    # axs[0, 0].plot(magn[:, 0, 0], c="b")
    # axs[0, 0].plot(magn[:, 1, 0], c="orange")
    # axs[1, 0].plot(new_mag[:, 0, 0], c="b")
    # axs[1, 0].plot(new_mag[:, 1, 0], c="orange")
    # axs[0, 1].plot(phase[:, 0, 0], c="b")
    # axs[0, 1].plot(phase[:, 1, 0], c="orange")
    # axs[1, 1].plot(new_phase[:, 0, 0], c="b")
    # axs[1, 1].plot(new_phase[:, 1, 0], c="orange")
    # fig.savefig("/tmp/plot.png")
    # plt.close(fig)

    global data
    if data.size == 0:
        data = np.expand_dims(new_data, axis=0)
    else:
        data = np.concat([data, np.expand_dims(new_data, axis=0)], axis=0)

        # Only keep latest 100 entries
        if data.shape[0] > 100:
            data = data[-100:]
    to_remove: list[Subscriber] = []
    for subscriber in subscriber_settings:
        try:
            subcarrier, rx, tx = subscriber_settings[subscriber]
            subscriber.send(new_data.real[subcarrier, rx, tx])
        except Exception as e:
            to_remove.append(subscriber)
            print(e)

    for subscriber in to_remove:
        del subscriber_settings[subscriber]


def make_heatmap():
    fig = plt.figure()
    ax = fig.add_axes([0, 0, 1, 1], polar=True)
    r = np.linspace(0, 3e-8, 100)  # Radius values
    theta = np.linspace(0, np.pi, 50)  # Angle values
    R, Theta = np.meshgrid(r, theta)  # Create a 2D grid of r and theta

    # Compute the function values
    Z = np.log(np.vectorize(aoa.evaluate)(Theta, R))

    ax.pcolormesh(Theta, R, Z, edgecolors="face")
    buf = io.BytesIO()
    fig.savefig(buf, format="jpeg")
    plt.close(fig)

    buf.seek(0)
    return buf


def gather_aoa():
    while True:
        # time.sleep(0.05)
        logger.info("Got AoA heatmap")
        buf = make_heatmap()
        yield (b"--frame\r\nContent-Type: image/jpeg\r\n\r\n" + buf.read() + b"\r\n")
        buf.close()


@app.route("/aoa_tof")
def aoa_tof():
    return Response(gather_aoa(), mimetype="multipart/x-mixed-replace; boundary=frame")


def start():
    app.run(debug=True, use_reloader=False, host="0.0.0.0")