#!/usr/bin/env python3
"""Render Gen2 collision boxes in the configured default standing pose."""

from __future__ import annotations

import argparse
import ast
import math
import re
import xml.etree.ElementTree as ET
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d.art3d import Poly3DCollection


REPO_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_URDF = REPO_ROOT / "source" / "gen2_lab" / "assets" / "robot_simplified_collision.urdf"
DEFAULT_OUTPUT = REPO_ROOT / "source" / "gen2_lab" / "assets" / "gen2_collision_visualization.png"
GEN2_CFG = REPO_ROOT / "source" / "engineai_lab" / "robots" / "gen2.py"


def _parse_floats(text: str | None, default: tuple[float, ...]) -> tuple[float, ...]:
    if not text:
        return default
    return tuple(float(value) for value in text.split())


def _rpy_matrix(rpy: tuple[float, float, float]) -> np.ndarray:
    roll, pitch, yaw = rpy
    cr, sr = math.cos(roll), math.sin(roll)
    cp, sp = math.cos(pitch), math.sin(pitch)
    cy, sy = math.cos(yaw), math.sin(yaw)
    return np.array(
        [
            [cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr],
            [sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr],
            [-sp, cp * sr, cp * cr],
        ],
        dtype=float,
    )


def _axis_angle(axis: tuple[float, float, float], angle: float) -> np.ndarray:
    vec = np.asarray(axis, dtype=float)
    norm = np.linalg.norm(vec)
    if norm < 1.0e-12:
        return np.eye(3)
    x, y, z = vec / norm
    c, s = math.cos(angle), math.sin(angle)
    one_c = 1.0 - c
    return np.array(
        [
            [c + x * x * one_c, x * y * one_c - z * s, x * z * one_c + y * s],
            [y * x * one_c + z * s, c + y * y * one_c, y * z * one_c - x * s],
            [z * x * one_c - y * s, z * y * one_c + x * s, c + z * z * one_c],
        ],
        dtype=float,
    )


def _transform(xyz=(0.0, 0.0, 0.0), rpy=(0.0, 0.0, 0.0)) -> np.ndarray:
    matrix = np.eye(4)
    matrix[:3, :3] = _rpy_matrix(rpy)
    matrix[:3, 3] = np.asarray(xyz, dtype=float)
    return matrix


def _load_default_joint_positions() -> dict[str, float]:
    if not GEN2_CFG.exists():
        return {}
    text = GEN2_CFG.read_text(encoding="utf-8")
    match = re.search(r"joint_pos=\{(.*?)\},\n\s*joint_vel=", text, flags=re.S)
    if match is None:
        return {}
    return ast.literal_eval("{" + match.group(1) + "}")


def _compute_link_world_transforms(root: ET.Element, joint_positions: dict[str, float]) -> dict[str, np.ndarray]:
    link_names = [link.attrib["name"] for link in root.findall("link")]
    children: dict[str, list[str]] = {}
    parent_of: dict[str, str] = {}
    edge = {}

    for joint in root.findall("joint"):
        parent = joint.find("parent").attrib["link"]
        child = joint.find("child").attrib["link"]
        origin = joint.find("origin")
        axis_elem = joint.find("axis")
        xyz = _parse_floats(origin.attrib.get("xyz"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
        rpy = _parse_floats(origin.attrib.get("rpy"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
        axis = (
            _parse_floats(axis_elem.attrib.get("xyz"), (1.0, 0.0, 0.0)) if axis_elem is not None else (1.0, 0.0, 0.0)
        )
        children.setdefault(parent, []).append(child)
        parent_of[child] = parent
        edge[(parent, child)] = (joint.attrib["name"], joint.attrib.get("type", "fixed"), xyz, rpy, axis)

    root_link = next(name for name in link_names if name not in parent_of)
    world = {root_link: np.eye(4)}
    stack = [root_link]
    while stack:
        parent = stack.pop()
        for child in children.get(parent, []):
            joint_name, joint_type, xyz, rpy, axis = edge[(parent, child)]
            joint_transform = _transform(xyz, rpy)
            if joint_type in {"revolute", "continuous"}:
                rotation = np.eye(4)
                rotation[:3, :3] = _axis_angle(axis, joint_positions.get(joint_name, 0.0))
                joint_transform = joint_transform @ rotation
            elif joint_type == "prismatic":
                translation = np.eye(4)
                translation[:3, 3] = np.asarray(axis, dtype=float) * joint_positions.get(joint_name, 0.0)
                joint_transform = joint_transform @ translation
            world[child] = world[parent] @ joint_transform
            stack.append(child)
    return world


def _box_vertices(transform: np.ndarray, size: tuple[float, float, float]) -> np.ndarray:
    sx, sy, sz = (np.asarray(size, dtype=float) * 0.5).tolist()
    local = np.array(
        [
            [-sx, -sy, -sz, 1.0],
            [sx, -sy, -sz, 1.0],
            [sx, sy, -sz, 1.0],
            [-sx, sy, -sz, 1.0],
            [-sx, -sy, sz, 1.0],
            [sx, -sy, sz, 1.0],
            [sx, sy, sz, 1.0],
            [-sx, sy, sz, 1.0],
        ]
    )
    return (transform @ local.T).T[:, :3]


def _box_faces(vertices: np.ndarray) -> list[np.ndarray]:
    indices = [(0, 1, 2, 3), (4, 5, 6, 7), (0, 1, 5, 4), (1, 2, 6, 5), (2, 3, 7, 6), (3, 0, 4, 7)]
    return [vertices[list(face)] for face in indices]


def _module_color(link_name: str) -> str:
    if link_name.startswith("left_leg"):
        return "#2a9d8f"
    if link_name.startswith("right_leg"):
        return "#e76f51"
    if link_name.startswith("left_arm"):
        return "#457b9d"
    if link_name.startswith("right_arm"):
        return "#f4a261"
    if link_name.startswith("waist"):
        return "#8d5a97"
    return "#5c677d"


def _collect_boxes(root: ET.Element) -> list[dict]:
    joint_positions = _load_default_joint_positions()
    link_world = _compute_link_world_transforms(root, joint_positions)
    boxes = []
    for link in root.findall("link"):
        link_name = link.attrib["name"]
        collision = link.find("collision")
        box = link.find("./collision/geometry/box")
        if collision is None or box is None:
            continue
        origin = collision.find("origin")
        xyz = _parse_floats(origin.attrib.get("xyz"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
        rpy = _parse_floats(origin.attrib.get("rpy"), (0.0, 0.0, 0.0)) if origin is not None else (0.0, 0.0, 0.0)
        size = _parse_floats(box.attrib["size"], (0.0, 0.0, 0.0))
        transform = link_world[link_name] @ _transform(xyz, rpy)
        boxes.append(
            {
                "name": link_name,
                "vertices": _box_vertices(transform, size),
                "center": transform[:3, 3],
                "size": size,
                "color": _module_color(link_name),
            }
        )
    return boxes


def _set_equal_axes_3d(ax, points: np.ndarray) -> None:
    mins = points.min(axis=0)
    maxs = points.max(axis=0)
    center = (mins + maxs) * 0.5
    radius = float((maxs - mins).max() * 0.56)
    ax.set_xlim(center[0] - radius, center[0] + radius)
    ax.set_ylim(center[1] - radius, center[1] + radius)
    ax.set_zlim(center[2] - radius, center[2] + radius)


def _draw_projection(ax, boxes: list[dict], dims: tuple[int, int], labels: tuple[str, str], title: str) -> None:
    edges = [(0, 1), (1, 2), (2, 3), (3, 0), (4, 5), (5, 6), (6, 7), (7, 4), (0, 4), (1, 5), (2, 6), (3, 7)]
    for box in boxes:
        vertices = box["vertices"]
        for start, end in edges:
            ax.plot(
                [vertices[start, dims[0]], vertices[end, dims[0]]],
                [vertices[start, dims[1]], vertices[end, dims[1]]],
                color=box["color"],
                linewidth=1.2,
            )
    ax.set_title(title)
    ax.set_xlabel(labels[0])
    ax.set_ylabel(labels[1])
    ax.set_aspect("equal", adjustable="box")
    ax.grid(True, alpha=0.25)


def _save_projection(output_path: Path, boxes: list[dict], dims: tuple[int, int], labels: tuple[str, str], title: str) -> None:
    fig, ax = plt.subplots(figsize=(8, 8), dpi=180)
    _draw_projection(ax, boxes, dims, labels, title)
    fig.tight_layout()
    fig.savefig(output_path)
    plt.close(fig)
    print(f"wrote {output_path}")


def _save_3d(output_path: Path, boxes: list[dict], all_points: np.ndarray) -> None:
    fig = plt.figure(figsize=(8, 8), dpi=180)
    ax3d = fig.add_subplot(1, 1, 1, projection="3d")
    for box in boxes:
        collection = Poly3DCollection(_box_faces(box["vertices"]), facecolor=box["color"], edgecolor="#202020", alpha=0.32)
        collection.set_linewidth(0.45)
        ax3d.add_collection3d(collection)
    _set_equal_axes_3d(ax3d, all_points)
    ax3d.set_title("Gen2 collision boxes - 3D")
    ax3d.set_xlabel("X forward (m)")
    ax3d.set_ylabel("Y left (m)")
    ax3d.set_zlabel("Z up (m)")
    ax3d.view_init(elev=18, azim=-58)
    fig.tight_layout()
    fig.savefig(output_path)
    plt.close(fig)
    print(f"wrote {output_path}")


def render(urdf_path: Path, output_path: Path) -> None:
    root = ET.parse(urdf_path).getroot()
    boxes = _collect_boxes(root)
    if not boxes:
        raise RuntimeError(f"No box collisions found in {urdf_path}")

    all_points = np.concatenate([box["vertices"] for box in boxes], axis=0)
    fig = plt.figure(figsize=(15, 11), dpi=170)

    ax3d = fig.add_subplot(2, 2, 1, projection="3d")
    for box in boxes:
        collection = Poly3DCollection(_box_faces(box["vertices"]), facecolor=box["color"], edgecolor="#202020", alpha=0.32)
        collection.set_linewidth(0.45)
        ax3d.add_collection3d(collection)
    _set_equal_axes_3d(ax3d, all_points)
    ax3d.set_title("Gen2 collision boxes - 3D")
    ax3d.set_xlabel("X forward (m)")
    ax3d.set_ylabel("Y left (m)")
    ax3d.set_zlabel("Z up (m)")
    ax3d.view_init(elev=18, azim=-58)

    _draw_projection(fig.add_subplot(2, 2, 2), boxes, (1, 2), ("Y left (m)", "Z up (m)"), "Front view")
    _draw_projection(fig.add_subplot(2, 2, 3), boxes, (0, 2), ("X forward (m)", "Z up (m)"), "Side view")
    _draw_projection(fig.add_subplot(2, 2, 4), boxes, (0, 1), ("X forward (m)", "Y left (m)"), "Top view")

    legend_items = [
        ("body", "#5c677d"),
        ("waist", "#8d5a97"),
        ("left leg", "#2a9d8f"),
        ("right leg", "#e76f51"),
        ("left arm", "#457b9d"),
        ("right arm", "#f4a261"),
    ]
    handles = [plt.Line2D([0], [0], color=color, lw=4, label=name) for name, color in legend_items]
    fig.legend(handles=handles, loc="lower center", ncol=6, frameon=False)
    fig.suptitle(f"{urdf_path.name}: {len(boxes)} collision boxes", fontsize=14)
    fig.tight_layout(rect=(0, 0.04, 1, 0.96))
    output_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output_path)
    plt.close(fig)
    print(f"wrote {output_path}")

    stem = output_path.with_suffix("")
    _save_3d(stem.with_name(f"{stem.name}_3d.png"), boxes, all_points)
    _save_projection(stem.with_name(f"{stem.name}_front.png"), boxes, (1, 2), ("Y left (m)", "Z up (m)"), "Front view")
    _save_projection(stem.with_name(f"{stem.name}_side.png"), boxes, (0, 2), ("X forward (m)", "Z up (m)"), "Side view")
    _save_projection(stem.with_name(f"{stem.name}_top.png"), boxes, (0, 1), ("X forward (m)", "Y left (m)"), "Top view")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--urdf", type=Path, default=DEFAULT_URDF)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    args = parser.parse_args()
    render(args.urdf.resolve(), args.output.resolve())


if __name__ == "__main__":
    main()