import sys
import socket
import json
import numpy as np
from collections import deque
import pyqtgraph as pg
from PyQt6.QtWidgets import QApplication, QMainWindow, QVBoxLayout, QHBoxLayout, QWidget, QLabel
from PyQt6.QtCore import QTimer, Qt


WIN            = 64
MASK_CALIB     = 50
BASELINE_CALIB = 150
NULL_FRACTION  = 0.20
HAMPEL_GAMMA   = 3.0
Z_THRESH       = 8.0
METRIC_SCALE   = 1000.0


class CSIProcessor:

    def __init__(self, name):
        self.name = name
        self.expected_len = None
        self.active_mask = None
        self._mask_sum = None
        self._mask_cnt = 0
        self._drop_first = False

        self.win = deque(maxlen=WIN)
        self._calib = deque(maxlen=BASELINE_CALIB)
        self.baseline_med = None
        self.baseline_mad = None

    @staticmethod
    def _amplitudes(raw):
        arr = np.asarray(raw, dtype=np.float64)
        n = arr.size - (arr.size & 1)
        iq = arr[:n].reshape(-1, 2)
        return np.sqrt(iq[:, 0] ** 2 + iq[:, 1] ** 2)

    def calibrating(self):
        return self.baseline_med is None

    def process(self, raw):
        if isinstance(raw, dict):
            csi = raw.get("csi")
            if raw.get("first_word_invalid"):
                if not self._drop_first:
                    self._drop_first = True
                    if self.active_mask is not None:
                        self.active_mask[0:2] = False
        else:
            csi = raw

        if not isinstance(csi, list) or len(csi) < 128:
            return None

        if self.expected_len is None:
            self.expected_len = len(csi)
        if len(csi) != self.expected_len:
            return None

        amp = self._amplitudes(csi)

        if self._drop_first and amp.size >= 2:
            amp[0] = 0.0
            amp[1] = 0.0

        if self.active_mask is None:
            if self._mask_sum is None:
                self._mask_sum = np.zeros_like(amp)
            self._mask_sum += amp
            self._mask_cnt += 1
            if self._mask_cnt < MASK_CALIB:
                return None
            mean_amp = self._mask_sum / self._mask_cnt
            ref = np.median(mean_amp[mean_amp > 0]) if np.any(mean_amp > 0) else 0.0
            self.active_mask = mean_amp > (NULL_FRACTION * ref)
            if self.active_mask.sum() < 4:
                self.active_mask = mean_amp > 0
        amp = amp[self.active_mask]

        m = amp.mean()
        if m <= 1e-9:
            return None
        amp_n = amp / m

        self.win.append(amp_n)
        if len(self.win) < WIN:
            return None
        W = np.stack(self.win)

        med = np.median(W, axis=0)
        mad = np.median(np.abs(W - med), axis=0)
        thr = HAMPEL_GAMMA * 1.4826 * mad + 1e-12
        Wc = np.where(np.abs(W - med) > thr, med, W)

        metric = float(np.mean(np.var(Wc, axis=0))) * METRIC_SCALE

        if self.baseline_med is None:
            self._calib.append(metric)
            if len(self._calib) >= BASELINE_CALIB:
                b = np.asarray(self._calib)
                self.baseline_med = float(np.median(b))
                self.baseline_mad = float(np.median(np.abs(b - self.baseline_med)) * 1.4826 + 1e-9)
            return (metric, 0.0)

        z = (metric - self.baseline_med) / self.baseline_mad
        return (metric, z)


class DualRadarUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("GhostWatch - Dual Zone Radar")
        self.resize(1000, 700)
        self.setStyleSheet("background-color: #0a0a0a; color: #00FF00; font-family: Consolas;")

        self.sock_a = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sock_a.bind(("127.0.0.1", 5005))
        self.sock_a.setblocking(False)

        self.sock_b = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sock_b.bind(("127.0.0.1", 5006))
        self.sock_b.setblocking(False)

        self.proc_a = CSIProcessor("ZONE A")
        self.proc_b = CSIProcessor("ZONE B")

        self.plot_data_a = np.zeros(200)
        self.plot_data_b = np.zeros(200)

        self.init_ui()

        self.timer = QTimer()
        self.timer.timeout.connect(self.update_radar)
        self.timer.start(30)

    def init_ui(self):
        main_layout = QVBoxLayout()
        graphs_layout = QHBoxLayout()

        pg.setConfigOption('background', '#0a0a0a')
        pg.setConfigOption('foreground', '#00FF00')

        self.plot_widget_a = pg.PlotWidget(title="СЕКТОР А: энергия движения")
        self.curve_a = self.plot_widget_a.plot(pen=pg.mkPen('#00FF00', width=2))
        self.thr_a = self.plot_widget_a.addLine(y=0, pen=pg.mkPen('#FFAA00', width=1, style=Qt.PenStyle.DashLine))
        graphs_layout.addWidget(self.plot_widget_a)

        self.plot_widget_b = pg.PlotWidget(title="СЕКТОР Б: энергия движения")
        self.curve_b = self.plot_widget_b.plot(pen=pg.mkPen('#00FF00', width=2))
        self.thr_b = self.plot_widget_b.addLine(y=0, pen=pg.mkPen('#FFAA00', width=1, style=Qt.PenStyle.DashLine))
        graphs_layout.addWidget(self.plot_widget_b)

        main_layout.addLayout(graphs_layout)
        map_layout = QHBoxLayout()

        self.alert_a = QLabel("ZONE A: CALIBRATING…")
        self.alert_a.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.alert_a.setStyleSheet("background-color: #002200; border: 2px solid #00FF00; font-size: 32px; font-weight: bold; padding: 20px;")

        self.alert_b = QLabel("ZONE B: CALIBRATING…")
        self.alert_b.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.alert_b.setStyleSheet("background-color: #002200; border: 2px solid #00FF00; font-size: 32px; font-weight: bold; padding: 20px;")

        map_layout.addWidget(self.alert_a)
        map_layout.addWidget(self.alert_b)

        main_layout.addLayout(map_layout)
        container = QWidget()
        container.setLayout(main_layout)
        self.setCentralWidget(container)

    def process_stream(self, sock, proc, plot_data, curve, thr_line, alert_label):
        try:
            while True:
                data, _ = sock.recvfrom(65536)
                try:
                    raw_csi = json.loads(data.decode('utf-8'))
                except (ValueError, UnicodeDecodeError):
                    continue

                res = proc.process(raw_csi)
                if res is None:
                    continue
                metric, z = res

                plot_data[:-1] = plot_data[1:]
                plot_data[-1] = metric
                curve.setData(plot_data)

                if proc.calibrating():
                    alert_label.setText(f"{proc.name}: CALIBRATING… ({len(proc._calib)}/{BASELINE_CALIB})")
                    alert_label.setStyleSheet("background-color: #111100; color: #FFAA00; border: 2px solid #FFAA00; font-size: 30px; font-weight: bold; padding: 20px;")
                    continue

                thr_val = proc.baseline_med + Z_THRESH * proc.baseline_mad
                thr_line.setValue(thr_val)

                if z > Z_THRESH:
                    alert_label.setText(f"!! {proc.name} BREACHED !!")
                    alert_label.setStyleSheet("background-color: #440000; color: #FF0000; border: 2px solid #FF0000; font-size: 32px; font-weight: bold; padding: 20px;")
                    curve.setPen(pg.mkPen('#FF0000', width=3))
                else:
                    alert_label.setText(f"{proc.name}: CLEAR")
                    alert_label.setStyleSheet("background-color: #002200; color: #00FF00; border: 2px solid #00FF00; font-size: 32px; font-weight: bold; padding: 20px;")
                    curve.setPen(pg.mkPen('#00FF00', width=2))
        except BlockingIOError:
            pass

    def update_radar(self):
        self.process_stream(self.sock_a, self.proc_a, self.plot_data_a, self.curve_a, self.thr_a, self.alert_a)
        self.process_stream(self.sock_b, self.proc_b, self.plot_data_b, self.curve_b, self.thr_b, self.alert_b)


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = DualRadarUI()
    window.show()
    sys.exit(app.exec())