"""
File: main_qt.py
Description:    stores custom widget class that builds Qt app for windows displaying 
                landmark and geom coordinates separately that have points that can be
                selected and mapped together
"""

# imports
import sys
import random
import json
import yaml # type: ignore

from PySide6 import QtCore, QtWidgets, QtGui # type: ignore # QT GUI framework
from .both_windows import get_both_canvases, mapping_canvas, event_handler_scale # returns two interactive 3D canvases
from wgpu.gui.qt import run # type: ignore # integrates QT event loop with WGPU rendering
import numpy as np
import pygfx as gfx # type: ignore
from tps import ThinPlateSpline # type: ignore

class MyWidget(QtWidgets.QWidget): # main widget class holding all UI components
    """
    Description:    custom widget class derived from QWidget class from Qt widgets
    Parameters:     None
    Returns:         None
    """

    def __init__(self, file1=None, file2=None):
        """
        Description:    intialize custom widget object from custom Qt Widget class
        Parameters:     None
        Returns:         None
        """

        # class body:

        super().__init__() # initializes Q widget

        self.file1 = file1
        self.file2 = file2

        # track clicked points
        self.selected_keypoint = None
        self.selected_geom = None
        self.dropdown_selection = 0
        self.dropdown_mapping_selection = 0
        self.kp_geometry = None
        self.gm_geometry = None

        self.mapping_pairs = []
        self.mapping_pairs_text = []
        self.mapping_pairs_id = []

        self.mapping_canvas = mapping_canvas()

        # QListWidget - on screen list displays mapping
        self.mapping_list = QtWidgets.QListWidget()
        self.mapping_list.installEventFilter(self)

        self.geoms_list = QtWidgets.QListWidget()
        self.keypoints_list = QtWidgets.QListWidget()

        # builds first canvas/widget -> keypoints and second canvas/widget -> geoms
        # select_kepoints/geom -> callbacks triggered on clicks
        self.canvas1, self.canvas2, self.geoms_to_body, self.kp_names = get_both_canvases(self.file1, self.file2, self.handle_selection, self.handle_hover)

        # === Set flexible size policies ===
        for w in [self.canvas1, self.canvas2, self.mapping_canvas["canvas"]]:
            w.setSizePolicy(QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Expanding)

        self.mapping_list.setSizePolicy(QtWidgets.QSizePolicy.Preferred, QtWidgets.QSizePolicy.Expanding)
        self.keypoints_list.setSizePolicy(QtWidgets.QSizePolicy.Preferred, QtWidgets.QSizePolicy.Expanding)
        self.geoms_list.setSizePolicy(QtWidgets.QSizePolicy.Preferred, QtWidgets.QSizePolicy.Expanding)

        self.build_qt_app()

    def handle_selection(self, name, index, typ, geometry, canvas, red_list, purple_list):
        position = geometry.positions.data[index].copy() # get 3d point
        if typ == "keypoint":
            self.selected_keypoint = (name, index, geometry, canvas, position, red_list, purple_list)
            self.kp_geometry = (geometry, canvas, red_list, purple_list)
        elif typ == "geom":
            self.selected_geom = (name, index, geometry, canvas, position, red_list, purple_list)
            self.gm_geometry = (geometry, canvas, red_list, purple_list)
        self.check_and_add_mapping()
    
    def handle_hover(self, name, typ):
        if typ == "keypoint":
            if self.keypoints_list.count() < 1:
                self.keypoints_list.addItem(name)
            else:
                self.keypoints_list.item(0).setText(name)
        elif typ == 'geom':
            if self.geoms_list.count() < 1:
                self.geoms_list.addItem(name)
            else:
                self.geoms_list.item(0).setText(name)
            

    def check_and_add_mapping(self):
        """
        Description:    maps a selected keypoint/geom pair then adds it to the mapping list widget
        Parameters:     self - class/object instance with attributes
        Returns:        None
        """

        # function body:

        if self.selected_keypoint and self.selected_geom:
            name_kp, idx_kp, geom_kp, canvas_kp, pos_kp, reds_kp, purples_kp = self.selected_keypoint
            name_gm, idx_gm, geom_gm, canvas_gm, pos_geom, reds_geom, purples_geom = self.selected_geom

            # Add mapping to list
            mapping_text = f"{name_kp} <--> {name_gm} (body: {self.geoms_to_body[name_gm]})"
            self.mapping_list.addItem(mapping_text)

            # Save mapping pair
            self.mapping_pairs.append((pos_kp, pos_geom))
            self.mapping_pairs_id.append((idx_kp, idx_gm))
            self.mapping_pairs_text.append((name_kp, name_gm))

            # Update the visualization with the full set
            self.update_mapping_visualization(geom_kp, geom_gm)

            def revert_colors():
                geom_kp.colors.data[idx_kp] = [1, 0, 0, 1] # changes to color red
                geom_gm.colors.data[idx_gm] = [1, 0, 0, 1]
                geom_kp.colors.update_range(idx_kp)
                geom_gm.colors.update_range(idx_gm)
                canvas_kp.request_draw()
                canvas_gm.request_draw()

                if idx_kp in purples_kp:
                    purples_kp.remove(idx_kp)
                if idx_gm in purples_geom:
                    purples_geom.remove(idx_gm)
                reds_kp.append(idx_kp)
                reds_geom.append(idx_gm)

            QtCore.QTimer.singleShot(750, revert_colors)

            self.selected_keypoint = None
            self.selected_geom = None
        
        elif self.selected_keypoint:
            idx_kp, geom_kp = self.selected_keypoint[1:3]
            if all(x == y for x, y in zip(geom_kp.colors.data[idx_kp], [0, 1, 1, 1])):
                self.selected_keypoint = None
            elif all(x == y for x, y in zip(geom_kp.colors.data[idx_kp], [1, 0, 0, 1])):
                self.selected_keypoint = None
            
        elif self.selected_geom:
            idx_gm, geom_gm = self.selected_geom[1:3]
            if all(x == y for x, y in zip(geom_gm.colors.data[idx_gm], [0, 1, 1, 1])):
                self.selected_geom = None
            elif all(x == y for x, y in zip(geom_gm.colors.data[idx_gm], [1, 0, 0, 1])):
                self.selected_geom = None

    def update_mapping_visualization(self, geom_kp, geom_gm):
        scene = self.mapping_canvas
        point_obj = scene["point_obj"]
        line_obj = scene["line_obj"]
        grid = scene['grid']

        Y_aligned = None
        X_centered = None

        if self.dropdown_selection == 0:
            Y_aligned, X_centered = self.apply_rigid_transformation(geom_kp, geom_gm)
        elif self.dropdown_selection == 1:
            Y_aligned, X_centered = self.apply_transformation(geom_kp, geom_gm)[0:2]

        all_points = []
        all_colors = []
        all_line_segments = []
        line_segment_ids = []
        index = -1
        kp_list = []
        gm_list = []

        for id_kp, id_geom in self.mapping_pairs_id:
            kp = X_centered[id_kp]
            geom = Y_aligned[id_geom]

            kp_list.append(id_kp)
            gm_list.append(id_geom)

            all_points.extend([kp, geom])
            all_colors.extend([[1, 0, 0, 1], [1, 0, 0, 1]])

            all_line_segments.append(kp)
            all_line_segments.append(geom)

            line_segment_ids.append(index + 1)
            line_segment_ids.append(index + 2)
            index += 2

            all_line_segments.append([np.nan, np.nan, np.nan])  # break line
        
        for i in range(len(X_centered)):
            kp = X_centered[i]
            if not (i in kp_list):
                all_points.append(kp)
                all_colors.append([0.8, 0, 1, 1])
        for i in range(len(Y_aligned)):
            gm = Y_aligned[i]
            if not (i in gm_list):
                all_points.append(gm)
                all_colors.append([0, 1, 0, 1])
        

        all_points = np.array(all_points, dtype=np.float32)
        center = all_points.mean(axis=0)

        all_colors = np.array(all_colors, dtype=np.float32)
        all_line_segments = np.array(all_line_segments, dtype=np.float32)

        point_obj.geometry = gfx.Geometry(
            positions=gfx.Buffer(all_points),
            colors=gfx.Buffer(all_colors)
        )
        line_obj.geometry = gfx.Geometry(
            positions=gfx.Buffer(all_line_segments)
        )

        # Remove NaN rows before computing min_y
        valid_points = all_points[~np.isnan(all_points).any(axis=1)]
        if len(valid_points) > 0:
            # Place grid just below the lowest point
            min_y = np.min(valid_points[:, 1])
            grid.local.y = min_y - 2
        
        scene['camera'].show_object(point_obj)

        event_handler_scale(all_points, scene['canvas'], scene['camera'], center, 
                            point_obj.geometry, point_obj.material, scene['renderer'], 
                            scene['scene'], line_obj, line_segment_ids
                            )

    def procrustes_transformation(self):
        if len(self.mapping_pairs) < 3:
            return 'NA', 'NA'
        
        kp_subset = []
        geom_subset = []
        for kp, geom in self.mapping_pairs:
            kp_subset.append(kp)
            geom_subset.append(geom)

        X = np.array(kp_subset)
        Y = np.array(geom_subset)

        X_centered = X - X.mean(axis=0)
        Y_centered = Y - Y.mean(axis=0)

        M = Y_centered.T @ X_centered

        U, S, Vt = np.linalg.svd(M)

        Q = U @ Vt

        k = np.sum(S) / np.trace(Y_centered.T @ Y_centered)

        return k, Q
    
    def apply_rigid_transformation(self, geom_kp, geom_gm):
        k, Q = self.procrustes_transformation()

        X = np.array(geom_kp.positions.data.copy())
        Y = np.array(geom_gm.positions.data.copy())

        if k == 'NA' and Q == 'NA':
            return Y, X

        Y_aligned = k * Y @ Q + X.mean(axis=0)

        return Y_aligned, X  # Return original X for direct comparison

    def apply_transformation(self, geom_kp, geom_gm):
        if len(self.mapping_pairs) < 3:
            return np.array(geom_gm.positions.data.copy()), np.array(geom_kp.positions.data.copy()), None

        X = np.array([kp for kp, _ in self.mapping_pairs])  # Keypoints (targets)
        Y = np.array([gm for _, gm in self.mapping_pairs])  # Geoms (sources)

        Y_full = np.array(geom_gm.positions.data.copy())

        tps_model = ThinPlateSpline()
        tps_model.fit(Y, X)

        # Only extract attributes that are not methods or private
        attrs = {
            k: v for k, v in vars(tps_model).items()
            if not k.startswith('_') and not callable(v)
        }

        TY = tps_model.transform(Y_full)

        return TY, np.array(geom_kp.positions.data.copy()), attrs
    
    def build_qt_app(self):

        # === First column: canvas1 and canvas2 stacked vertically ===
        left_splitter = QtWidgets.QSplitter(QtCore.Qt.Vertical)
        kp_label = QtWidgets.QLabel("Motion Capture Landmarks")
        gm_label = QtWidgets.QLabel("Virtual Body Model")
        left_splitter.addWidget(kp_label)
        left_splitter.addWidget(self.canvas1)
        left_splitter.addWidget(gm_label)
        left_splitter.addWidget(self.canvas2)
        #left_splitter.setSizes([0.1, 0.4, 0.1, 0.4])  # equal height

        # === Second column: mapping list ===
        mapping_list_layout = QtWidgets.QVBoxLayout()
        mapping_list_container = QtWidgets.QWidget()
        mapping_list_container.setLayout(mapping_list_layout)

        # Motion Capture Landmarks
        mapping_list_layout.addWidget(QtWidgets.QLabel("Motion Capture Landmark"))
        self.keypoints_list.setFixedHeight(25)
        mapping_list_layout.addWidget(self.keypoints_list)

        # Virtual Body Geoms
        mapping_list_layout.addWidget(QtWidgets.QLabel("Virtual Body Geom"))
        self.geoms_list.setFixedHeight(25)
        mapping_list_layout.addWidget(self.geoms_list)

        # Correspondences Title
        mapping_list_layout.addWidget(QtWidgets.QLabel("Correspondences"))
        # Correspondences list (expand to fill)
        mapping_list_layout.addWidget(self.mapping_list, stretch=1)

        button = QtWidgets.QPushButton("Export Correspondences")
        mapping_list_layout.addWidget(button)
        button.clicked.connect(self.export_correspondences)

        # dropdown to chose geoms or body for mapping to keypoints
        correspondence_dropdown = QtWidgets.QComboBox()
        correspondence_dropdown.addItems(["Map to Geoms", "Map to Body"])
        mapping_list_layout.addWidget(correspondence_dropdown)
        correspondence_dropdown.currentIndexChanged.connect(self.selected_mapping_option)

        # === Third column: mapping canvas ===
        mapping_canvas_container = QtWidgets.QWidget()
        mapping_canvas_layout = QtWidgets.QVBoxLayout(mapping_canvas_container)

        mapping_canvas_widget = self.mapping_canvas["canvas"]
        mapping_canvas_layout.addWidget(mapping_canvas_widget)

        # Legend
        legend_widget = QtWidgets.QWidget()
        legend_layout = QtWidgets.QHBoxLayout(legend_widget)
        legend_layout.setContentsMargins(0, 0, 0, 0)
        legend_layout.setSpacing(8)

        # Purple dot = motion capture keypoint
        purple_label = QtWidgets.QLabel("Keypoint")
        purple_color = QtWidgets.QLabel()
        purple_color.setFixedSize(12, 12)
        purple_color.setStyleSheet("background-color: rgb(204, 0, 255); border: 1px solid #333;")
        legend_layout.addWidget(purple_color)
        legend_layout.addWidget(purple_label)

        # Green dot = virtual rodent geom
        green_label = QtWidgets.QLabel("Geom")
        green_color = QtWidgets.QLabel()
        green_color.setFixedSize(12, 12)
        green_color.setStyleSheet("background-color: rgb(0, 255, 0); border: 1px solid #333;")
        legend_layout.addWidget(green_color)
        legend_layout.addWidget(green_label)

        # Dropdown box to select rigid or non rigid transformation
        dropdown = QtWidgets.QComboBox()
        dropdown.addItems(["Rigid Transformation", 'Non-rigid Transformation'])
        legend_layout.addWidget(dropdown)
        dropdown.currentIndexChanged.connect(self.selected_transformation)

        # button to export current transformation parameters
        transformation_button = QtWidgets.QPushButton("Export Transformation Parameters")
        legend_layout.addWidget(transformation_button)
        transformation_button.clicked.connect(self.export_transformation_params)

        # Add to layout
        mapping_canvas_layout.addWidget(legend_widget)

        # === Split main horizontally across 3 regions ===
        main_splitter = QtWidgets.QSplitter(QtCore.Qt.Horizontal)
        main_splitter.addWidget(left_splitter)
        main_splitter.addWidget(mapping_list_container)
        main_splitter.addWidget(mapping_canvas_container)

        # Set initial pixel sizes for each splitter panel
        total_width = QtWidgets.QApplication.primaryScreen().availableGeometry().width()
        main_splitter.setSizes([
            int(total_width * 0.45),  # left column (canvas1 + canvas2 stacked)
            int(total_width * 0.15),  # middle column (mapping list)
            int(total_width * 0.40)   # right column (mapping canvas)
        ])

        # === Final layout ===
        main_layout = QtWidgets.QHBoxLayout()
        main_layout.addWidget(main_splitter)
        self.setLayout(main_layout)

        # Trigger initial draw
        self.canvas1.draw_frame()
        self.canvas2.draw_frame()
    
    def export_correspondences(self):
        kp_to_gm = {}
        kp_names = []

        write_file = 'kp_to_gm_correspondences.yml'
        if self.dropdown_mapping_selection == 1:
            write_file = "kp_to_bdy_correspondences.yml"
        
        for kp, gm in self.mapping_pairs_text:
            if self.dropdown_mapping_selection == 0:
                kp_to_gm[kp] = gm
            else:
                kp_to_gm[kp] = self.geoms_to_body[gm]

        for each in self.kp_names:
            kp_names.append(str(each))

        with open(write_file, 'w') as file:
            yaml.dump({"KP_NAMES":kp_names, "KEYPOINT_MODEL_PAIRS":kp_to_gm}, file, default_flow_style=False)

    def selected_transformation(self, index):
        # rigid transformation index = 0
        # non rigid transformation index = 1
        self.dropdown_selection = index
        if self.gm_geometry[0] and self.kp_geometry[0]:
            self.update_mapping_visualization(self.kp_geometry[0], self.gm_geometry[0])
    
    def selected_mapping_option(self, index):
        self.dropdown_mapping_selection = index
    
    def eventFilter(self, source, event):
        if source == self.mapping_list and event.type() == QtCore.QEvent.KeyPress:
            if event.key() == QtCore.Qt.Key_Backspace:
                for item in self.mapping_list.selectedItems():
                    remove_item_index = self.mapping_list.row(item)
                    kp_id, gm_id = self.mapping_pairs_id[remove_item_index]

                    geom_kp, canvas_kp, reds_kp, purples_kp = self.kp_geometry
                    geom_gm, canvas_gm, reds_gm, purples_gm = self.gm_geometry

                    geom_kp.colors.data[kp_id] = [0, 1, 1, 1] # changes to color cyan
                    geom_gm.colors.data[gm_id] = [0, 1, 1, 1]
                    geom_kp.colors.update_range(kp_id)
                    geom_gm.colors.update_range(gm_id )
                    canvas_kp.request_draw()
                    canvas_gm.request_draw()

                    if kp_id in reds_kp:
                        reds_kp.remove(kp_id)
                    if kp_id in purples_kp:
                        purples_kp.remove(kp_id)

                    if gm_id in reds_gm:
                        reds_gm.remove(gm_id)
                    if gm_id in purples_gm:
                        purples_gm.remove(gm_id)

                    del self.mapping_pairs[remove_item_index]
                    del self.mapping_pairs_text[remove_item_index]
                    del self.mapping_pairs_id[remove_item_index]

                    self.mapping_list.takeItem(remove_item_index)
                    if len(self.mapping_pairs) >= 1:
                        self.update_mapping_visualization(geom_kp, geom_gm)
                return True
        return super().eventFilter(source, event)
    
    def export_transformation_params(self):
        if self.dropdown_selection == 0:
            k, Q = self.procrustes_transformation()
            rigid_data = {'scale': float(k), 'rotation': Q.tolist()}

            with open('rigid_transf_params.json', 'w') as file:
                json.dump(rigid_data, file)
        
        elif self.dropdown_selection == 1:
            tps_config = self.apply_transformation(self.kp_geometry[0], self.gm_geometry[0])[2]

            with open("tps_config.json", 'w') as file:
                json.dump({key: value.tolist() 
                           if isinstance(value, np.ndarray) else value for key, value in tps_config.items()
                           }, file)

def run_gui(file1: str, file2: str):
    app = QtWidgets.QApplication([])

    widget = MyWidget(file1, file2)

    # Dynamically size to screen
    screen = QtWidgets.QApplication.primaryScreen()
    available = screen.availableGeometry()

    width = int(available.width() * 0.95)
    height = int(available.height() * 0.8)
    widget.resize(width, height)

    # Center on screen
    widget.move(
        available.left() + (available.width() - width) // 2,
        available.top() + (available.height() - height) // 2
    )

    # Limit max size to avoid overflow
    widget.setMaximumSize(available.width(), available.height())

    widget.show()
    run()