#!/usr/bin/env python3
"""H1 calibration-v3 branch-crossing and differential evidence contracts."""

from __future__ import annotations

from copy import deepcopy
from pathlib import Path
import sys
import unittest

import numpy as np


CODE_ROOT = Path(__file__).resolve().parents[1]
if str(CODE_ROOT) not in sys.path:
    sys.path.insert(0, str(CODE_ROOT))

from analysis.metrics import MetricError, derive_trial_metrics  # noqa: E402
from core.model_contract import (  # noqa: E402
    MASTER_JOINT_NAMES,
    finite_joint_limits,
    load_models,
)
from core.retargeting_baselines import (  # noqa: E402
    build_canonical_sew_target_baselines,
)
from experiments.executors import (  # noqa: E402
    H1ValidityReason,
    _master_trajectory,
    execute_h1_retargeting,
)
from experiments.plan import build_trial_plan, load_document  # noqa: E402


CONFIG_PATH = (
    CODE_ROOT / "config" / "experiments" / "h1_calibration_v3.json"
)
METRIC_CONFIG_PATH = (
    CODE_ROOT
    / "config"
    / "experiments"
    / "metrics_h1_calibration_v3.json"
)


class H1CalibrationV3Test(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.models = load_models(add_simulated_tcp=True)
        cls.lower, cls.upper = finite_joint_limits(
            cls.models.master, MASTER_JOINT_NAMES
        )
        cls.plan = build_trial_plan(load_document(CONFIG_PATH))
        cls.metric_configuration = load_document(METRIC_CONFIG_PATH)
        _, cls.sew_method = build_canonical_sew_target_baselines(cls.models)

        cls.sew_payloads = {}
        for path_type in ("linear", "cosine_roundtrip"):
            cls.sew_payloads[path_type] = execute_h1_retargeting(
                cls._trial("sew", path_type)
            )
        cls.na_payload = execute_h1_retargeting(
            cls._trial("bounded_dls_ik", "linear")
        )

    @classmethod
    def _trial(cls, method_id: str, path_type: str):
        return next(
            trial
            for trial in cls.plan["trials"]
            if trial["method"]["method_id"] == method_id
            and trial["trajectory"]["path_type"] == path_type
        )

    @classmethod
    def _trajectory(cls, trial):
        return _master_trajectory(
            trial, lower=cls.lower, upper=cls.upper
        )

    def test_explicit_linear_and_roundtrip_paths_are_strictly_paired(self) -> None:
        self.assertEqual(self.plan["pair_count"], 2)
        self.assertEqual(self.plan["trial_count"], 8)
        by_pair = {}
        for trial in self.plan["trials"]:
            trajectory = self._trajectory(trial)
            by_pair.setdefault(trial["pair_id"], []).append(trajectory)
            self.assertLess(
                np.max(np.linalg.norm(np.diff(trajectory, axis=0), axis=1)),
                0.05,
            )

        for trajectories in by_pair.values():
            self.assertEqual(len(trajectories), 4)
            for candidate in trajectories[1:]:
                np.testing.assert_array_equal(candidate, trajectories[0])

        linear_trial = self._trial("sew", "linear")
        linear = self._trajectory(linear_trial)
        np.testing.assert_array_equal(
            linear[0], np.asarray(linear_trial["trajectory"]["start"])
        )
        np.testing.assert_array_equal(
            linear[-1], np.asarray(linear_trial["trajectory"]["end"])
        )

        roundtrip_trial = self._trial("sew", "cosine_roundtrip")
        roundtrip = self._trajectory(roundtrip_trial)
        np.testing.assert_array_equal(
            roundtrip[0], np.asarray(roundtrip_trial["trajectory"]["start"])
        )
        np.testing.assert_allclose(
            roundtrip[len(roundtrip) // 2],
            np.asarray(roundtrip_trial["trajectory"]["end"]),
            atol=1e-15,
            rtol=0.0,
        )
        np.testing.assert_array_equal(roundtrip[-1], roundtrip[0])

    def test_explicit_path_contract_rejects_partial_or_unknown_paths(self) -> None:
        trial = deepcopy(self._trial("sew", "linear"))
        del trial["trajectory"]["end"]
        with self.assertRaisesRegex(ValueError, "both start and end"):
            self._trajectory(trial)

        trial = deepcopy(self._trial("sew", "linear"))
        trial["trajectory"]["path_type"] = "triangle"
        with self.assertRaisesRegex(ValueError, "path_type"):
            self._trajectory(trial)

    def test_target_debug_has_phi_reference_and_reach_margins(self) -> None:
        linear = self.sew_payloads["linear"].samples
        crossing = np.flatnonzero(
            np.abs(np.diff(linear["map_sew_phi_rad"])) > np.pi
        )
        self.assertEqual(crossing.size, 1)
        debug = self.sew_method.mapper._target_from_master(
            linear["q_master"][int(crossing[0]) + 1]
        )
        for field in (
            "phi_rad",
            "reference_axis_norm",
            "reach_lower_margin_m",
            "reach_upper_margin_m",
            "master_arm_normal_norm",
        ):
            self.assertIn(field, debug)
            self.assertTrue(np.isfinite(debug[field]))
        self.assertGreater(debug["reference_axis_norm"], 0.0)
        self.assertGreater(debug["reach_lower_margin_m"], 0.0)
        self.assertGreater(debug["reach_upper_margin_m"], 0.0)

        for payload in self.sew_payloads.values():
            samples = payload.samples
            self.assertTrue(np.all(samples["map_reach_clip_code"] == 0))
            self.assertTrue(
                np.all(samples["map_reach_lower_margin_m"] > 0.0)
            )
            self.assertTrue(
                np.all(samples["map_reach_upper_margin_m"] > 0.0)
            )
            self.assertTrue(np.all(samples["map_reference_axis_norm"] > 0.0))

    def test_actual_sew_differential_and_branch_metrics_are_reconstructable(
        self,
    ) -> None:
        expected_crossings = {"linear": 1, "cosine_roundtrip": 2}
        for path_type, payload in self.sew_payloads.items():
            with self.subTest(path_type=path_type):
                samples = payload.samples
                self.assertEqual(
                    samples["map_differential_A"].shape, (81, 7, 7)
                )
                self.assertTrue(
                    np.all(samples["map_differential_applicable"] == 1)
                )
                self.assertTrue(np.all(samples["map_branch_smooth"] == 1))
                self.assertTrue(
                    np.all(samples["map_differential_valid"] == 1)
                )
                self.assertTrue(
                    np.all(np.isfinite(samples["map_differential_A"]))
                )
                self.assertTrue(
                    np.all(samples["map_differential_runtime_s"] >= 0.0)
                )
                self.assertTrue(
                    np.all(
                        np.isfinite(
                            samples[
                                "map_differential_max_"
                                "one_sided_consistency"
                            ]
                        )
                    )
                )

                metrics = derive_trial_metrics(
                    samples, self.metric_configuration
                )
                self.assertEqual(metrics["h1_F_r"], 0)
                self.assertEqual(metrics["h1_D_r"], 0)
                self.assertEqual(metrics["h1_C_r"], 0)
                self.assertEqual(
                    metrics["h1_phi_raw_wrap_crossing_count"],
                    expected_crossings[path_type],
                )
                self.assertTrue(metrics["h1_phi_wrap_metric_valid"])
                self.assertLess(
                    metrics[
                        "h1_phi_wrap_crossing_max_slave_joint_step_rad"
                    ],
                    0.25,
                )
                self.assertEqual(
                    metrics["h1_differential_valid_fraction"], 1.0
                )
                for name in (
                    "h1_pose_runtime_p50_ms",
                    "h1_pose_runtime_p95_ms",
                    "h1_pose_runtime_p99_ms",
                    "h1_pose_runtime_max_ms",
                    "h1_pose_runtime_warm_p95_ms",
                    "h1_differential_runtime_p50_ms",
                    "h1_differential_runtime_p95_ms",
                    "h1_differential_runtime_p99_ms",
                    "h1_differential_runtime_max_ms",
                    "h1_feedback_ready_runtime_p95_ms",
                ):
                    self.assertIn(name, metrics)
                    self.assertIsNotNone(metrics[name])
                    self.assertGreaterEqual(metrics[name], 0.0)

    def test_non_sew_actual_differential_is_explicit_na_not_failure(self) -> None:
        payload = self.na_payload
        samples = payload.samples
        self.assertTrue(
            np.all(samples["map_differential_applicable"] == 0)
        )
        self.assertTrue(np.all(np.isnan(samples["map_differential_A"])))
        np.testing.assert_array_equal(
            samples["map_differential_valid"],
            samples["map_branch_smooth"],
        )
        self.assertIsNotNone(payload.metadata["differential_n_a_reason"])

        metrics = derive_trial_metrics(samples, self.metric_configuration)
        self.assertEqual(metrics["h1_F_r"], 0)
        self.assertEqual(metrics["h1_D_r"], 0)
        self.assertEqual(metrics["h1_differential_applicable_fraction"], 0.0)
        self.assertIsNone(metrics["h1_differential_valid_fraction"])
        self.assertIsNone(metrics["h1_differential_runtime_p95_ms"])

    def test_v3_metrics_reject_missing_or_nonbinary_evidence(self) -> None:
        original = self.sew_payloads["linear"].samples

        missing = dict(original)
        missing.pop("map_differential_applicable")
        with self.assertRaisesRegex(
            MetricError, "requires v3 evidence fields"
        ):
            derive_trial_metrics(missing, self.metric_configuration)

        nonbinary = {
            name: value.copy() for name, value in original.items()
        }
        nonbinary["map_pose_success"] = np.asarray(
            nonbinary["map_pose_success"], dtype=float
        )
        nonbinary["map_pose_success"][0] = np.nan
        with self.assertRaisesRegex(MetricError, "finite 0/1 flags"):
            derive_trial_metrics(nonbinary, self.metric_configuration)

    def test_empty_or_missing_wrap_evidence_cannot_look_perfect(self) -> None:
        original = self.sew_payloads["linear"].samples
        no_crossing = {
            name: value.copy() for name, value in original.items()
        }
        no_crossing["map_sew_phi_rad"] = np.zeros_like(
            no_crossing["map_sew_phi_rad"]
        )
        no_crossing_metrics = derive_trial_metrics(
            no_crossing, self.metric_configuration
        )
        self.assertFalse(
            no_crossing_metrics["h1_phi_wrap_metric_valid"]
        )
        self.assertIsNone(
            no_crossing_metrics[
                "h1_phi_wrap_crossing_max_slave_joint_step_rad"
            ]
        )

        no_accepted = {
            name: value.copy() for name, value in original.items()
        }
        no_accepted["map_accepted"] = np.zeros_like(
            no_accepted["map_accepted"]
        )
        empty_metrics = derive_trial_metrics(
            no_accepted, self.metric_configuration
        )
        self.assertFalse(empty_metrics["h1_composite_metric_valid"])
        self.assertIsNone(empty_metrics["h1_F_r"])
        self.assertIsNone(empty_metrics["h1_D_r"])
        self.assertIsNone(empty_metrics["h1_C_r"])

    def test_metric_reason_labels_match_executor_enum(self) -> None:
        self.assertEqual(
            self.metric_configuration["h1"]["validity_reason_labels"],
            [member.value for member in H1ValidityReason],
        )


if __name__ == "__main__":
    unittest.main()
