import cv2
import numpy

from Tbd.helper import  CameraObject, CharucoDetection, CalibrationSample
from UI.CameraSetupWidget import CameraSetupWidget
from Helpers import saveManager as sm

from PySide6.QtCore import Qt, QTimer
from PySide6.QtGui import QImage, QPixmap

from PySide6.QtWidgets import QWidget, QLabel, QVBoxLayout, QSizePolicy
from typing import List
from PySide6.QtWidgets import QScrollArea, QFrame, QPushButton

from UI.UILogger import UILogger

MIN_CHARUCO_CORNERS = 15   # absolute minimum
#GOOD_CHARUCO_CORNERS = 25  # ideal

DICT_ID = cv2.aruco.DICT_4X4_50
SQUARES_X = 6                    # number of chessboard squares in X
SQUARES_Y = 9                    # number of chessboard squares in Y
SQUARE_LEN_MM = 25.0             # square size in mm
MARKER_LEN_MM = 18.0             # marker size in mm (must be < square)

#FRAME_SIZE = [3264, 2448]
#FRAME_SIZE = [1536, 2048]
FRAME_SIZE = [2048, 1536]

REFERENCE_CAMERA_INDEX = 1

class HomeWindow(QWidget):

    def __init__(self, cameraList: List[CameraObject], on_cameraList_changed, P_mats, logger : UILogger):
        super().__init__()

        # Initiating variables
        self.P_mats = P_mats
        self.on_cameraList_changed = on_cameraList_changed
        self.cameraList = cameraList
        self.logger = logger
        self.calibration_samples: list[CalibrationSample] = []
        self.camera_Setup_Widget_List = []


        # Charcuo detection stuff
        dictionary = cv2.aruco.getPredefinedDictionary(DICT_ID)
        self.board = cv2.aruco.CharucoBoard(
            (SQUARES_X, SQUARES_Y),
            SQUARE_LEN_MM,
            MARKER_LEN_MM,
            dictionary
        )
        self.charuco_detector = cv2.aruco.CharucoDetector(self.board)
        # Detector parameters
        detector_params = cv2.aruco.DetectorParameters()
            # You can tweak these if detection is unstable:
            # detector_params.adaptiveThreshWinSizeMin = 3
            # detector_params.adaptiveThreshWinSizeMax = 23
            # detector_params.adaptiveThreshWinSizeStep = 10

        self.arcuo_detector = cv2.aruco.ArucoDetector(dictionary, detector_params)

        # Creating the main_Layout and other Widgets 
        self.main_Layout = QVBoxLayout(self)
        self.scroll_Box = QScrollArea()
        self.calibration_button = QPushButton("Calibrate")
        self.caputre_calibration_sample_button = QPushButton("Capture calibration sample")
        self.clear_calibration_samples_button = QPushButton("Clear calibration samples")
        self.load_calibration_button = QPushButton("Load calibration")
        self.save_calibration_button = QPushButton("Save calibration")
        self.connect_cameras_button = QPushButton("Connect cameras")
        self.disconnect_cameras_button = QPushButton("Disconnect cameras")


        # Creating the Scrollable area
        self.scroll_Box.setWidgetResizable(True)
        self.scroll_Content = QWidget()
        self.scroll_Layout = QVBoxLayout(self.scroll_Content)
        self.scroll_Layout.setAlignment(Qt.AlignTop)
        self.scroll_Box.setWidget(self.scroll_Content)

        # Adding Primary show elements
        self.main_Layout.addWidget(self.connect_cameras_button)
        self.main_Layout.addWidget(self.disconnect_cameras_button)
        self.main_Layout.addWidget(self.clear_calibration_samples_button)
        self.main_Layout.addWidget(self.load_calibration_button)
        self.main_Layout.addWidget(self.save_calibration_button)
        self.main_Layout.addWidget(self.caputre_calibration_sample_button)
        self.main_Layout.addWidget(self.calibration_button)
        self.main_Layout.addWidget(self.scroll_Box)

        # Connecting Buttons to functions
        self.calibration_button.clicked.connect(self._calibrate_camera)
        self.caputre_calibration_sample_button.clicked.connect(self._caputre_calibration_sample)
        self.clear_calibration_samples_button.clicked.connect(self._clear_calibration_samples)
        self.load_calibration_button.clicked.connect(self._load_calibration)
        self.save_calibration_button.clicked.connect(self._save_calibration)
        self.connect_cameras_button.clicked.connect(self._connect_cameras)
        self.disconnect_cameras_button.clicked.connect(self._disconnect_cameras)


    def _connect_cameras(self):
        # Keep all cameras open simultaneously during calibration - this hardware
        # has been confirmed to sustain 6 concurrent MSMF streams once split
        # across separate USB controllers, so per-sample open/close is no longer
        # needed and was just adding latency to every capture.
        for camera in self.cameraList:
            if camera.capture is not None:
                continue
            if camera.open():
                self.logger.log(f"cam{camera.index}: connected.")
            else:
                self.logger.log(f"cam{camera.index}: failed to connect.")

    def _disconnect_cameras(self):
        for camera in self.cameraList:
            camera.close()
        self.logger.log("All cameras disconnected.")

    def _clear_calibration_samples(self):
        self.calibration_samples.clear()
        self.logger.log("Calibration samples where cleared.")

    def _load_calibration(self):
        loaded = sm.load_calibration_json("C:\\git\\Table\\Test\\test.json", self.logger)
        self.cameraList.clear()
        self.cameraList.extend(loaded)
        self.P_mats.clear()
        self.P_mats.update({
            camera.index: camera.camera_projection_matrix
            for camera in self.cameraList
            if camera.camera_projection_matrix is not None
        })
        self.on_cameraList_changed()

    def _save_calibration(self):
        sm.save_calibration_json(self.cameraList, "C:\\git\\Table\\Test\\test.json", self.logger) 

    def _calibrate_camera(self):
        self.logger.log("Fuck you")

        if(len(self.cameraList) < 2):
            self.logger.log("You need at least 2 cameras for calibration")
            return

        self._clear_camera_calibrations()

        REFERENCE_CAMERA_INDEX = self.cameraList[0].index

        for camera in self.cameraList:
            reprojection_error_rms, camera_matrix, distortion_coefficients = self._calibrate_intrinsics_characuo_for_camera(
                sample_list = self.calibration_samples,
                camera_index = camera.index,
                board = self.board
            )

            camera.camera_matrix = camera_matrix
            camera.distortion_coefficients = distortion_coefficients

            if camera_matrix is None:
                continue

            self.logger.log(f"[Intrinsics] cam{camera.index}: RMS={reprojection_error_rms:.4f} fx={camera_matrix[0,0]:.2f} fy={camera_matrix[1,1]:.2f} cx={camera_matrix[0,2]:.2f} cy={camera_matrix[1,2]:.2f}")

        reference_camera = next(camera for camera in self.cameraList if camera.index == REFERENCE_CAMERA_INDEX)

        if reference_camera.camera_matrix is None:
            self.logger.log(f"[Calibration] Reference camera cam{reference_camera.index} has no intrinsics, aborting calibration.")
            return

        self.logger.log("Intrinsics done")

        self.logger.log("Start extrinsics")
        # Reference camera projection (world origin)
        reference_camera.rotation_matrix_world_to_camera = numpy.eye(3)
        reference_camera.translation_vector_world_to_camera = numpy.zeros((3,1))
        reference_camera.camera_projection_matrix = (
            reference_camera.camera_matrix
            @ numpy.hstack([numpy.eye(3), numpy.zeros((3,1))])
        )
        self.P_mats.clear()
        self.P_mats[reference_camera.index] = reference_camera.camera_projection_matrix

        # Not every camera necessarily overlaps with the reference camera directly
        # (e.g. cameras on opposite sides of the table) - chain extrinsics through
        # whichever already-posed camera has the most shared samples with each
        # remaining camera instead of requiring a direct link to the reference.
        cameras_by_index = {camera.index: camera for camera in self.cameraList}
        posed_indices = {reference_camera.index}
        remaining_indices = set(cameras_by_index.keys()) - posed_indices

        while remaining_indices:
            best_known_index, best_target_index, best_count = None, None, 0

            for known_index in posed_indices:
                for target_index in remaining_indices:
                    count = self._count_common_samples(known_index, target_index)
                    if count > best_count:
                        best_known_index, best_target_index, best_count = known_index, target_index, count

            if best_known_index is None:
                for index in remaining_indices:
                    self.logger.log(f"cam{index}: no chain of overlapping samples back to the reference camera, skipping extrinsics.")
                break

            known_camera = cameras_by_index[best_known_index]
            target_camera = cameras_by_index[best_target_index]

            stereo_rms, R_known_to_target, t_known_to_target, E, F, P_known, P_target = self._stereo_calibrate_from_charuco_samples(
                sample_list=self.calibration_samples,
                board=self.board,
                camera_refernece=known_camera,
                camera_to_calibrate=target_camera
            )

            if R_known_to_target is None:
                self.logger.log(f"cam{best_known_index} <-> cam{best_target_index}: stereo calibration failed despite {best_count} shared samples, skipping.")
                remaining_indices.discard(best_target_index)
                continue

            # Compose known camera's world pose with the known->target relative pose
            R_world_known = known_camera.rotation_matrix_world_to_camera
            t_world_known = known_camera.translation_vector_world_to_camera

            R_world_target = R_known_to_target @ R_world_known
            t_world_target = R_known_to_target @ t_world_known + t_known_to_target

            target_camera.rotation_matrix_world_to_camera = R_world_target
            target_camera.translation_vector_world_to_camera = t_world_target
            target_camera.camera_projection_matrix = target_camera.camera_matrix @ numpy.hstack([R_world_target, t_world_target])

            self.P_mats[target_camera.index] = target_camera.camera_projection_matrix
            self.logger.log(f"[Extrinsics] cam{target_camera.index} chained via cam{known_camera.index} (stereo_rms={stereo_rms:.4f}, shared_samples={best_count})")

            posed_indices.add(best_target_index)
            remaining_indices.discard(best_target_index)

        self.logger.log("Extrinsics done")
        self.logger.log("Amount of P_mats:" + str(len(self.P_mats)))
        self.logger.log("Amount of cameras:" + str(len(self.cameraList)))
        self.log_all_camera_calibration(self.cameraList)


    def rotation_matrix_to_rpy_deg(self, R: numpy.ndarray):
        """
        Convert rotation matrix to roll, pitch, yaw in degrees.
        Assumes right-handed, OpenCV convention.
        """
        sy = numpy.sqrt(R[0,0]*R[0,0] + R[1,0]*R[1,0])
        singular = sy < 1e-6

        if not singular:
            roll  = numpy.arctan2(R[2,1], R[2,2])
            pitch = numpy.arctan2(-R[2,0], sy)
            yaw   = numpy.arctan2(R[1,0], R[0,0])
        else:
            roll  = numpy.arctan2(-R[1,2], R[1,1])
            pitch = numpy.arctan2(-R[2,0], sy)
            yaw   = 0.0

        return numpy.degrees([roll, pitch, yaw])
    
    def log_all_camera_calibration(self, camera_list):
        self.logger.log("========== CAMERA CALIBRATION SUMMARY ==========")

        for cam in camera_list:
            self.logger.log(f"--- Camera {cam.index} ---")

            # -------- Intrinsics --------
            if cam.camera_matrix is not None:
                K = cam.camera_matrix
                fx, fy = K[0,0], K[1,1]
                cx, cy = K[0,2], K[1,2]

                self.logger.log(
                    f"Intrinsics:"
                    f" fx={fx:.2f}, fy={fy:.2f},"
                    f" cx={cx:.2f}, cy={cy:.2f}"
                )

                if cam.distortion_coefficients is not None:
                    d = cam.distortion_coefficients.flatten()
                    d_short = ", ".join(f"{v:.4f}" for v in d[:5])
                    self.logger.log(f"Distortion: [{d_short}{'...' if len(d) > 5 else ''}]")
            else:
                self.logger.log("Intrinsics: NOT SET")

            # -------- Extrinsics --------
            if cam.rotation_matrix_world_to_camera is not None and cam.translation_vector_world_to_camera is not None:
                R = cam.rotation_matrix_world_to_camera
                t = cam.translation_vector_world_to_camera.reshape(3)

                roll, pitch, yaw = self.rotation_matrix_to_rpy_deg(R)
                dist = numpy.linalg.norm(t)

                self.logger.log(
                    f"Extrinsics (world → cam):"
                    f" t=({t[0]:.1f}, {t[1]:.1f}, {t[2]:.1f})"
                    f" | |t|={dist:.1f}"
                )
                self.logger.log(
                    f"Rotation (deg):"
                    f" roll={roll:.2f}, pitch={pitch:.2f}, yaw={yaw:.2f}"
                )
            else:
                self.logger.log("Extrinsics: NOT SET")

            # -------- Projection --------
            if cam.camera_projection_matrix is not None:
                P = cam.camera_projection_matrix
                self.logger.log(f"Projection matrix: shape={P.shape}")
            else:
                self.logger.log("Projection matrix: NOT SET")

        self.logger.log("==============================================")

    def _build_stereo_correspondences_from_samples(
            self,
            sample_list,
            board,
            camera_index_reference,
            camera_index_to_calibrate
    ):
        object_points_per_frame  = []
        frame_points_per_frame_reference  = []
        frame_points_per_frame_to_calibrate  = []

        for sample in sample_list:

            # Get detections for both cameras and check if valide
            detection_reference = sample.detections.get(camera_index_reference)
            detection_to_calibrate = sample.detections.get(camera_index_to_calibrate)

            if detection_reference is None or detection_to_calibrate is None:
                continue

            if detection_reference.charuco_ids is None or detection_reference.charuco_corners is None:
                continue

            if detection_to_calibrate.charuco_ids is None or detection_to_calibrate.charuco_corners is None:
                continue

            # Find corner IDs that both cameras can see and check if the amount is enough
            ids_reference = detection_reference.charuco_ids.reshape(-1)
            ids_to_calibrate = detection_to_calibrate.charuco_ids.reshape(-1)

            common_ids = numpy.intersect1d(ids_reference, ids_to_calibrate)

            if len(common_ids) < MIN_CHARUCO_CORNERS:
                continue

            # Build ordered correspondences by common_ids
            # Map id -> corner for each cam
            map_a = {int(i): detection_reference.charuco_corners[idx] for idx, i in enumerate(ids_reference)}
            map_b = {int(i): detection_to_calibrate.charuco_corners[idx] for idx, i in enumerate(ids_to_calibrate)}

            # Assemble corners/ids arrays in matching order
            corners_reference = numpy.array([map_a[int(i)] for i in common_ids], dtype=numpy.float32).reshape(-1, 1, 2)
            corners_to_calibrate = numpy.array([map_b[int(i)] for i in common_ids], dtype=numpy.float32).reshape(-1, 1, 2)
            ids_common = common_ids.astype(numpy.int32).reshape(-1, 1)

             # Convert ChArUco corners+ids -> (objectPts, imagePts) for the board
            object_points_reference, frame_pts_reference = board.matchImagePoints(corners_reference, ids_common)
            _, frame_pts_to_calibrate = board.matchImagePoints(corners_to_calibrate, ids_common)

            if object_points_reference is None or frame_pts_reference is None or frame_pts_to_calibrate is None:
                continue

            object_points_reference = self._as_np_float32(object_points_reference)
            frame_pts_reference = self._as_np_float32(frame_pts_reference)
            frame_pts_to_calibrate = self._as_np_float32(frame_pts_to_calibrate)

            if len(object_points_reference) < MIN_CHARUCO_CORNERS:
                continue

            object_points_per_frame.append(object_points_reference)
            frame_points_per_frame_reference.append(frame_pts_reference)
            frame_points_per_frame_to_calibrate.append(frame_pts_to_calibrate)
            
        return object_points_per_frame, frame_points_per_frame_reference, frame_points_per_frame_to_calibrate

    def _as_np_float32(self, x):
        return numpy.asarray(x, dtype=numpy.float32)

    def _count_common_samples(self, camera_index_a: int, camera_index_b: int) -> int:
        object_points_per_frame, _, _ = self._build_stereo_correspondences_from_samples(
            sample_list=self.calibration_samples,
            board=self.board,
            camera_index_reference=camera_index_a,
            camera_index_to_calibrate=camera_index_b
        )
        return len(object_points_per_frame)

    def _stereo_calibrate_from_charuco_samples(
            self,
            sample_list,
            board,
            camera_refernece,
            camera_to_calibrate
    ):
        
        if camera_refernece.camera_matrix is None or camera_refernece.distortion_coefficients is None:
            self.logger.log(f"cam{camera_refernece.index} missing intrinsics, skipping extrinsics.")
            return None, None, None, None, None, None, None
        if camera_to_calibrate.camera_matrix is None or camera_to_calibrate.distortion_coefficients is None:
            self.logger.log(f"cam{camera_to_calibrate.index} missing intrinsics, skipping extrinsics.")
            return None, None, None, None, None, None, None

        object_points_per_frame, frame_points_per_frame_reference, frame_points_per_frame_to_calibrate = self._build_stereo_correspondences_from_samples(
            sample_list=sample_list,
            board=board,
            camera_index_reference=camera_refernece.index,
            camera_index_to_calibrate=camera_to_calibrate.index
        )

        if len(object_points_per_frame) == 0:
            self.logger.log(f"cam{camera_refernece.index} <-> cam{camera_to_calibrate.index}: no overlapping calibration samples, skipping extrinsics.")
            return None, None, None, None, None, None, None

        frame_width, frame_height = FRAME_SIZE

        # Keep intrinsics fixed (recommended since you already calibrated them)
        flags = cv2.CALIB_FIX_INTRINSIC

        # Termination criteria for stereo calibration optimizer
        criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 1e-6)

        stereo_rms_reprojection_error_px, camera_matrix_reference, dist_reference, camera_matrix_to_calibrate, dist_to_calibrate, \
            rotation_matrix_reference_to_calibrate, translation_vector_reference_to_calibrate, \
            essential_matrix, fundamental_matrix = cv2.stereoCalibrate(
            objectPoints=object_points_per_frame,
            imagePoints1=frame_points_per_frame_reference,
            imagePoints2=frame_points_per_frame_to_calibrate,
            cameraMatrix1=camera_refernece.camera_matrix,
            distCoeffs1=camera_refernece.distortion_coefficients,
            cameraMatrix2=camera_to_calibrate.camera_matrix,
            distCoeffs2=camera_to_calibrate.distortion_coefficients,
            imageSize=(frame_width, frame_height),
            criteria=criteria,
            flags=flags
        )

        projection_matrix_camA = camera_refernece.camera_matrix @ numpy.hstack([numpy.eye(3, dtype=numpy.float64), numpy.zeros((3,1), dtype=numpy.float64)])
        projection_matrix_camB = camera_to_calibrate.camera_matrix @ numpy.hstack([rotation_matrix_reference_to_calibrate, translation_vector_reference_to_calibrate])



        return (
            float(stereo_rms_reprojection_error_px),
            rotation_matrix_reference_to_calibrate,
            translation_vector_reference_to_calibrate,
            essential_matrix,
            fundamental_matrix,
            projection_matrix_camA,
            projection_matrix_camB,
        )

    def _clear_camera_calibrations(self):
        for camera in self.cameraList:
            camera.camera_matrix = None
            camera.distortion_coefficients = None
            camera.rotation_matrix_world_to_camera = None
            camera.translation_vector_world_to_camera = None

    def _calibrate_intrinsics_characuo_for_camera(
            self,
            sample_list: CalibrationSample,
            camera_index: int,
            board
    ):
        object_points_per_frame  = []
        frame_points_per_frame  = []

        for sample in sample_list:
            detection = sample.detections.get(camera_index)

            if detection is None:
                continue

            if detection.charuco_ids is None or detection.charuco_corners is None:
                continue

            if len(detection.charuco_ids) < MIN_CHARUCO_CORNERS:
                continue

            # map detected 2D corners + ids to corresponding 3D board points
            object_points, frame_points = board.matchImagePoints(detection.charuco_corners, detection.charuco_ids)

            if object_points is None or frame_points is None:
                continue

            object_points_per_frame.append(object_points)
            frame_points_per_frame.append(frame_points)

        if len(object_points_per_frame) == 0:
            self.logger.log(f"cam{camera_index}: no valid calibration samples, skipping intrinsics.")
            return None, None, None


        frame_width, frame_height = FRAME_SIZE
        self.logger.log(f"Calibrating cam{camera_index} with imageSize={FRAME_SIZE} (w,h)")
        self.logger.log(f"H= {frame_width} W= {frame_height}")

        reprojection_error_rms, camera_matrix, distortion_coefficients, rvecs, tvecs = cv2.calibrateCamera(
            objectPoints=object_points_per_frame,
            imagePoints=frame_points_per_frame,
            imageSize=(frame_width, frame_height),
            cameraMatrix=None,
            distCoeffs=None
        )

        return float(reprojection_error_rms), camera_matrix, distortion_coefficients

    def _caputre_calibration_sample(self):
        detections = {}

        for camera in self.cameraList: # Loop all cameras - expects them already connected via "Connect cameras"
            if camera.capture is None:
                self.logger.log(f"cam{camera.index}: not connected, click 'Connect cameras' first. Skipping.")
                continue

            ok, frame = camera.capture.read() #Capture a frame

            if not ok: # A single failed grab shouldn't discard what other cameras saw
                self.logger.log(f"Could not grab frame for calibration sample, from camera {camera.index}.")
                continue

            h, w = frame.shape[:2]
            self.logger.log(f"cam{camera.index} capture size = {w}x{h}")

            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) # Turn the frame black and white
            gray = 255 - gray # Invert the color of the frame, as the ChArCuo board I use is inverted (to save printer Ink)

            #This is apparently a fast check, bevor i detect the actuall board
            corner, ids, _ = self.arcuo_detector.detectMarkers(gray) #Detect ChArCuo Markers
            if ids is None or len(ids) == 0: #Board not visible from this camera this time, skip it, not the whole sample
                self.logger.log(f"Not enought markers found for calibration sample, from camera {camera.index}.")
                continue

            # Detecting the Charcuo board and evaluating if its good as a sample
            result = self.charuco_detector.detectBoard(gray)

            charuco_corners = result[0]
            charuco_ids = result[1]

            if charuco_ids is None:
                self.logger.log(f"Not enought markers found for calibration sample, from camera {camera.index}.")
                continue

            if len(charuco_ids) < MIN_CHARUCO_CORNERS:
                self.logger.log(f"Not enought markers found for calibration sample, from camera {camera.index}.")
                continue

            # Adding the sample for the current camera to detections
            detections[camera.index] = CharucoDetection(
                camera_id = camera.index,
                charuco_corners = charuco_corners,
                charuco_ids=charuco_ids,
                gray_frame = gray
            )

        # Need at least two cameras' worth of detections for this sample to be
        # useful for stereo pairing (a single-camera detection can still help
        # that camera's intrinsics, but can't contribute to any camera pair).
        if len(detections) < 2:
            self.logger.log("Not enough cameras saw the board in this sample, discarding.")
            return

        self.calibration_samples.append(
            CalibrationSample(detections=detections)
        )

        self.logger.log(f"Added calibration sample ({len(detections)}/{len(self.cameraList)} cameras saw the board).")

    def updateCameras(self):
        self._clearLayout()

        for camera_Object in self.cameraList:
            camera_Setup_Widget = CameraSetupWidget(camera_Object)

            self.scroll_Layout.addWidget(camera_Setup_Widget)
            self.camera_Setup_Widget_List.append(camera_Setup_Widget)
    
    def _clearLayout(self):
        self.camera_Setup_Widget_List.clear()
        
        while self.scroll_Layout.count():
            item = self.scroll_Layout.takeAt(0)
            widget = item.widget()
            if widget is not None:
                widget.setParent(None)
                widget.deleteLater()

