from __future__ import annotations

import csv
import json
from dataclasses import dataclass
from datetime import datetime, timezone
from math import radians
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import mujoco
import numpy as np
import pinocchio as pin
from PIL import Image

from .dual_arm import DualArmAssembly
from .kinematics import RM75Kinematics, pose_errors
from .mujoco_backend import CONTROLLED_HOME_Q_RAD, DualArmMuJoCo
from .mujoco_trajectories import (
    DEMO_TRAJECTORIES,
    cartesian_demo_targets,
    joint_demo_trajectory,
    solve_pose_trajectory,
)
from .realman_reference import RealManFkReference
from .solver import RM75IkSolver
from .types import IkOptions, physical_joint_limits, teleop_joint_limits


MUJOCO_PIN_POSITION_LIMIT_M = 1e-9
MUJOCO_PIN_ORIENTATION_LIMIT_RAD = 1e-9
ALGO_FK_POSITION_LIMIT_M = 1e-4
ALGO_FK_ORIENTATION_LIMIT_RAD = radians(0.01)
IK_POSITION_LIMIT_M = 1e-3
IK_ORIENTATION_LIMIT_RAD = radians(0.1)
NEAR_IK_RATE_LIMIT = 0.995
CONTINUOUS_IK_RATE_LIMIT = 0.999
MAX_JOINT_STEP_RAD = radians(2.0)
INACTIVE_ARM_DELTA_LIMIT_RAD = 1e-12


@dataclass(frozen=True)
class Stage2Settings:
    seed: int = 20260630
    fk_samples: int = 10_000
    ik_samples: int = 1_000
    trajectories: int = 20
    trajectory_points: int = 500
    render_width: int = 1280
    render_height: int = 720

    @classmethod
    def quick(cls, seed: int = 20260630) -> "Stage2Settings":
        return cls(
            seed=seed,
            fk_samples=100,
            ik_samples=30,
            trajectories=2,
            trajectory_points=25,
            render_width=640,
            render_height=360,
        )


def _sample_configurations(
    rng: np.random.Generator,
    lower: np.ndarray,
    upper: np.ndarray,
    count: int,
    margin: Optional[np.ndarray] = None,
) -> np.ndarray:
    selected_margin = np.zeros(7) if margin is None else np.asarray(margin)
    return rng.uniform(
        lower + selected_margin,
        upper - selected_margin,
        size=(count, 7),
    )


def _percentile(values: List[float], percentile: float) -> float:
    return float(np.percentile(values, percentile)) if values else float("nan")


class Stage2Validator:
    def __init__(
        self,
        reference: RealManFkReference,
        controlled_arm: str = "left",
        settings: Stage2Settings = Stage2Settings(),
    ) -> None:
        self.reference = reference
        self.controlled_arm = controlled_arm
        self.settings = settings
        self.rng = np.random.default_rng(settings.seed)
        self.scene = DualArmMuJoCo(controlled_arm=controlled_arm)
        self.assembly = DualArmAssembly.from_source_urdf()
        self.teleop_kinematics = RM75Kinematics(limits=teleop_joint_limits())
        self.solver = RM75IkSolver(self.teleop_kinematics)
        self.mount = (
            self.assembly.mounts.left_base
            if controlled_arm == "left"
            else self.assembly.mounts.right_base
        )
        self._inactive_initial = self.scene.get_arm_configuration(self.scene.inactive_arm)
        self._max_inactive_delta = 0.0
        self.checks: Dict[str, Dict[str, Any]] = {}
        self.failures: List[Dict[str, Any]] = []
        self.images: Dict[str, Tuple[np.ndarray, np.ndarray]] = {}

    def _check_inactive_arm(self) -> None:
        delta = float(
            np.max(
                np.abs(
                    self.scene.get_arm_configuration(self.scene.inactive_arm)
                    - self._inactive_initial
                )
            )
        )
        self._max_inactive_delta = max(self._max_inactive_delta, delta)

    def _record_failure(
        self,
        category: str,
        index: int,
        reason: str,
        q: Optional[np.ndarray] = None,
        position_error_m: float = float("nan"),
        orientation_error_rad: float = float("nan"),
    ) -> None:
        if len(self.failures) >= 1000:
            return
        self.failures.append(
            {
                "category": category,
                "sample": index,
                "reason": reason,
                "position_error_m": position_error_m,
                "orientation_error_rad": orientation_error_rad,
                "q_rad": json.dumps(q.tolist()) if q is not None else "",
            }
        )

    def _add_check(self, name: str, passed: bool, metrics: Dict[str, Any]) -> None:
        self.checks[name] = {"passed": bool(passed), "required": True, **metrics}

    def _local_mujoco_pose(self) -> pin.SE3:
        return self.mount.actInv(self.scene.get_flange_pose(self.controlled_arm))

    def run(self) -> Dict[str, Any]:
        self._model_structure_check()
        self._fk_check()
        self._single_point_ik_check()
        self._continuous_ik_check()
        self._trajectory_scenario_check()
        self._visual_check()
        self._add_check(
            "inactive_arm_fixed",
            self._max_inactive_delta < INACTIVE_ARM_DELTA_LIMIT_RAD,
            {
                "arm": self.scene.inactive_arm,
                "max_qpos_delta_rad": self._max_inactive_delta,
                "limit_rad": INACTIVE_ARM_DELTA_LIMIT_RAD,
            },
        )
        passed = all(check["passed"] for check in self.checks.values())
        return {
            "schema_version": 1,
            "generated_at": datetime.now(timezone.utc).isoformat(),
            "seed": self.settings.seed,
            "controlled_arm": self.controlled_arm,
            "inactive_arm": self.scene.inactive_arm,
            "realman_api_version": self.reference.api_version,
            "mujoco_version": mujoco.__version__,
            "passed": passed,
            "checks": self.checks,
            "failure_count": len(self.failures),
        }

    def _model_structure_check(self) -> None:
        model = self.scene.model
        joint_names = [model.joint(index).name for index in range(model.njnt)]
        expected = [
            f"{arm}_joint_{joint_index}"
            for arm in ("left", "right")
            for joint_index in range(1, 8)
        ]
        obj_assets = [key for key in self.scene.assets if key.endswith(".obj")]
        stl_assets = [key for key in self.scene.assets if key.endswith(".STL")]
        passed = (
            model.nq == model.nv == model.njnt == 14
            and model.nu == 0
            and model.nsite == 2
            and model.nmesh == 27
            and joint_names == expected
            and len(obj_assets) == 19
            and len(stl_assets) == 8
        )
        self._add_check(
            "model_structure",
            passed,
            {
                "nq": model.nq,
                "nv": model.nv,
                "njnt": model.njnt,
                "nu": model.nu,
                "nmesh": model.nmesh,
                "obj_assets": len(obj_assets),
                "stl_assets": len(stl_assets),
            },
        )

    def _fk_check(self) -> None:
        limits = physical_joint_limits()
        samples = _sample_configurations(
            self.rng, limits.lower, limits.upper, self.settings.fk_samples
        )
        pin_position_errors: List[float] = []
        pin_orientation_errors: List[float] = []
        algo_position_errors: List[float] = []
        algo_orientation_errors: List[float] = []
        for index, q in enumerate(samples):
            self.scene.set_arm_configuration(self.controlled_arm, q)
            self._check_inactive_arm()
            mujoco_world = self.scene.get_flange_pose(self.controlled_arm)
            pin_world = self.assembly.forward(self.controlled_arm, q)
            pin_position, pin_orientation = pose_errors(mujoco_world, pin_world)
            local_mujoco = self.mount.actInv(mujoco_world)
            algo_position, algo_orientation = pose_errors(
                local_mujoco, self.reference.forward(q)
            )
            pin_position_errors.append(pin_position)
            pin_orientation_errors.append(pin_orientation)
            algo_position_errors.append(algo_position)
            algo_orientation_errors.append(algo_orientation)
            if (
                pin_position >= MUJOCO_PIN_POSITION_LIMIT_M
                or pin_orientation >= MUJOCO_PIN_ORIENTATION_LIMIT_RAD
                or algo_position >= ALGO_FK_POSITION_LIMIT_M
                or algo_orientation >= ALGO_FK_ORIENTATION_LIMIT_RAD
            ):
                self._record_failure(
                    "fk",
                    index,
                    "MuJoCo FK residual exceeded an acceptance limit",
                    q,
                    algo_position,
                    algo_orientation,
                )
        max_pin_position = max(pin_position_errors, default=float("inf"))
        max_pin_orientation = max(pin_orientation_errors, default=float("inf"))
        max_algo_position = max(algo_position_errors, default=float("inf"))
        max_algo_orientation = max(algo_orientation_errors, default=float("inf"))
        self._add_check(
            "fk",
            max_pin_position < MUJOCO_PIN_POSITION_LIMIT_M
            and max_pin_orientation < MUJOCO_PIN_ORIENTATION_LIMIT_RAD
            and max_algo_position < ALGO_FK_POSITION_LIMIT_M
            and max_algo_orientation < ALGO_FK_ORIENTATION_LIMIT_RAD,
            {
                "samples": len(samples),
                "max_mujoco_pin_position_error_m": max_pin_position,
                "max_mujoco_pin_orientation_error_rad": max_pin_orientation,
                "max_mujoco_algo_position_error_m": max_algo_position,
                "max_mujoco_algo_orientation_error_rad": max_algo_orientation,
            },
        )

    @staticmethod
    def _ik_options(max_iterations: int) -> IkOptions:
        return IkOptions(
            position_tolerance_m=0.9 * IK_POSITION_LIMIT_M,
            orientation_tolerance_rad=0.9 * IK_ORIENTATION_LIMIT_RAD,
            max_iterations=max_iterations,
        )

    def _single_point_ik_check(self) -> None:
        limits = teleop_joint_limits()
        margin = np.deg2rad([5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 10.0])
        targets_q = _sample_configurations(
            self.rng,
            limits.lower,
            limits.upper,
            self.settings.ik_samples,
            margin,
        )
        successes = 0
        solve_times: List[float] = []
        max_position = 0.0
        max_orientation = 0.0
        for index, target_q in enumerate(targets_q):
            target = self.reference.forward(target_q)
            seed = np.clip(
                target_q + self.rng.uniform(-radians(10), radians(10), 7),
                limits.lower,
                limits.upper,
            )
            result = self.solver.solve(target, seed, self._ik_options(200))
            solve_times.append(result.solve_time_sec)
            if result.success and result.q is not None:
                self.scene.set_arm_configuration(self.controlled_arm, result.q)
                self._check_inactive_arm()
                position_error, orientation_error = pose_errors(
                    self._local_mujoco_pose(), target
                )
                max_position = max(max_position, position_error)
                max_orientation = max(max_orientation, orientation_error)
                if (
                    position_error <= IK_POSITION_LIMIT_M
                    and orientation_error <= IK_ORIENTATION_LIMIT_RAD
                ):
                    successes += 1
                    continue
            else:
                position_error = result.position_error_m
                orientation_error = result.orientation_error_rad
            self._record_failure(
                "single_point_ik",
                index,
                f"status={result.status.value}; {result.message}",
                seed,
                position_error,
                orientation_error,
            )
        rate = successes / max(len(targets_q), 1)
        self._add_check(
            "single_point_ik",
            rate >= NEAR_IK_RATE_LIMIT,
            {
                "samples": len(targets_q),
                "successes": successes,
                "success_rate": rate,
                "max_position_error_m": max_position,
                "max_orientation_error_rad": max_orientation,
                "p99_solve_time_sec": _percentile(solve_times, 99),
                "max_solve_time_sec": max(solve_times, default=float("nan")),
            },
        )

    def _continuous_ik_check(self) -> None:
        limits = teleop_joint_limits()
        successes = 0
        total = 0
        max_joint_step = 0.0
        max_position = 0.0
        max_orientation = 0.0
        for trajectory_index in range(self.settings.trajectories):
            span = limits.upper - limits.lower
            center = self.rng.uniform(
                limits.lower + 0.3 * span,
                limits.upper - 0.3 * span,
            )
            amplitude = self.rng.uniform(0.015, 0.04, 7) * span
            frequency = self.rng.uniform(0.03, 0.08, 7)
            phase = self.rng.uniform(-np.pi, np.pi, 7)
            times = np.arange(self.settings.trajectory_points) / 90.0
            target_path = center + amplitude * np.sin(
                2.0 * np.pi * times[:, None] * frequency + phase
            )
            seed = target_path[0].copy()
            previous = seed.copy()
            for point_index, target_q in enumerate(target_path):
                total += 1
                target = self.reference.forward(target_q)
                result = self.solver.solve(target, seed, self._ik_options(100))
                if result.success and result.q is not None:
                    self.scene.set_arm_configuration(self.controlled_arm, result.q)
                    self._check_inactive_arm()
                    position_error, orientation_error = pose_errors(
                        self._local_mujoco_pose(), target
                    )
                    joint_step = float(np.max(np.abs(result.q - previous)))
                    max_position = max(max_position, position_error)
                    max_orientation = max(max_orientation, orientation_error)
                    max_joint_step = max(max_joint_step, joint_step)
                    if (
                        position_error <= IK_POSITION_LIMIT_M
                        and orientation_error <= IK_ORIENTATION_LIMIT_RAD
                    ):
                        successes += 1
                        seed = result.q
                        previous = result.q
                        continue
                else:
                    position_error = result.position_error_m
                    orientation_error = result.orientation_error_rad
                self._record_failure(
                    "continuous_ik",
                    trajectory_index * self.settings.trajectory_points + point_index,
                    f"status={result.status.value}; {result.message}",
                    seed,
                    position_error,
                    orientation_error,
                )
        rate = successes / max(total, 1)
        self._add_check(
            "continuous_ik",
            rate >= CONTINUOUS_IK_RATE_LIMIT
            and max_joint_step < MAX_JOINT_STEP_RAD,
            {
                "trajectories": self.settings.trajectories,
                "points": total,
                "successes": successes,
                "success_rate": rate,
                "max_position_error_m": max_position,
                "max_orientation_error_rad": max_orientation,
                "max_joint_step_rad": max_joint_step,
                "joint_step_limit_rad": MAX_JOINT_STEP_RAD,
            },
        )

    def _trajectory_scenario_check(self) -> None:
        scenario_results: Dict[str, bool] = {}
        start = CONTROLLED_HOME_Q_RAD.copy()
        for kind in DEMO_TRAJECTORIES:
            try:
                if kind == "joint":
                    q_trajectory = joint_demo_trajectory(start, 120)
                else:
                    targets = cartesian_demo_targets(
                        kind, self.teleop_kinematics.forward(start), 120
                    )
                    q_trajectory = solve_pose_trajectory(self.solver, targets, start)
                self.scene.play_trajectory(q_trajectory)
                self._check_inactive_arm()
                scenario_results[kind] = True
            except (RuntimeError, ValueError) as exc:
                scenario_results[kind] = False
                self._record_failure("trajectory_scenario", 0, f"{kind}: {exc}")

        limits = teleop_joint_limits()
        scenario_q = {
            "limit_near": 0.8 * limits.upper + 0.2 * CONTROLLED_HOME_Q_RAD,
            "singularity_near": np.deg2rad([0, 0, 0, 90, 0, 0, 0]),
        }
        blend = 0.5 - 0.5 * np.cos(np.linspace(0.0, np.pi, 120))
        for name, endpoint in scenario_q.items():
            path = start[None, :] + blend[:, None] * (endpoint - start)[None, :]
            try:
                targets = [self.reference.forward(q) for q in path]
                solutions = solve_pose_trajectory(self.solver, targets, start)
                self.scene.play_trajectory(solutions)
                self._check_inactive_arm()
                scenario_results[name] = True
            except (RuntimeError, ValueError) as exc:
                scenario_results[name] = False
                self._record_failure("trajectory_scenario", 0, f"{name}: {exc}")
        self._add_check(
            "trajectory_scenarios",
            all(scenario_results.values()),
            {"scenarios": scenario_results},
        )

    def _visible_geom_names(self, segmentation: np.ndarray) -> set[str]:
        pairs = np.unique(segmentation.reshape(-1, 2), axis=0)
        names = set()
        for object_id, object_type in pairs:
            if object_type != mujoco.mjtObj.mjOBJ_GEOM or object_id < 0:
                continue
            name = mujoco.mj_id2name(
                self.scene.model, mujoco.mjtObj.mjOBJ_GEOM, int(object_id)
            )
            if name is not None:
                names.add(name)
        return names

    @staticmethod
    def _segmentation_preview(segmentation: np.ndarray) -> np.ndarray:
        object_ids = segmentation[:, :, 0]
        preview = np.zeros((*object_ids.shape, 3), dtype=np.uint8)
        foreground = object_ids >= 0
        ids = object_ids[foreground].astype(np.uint32)
        preview[foreground, 0] = (37 * ids + 53) % 255
        preview[foreground, 1] = (97 * ids + 101) % 255
        preview[foreground, 2] = (17 * ids + 211) % 255
        return preview

    def _visual_check(self) -> None:
        start = CONTROLLED_HOME_Q_RAD.copy()
        targets = cartesian_demo_targets(
            "combined", self.teleop_kinematics.forward(start), 121
        )
        q_trajectory = solve_pose_trajectory(self.solver, targets, start)
        all_frames_pass = True
        frame_metrics: Dict[str, Dict[str, Any]] = {}
        for label, index in (("start", 0), ("middle", 60), ("end", 120)):
            self.scene.set_arm_configuration(self.controlled_arm, q_trajectory[index])
            self.scene.set_target_marker(self.mount * targets[index])
            self._check_inactive_arm()
            rgb = self.scene.render(
                self.settings.render_width, self.settings.render_height
            )
            segmentation = self.scene.render(
                self.settings.render_width,
                self.settings.render_height,
                segmentation=True,
            )
            visible_names = self._visible_geom_names(segmentation)
            left_visible = any(name.startswith("left_") for name in visible_names)
            right_visible = any(name.startswith("right_") for name in visible_names)
            arm_ids = {
                mujoco.mj_name2id(
                    self.scene.model, mujoco.mjtObj.mjOBJ_GEOM, name
                )
                for name in visible_names
                if name.startswith(("left_", "right_"))
            }
            mask = np.isin(segmentation[:, :, 0], list(arm_ids))
            rows, columns = np.where(mask)
            not_touching_border = bool(
                len(rows)
                and rows.min() > 1
                and columns.min() > 1
                and rows.max() < rgb.shape[0] - 2
                and columns.max() < rgb.shape[1] - 2
            )
            frame_pass = (
                float(rgb.std()) > 5.0
                and left_visible
                and right_visible
                and not_touching_border
            )
            all_frames_pass &= frame_pass
            frame_metrics[label] = {
                "rgb_std": float(rgb.std()),
                "left_visible": left_visible,
                "right_visible": right_visible,
                "not_touching_border": not_touching_border,
                "visible_geom_count": len(visible_names),
            }
            self.images[label] = (rgb, self._segmentation_preview(segmentation))
        self._add_check("visual_rendering", all_frames_pass, {"frames": frame_metrics})


def write_stage2_report(
    output_dir: Path | str,
    summary: Dict[str, Any],
    failures: List[Dict[str, Any]],
    images: Dict[str, Tuple[np.ndarray, np.ndarray]],
) -> Tuple[Path, Path, Path, List[Path]]:
    directory = Path(output_dir)
    directory.mkdir(parents=True, exist_ok=True)
    json_path = directory / "stage2_summary.json"
    csv_path = directory / "stage2_failures.csv"
    markdown_path = directory / "stage2_report.md"

    with json_path.open("w", encoding="utf-8") as stream:
        json.dump(summary, stream, ensure_ascii=True, indent=2, sort_keys=True)
        stream.write("\n")

    fields = [
        "category",
        "sample",
        "reason",
        "position_error_m",
        "orientation_error_rad",
        "q_rad",
    ]
    with csv_path.open("w", encoding="utf-8", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=fields)
        writer.writeheader()
        writer.writerows(failures)

    image_paths: List[Path] = []
    for label, (rgb, segmentation) in images.items():
        for suffix, image in (("rgb", rgb), ("segmentation", segmentation)):
            path = directory / f"stage2_{label}_{suffix}.png"
            Image.fromarray(image).save(path)
            image_paths.append(path)

    lines = [
        "# RM75-B Stage 2 MuJoCo Validation",
        "",
        f"- Overall: **{'PASS' if summary['passed'] else 'FAIL'}**",
        f"- Controlled arm: `{summary['controlled_arm']}`",
        f"- Seed: `{summary['seed']}`",
        f"- MuJoCo: `{summary['mujoco_version']}`",
        f"- RealMan API: `{summary['realman_api_version']}`",
        f"- Failures recorded: `{summary['failure_count']}`",
        "",
        "| Check | Result | Metrics |",
        "|---|---:|---|",
    ]
    for name, check in summary["checks"].items():
        metrics = {
            key: value
            for key, value in check.items()
            if key not in {"passed", "required"}
        }
        lines.append(
            f"| `{name}` | {'PASS' if check['passed'] else 'FAIL'} | "
            f"`{json.dumps(metrics, ensure_ascii=True, sort_keys=True)}` |"
        )
    markdown_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return json_path, csv_path, markdown_path, image_paths

