"""
Utility script to visualize camera extrinsics from a JSON file using Plotly.
"""

import json
import click
import numpy as np
import plotly.graph_objects as go
from typing import Any, Dict, Optional, List
import configparser
from pathlib import Path
import re


RESOLUTION_MAP = {
    "FHD1200": "FHD1200",
    "FHD": "FHD",
    "2K": "2K",
    "HD": "HD",
    "SVGA": "SVGA",
    "VGA": "VGA",
}


def parse_pose(pose_str: str) -> np.ndarray:
    """Parses a 16-float pose string into a 4x4 matrix."""
    try:
        vals = [float(x) for x in pose_str.split()]
        if len(vals) != 16:
            raise ValueError(f"Expected 16 values, got {len(vals)}")
        return np.array(vals).reshape((4, 4))
    except Exception as e:
        raise ValueError(f"Failed to parse pose string: {e}")


def world_to_plot(points: np.ndarray, world_basis: str = "cv") -> np.ndarray:
    """
    Transforms world-space points to plot-space.
    'cv' basis: +X right, +Y down, +Z forward (no-op).
    'opengl' basis: +X right, +Y up, +Z backward.

    Args:
        points: (N, 3) array of points in world coordinates.
        world_basis: 'cv' or 'opengl'.

    Returns:
        (N, 3) array of points.
    """
    if world_basis == "opengl":
        # CV -> OpenGL: Y = -Y, Z = -Z
        pts_plot = points.copy()
        pts_plot[:, 1] *= -1
        pts_plot[:, 2] *= -1
        return pts_plot
    return points


def load_zed_configs(
    paths: List[str], resolution: str, eye: str
) -> Dict[str, Dict[str, float]]:
    """
    Loads ZED intrinsics from config files.
    Returns a mapping from serial (string) to intrinsics dict.
    """
    configs = {}
    eye_prefix = eye.upper()

    # Map resolution to section suffix
    res_map = {
        "1200": "FHD1200",
        "fhd": "FHD",
        "2k": "2K",
        "hd": "HD",
        "svga": "SVGA",
        "vga": "VGA",
    }
    res_suffix = res_map.get(resolution.lower(), resolution.upper())
    section_name = f"{eye_prefix}_CAM_{res_suffix}"

    all_files = []
    for p in paths:
        path = Path(p)
        if path.is_dir():
            all_files.extend(list(path.glob("SN*.conf")))
        else:
            all_files.append(path)

    for f in all_files:
        # Extract serial from filename SN<serial>.conf
        match = re.search(r"SN(\d+)", f.name)
        serial = match.group(1) if match else None

        parser = configparser.ConfigParser()
        try:
            parser.read(f)
            if section_name in parser:
                sect = parser[section_name]
                intrinsics = {
                    "fx": float(sect.get("fx", 0)),
                    "fy": float(sect.get("fy", 0)),
                    "cx": float(sect.get("cx", 0)),
                    "cy": float(sect.get("cy", 0)),
                }
                if serial:
                    configs[serial] = intrinsics

                # Always store as default in case it's the only file
                configs["default"] = intrinsics
        except Exception as e:
            print(f"Warning: Failed to parse config {f}: {e}")

    # If only one config was provided, apply to all
    if len(all_files) == 1 and "default" in configs:
        return {"all": configs["default"]}

    return configs


def get_frustum_points(
    intrinsics: Optional[Dict[str, float]],
    frustum_scale: float,
    fov_deg: float,
) -> np.ndarray:
    """
    Returns 5 points in local camera coordinates: center + 4 corners of the far plane.
    Local coordinates: forward is +Z, right is +X, down is +Y (OpenCV convention).
    """
    if intrinsics and all(k in intrinsics for k in ["fx", "fy", "cx", "cy"]):
        fx, fy = intrinsics["fx"], intrinsics["fy"]
        cx, cy = intrinsics["cx"], intrinsics["cy"]
        # We assume the frustum plane is at Z = frustum_scale
        # x = (u - cx) * Z / fx
        # y = (v - cy) * Z / fy
        # We'll assume a standard aspect ratio and center cx/cy for visualization
        # if we don't have image dimensions.
        # Let's approximate image size from principal point (assuming it's roughly center)
        w_half = (cx / fx) * frustum_scale
        h_half = (cy / fy) * frustum_scale
        w, h = w_half, h_half
    else:
        fov_rad = np.radians(fov_deg)
        # Assuming horizontal FOV
        w = frustum_scale * np.tan(fov_rad / 2.0)
        h = w * 0.75  # 4:3 aspect ratio assumption

    # 5 points: center + 4 corners of the far plane
    # OpenCV: +Z forward, +X right, +Y down
    pts_local = np.array(
        [
            [0, 0, 0],  # Center
            [
                -w,
                -h,
                frustum_scale,
            ],  # Top-Left (if Y down is positive, -h is up) -> Wait.
            # OpenCV: Y is down. So -h is UP in 3D space if we map Y->Y.
            # But usually we want to visualize it.
            # Let's stick to:
            # +X right
            # +Y down
            # +Z forward
            [w, -h, frustum_scale],  # Top-Right
            [w, h, frustum_scale],  # Bottom-Right
            [-w, h, frustum_scale],  # Bottom-Left
        ]
    )
    return pts_local


def add_camera_trace(
    fig: go.Figure,
    pose: np.ndarray,
    label: str,
    scale: float = 0.2,
    frustum_scale: float = 0.5,
    fov_deg: float = 60.0,
    intrinsics: Optional[Dict[str, float]] = None,
    color: str = "blue",
    world_basis: str = "cv",
):
    """
    Adds a camera frustum and axes to the Plotly figure.
    """
    R = pose[:3, :3]
    t = pose[:3, 3]

    # world_from_cam (Standard convention for calibrate_extrinsics.py)
    # calibrate_extrinsics.py inverts the solvePnP result before saving.
    center = t
    R_world = R

    # OpenCV convention: X right, Y down, Z forward
    x_axis_local = np.array([1, 0, 0])
    y_axis_local = np.array([0, 1, 0])
    z_axis_local = np.array([0, 0, 1])

    # Transform local axes to world
    x_axis_world = R_world @ x_axis_local
    y_axis_world = R_world @ y_axis_local
    z_axis_world = R_world @ z_axis_local

    # Frustum points in local coordinates (OpenCV: +Z fwd, +X right, +Y down)
    pts_local = get_frustum_points(intrinsics, frustum_scale, fov_deg)

    # Transform frustum to world
    pts_world = (R_world @ pts_local.T).T + center

    # --- Apply Global Basis Transform ---
    # Transform everything from World Space -> Plot Space
    center_plot = world_to_plot(center[None, :], world_basis=world_basis)[0]

    x_end_world = center + x_axis_world * scale
    y_end_world = center + y_axis_world * scale
    z_end_world = center + z_axis_world * scale

    x_end_plot = world_to_plot(x_end_world[None, :], world_basis=world_basis)[0]
    y_end_plot = world_to_plot(y_end_world[None, :], world_basis=world_basis)[0]
    z_end_plot = world_to_plot(z_end_world[None, :], world_basis=world_basis)[0]

    pts_plot = world_to_plot(pts_world, world_basis=world_basis)

    # Create lines for frustum
    # Edges: 0-1, 0-2, 0-3, 0-4 (pyramid sides)
    #        1-2, 2-3, 3-4, 4-1 (base)
    x_lines = []
    y_lines = []
    z_lines = []

    def add_line(i, j):
        x_lines.extend([pts_plot[i, 0], pts_plot[j, 0], None])
        y_lines.extend([pts_plot[i, 1], pts_plot[j, 1], None])
        z_lines.extend([pts_plot[i, 2], pts_plot[j, 2], None])

    # Pyramid sides
    for i in range(1, 5):
        add_line(0, i)
    # Base
    add_line(1, 2)
    add_line(2, 3)
    add_line(3, 4)
    add_line(4, 1)

    # Add frustum trace
    fig.add_trace(
        go.Scatter3d(
            x=x_lines,
            y=y_lines,
            z=z_lines,
            mode="lines",
            line=dict(color=color, width=2),
            name=f"{label} Frustum",
            showlegend=False,
            hoverinfo="skip",
        )
    )

    # Add center point with label
    fig.add_trace(
        go.Scatter3d(
            x=[center_plot[0]],
            y=[center_plot[1]],
            z=[center_plot[2]],
            mode="markers+text",
            marker=dict(size=4, color="black"),
            text=[label],
            textposition="top center",
            name=label,
            showlegend=True,
        )
    )

    # Add axes (RGB = XYZ)
    # X axis (Red)
    fig.add_trace(
        go.Scatter3d(
            x=[center_plot[0], x_end_plot[0]],
            y=[center_plot[1], x_end_plot[1]],
            z=[center_plot[2], x_end_plot[2]],
            mode="lines",
            line=dict(color="red", width=3),
            showlegend=False,
            hoverinfo="skip",
        )
    )
    # Y axis (Green)
    fig.add_trace(
        go.Scatter3d(
            x=[center_plot[0], y_end_plot[0]],
            y=[center_plot[1], y_end_plot[1]],
            z=[center_plot[2], y_end_plot[2]],
            mode="lines",
            line=dict(color="green", width=3),
            showlegend=False,
            hoverinfo="skip",
        )
    )
    # Z axis (Blue)
    fig.add_trace(
        go.Scatter3d(
            x=[center_plot[0], z_end_plot[0]],
            y=[center_plot[1], z_end_plot[1]],
            z=[center_plot[2], z_end_plot[2]],
            mode="lines",
            line=dict(color="blue", width=3),
            showlegend=False,
            hoverinfo="skip",
        )
    )


@click.command()
@click.option("--input", "-i", required=True, help="Path to input JSON file.")
@click.option(
    "--output", "-o", help="Path to save the output visualization (HTML or PNG)."
)
@click.option("--show", is_flag=True, help="Show the plot interactively.")
@click.option("--scale", type=float, default=0.2, help="Scale of the camera axes.")
@click.option(
    "--birdseye",
    is_flag=True,
    help="Show a top-down bird-eye view (X-Z plane).",
)
@click.option(
    "--frustum-scale", type=float, default=0.5, help="Scale of the camera frustum."
)
@click.option(
    "--fov",
    type=float,
    default=60.0,
    help="Horizontal FOV in degrees for frustum visualization.",
)
@click.option(
    "--zed-configs",
    multiple=True,
    help="Path to ZED config file(s) or directory containing SN*.conf files.",
)
@click.option(
    "--resolution",
    type=click.Choice(RESOLUTION_MAP.keys()),
    default="FHD1200",
    help="Resolution suffix to use from ZED config.",
)
@click.option(
    "--eye",
    type=click.Choice(["left", "right"]),
    default="left",
    help="Which eye's intrinsics to use from ZED config.",
)
@click.option(
    "--show-ground/--no-show-ground",
    default=False,
    help="Show a ground plane at Y=ground-y.",
)
@click.option(
    "--ground-y",
    type=float,
    default=0.0,
    help="Y height of the ground plane.",
)
@click.option(
    "--ground-size",
    type=float,
    default=8.0,
    help="Size of the ground plane (side length in meters).",
)
@click.option(
    "--show-origin-axes/--no-show-origin-axes",
    default=True,
    help="Show a world-origin axis triad (X:red, Y:green, Z:blue).",
)
@click.option(
    "--origin-axes-scale",
    type=float,
    help="Scale of the world-origin axes triad. Defaults to --scale if not provided.",
)
@click.option(
    "--world-basis",
    type=click.Choice(["cv", "opengl"]),
    default="cv",
    help="World coordinate basis convention. 'cv' is +Y down, +Z forward. 'opengl' is +Y up, +Z backward.",
)
def main(
    input: str,
    output: Optional[str],
    show: bool,
    scale: float,
    birdseye: bool,
    frustum_scale: float,
    fov: float,
    zed_configs: List[str],
    resolution: str,
    eye: str,
    show_ground: bool,
    ground_y: float,
    ground_size: float,
    show_origin_axes: bool,
    origin_axes_scale: Optional[float],
    world_basis: str,
):
    """Visualize camera extrinsics from JSON using Plotly."""

    try:
        with open(input, "r") as f:
            data = json.load(f)
    except Exception as e:
        print(f"Error reading input file: {e}")
        return

    # Parse poses
    poses = {}
    for serial, cam_data in data.items():
        if not isinstance(cam_data, dict) or "pose" not in cam_data:
            continue
        try:
            poses[serial] = parse_pose(str(cam_data["pose"]))
        except ValueError as e:
            print(f"Warning: Skipping camera {serial} due to error: {e}")

    if not poses:
        print("No valid camera poses found in the input file.")
        return

    # Load ZED configs if provided
    zed_intrinsics = {}
    if zed_configs:
        zed_intrinsics = load_zed_configs(list(zed_configs), resolution, eye)
        matched_count = 0
        for serial in poses.keys():
            if "all" in zed_intrinsics or serial in zed_intrinsics:
                matched_count += 1
        print(
            f"ZED Configs: matched {matched_count}/{len(poses)} cameras (fallback: {len(poses) - matched_count})"
        )

    # Create Plotly figure
    fig = go.Figure()

    for serial, pose in poses.items():
        cam_intrinsics = zed_intrinsics.get("all") or zed_intrinsics.get(str(serial))
        add_camera_trace(
            fig,
            pose,
            str(serial),
            scale=scale,
            frustum_scale=frustum_scale,
            fov_deg=fov,
            intrinsics=cam_intrinsics,
            world_basis=world_basis,
        )

    if show_origin_axes:
        origin = np.zeros(3)
        axis_len = origin_axes_scale if origin_axes_scale is not None else scale

        # Define world axes points
        x_end = np.array([axis_len, 0, 0])
        y_end = np.array([0, axis_len, 0])
        z_end = np.array([0, 0, axis_len])

        # Transform to plot space
        origin_plot = world_to_plot(origin[None, :], world_basis=world_basis)[0]
        x_end_plot = world_to_plot(x_end[None, :], world_basis=world_basis)[0]
        y_end_plot = world_to_plot(y_end[None, :], world_basis=world_basis)[0]
        z_end_plot = world_to_plot(z_end[None, :], world_basis=world_basis)[0]

        fig.add_trace(
            go.Scatter3d(
                x=[origin_plot[0], x_end_plot[0]],
                y=[origin_plot[1], x_end_plot[1]],
                z=[origin_plot[2], x_end_plot[2]],
                mode="lines",
                line=dict(color="red", width=4),
                name="World X",
                legendgroup="Origin",
                showlegend=True,
                hoverinfo="text",
                text="World X",
            )
        )
        fig.add_trace(
            go.Scatter3d(
                x=[origin_plot[0], y_end_plot[0]],
                y=[origin_plot[1], y_end_plot[1]],
                z=[origin_plot[2], y_end_plot[2]],
                mode="lines",
                line=dict(color="green", width=4),
                name="World Y",
                legendgroup="Origin",
                showlegend=True,
                hoverinfo="text",
                text="World Y",
            )
        )
        fig.add_trace(
            go.Scatter3d(
                x=[origin_plot[0], z_end_plot[0]],
                y=[origin_plot[1], z_end_plot[1]],
                z=[origin_plot[2], z_end_plot[2]],
                mode="lines",
                line=dict(color="blue", width=4),
                name="World Z",
                legendgroup="Origin",
                showlegend=True,
                hoverinfo="text",
                text="World Z",
            )
        )

    if show_ground:
        half_size = ground_size / 2.0
        x_grid = np.linspace(-half_size, half_size, 2)
        z_grid = np.linspace(-half_size, half_size, 2)
        x_mesh, z_mesh = np.meshgrid(x_grid, z_grid)
        y_mesh = np.full_like(x_mesh, ground_y)

        # Flatten for transformation
        pts_ground = np.stack(
            [x_mesh.flatten(), y_mesh.flatten(), z_mesh.flatten()], axis=1
        )
        pts_ground_plot = world_to_plot(pts_ground, world_basis=world_basis)

        # Reshape back
        x_mesh_plot = pts_ground_plot[:, 0].reshape(x_mesh.shape)
        y_mesh_plot = pts_ground_plot[:, 1].reshape(y_mesh.shape)
        z_mesh_plot = pts_ground_plot[:, 2].reshape(z_mesh.shape)

        fig.add_trace(
            go.Surface(
                x=x_mesh_plot,
                y=y_mesh_plot,
                z=z_mesh_plot,
                showscale=False,
                opacity=0.15,
                colorscale=[[0, "gray"], [1, "gray"]],
                name="Ground Plane",
                hoverinfo="skip",
            )
        )

    # Configure layout
    # CV basis: +Y down, +Z forward
    if world_basis == "cv":
        scene_dict: Dict[str, Any] = dict(
            xaxis_title="X (Right)",
            yaxis_title="Y (Down)",
            zaxis_title="Z (Forward)",
            aspectmode="data",
            camera=dict(
                up=dict(
                    x=0, y=-1, z=0
                ),  # In Plotly's default view, +Y is up. To show +Y down, we set up to -Y.
                eye=dict(x=1.25, y=-1.25, z=1.25),
            ),
        )
    else:
        scene_dict: Dict[str, Any] = dict(
            xaxis_title="X (Right)",
            yaxis_title="Y (Up)",
            zaxis_title="Z (Backward)",
            aspectmode="data",
            camera=dict(
                up=dict(x=0, y=1, z=0),
                eye=dict(x=1.25, y=1.25, z=1.25),
            ),
        )

    if birdseye:
        # For birdseye, we force top-down view (looking down +Y towards X-Z plane)
        scene_dict["camera"] = dict(
            projection=dict(type="orthographic"),
            up=dict(x=0, y=0, z=1),  # World +Z is 'up' on screen
            eye=dict(x=0, y=2.5, z=0),
        )

    fig.update_layout(
        title=f"Camera Extrinsics<br><sup>World Basis: {world_basis.upper()} ({' +Y down, +Z fwd' if world_basis == 'cv' else '+Y up, +Z backward'})</sup>",
        scene=scene_dict,
        margin=dict(l=0, r=0, b=0, t=60),
        legend=dict(x=0, y=1),
    )

    if output:
        if output.endswith(".html"):
            fig.write_html(output)
            print(f"Saved interactive plot to {output}")
        elif (
            output.endswith(".png")
            or output.endswith(".jpg")
            or output.endswith(".jpeg")
        ):
            try:
                # Requires kaleido
                fig.write_image(output)
                print(f"Saved static image to {output}")
            except Exception as e:
                print(f"Error saving image (ensure kaleido is installed): {e}")
        else:
            # Default to HTML if unknown extension
            out_path = output + ".html"
            fig.write_html(out_path)
            print(f"Saved interactive plot to {out_path}")

    if show:
        fig.show()
    elif not output:
        print(
            "No output path specified and --show not passed. Plot not saved or shown."
        )


if __name__ == "__main__":
    # pylint: disable=no-value-for-parameter
    main()