#!/usr/bin/env python3
"""Static readiness checks for the Gen2 IsaacLab locomotion asset."""

from __future__ import annotations

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

import numpy as np


REPO_ROOT = Path(__file__).resolve().parents[1]
URDF_PATH = (
    REPO_ROOT
    / "source"
    / "engineai_lab"
    / "assets"
    / "gen2"
    / "urdf"
    / "robot_simplified_collision.urdf"
)
GEN2_ROBOT_CFG = REPO_ROOT / "source" / "engineai_lab" / "robots" / "gen2.py"
GEN2_ENV_CFG = (
    REPO_ROOT
    / "source"
    / "engineai_lab"
    / "tasks"
    / "velocity"
    / "config"
    / "gen2"
    / "common_env_cfg.py"
)
FOOT_LINKS = ("left_leg_link_6", "right_leg_link_6")


@dataclass
class Box:
    name: str
    center: np.ndarray
    axes: np.ndarray
    half: np.ndarray
    vertices: np.ndarray


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]:
    text = GEN2_ROBOT_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 _load_base_init_z() -> float | None:
    text = GEN2_ROBOT_CFG.read_text(encoding="utf-8")
    match = re.search(r"pos=\(([^)]*)\)", text)
    if match is None:
        return None
    return float(match.group(1).split(",")[2].strip())


def _load_height_targets() -> dict[str, float]:
    if not GEN2_ENV_CFG.exists():
        return {}
    text = GEN2_ENV_CFG.read_text(encoding="utf-8")
    targets = {}
    for name in ("GEN2_PELVIS_HEIGHT_TARGET", "GEN2_TORSO_HEIGHT_TARGET"):
        match = re.search(rf"{name}\s*=\s*([0-9.]+)", text)
        if match is not None:
            targets[name] = float(match.group(1))
    return targets


def _link_tree(root: ET.Element):
    links = {link.attrib["name"]: link 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)
    return links, children, parent_of, edge


def _link_world_transforms(root: ET.Element, joint_positions: dict[str, float]) -> tuple[dict[str, np.ndarray], dict[str, str], dict]:
    links, children, parent_of, edge = _link_tree(root)
    root_link = next(name for name in links 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, parent_of, edge


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],
        ],
        dtype=float,
    )
    return (transform @ local.T).T[:, :3]


def _collect_boxes(root: ET.Element) -> tuple[list[Box], dict[str, str], dict]:
    joint_positions = _load_default_joint_positions()
    link_world, parent_of, edge = _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(Box(link_name, transform[:3, 3], transform[:3, :3], np.asarray(size, dtype=float) * 0.5, _box_vertices(transform, size)))
    return boxes, parent_of, edge


def _is_adjacent(parent_of: dict[str, str], a: str, b: str) -> bool:
    return parent_of.get(a) == b or parent_of.get(b) == a


def _obb_penetration(a: Box, b: Box) -> float | None:
    axes = [a.axes[:, idx] for idx in range(3)] + [b.axes[:, idx] for idx in range(3)]
    for i in range(3):
        for j in range(3):
            cross = np.cross(a.axes[:, i], b.axes[:, j])
            norm = np.linalg.norm(cross)
            if norm > 1.0e-8:
                axes.append(cross / norm)

    min_margin = float("inf")
    for axis in axes:
        axis = axis / np.linalg.norm(axis)
        distance = abs(np.dot(b.center - a.center, axis))
        radius_a = sum(a.half[i] * abs(np.dot(a.axes[:, i], axis)) for i in range(3))
        radius_b = sum(b.half[i] * abs(np.dot(b.axes[:, i], axis)) for i in range(3))
        margin = radius_a + radius_b - distance
        if margin <= 0.0:
            return None
        min_margin = min(min_margin, margin)
    return min_margin


def _check_inertials(root: ET.Element):
    total_mass = 0.0
    missing_inertial = []
    bad_mass = []
    bad_inertia = []
    for link in root.findall("link"):
        name = link.attrib["name"]
        inertial = link.find("inertial")
        if inertial is None:
            missing_inertial.append(name)
            continue
        mass = float(inertial.find("mass").attrib["value"])
        total_mass += mass
        if mass <= 0.0:
            bad_mass.append((name, mass))
        inertia = inertial.find("inertia").attrib
        matrix = np.array(
            [
                [float(inertia.get("ixx", 0.0)), float(inertia.get("ixy", 0.0)), float(inertia.get("ixz", 0.0))],
                [float(inertia.get("ixy", 0.0)), float(inertia.get("iyy", 0.0)), float(inertia.get("iyz", 0.0))],
                [float(inertia.get("ixz", 0.0)), float(inertia.get("iyz", 0.0)), float(inertia.get("izz", 0.0))],
            ]
        )
        eig = np.linalg.eigvalsh(matrix)
        if eig.min() <= 0.0:
            bad_inertia.append((name, eig.tolist()))
    return total_mass, missing_inertial, bad_mass, bad_inertia


def _check_joint_limits(root: ET.Element):
    movable = []
    missing = []
    bad = []
    for joint in root.findall("joint"):
        if joint.attrib.get("type") == "fixed":
            continue
        name = joint.attrib["name"]
        limit = joint.find("limit")
        if limit is None:
            missing.append(name)
            continue
        lower = float(limit.attrib.get("lower", "nan"))
        upper = float(limit.attrib.get("upper", "nan"))
        effort = float(limit.attrib.get("effort", "nan"))
        velocity = float(limit.attrib.get("velocity", "nan"))
        movable.append((name, lower, upper, effort, velocity))
        if not (lower < upper and effort > 0.0 and velocity > 0.0):
            bad.append((name, lower, upper, effort, velocity))
    return movable, missing, bad


def main() -> None:
    root = ET.parse(URDF_PATH).getroot()
    boxes, parent_of, edge = _collect_boxes(root)
    movable, missing_limits, bad_limits = _check_joint_limits(root)
    total_mass, missing_inertial, bad_mass, bad_inertia = _check_inertials(root)

    adjacent_overlaps = []
    non_adjacent_overlaps = []
    for i, box_a in enumerate(boxes):
        for box_b in boxes[i + 1 :]:
            penetration = _obb_penetration(box_a, box_b)
            if penetration is None:
                continue
            if _is_adjacent(parent_of, box_a.name, box_b.name):
                joint_name = edge.get((box_a.name, box_b.name), edge.get((box_b.name, box_a.name), ("?",)))[0]
                adjacent_overlaps.append((penetration, box_a.name, box_b.name, joint_name))
            else:
                non_adjacent_overlaps.append((penetration, box_a.name, box_b.name))

    adjacent_overlaps.sort(reverse=True)
    non_adjacent_overlaps.sort(reverse=True)

    base_z = _load_base_init_z()
    height_targets = _load_height_targets()
    foot_min_z = {box.name: float(box.vertices[:, 2].min()) for box in boxes if box.name in FOOT_LINKS}
    lowest_foot_z = min(foot_min_z.values())
    suggested_base_z = -lowest_foot_z + 0.01

    print("GEN2 RL READINESS CHECK")
    print(f"urdf={URDF_PATH}")
    print(f"links={len(root.findall('link'))} joints={len(root.findall('joint'))} collision_boxes={len(boxes)}")
    print(f"total_mass_kg={total_mass:.6f}")
    print(f"inertial_missing={len(missing_inertial)} bad_mass={len(bad_mass)} bad_inertia={len(bad_inertia)}")
    print(f"movable_joints={len(movable)} missing_limits={len(missing_limits)} bad_limits={len(bad_limits)}")
    if movable:
        print(f"effort_range_nm={min(item[3] for item in movable):.3f}..{max(item[3] for item in movable):.3f}")
        print(f"velocity_range_rad_s={min(item[4] for item in movable):.3f}..{max(item[4] for item in movable):.3f}")
    print(f"non_adjacent_collision_overlaps={len(non_adjacent_overlaps)}")
    for penetration, link_a, link_b in non_adjacent_overlaps:
        print(f"  NON_ADJ {link_a} <-> {link_b}: penetration_m={penetration:.6f}")
    print(f"adjacent_collision_overlaps={len(adjacent_overlaps)}")
    for penetration, link_a, link_b, joint_name in adjacent_overlaps:
        print(f"  ADJ {link_a} <-> {link_b} via {joint_name}: penetration_m={penetration:.6f}")
    print("foot_collision_bottom_relative_to_body_root_m:")
    for name, min_z in foot_min_z.items():
        world_min = min_z + base_z if base_z is not None else float("nan")
        print(f"  {name}: relative={min_z:.6f} world_at_init={world_min:.6f}")
    print(f"base_init_z={base_z:.6f}" if base_z is not None else "base_init_z=unknown")
    for target_name in ("GEN2_PELVIS_HEIGHT_TARGET", "GEN2_TORSO_HEIGHT_TARGET"):
        target = height_targets.get(target_name)
        print(f"{target_name.lower()}={target:.6f}" if target is not None else f"{target_name.lower()}=unknown")
    print(f"suggested_base_z_for_1cm_foot_clearance={suggested_base_z:.6f}")

    missing_height_targets = len(height_targets) != 2
    hard_fail = bool(
        missing_inertial
        or bad_mass
        or bad_inertia
        or missing_limits
        or bad_limits
        or non_adjacent_overlaps
        or missing_height_targets
    )
    ground_penetration = base_z is not None and lowest_foot_z + base_z < -0.005
    if hard_fail:
        print("RESULT=FAIL")
    elif ground_penetration:
        print("RESULT=WARN_RESET_HEIGHT")
    elif adjacent_overlaps:
        print("RESULT=WARN_ADJACENT_OVERLAPS")
    else:
        print("RESULT=PASS")


if __name__ == "__main__":
    main()