#!/usr/bin/env python3
"""H1 calibration-v2 seeded pairing and audit-semantic 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 derive_trial_metrics  # noqa: E402
from core.model_contract import (  # noqa: E402
    MASTER_JOINT_NAMES,
    finite_joint_limits,
    load_models,
)
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_v2.json"
)
METRIC_CONFIG_PATH = (
    CODE_ROOT / "config" / "experiments" / "metrics_h1_calibration_v2.json"
)


class H1CalibrationV2Test(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)

    def _trajectory(self, trial):
        return _master_trajectory(
            trial, lower=self.lower, upper=self.upper
        )

    def _sew_trial(self, family: str, sample_count: int = 7):
        trial = next(
            item
            for item in self.plan["trials"]
            if item["method"]["method_id"] == "sew"
            and item["trajectory"]["family"] == family
            and item["replicate"] == 0
        )
        short_trial = deepcopy(trial)
        short_trial["trajectory"]["sample_count"] = sample_count
        return short_trial

    def test_v2_plan_is_bounded_seeded_and_strictly_paired(self) -> None:
        self.assertEqual(self.plan["pair_count"], 15)
        self.assertEqual(self.plan["trial_count"], 60)
        by_pair = {}
        by_trajectory = {}
        for trial in self.plan["trials"]:
            trajectory = self._trajectory(trial)
            self.assertTrue(np.all(trajectory > self.lower))
            self.assertTrue(np.all(trajectory < self.upper))
            by_pair.setdefault(trial["pair_id"], []).append(trajectory)
            by_trajectory.setdefault(
                trial["trajectory"]["trajectory_id"], {}
            ).setdefault(trial["replicate"], trajectory)

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

        for instances in by_trajectory.values():
            self.assertEqual(set(instances), {0, 1, 2})
            self.assertFalse(
                np.array_equal(instances[0], instances[1]),
                "replicates must not be deterministic pseudo-replicates",
            )
            self.assertFalse(np.array_equal(instances[1], instances[2]))

    def test_valid_low_manipulability_is_not_reach_clipped(self) -> None:
        payload = execute_h1_retargeting(
            self._sew_trial("low_manipulability_valid")
        )
        samples = payload.samples
        self.assertTrue(np.all(samples["map_pose_success"]))
        self.assertTrue(np.all(samples["map_differential_valid"]))
        self.assertTrue(np.any(samples["map_low_manipulability"]))
        self.assertTrue(np.all(samples["map_reach_clip_code"] == 0))
        self.assertTrue(
            np.all(
                samples["map_validity_reason_code"]
                == list(H1ValidityReason).index(H1ValidityReason.NONE)
            )
        )

    def test_reach_clip_has_reason_despite_pose_success(self) -> None:
        payload = execute_h1_retargeting(
            self._sew_trial("reach_clip_upper")
        )
        samples = payload.samples
        upper_code = list(H1ValidityReason).index(
            H1ValidityReason.REACH_CLIPPED_UPPER
        )
        self.assertTrue(np.all(samples["map_pose_success"]))
        self.assertTrue(np.all(samples["map_differential_valid"] == 0))
        self.assertTrue(np.all(samples["map_failure_code"] == 0))
        self.assertTrue(np.all(samples["map_reach_clip_code"] == 1))
        self.assertTrue(
            np.all(samples["map_validity_reason_code"] == upper_code)
        )
        invalid_events = [
            event
            for event in payload.events
            if event["event"] == "differential_invalid"
        ]
        self.assertEqual(len(invalid_events), len(samples["sample_index"]))
        self.assertTrue(
            all(
                event["validity_reason"]
                == H1ValidityReason.REACH_CLIPPED_UPPER.value
                for event in invalid_events
            )
        )
        metrics = derive_trial_metrics(
            samples, self.metric_configuration
        )
        self.assertEqual(metrics["h1_reach_clip_upper_fraction"], 1.0)
        self.assertEqual(metrics["h1_reach_clip_fraction"], 1.0)
        self.assertEqual(
            metrics["h1_primary_invalid_reason"],
            H1ValidityReason.REACH_CLIPPED_UPPER.value,
        )
        self.assertEqual(
            metrics["h1_validity_reason_histogram"],
            {H1ValidityReason.REACH_CLIPPED_UPPER.value: 7},
        )
        self.assertEqual(metrics["h1_unexplained_invalid_fraction"], 0.0)
        self.assertGreater(
            metrics["h1_min_slave_min_singular_value"], 0.0
        )

    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],
        )

    def test_every_invalid_sample_has_a_nonzero_audit_reason(self) -> None:
        payload = execute_h1_retargeting(self._sew_trial("joint_limit"))
        samples = payload.samples
        invalid = np.asarray(samples["map_differential_valid"]) == 0
        self.assertTrue(np.any(invalid))
        self.assertTrue(
            np.all(np.asarray(samples["map_validity_reason_code"])[invalid] != 0)
        )


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