import numpy as np
import cv2
import os
from PySide6.QtCore import QObject, QTimer
import mediapipe as mp
from mediapipe.tasks.python import BaseOptions
from mediapipe.tasks.python.vision import HandLandmarker, HandLandmarkerOptions, RunningMode
from Tbd.helper import Hand2D
from collections import defaultdict
from Helpers.HandUdpSender import HandUdpSender

import struct
import time

MODEL_PATH = os.path.join(os.path.dirname(__file__), "models", "hand_landmarker.task")

class TrackingController(QObject):
    def __init__(self, cameraList, P_mats, logger=None, tick_ms=30):
        super().__init__()

        self.cameraList = cameraList
        self.P_mats = P_mats
        self.logger = logger

        self.udp = HandUdpSender(host="127.0.0.1", port=9000)

        self.timer = QTimer(self)
        self.timer.timeout.connect(self._tick)
        self.tick_ms = tick_ms

        self.running = False

        # MediaPipe's legacy solutions.hands API no longer exists on the
        # installed mediapipe version, so hand detection runs on the newer
        # Tasks API instead. Each camera gets its own HandLandmarker running
        # in VIDEO mode, fed only that camera's own frames with a monotonically
        # increasing per-camera timestamp - this lets MediaPipe track hands
        # frame-to-frame per view (much less jitter than re-detecting blind
        # every frame), which wouldn't be valid if frames from different
        # camera viewpoints were interleaved through a single VIDEO-mode
        # detector.
        self._frame_counter = defaultdict(int)
        self.detectors = {}
        for camera in self.cameraList:
            options = HandLandmarkerOptions(
                base_options=BaseOptions(model_asset_path=MODEL_PATH),
                running_mode=RunningMode.VIDEO,
                num_hands=12,
                min_hand_detection_confidence=0.5,
                min_hand_presence_confidence=0.5,
                min_tracking_confidence=0.5,
            )
            self.detectors[camera.index] = HandLandmarker.create_from_options(options)

    def start(self):
        # Guard: need projections
        self.logger.log("Amount of P_mats:" + str(len(self.P_mats)))
        if not self.P_mats or len(self.P_mats) < 2:
            if self.logger: self.logger.log("[Tracking] Missing P_mats (need >=2 cameras calibrated)")
            return

        if self.running:
            return

        for camera in self.cameraList:
            if not camera.open():
                if self.logger: self.logger.log(f"[Tracking] cam{camera.index}: failed to open, will be skipped.")

        self.running = True
        self.timer.start(self.tick_ms)
        if self.logger: self.logger.log("[Tracking] Started")

    def stop(self):
        if not self.running:
            return
        self.running = False
        self.timer.stop()

        for camera in self.cameraList:
            camera.close()

        for detector in self.detectors.values():
            detector.close()

        if self.logger: self.logger.log("[Tracking] Stopped")

    def _tick(self):
        hands2d = self._detectHands_all_cameras()
        if not hands2d:
            return

        # Group hands across cameras (you already have _allocateDetectedHands logic)
        hand_groups = self._allocateDetectedHands(hands2d)

        # Triangulate (your existing triangulation approach)
        hands3d = self._triangulateHands(hand_groups)

        if len(hands3d) > 0:
            self.udp.send_hands3d(hands3d)

        print(
            f"hands2d={len(hands2d)} groups={len(hand_groups)} hands3d={len(hands3d)}",
            flush=True
        )

        if len(hands2d) > 0:
            print("cam_ids:", sorted(set(h.camera_id for h in hands2d)), flush=True)

    def _detectHands_all_cameras(self):
        out = []
        for camera in self.cameraList:
            if camera.capture is None:
                continue

            detector = self.detectors.get(camera.index)
            if detector is None:
                continue

            ok, frameBGR = camera.capture.read()
            if not ok or frameBGR is None:
                continue

            h, w = frameBGR.shape[:2]
            frameRGB = cv2.cvtColor(frameBGR, cv2.COLOR_BGR2RGB)

            mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=frameRGB)

            self._frame_counter[camera.index] += 1
            timestamp_ms = self._frame_counter[camera.index] * self.tick_ms
            result = detector.detect_for_video(mp_image, timestamp_ms)

            if not result.hand_landmarks:
                continue

            for hi, lm_list in enumerate(result.hand_landmarks):
                landmarks_px = np.array(
                    [[lm.x * w, lm.y * h] for lm in lm_list],
                    dtype=np.float32
                )

                label, score = "Unknown", 0.0
                if hi < len(result.handedness) and result.handedness[hi]:
                    label = result.handedness[hi][0].category_name
                    score = result.handedness[hi][0].score

                out.append(Hand2D(
                    camera_id=camera.index,
                    handedness=label,
                    score=score,
                    landmarks_px=landmarks_px
                ))
        return out

    def _allocateDetectedHands(self, hands2d):
        frames_by_camera = defaultdict(list)
        for hand in hands2d:
            frames_by_camera[hand.camera_id].append(hand)

        used = set()
        hand_groups = []
        reproj_threshold = 50.0

        camera_ids = sorted(frames_by_camera.keys())

        for cam_a in camera_ids:
            if cam_a not in self.P_mats:
                continue
            for hand_a in frames_by_camera[cam_a]:
                if id(hand_a) in used:
                    continue

                group = {cam_a: hand_a}
                used.add(id(hand_a))

                P_a = self.P_mats[cam_a]

                for cam_b in camera_ids:
                    if cam_b == cam_a:
                        continue
                    if cam_b not in self.P_mats:
                        continue

                    P_b = self.P_mats[cam_b]
                    best_hand = None
                    best_err = np.inf

                    for hand_b in frames_by_camera[cam_b]:
                        if id(hand_b) in used:
                            continue

                        err = self.pair_reprojection_error(
                            hand_a, hand_b, P_a, P_b, key_idxs=[0, 9]
                        )
                        if err < best_err:
                            best_err = err
                            best_hand = hand_b

                    if best_hand is not None and best_err < reproj_threshold:
                        group[cam_b] = best_hand
                        used.add(id(best_hand))

                hand_groups.append(group)

        return hand_groups

    def triangulate_hands_with_metadata(self, hand_groups):
        hands3d = []
        for group in hand_groups:
            if len(group) < 2:
                continue

            handedness, confidence = self.choose_metadata_from_group(group)

            cams = list(group.keys())
            P1 = self.P_mats[cams[0]]
            P2 = self.P_mats[cams[1]]
            lm1 = group[cams[0]].landmarks_px
            lm2 = group[cams[1]].landmarks_px

            pts3d = []
            for i in range(lm1.shape[0]):
                X = self.triangulate_point(P1, P2, lm1[i], lm2[i])
                pts3d.append(X)

            hands3d.append({
                "handedness": handedness,          # "Left"/"Right"/"Unknown"
                "confidence": confidence,          # 0..1
                "landmarks": np.asarray(pts3d, dtype=np.float32)  # (21,3)
            })
        return hands3d

    # --- math helpers (copy from your existing code) ---
    def triangulate_point(self, P1, P2, x1, x2):
        A = np.zeros((4, 4), dtype=np.float32)
        A[0] = x1[0] * P1[2] - P1[0]
        A[1] = x1[1] * P1[2] - P1[1]
        A[2] = x2[0] * P2[2] - P2[0]
        A[3] = x2[1] * P2[2] - P2[1]
        _, _, Vt = np.linalg.svd(A)
        X_h = Vt[-1]
        X_h /= X_h[3]
        
        return X_h[:3]

    def project_point(self, P, X):
        X_h = np.array([X[0], X[1], X[2], 1.0], dtype=np.float32)
        x = P @ X_h
        return np.array([x[0] / x[2], x[1] / x[2]], dtype=np.float32)

    def pair_reprojection_error(self, hand_a, hand_b, P_a, P_b, key_idxs):
        errors = []
        for idx in key_idxs:
            x1 = hand_a.landmarks_px[idx]
            x2 = hand_b.landmarks_px[idx]
            X = self.triangulate_point(P_a, P_b, x1, x2)
            x1_hat = self.project_point(P_a, X)
            x2_hat = self.project_point(P_b, X)
            errors.append(np.linalg.norm(x1_hat - x1))
            errors.append(np.linalg.norm(x2_hat - x2))
        return float(np.mean(errors))

    def choose_metadata_from_group(group: dict):
        # group: {cam_id: Hand2D, ...}
        best = max(group.values(), key=lambda h: float(getattr(h, "score", 0.0)))
        handedness = getattr(best, "handedness", "Unknown")
        confidence = float(getattr(best, "score", 0.0))
        return handedness, confidence
    
    def _triangulateHands(self, hand_groups):
        hands3d = []

        for group in hand_groups:
            cams = list(group.keys())
            if len(cams) < 2:
                continue

            # Pick metadata from best 2D view
            best = max(group.values(), key=lambda hh: float(getattr(hh, "score", 0.0)))
            handedness = getattr(best, "handedness", "Unknown")
            confidence = float(getattr(best, "score", 0.0))

            P1 = self.P_mats[cams[0]]
            P2 = self.P_mats[cams[1]]
            lm1 = group[cams[0]].landmarks_px
            lm2 = group[cams[1]].landmarks_px

            pts3d = []
            for i in range(lm1.shape[0]):
                X = self.triangulate_point(P1, P2, lm1[i], lm2[i])
                pts3d.append(X)

            hands3d.append({
                "landmarks": np.asarray(pts3d, dtype=np.float32),
                "handedness": handedness,
                "confidence": confidence,
            })

        return hands3d