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, Iterable, List, Optional, Tuple

import numpy as np
import pinocchio as pin
import yaml

from .dual_arm import DualArmAssembly
from .kinematics import RM75Kinematics, pose_errors
from .realman_reference import RealManFkReference
from .solver import RM75IkSolver, deterministic_recovery_seeds
from .types import IkOptions, IkStatus, JointLimits, physical_joint_limits, teleop_joint_limits


FK_POSITION_LIMIT_M = 1e-4
FK_ORIENTATION_LIMIT_RAD = radians(0.01)
IK_POSITION_LIMIT_M = 1e-3
IK_ORIENTATION_LIMIT_RAD = radians(0.1)
JACOBIAN_RELATIVE_LIMIT = 1e-3
JACOBIAN_ABSOLUTE_LIMIT = 5e-4
NEAR_IK_RATE_LIMIT = 0.995
CONTINUOUS_IK_RATE_LIMIT = 0.999
GLOBAL_RECOVERY_RATE_LIMIT = 0.85
NEAR_IK_P99_LIMIT_SEC = 0.008
CONTROL_PERIOD_SEC = 1.0 / 90.0
MAX_CONTINUOUS_JOINT_STEP_RAD = radians(2.0)


def _validation_ik_options(max_iterations: int) -> IkOptions:
    # Keep a 10% convergence guard band for independent Algo verification.
    return IkOptions(
        position_tolerance_m=0.9 * IK_POSITION_LIMIT_M,
        orientation_tolerance_rad=0.9 * IK_ORIENTATION_LIMIT_RAD,
        max_iterations=max_iterations,
    )


@dataclass(frozen=True)
class ValidationSettings:
    seed: int = 20260629
    fk_samples: int = 10_000
    jacobian_samples: int = 200
    near_ik_samples: int = 1_000
    global_samples: int = 200
    continuous_trajectories: int = 20
    continuous_points: int = 500
    tool_samples: int = 100
    dual_samples: int = 100
    strict: bool = True

    @classmethod
    def quick(cls, seed: int = 20260629, strict: bool = False) -> "ValidationSettings":
        return cls(
            seed=seed,
            fk_samples=100,
            jacobian_samples=10,
            near_ik_samples=30,
            global_samples=10,
            continuous_trajectories=2,
            continuous_points=25,
            tool_samples=10,
            dual_samples=10,
            strict=strict,
        )


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


def _sample_configurations(
    rng: np.random.Generator,
    limits: JointLimits,
    count: int,
    margin_rad: Optional[np.ndarray] = None,
) -> np.ndarray:
    margin = np.zeros(7) if margin_rad is None else np.asarray(margin_rad, dtype=float)
    lower = limits.lower + margin
    upper = limits.upper - margin
    if np.any(lower >= upper):
        raise ValueError(f"sampling margin is too large for {limits.name} limits")
    return rng.uniform(lower, upper, size=(count, 7))


def _tool_pose_from_values(values: Iterable[float]) -> pin.SE3:
    pose = np.asarray(list(values), dtype=float)
    if pose.shape != (7,) or not np.all(np.isfinite(pose)):
        raise ValueError("tool pose must be [x,y,z,qx,qy,qz,qw]")
    quaternion = pin.Quaternion(pose[6], pose[3], pose[4], pose[5])
    if quaternion.norm() <= 0.0:
        raise ValueError("tool quaternion must be non-zero")
    quaternion.normalize()
    return pin.SE3(quaternion.matrix(), pose[:3])


def load_project_tools(config_path: Path | str) -> Dict[str, pin.SE3]:
    with Path(config_path).open("r", encoding="utf-8") as stream:
        data = yaml.safe_load(stream)
    tools = data.get("tools_in_ee", {})
    selected: Dict[str, pin.SE3] = {}
    for name in ("scissor", "omnipic", "minisci"):
        if name not in tools or "pose" not in tools[name]:
            raise ValueError(f"missing tool pose for {name!r}")
        selected[name] = _tool_pose_from_values(tools[name]["pose"])
    return selected


class Stage1Validator:
    def __init__(
        self,
        reference: RealManFkReference,
        settings: ValidationSettings = ValidationSettings(),
        tools: Optional[Dict[str, pin.SE3]] = None,
    ) -> None:
        self.reference = reference
        self.settings = settings
        self.tools = tools or {}
        self.rng = np.random.default_rng(settings.seed)
        self.checks: Dict[str, Dict[str, Any]] = {}
        self.failures: List[Dict[str, Any]] = []

    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"),
        profile: str = "",
    ) -> None:
        if len(self.failures) >= 1000:
            return
        self.failures.append(
            {
                "category": category,
                "profile": profile,
                "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],
        *,
        required: bool = True,
    ) -> None:
        self.checks[name] = {
            "passed": bool(passed),
            "required": required,
            **metrics,
        }

    def run(self) -> Dict[str, Any]:
        self._model_checks()
        self._fk_checks()
        self._jacobian_checks()
        self._near_ik_checks()
        self._continuous_ik_checks()
        self._global_recovery_checks()
        self._singularity_checks()
        self._dual_arm_checks()
        self._tool_checks()
        required_checks = [
            check["passed"]
            for check in self.checks.values()
            if check.get("required", True)
        ]
        return {
            "schema_version": 1,
            "generated_at": datetime.now(timezone.utc).isoformat(),
            "seed": self.settings.seed,
            "strict": self.settings.strict,
            "realman_api_version": self.reference.api_version,
            "passed": bool(all(required_checks)),
            "checks": self.checks,
            "failure_count": len(self.failures),
        }

    def _model_checks(self) -> None:
        physical = RM75Kinematics(limits=physical_joint_limits())
        teleop = RM75Kinematics(limits=teleop_joint_limits())
        assembly = DualArmAssembly.from_source_urdf(limits=physical_joint_limits())
        passed = (
            physical.model.nq == physical.model.nv == 7
            and teleop.model.nq == teleop.model.nv == 7
            and assembly.dof == 14
            and physical.limits.contains(np.zeros(7))
            and teleop.limits.contains(np.zeros(7))
        )
        self._add_check(
            "model_structure",
            passed,
            {
                "single_arm_nq": physical.model.nq,
                "single_arm_nv": physical.model.nv,
                "dual_arm_dof": assembly.dof,
                "right_visual_origin_delta_m": assembly.mounts.right_visual_origin_delta_m,
            },
        )

    def _fk_checks(self) -> None:
        for limits in (physical_joint_limits(), teleop_joint_limits()):
            kinematics = RM75Kinematics(limits=limits)
            samples = _sample_configurations(
                self.rng, limits, self.settings.fk_samples
            )
            position_errors: List[float] = []
            orientation_errors: List[float] = []
            for index, q in enumerate(samples):
                pin_pose = kinematics.forward(q)
                reference_pose = self.reference.forward(q)
                position_error, orientation_error = pose_errors(pin_pose, reference_pose)
                position_errors.append(position_error)
                orientation_errors.append(orientation_error)
                if (
                    position_error >= FK_POSITION_LIMIT_M
                    or orientation_error >= FK_ORIENTATION_LIMIT_RAD
                ):
                    self._record_failure(
                        "fk",
                        index,
                        "FK residual exceeded limit",
                        q,
                        position_error,
                        orientation_error,
                        limits.name,
                    )
            max_position = max(position_errors, default=float("inf"))
            max_orientation = max(orientation_errors, default=float("inf"))
            self._add_check(
                f"fk_{limits.name}",
                max_position < FK_POSITION_LIMIT_M
                and max_orientation < FK_ORIENTATION_LIMIT_RAD,
                {
                    "samples": len(samples),
                    "max_position_error_m": max_position,
                    "p99_position_error_m": _percentile(position_errors, 99),
                    "max_orientation_error_rad": max_orientation,
                    "p99_orientation_error_rad": _percentile(
                        orientation_errors, 99
                    ),
                },
            )

    def _numeric_reference_jacobian(self, q: np.ndarray, step: float = 2e-3) -> np.ndarray:
        center = self.reference.forward(q)
        numeric = np.zeros((6, 7))
        for joint_index in range(7):
            delta = np.zeros(7)
            delta[joint_index] = step
            plus = self.reference.forward(q + delta)
            minus = self.reference.forward(q - delta)
            plus_twist = pin.log6(center.actInv(plus)).vector
            minus_twist = pin.log6(center.actInv(minus)).vector
            numeric[:, joint_index] = (plus_twist - minus_twist) / (2.0 * step)
        return numeric

    def _jacobian_checks(self) -> None:
        limits = physical_joint_limits()
        kinematics = RM75Kinematics(limits=limits)
        # Algo FK is returned as float32. A 2e-3 rad central-difference step
        # keeps quantization below the analytic-Jacobian acceptance limit.
        margin = np.full(7, 3e-3)
        samples = _sample_configurations(
            self.rng, limits, self.settings.jacobian_samples, margin
        )
        relative_errors: List[float] = []
        absolute_errors: List[float] = []
        for index, q in enumerate(samples):
            analytic = kinematics.jacobian(q)
            numeric = self._numeric_reference_jacobian(q)
            difference = analytic - numeric
            relative = float(
                np.linalg.norm(difference) / max(np.linalg.norm(numeric), 1e-12)
            )
            absolute = float(np.max(np.abs(difference)))
            relative_errors.append(relative)
            absolute_errors.append(absolute)
            if relative >= JACOBIAN_RELATIVE_LIMIT or absolute >= JACOBIAN_ABSOLUTE_LIMIT:
                self._record_failure(
                    "jacobian",
                    index,
                    f"relative={relative:.6g}, absolute={absolute:.6g}",
                    q,
                    profile=limits.name,
                )
        max_relative = max(relative_errors, default=float("inf"))
        max_absolute = max(absolute_errors, default=float("inf"))
        self._add_check(
            "jacobian",
            max_relative < JACOBIAN_RELATIVE_LIMIT
            and max_absolute < JACOBIAN_ABSOLUTE_LIMIT,
            {
                "samples": len(samples),
                "max_relative_error": max_relative,
                "max_absolute_error": max_absolute,
            },
        )

    def _externally_accept_solution(
        self,
        target: pin.SE3,
        result_q: Optional[np.ndarray],
        limits: JointLimits,
    ) -> Tuple[bool, float, float]:
        if result_q is None or not limits.contains(result_q):
            return False, float("inf"), float("inf")
        verified = self.reference.forward(result_q)
        position_error, orientation_error = pose_errors(verified, target)
        return (
            position_error <= IK_POSITION_LIMIT_M
            and orientation_error <= IK_ORIENTATION_LIMIT_RAD,
            position_error,
            orientation_error,
        )

    def _near_ik_checks(self) -> None:
        all_times: List[float] = []
        profile_rates: Dict[str, float] = {}
        for limits in (physical_joint_limits(), teleop_joint_limits()):
            kinematics = RM75Kinematics(limits=limits)
            solver = RM75IkSolver(kinematics)
            margin = np.deg2rad([5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 10.0])
            targets_q = _sample_configurations(
                self.rng, limits, self.settings.near_ik_samples, margin
            )
            successes = 0
            profile_times: List[float] = []
            options = _validation_ik_options(max_iterations=200)
            for index, target_q in enumerate(targets_q):
                seed = np.clip(
                    target_q + self.rng.uniform(-radians(10), radians(10), 7),
                    limits.lower,
                    limits.upper,
                )
                target = self.reference.forward(target_q)
                result = solver.solve(target, seed, options)
                profile_times.append(result.solve_time_sec)
                all_times.append(result.solve_time_sec)
                accepted, position_error, orientation_error = self._externally_accept_solution(
                    target, result.q, limits
                )
                if result.success and accepted:
                    successes += 1
                else:
                    self._record_failure(
                        "near_ik",
                        index,
                        f"status={result.status.value}; {result.message}",
                        seed,
                        position_error if np.isfinite(position_error) else result.position_error_m,
                        orientation_error if np.isfinite(orientation_error) else result.orientation_error_rad,
                        limits.name,
                    )
            rate = successes / max(len(targets_q), 1)
            profile_rates[limits.name] = rate
            self._add_check(
                f"near_ik_{limits.name}",
                rate >= NEAR_IK_RATE_LIMIT,
                {
                    "samples": len(targets_q),
                    "successes": successes,
                    "success_rate": rate,
                    "p99_time_sec": _percentile(profile_times, 99),
                    "max_time_sec": max(profile_times, default=float("nan")),
                },
            )

        p99_time = _percentile(all_times, 99)
        max_time = max(all_times, default=float("inf"))
        self._add_check(
            "near_ik_performance",
            p99_time < NEAR_IK_P99_LIMIT_SEC and max_time < CONTROL_PERIOD_SEC,
            {
                "p99_time_sec": p99_time,
                "max_time_sec": max_time,
                "p99_limit_sec": NEAR_IK_P99_LIMIT_SEC,
                "control_period_sec": CONTROL_PERIOD_SEC,
            },
            required=self.settings.strict,
        )

    def _continuous_ik_checks(self) -> None:
        limits = teleop_joint_limits()
        kinematics = RM75Kinematics(limits=limits)
        solver = RM75IkSolver(kinematics)
        options = _validation_ik_options(max_iterations=100)
        successes = 0
        total = 0
        max_joint_step = 0.0
        for trajectory_index in range(self.settings.continuous_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.continuous_points) / 90.0
            path = center + amplitude * np.sin(
                2.0 * np.pi * times[:, None] * frequency + phase
            )
            seed = path[0].copy()
            previous_solution = seed.copy()
            for point_index, target_q in enumerate(path):
                total += 1
                target = self.reference.forward(target_q)
                result = solver.solve(target, seed, options)
                accepted, position_error, orientation_error = self._externally_accept_solution(
                    target, result.q, limits
                )
                if result.success and accepted and result.q is not None:
                    joint_step = float(np.max(np.abs(result.q - previous_solution)))
                    max_joint_step = max(max_joint_step, joint_step)
                    previous_solution = result.q
                    seed = result.q
                    successes += 1
                else:
                    self._record_failure(
                        "continuous_ik",
                        trajectory_index * self.settings.continuous_points + point_index,
                        f"status={result.status.value}; {result.message}",
                        seed,
                        position_error,
                        orientation_error,
                        limits.name,
                    )
        rate = successes / max(total, 1)
        self._add_check(
            "continuous_ik",
            rate >= CONTINUOUS_IK_RATE_LIMIT
            and max_joint_step <= MAX_CONTINUOUS_JOINT_STEP_RAD,
            {
                "trajectories": self.settings.continuous_trajectories,
                "points": total,
                "successes": successes,
                "success_rate": rate,
                "max_joint_step_rad": max_joint_step,
                "joint_step_limit_rad": MAX_CONTINUOUS_JOINT_STEP_RAD,
            },
        )

    def _global_recovery_checks(self) -> None:
        limits = physical_joint_limits()
        kinematics = RM75Kinematics(limits=limits)
        solver = RM75IkSolver(kinematics)
        options = _validation_ik_options(max_iterations=500)
        margin = np.deg2rad([5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 10.0])
        target_configurations = _sample_configurations(
            self.rng, limits, self.settings.global_samples, margin
        )
        recovery_seeds = deterministic_recovery_seeds(limits)
        single_successes = 0
        recovery_successes = 0
        recovery_times: List[float] = []
        for index, target_q in enumerate(target_configurations):
            target = self.reference.forward(target_q)
            random_seed = self.rng.uniform(limits.lower, limits.upper)
            single = solver.solve(target, random_seed, options)
            single_accepted, _, _ = self._externally_accept_solution(
                target, single.q, limits
            )
            single_successes += int(single.success and single_accepted)

            recovered = solver.solve_multistart(target, recovery_seeds, options)
            recovery_times.append(recovered.solve_time_sec)
            accepted, position_error, orientation_error = self._externally_accept_solution(
                target, recovered.q, limits
            )
            if recovered.success and accepted:
                recovery_successes += 1
            else:
                self._record_failure(
                    "global_recovery",
                    index,
                    f"status={recovered.status.value}; {recovered.message}",
                    target_q,
                    position_error,
                    orientation_error,
                    limits.name,
                )
        count = max(len(target_configurations), 1)
        recovery_rate = recovery_successes / count
        self._add_check(
            "global_recovery",
            recovery_rate >= GLOBAL_RECOVERY_RATE_LIMIT,
            {
                "samples": len(target_configurations),
                "single_seed_success_rate": single_successes / count,
                "recovery_successes": recovery_successes,
                "recovery_success_rate": recovery_rate,
                "recovery_p95_time_sec": _percentile(recovery_times, 95),
                "recovery_max_time_sec": max(recovery_times, default=float("nan")),
            },
        )

    def _singularity_checks(self) -> None:
        limits = physical_joint_limits()
        solver = RM75IkSolver(RM75Kinematics(limits=limits))
        singular_degrees = np.asarray(
            [
                [0, 0, 0, 90, 0, 0, 0],
                [0, 60, 0, 0, 0, 90, 0],
                [0, 0, 90, 90, 0, 90, 0],
                [0, 90, 90, 90, 90, 0, 0],
            ],
            dtype=float,
        )
        invalid_results = 0
        total = 0
        statuses: Dict[str, int] = {}
        for case_index, singular_q in enumerate(np.deg2rad(singular_degrees)):
            for perturbation in (-radians(0.1), 0.0, radians(0.1)):
                total += 1
                target_q = singular_q.copy()
                target_q[case_index % 7] += perturbation
                target = self.reference.forward(target_q)
                seed = np.clip(
                    target_q + self.rng.uniform(-radians(0.5), radians(0.5), 7),
                    limits.lower,
                    limits.upper,
                )
                result = solver.solve(
                    target,
                    seed,
                    _validation_ik_options(max_iterations=200),
                )
                statuses[result.status.value] = statuses.get(result.status.value, 0) + 1
                finite_diagnostics = np.isfinite(result.position_error_m) and np.isfinite(
                    result.orientation_error_rad
                )
                accepted, _, _ = self._externally_accept_solution(target, result.q, limits)
                pseudo_success = result.status is IkStatus.SUCCESS and not accepted
                if not finite_diagnostics or pseudo_success:
                    invalid_results += 1
                    self._record_failure(
                        "singularity",
                        total - 1,
                        "non-finite diagnostic or false success",
                        seed,
                        result.position_error_m,
                        result.orientation_error_rad,
                        limits.name,
                    )
        self._add_check(
            "singularity_behavior",
            invalid_results == 0,
            {
                "samples": total,
                "invalid_results": invalid_results,
                "statuses": statuses,
            },
        )

    def _dual_arm_checks(self) -> None:
        limits = physical_joint_limits()
        assembly = DualArmAssembly.from_source_urdf(limits=limits)
        samples = _sample_configurations(
            self.rng, limits, self.settings.dual_samples
        )
        max_position = 0.0
        max_orientation = 0.0
        failures = 0
        for arm in ("left", "right"):
            mount = (
                assembly.mounts.left_base if arm == "left" else assembly.mounts.right_base
            )
            for index, q in enumerate(samples):
                world_pose = assembly.forward(arm, q)
                local_pose = mount.actInv(world_pose)
                reference_pose = self.reference.forward(q)
                position_error, orientation_error = pose_errors(local_pose, reference_pose)
                max_position = max(max_position, position_error)
                max_orientation = max(max_orientation, orientation_error)
                if (
                    position_error >= FK_POSITION_LIMIT_M
                    or orientation_error >= FK_ORIENTATION_LIMIT_RAD
                ):
                    failures += 1
                    self._record_failure(
                        "dual_arm",
                        index,
                        f"{arm} local FK residual exceeded limit",
                        q,
                        position_error,
                        orientation_error,
                        arm,
                    )
        self._add_check(
            "dual_arm_assembly",
            failures == 0,
            {
                "samples_per_arm": len(samples),
                "max_position_error_m": max_position,
                "max_orientation_error_rad": max_orientation,
                "right_visual_origin_delta_m": assembly.mounts.right_visual_origin_delta_m,
            },
        )

    def _tool_checks(self) -> None:
        if not self.tools:
            self._add_check(
                "tool_frames",
                True,
                {"samples": 0, "message": "no tool configuration supplied"},
                required=False,
            )
            return
        limits = teleop_joint_limits()
        kinematics = RM75Kinematics(limits=limits)
        samples = _sample_configurations(
            self.rng, limits, self.settings.tool_samples
        )
        max_position = 0.0
        max_orientation = 0.0
        failures = 0
        for tool_name, tool in self.tools.items():
            for index, q in enumerate(samples):
                pin_pose = kinematics.forward(q, tool)
                reference_pose = self.reference.forward(q, tool)
                position_error, orientation_error = pose_errors(pin_pose, reference_pose)
                max_position = max(max_position, position_error)
                max_orientation = max(max_orientation, orientation_error)
                if (
                    position_error >= FK_POSITION_LIMIT_M
                    or orientation_error >= FK_ORIENTATION_LIMIT_RAD
                ):
                    failures += 1
                    self._record_failure(
                        "tool_frame",
                        index,
                        f"{tool_name} residual exceeded limit",
                        q,
                        position_error,
                        orientation_error,
                        tool_name,
                    )
        self._add_check(
            "tool_frames",
            failures == 0,
            {
                "tools": sorted(self.tools),
                "samples_per_tool": len(samples),
                "max_position_error_m": max_position,
                "max_orientation_error_rad": max_orientation,
            },
        )


def write_validation_report(
    output_dir: Path | str,
    summary: Dict[str, Any],
    failures: List[Dict[str, Any]],
) -> Tuple[Path, Path, Path]:
    directory = Path(output_dir)
    directory.mkdir(parents=True, exist_ok=True)
    json_path = directory / "stage1_summary.json"
    csv_path = directory / "stage1_failures.csv"
    markdown_path = directory / "stage1_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")

    fieldnames = [
        "category",
        "profile",
        "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=fieldnames)
        writer.writeheader()
        writer.writerows(failures)

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