"""
File: both_windows.py
Description:    creates two canvases both having point clouds with the first one coming from
                motion capture landmarks and the second coming from mujoco virtual rodent body
                geoms
"""

# imports
import numpy as np
from wgpu.gui.qt import QWgpuCanvas # type: ignore
import pygfx as gfx # type: ignore
from .parse_data import get_both_coordinates

def plotting_keypoints(positions, kp_names, on_select=None, on_hover=None):
    """
    Description:    create canvas that shows point cloud for 3d motion landmark coordinates
                    points can be clicked and mapped to virtual rodent mujoco model geoms
    Parameters:     positions - list of coordinates list for each 3d landmark
                    kp_names - corresponding landmark names of each coordinate 
    """
    # # Flip the Y axis (invert vertical axis)
    # positions[:, 1] *= -1

    # # # flip the x axis as well (apply -1 scalar to x coordinates)

    # positions[:, 0] *= -1

    positions = positions[:, [0, 2, 1]]
    positions[:, 2] *= -1
    

    center = positions.mean(axis=0) # get center of points to set up camera orientation
    canvas = QWgpuCanvas(title="Motion Capture Landmarks") # intialize 3d canvas object
    renderer = gfx.renderers.WgpuRenderer(canvas, blend_mode="weighted") # intialize gpu renderer
    scene = gfx.Scene() # initialize scene and background color
    scene.add(gfx.Background.from_color("#000"))

    kp_canvas = widget_builder(scene, center, positions, renderer, canvas, "keypoint", kp_names, on_select, on_hover)

    return kp_canvas

def plotting_geoms(positions, geom_names, on_select=None, on_hover=None):
    """
    Description:    return canvas that contains mujococ virtual body geoms plotted as 3d points
    Parameters:     positions - 3d coordinates of the geom
                    geom_names - list of names corresponding to geom positions
                    on_select - callback stored for when points are pressed
    Returns:        Canvas for geoms to be used in Qt app
    """

    positions = positions * 500 # positions are scaled up since the point values were small

    # set up to rotate points across the z-axis
    theta = np.radians(-90)
    rotation_matrix_x = np.array([
        [1, 0, 0],
        [0, np.cos(theta), -np.sin(theta)],
        [0, np.sin(theta),  np.cos(theta)],
    ], dtype=np.float32)

    # Apply rotation
    positions = positions @ rotation_matrix_x.T

    # Compute center of the points
    center = positions.mean(axis=0)

    canvas = QWgpuCanvas(title="Virtual Rodent Skeleton Points")
    renderer = gfx.renderers.WgpuRenderer(canvas, blend_mode="weighted")
    scene = gfx.Scene()
    scene.add(gfx.Background.from_color("#000"))

    geom_canvas = widget_builder(scene, center, positions, renderer, canvas, "geom", geom_names, on_select, on_hover)

    return geom_canvas


def mapping_canvas():
    canvas = QWgpuCanvas(title="Correspondences Visualizer")
    renderer = gfx.renderers.WgpuRenderer(canvas, blend_mode="weighted")
    scene = gfx.Scene()
    scene.add(gfx.Background.from_color("#000"))

    # Initialize with empty point/line geometry at origin
    positions = np.zeros((1, 3), dtype=np.float32)
    colors = np.array([[1, 0, 0, 1]], dtype=np.float32)

    point_geom = gfx.Geometry(positions=positions, colors=colors)
    point_mat = gfx.PointsMaterial(color_mode='vertex', size=12)
    point_obj = gfx.Points(point_geom, point_mat)
    scene.add(point_obj)

    line_geom = gfx.Geometry(positions=positions.copy())
    line_mat = gfx.LineMaterial(thickness=3.0)
    line_obj = gfx.Line(line_geom, line_mat)
    scene.add(line_obj)

    # Add grid at y=0
    grid = gfx.Grid(
        None,
        gfx.GridMaterial(
            major_step=100,
            minor_step=10,
            thickness_space="world",
            major_thickness=2,
            minor_thickness=0.1,
            infinite=True,
        ),
        orientation="xz",
    )
    grid.local.y = 0  # keep grid on ground plane
    scene.add(grid)

    camera = gfx.PerspectiveCamera(70, 16 / 9)
    controller = gfx.OrbitController(camera)
    controller.register_events(renderer)

    # Camera should initially look at origin
    camera.show_object((0, 0, 0, 500))  # use grid to initialize framing
    camera.look_at((0, 0, 0))

    def draw():
        renderer.render(scene, camera)
    canvas.request_draw(draw)

    return {
        "canvas": canvas,
        "scene": scene,
        "point_obj": point_obj,
        "line_obj": line_obj,
        "renderer": renderer,
        "camera": camera,
        "grid": grid
    }

def widget_builder(scene, center, positions, renderer, canvas, type, names, on_select, on_hover):
   # initialize a grid to be applied to the scene; grid will be on the x, z axis
    grid = gfx.Grid(
        None,
        gfx.GridMaterial(
            major_step=100,
            minor_step=10,
            thickness_space="world",
            major_thickness=2,
            minor_thickness=0.1,
            infinite=True,
        ),
        orientation="xz",
    )

    # set up the grid in terms of y-axis (height of grid) and add to scene
    # grid.local.y = scalar * (center[1])
    # scene.add(grid)

    # Place grid just beneath lowest point
    min_y = np.min(positions[:, 1])
    grid.local.y = min_y - 2
    scene.add(grid)

    # Set initial colors (cyan)
    colors = np.ones((len(positions), 4), dtype=np.float32)
    colors[:, :3] = [0, 1, 1]  # RGB = cyan

    # set up Geometry object for all points and appy cyan color to all
    geometry = gfx.Geometry(positions=positions, colors=colors)
    # Enable per-vertex color material and enable picking
    material = gfx.PointsMaterial(color_mode="vertex", size=10, pick_write=True)
    ob = gfx.Points(geometry, material) # create the points initializing Points object
    scene.add(ob)

    # # Draw line between points to create skeleton
    # line_segments = []
    # for i in range(len(positions) - 1):
    #     line_segments.extend([positions[i], positions[i+1], [np.nan, np.nan, np.nan]])
    # line_segments = np.array(line_segments, dtype=np.float32)

    # line_geom = gfx.Geometry(positions=gfx.Buffer(line_segments))
    # line_mat = gfx.LineMaterial(thickness=3.0)
    # line_obj = gfx.Line(line_geom, line_mat)
    # scene.add(line_obj)

    # Set camera based on data
    camera = gfx.PerspectiveCamera(50, 16 / 9)
    camera.show_object(ob)
    camera.look_at(center)

    # set up controller to be able to Orbit and view from many directions
    controller = gfx.OrbitController(camera)
    controller.register_events(renderer)

    red_list = []
    green_list = []
    purple_list = []

    @ob.add_event_handler("pointer_down")
    def offset_point(event):
        # function handling clicking on points
        event_handler_function(event, geometry, canvas, on_select, names, type, red_list, purple_list)
    
    @ob.add_event_handler("pointer_move")
    def hovering_over_point(event):
        # function handling hovering over points
        event_handler_hover(event, on_hover, names, type, geometry, canvas, green_list, red_list, purple_list)

    event_handler_scale(positions, canvas, camera, center, geometry, material, renderer, scene)
    
    return canvas

def event_handler_function(event, geometry, canvas, on_select, names, type, red_list, purple_list):
    info = event.pick_info # grab info
    if "vertex_index" in info: # if vertex index exists for point
        i = int(info["vertex_index"]) # grab the actual index

        non_cyan_colors = []
        cyan = [0, 1, 1, 1]
        red = [1, 0, 0, 1]
        purple = [0.5, 0, 0.5, 1]

        for k in range(len(geometry.colors.data)):
            for j in range(len(geometry.colors.data[k])):
                color_check = int(geometry.colors.data[k][j])
                if color_check != cyan[j] and k != i:
                    non_cyan_colors.append(k)

        for index in non_cyan_colors:
            if all(x == y for x, y in zip(geometry.colors.data[index], purple)):
                if index in red_list:
                    geometry.colors.data[index] = red
                else:
                    geometry.colors.data[index] = cyan
                purple_list.remove(index)
                
                geometry.colors.update_range(index)
                canvas.request_draw()
        
        if all(x == y for x, y in zip(geometry.colors.data[i], purple)): # if point color is purple switch to cyan or red if already mapped
            if i in red_list:
                geometry.colors.data[i] = red
            else:
                geometry.colors.data[i] = cyan  # RGBA = CYAN
            purple_list.remove(i)
        elif all(x == y for x, y in zip(geometry.colors.data[i], red)):
            red_list.append(i)
            purple_list.append(i)
            geometry.colors.data[i] = purple
        else:
            purple_list.append(i)
            geometry.colors.data[i] = purple  # RGBA = purple

        on_select(names[i], i, type, geometry, canvas, red_list, purple_list) # callback
        geometry.colors.update_range(i)
        canvas.request_draw()

def event_handler_hover(event, on_hover, names, type, geometry, canvas, green_list, red_list, purple_list):
    info = event.pick_info
    green = [0, 1, 0, 1]
    cyan = [0, 1, 1, 1]
    red = [1, 0, 0, 1]
    if "vertex_index" in info:
        i = int(info['vertex_index'])
        # process for grabbing name of point based on the index of the point
        if on_hover and (not (i in purple_list)):
            if len(green_list) == 1:
                index = green_list[0]
                if not (index in purple_list):
                    if index in red_list:
                        geometry.colors.data[index] = red
                    else:
                        geometry.colors.data[index] = cyan
                    geometry.colors.update_range(index)
                    canvas.request_draw()
                green_list.remove(index)
            geometry.colors.data[i] = green
            green_list.append(i)
            on_hover(names[i], type)
            geometry.colors.update_range(i)
            canvas.request_draw()

def event_handler_scale(positions, canvas, camera, center, geometry, material, renderer, scene, line_obj=None, segment_ids=None):
    accumulated_zoom = 0
    pending_zoom = False
    scale_factor = {"scale": 1.0}
    original_positions = positions.copy()

    @canvas.add_event_handler("wheel")
    def on_wheel(event):
        nonlocal accumulated_zoom, pending_zoom

        delta = event['dy']
        if delta == 0:
            return

        # Accumulate zoom for smoother control
        accumulated_zoom += delta
        pending_zoom = True

    # Draw function
    def draw():
        nonlocal pending_zoom, accumulated_zoom

        if pending_zoom:
            zoom_steps = abs(accumulated_zoom)

            # Camera vector & distances
            camera_pos = camera.world.position
            cam_to_center = center - camera_pos
            cam_to_center /= np.linalg.norm(cam_to_center)
            dist_to_center = np.linalg.norm(center - camera_pos)

            # Bounding sphere radius + margin
            max_extent = np.max(np.linalg.norm(original_positions - center, axis=1))
            margin = 1.0250
            min_visible_dist = max_extent * margin / np.tan(np.radians(camera.fov / 2))

            for _ in range(int(zoom_steps)):
                if accumulated_zoom > 0 and dist_to_center >= min_visible_dist:
                    scale_factor["scale"] *= 1.00075
                    #scale_factor["scale"] = min(25.0, scale_factor["scale"])
                elif accumulated_zoom < 0:
                    scale_factor["scale"] /= 1.0007
                    scale_factor["scale"] = max(1, scale_factor["scale"])

            # Apply scale to geometry
            new_positions = (original_positions - center) * scale_factor["scale"] + center
            geometry.positions.data[:] = new_positions
            geometry.positions.update_range(0, len(new_positions))

            # Update material size smoothly and clamped
            material.size = np.clip(9 * scale_factor["scale"], 10, 12)

            if line_obj and segment_ids:
                all_line_segments = []

                for i in range(0, len(segment_ids), 2):
                    kp = geometry.positions.data[segment_ids[i]]
                    geom = geometry.positions.data[segment_ids[i+1]]

                    all_line_segments.append(kp)
                    all_line_segments.append(geom)
                    all_line_segments.append([np.nan, np.nan, np.nan])  # break line
                
                all_line_segments = np.array(all_line_segments, dtype=np.float32)
                line_obj.geometry = gfx.Geometry(
                    positions=gfx.Buffer(all_line_segments)
                )

            canvas.request_draw()
            accumulated_zoom = 0
            pending_zoom = False

        renderer.render(scene, camera)

    canvas.request_draw(draw)

def get_both_canvases(file1, file2, on_select, on_hover):
    """
    Description:    combines two 3D point plotting functions into a single interface
                    retrieves both canvases to be displayed in widget app

    Parameters:     on_select_kepoint - callback when keypoints are selected in the first window
                    on_select_geom - callback when geoms are selected in second window

    Returns:        canvas1 - first window/widget/canvas for 3D motion capture keypoints
                    canvas2 - second window/widget/canvas for virtual body geoms from mujoco
    """

    # function body: 

    positions, model, kp_names, geom_names, geoms_to_body = get_both_coordinates(file1, file2) # retrieve positions data and their respective names

    canvas1 = plotting_keypoints(positions, kp_names, on_select=on_select, on_hover=on_hover) # grab canvas from pygfx function for keypoints

    canvas2 = plotting_geoms(model, geom_names, on_select=on_select, on_hover=on_hover) # grab canvas from pygfx function for geoms

    return canvas1, canvas2, geoms_to_body, kp_names