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

import matplotlib.pyplot as plt
import io
import logging

from ..aoa import AoA
from .. import config

import matplotlib

matplotlib.use("agg")

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

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]
):
    if config.VISUALISE_RAW:
        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(aoa: AoA, max_tof: float):
    logger.info(f"Making heatmap with aoa of {aoa.timestamp}")
    fig = plt.figure()
    ax = fig.add_axes([0, 0, 1, 1], polar=True)
    r = np.linspace(0, max_tof, 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(max_tof: float):
    assert aoa_queue is not None

    prev_frame = datetime.now()
    while True:
        while (datetime.now() - prev_frame).total_seconds() < 1 / config.HEATMAP_FPS:
            time.sleep(0.01)
        while not aoa_queue.empty():
            logger.debug("Receiving from aoa pipe")
            aoa = aoa_queue.get()
        prev_frame = datetime.now()
        logger.debug(f"Generating heatmap of time {aoa.timestamp}")
        buf = make_heatmap(aoa, max_tof)
        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():
    max_tof_str = request.args.get("max_tof")
    try:
        max_tof = float(max_tof_str)
    except Exception:
        max_tof = 5e-8
    return Response(
        gather_aoa(max_tof), mimetype="multipart/x-mixed-replace; boundary=frame"
    )


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