#!/usr/bin/env python3
"""Bilateral v3 environment, force-audit, and stable-contact contracts."""

from pathlib import Path
import sys
import unittest

import numpy as np


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

from analysis.metrics import MetricError, derive_trial_metrics  # noqa: E402
from experiments.executors import (  # noqa: E402
    _bilateral_scenario_and_config,
    execute_bilateral_simulation,
)
from experiments.plan import build_trial_plan, load_document  # noqa: E402
from experiments.rng import named_seed_record  # noqa: E402
from simulate_closed_loop import Wall  # noqa: E402


CONFIG_ROOT = CODE_ROOT / "config" / "experiments"


def propagation_trial():
    return {
        "method": {"method_id": "proposed_energy"},
        "trajectory": {
            "trajectory_id": "unit_contact_v3",
            "family": "contact_roundtrip",
            "duration_s": 1.0,
            "contact_probe_fraction": 0.03,
        },
        "factors": {
            "map_policy": "source_stamped",
            "environment_profile": {
                "duration": 6.0,
                "mapping_hz": 40.0,
                "wall_fraction": 0.5,
                "stiffness": 3200.0,
                "damping": 25.0,
                "force_limit": 20.0,
                "transition_depth": 0.001,
                "probe_fraction": 0.0,
                "contact_probe_cycles": 0.0,
            },
            "haptic_profile": {
                "feedback_strength": 0.35,
                "energy_min": 0.0,
                "energy_max": 2.0,
                "energy_initial": 1.0,
            },
            # Direct aliases must override the coupled environment profile.
            "duration_s": 5.0,
            "wall_stiffness": 6400.0,
            "probe_fraction": 0.02,
        },
        "seeds": named_seed_record(11, {"test": "bilateral-v3"}),
    }


class BilateralCalibrationV3Test(unittest.TestCase):
    def test_environment_profile_and_direct_factors_propagate(self):
        _, config = _bilateral_scenario_and_config(propagation_trial())
        self.assertEqual(config.duration, 5.0)
        self.assertEqual(config.mapping_hz, 40.0)
        self.assertEqual(config.wall_fraction, 0.5)
        self.assertEqual(config.wall_stiffness, 6400.0)
        self.assertEqual(config.wall_damping, 25.0)
        self.assertEqual(config.wall_force_limit, 20.0)
        self.assertEqual(config.wall_transition_depth, 0.001)
        self.assertEqual(config.contact_probe_fraction, 0.02)
        self.assertEqual(config.contact_probe_cycles, 0.0)

    def test_wall_reports_raw_applied_and_saturation_compatibly(self):
        wall = Wall(
            point=np.zeros(3),
            normal=np.array([1.0, 0.0, 0.0]),
            stiffness=100.0,
            damping=10.0,
            force_limit=5.0,
        )
        contact = wall.contact(
            np.array([0.1, 0.0, 0.0]),
            np.array([2.0, 0.0, 0.0]),
        )
        self.assertAlmostEqual(contact.penetration, 0.1)
        self.assertAlmostEqual(contact.force_raw_N, 30.0)
        self.assertAlmostEqual(contact.force_applied_N, 5.0)
        self.assertTrue(contact.saturation_active)
        np.testing.assert_allclose(
            contact.wrench_applied[:3], [-5.0, 0.0, 0.0]
        )

        legacy_wrench, legacy_penetration = wall.wrench(
            np.array([0.1, 0.0, 0.0]),
            np.array([2.0, 0.0, 0.0]),
        )
        np.testing.assert_allclose(legacy_wrench, contact.wrench_applied)
        self.assertEqual(legacy_penetration, contact.penetration)

    def test_v3_grids_are_separate_and_h3_challenge_is_disabled(self):
        stable = build_trial_plan(
            load_document(
                CONFIG_ROOT
                / "bilateral_calibration_v3_stable_contact.json"
            )
        )
        challenge_spec = load_document(
            CONFIG_ROOT
            / "bilateral_calibration_v3_energy_challenge.json"
        )
        challenge = build_trial_plan(challenge_spec)
        challenge_metrics = load_document(
            CONFIG_ROOT
            / "metrics_bilateral_v3_energy_challenge.json"
        )

        self.assertEqual(stable["pair_count"], 18)
        self.assertEqual(stable["trial_count"], 18)
        self.assertEqual(challenge["pair_count"], 3)
        self.assertEqual(challenge["trial_count"], 3)
        self.assertFalse(challenge_spec["h3_eligible"])
        self.assertNotIn("h3", challenge_metrics["enabled"])
        self.assertIn("not frozen", stable["specification"]["status"])

    def test_haptic_profiles_and_challenge_share_exogenous_seed_groups(self):
        stable = build_trial_plan(
            load_document(
                CONFIG_ROOT
                / "bilateral_calibration_v3_stable_contact.json"
            )
        )
        challenge = build_trial_plan(
            load_document(
                CONFIG_ROOT
                / "bilateral_calibration_v3_energy_challenge.json"
            )
        )

        stable_k3200_rep0 = [
            trial
            for trial in stable["trials"]
            if trial["replicate"] == 0
            and trial["factors"]["environment_profile"]["profile_id"]
            == "stable_candidate_k3200"
        ]
        self.assertEqual(len(stable_k3200_rep0), 3)
        stable_seeds = {
            _bilateral_scenario_and_config(trial)[1].seed
            for trial in stable_k3200_rep0
        }
        self.assertEqual(len(stable_seeds), 1)

        challenge_rep0 = next(
            trial
            for trial in challenge["trials"]
            if trial["replicate"] == 0
        )
        challenge_seed = _bilateral_scenario_and_config(
            challenge_rep0
        )[1].seed
        self.assertEqual(challenge_seed, next(iter(stable_seeds)))

        stable_k6400_rep0 = next(
            trial
            for trial in stable["trials"]
            if trial["replicate"] == 0
            and trial["factors"]["environment_profile"]["profile_id"]
            == "stiff_candidate_k6400"
        )
        self.assertNotEqual(
            _bilateral_scenario_and_config(stable_k6400_rep0)[1].seed,
            challenge_seed,
        )

    def test_stable_six_second_smoke_has_contact_without_force_limiting(self):
        plan = build_trial_plan(
            load_document(
                CONFIG_ROOT
                / "bilateral_calibration_v3_stable_contact.json"
            )
        )
        trial = next(
            entry
            for entry in plan["trials"]
            if entry["factors"]["environment_profile"]["profile_id"]
            == "stable_candidate_k3200"
            and entry["factors"]["haptic_profile"]["profile_id"]
            == "gain_035_wide_tank"
        )
        payload = execute_bilateral_simulation(trial)
        applied = payload.samples["wall_force_applied_N"]

        self.assertEqual(applied.shape[0], 3000)
        self.assertTrue(
            np.all(payload.samples["wall_force_saturation_active"] == 0)
        )
        self.assertGreaterEqual(float(np.max(applied)), 0.5)
        self.assertLessEqual(float(np.max(applied)), 5.0)
        np.testing.assert_allclose(
            payload.samples["wall_force_raw_N"],
            payload.samples["wall_force_applied_N"],
        )

        metric_config = load_document(
            CONFIG_ROOT
            / "metrics_bilateral_v3_stable_contact.json"
        )
        metrics = derive_trial_metrics(
            payload.samples,
            metric_config,
        )
        self.assertTrue(metrics["bilateral_stable_contact_gate_pass"])
        self.assertEqual(metrics["bilateral_force_limit_hit_fraction"], 0.0)
        self.assertEqual(metrics["bilateral_limit_active_fraction"], 0.0)
        self.assertGreaterEqual(
            metrics["bilateral_contact_force_rms_N"], 0.1
        )
        incomplete = dict(payload.samples)
        incomplete.pop("haptic_rate_limit_active")
        with self.assertRaisesRegex(
            MetricError, "missing required bilateral audit fields"
        ):
            derive_trial_metrics(incomplete, metric_config)

    def test_energy_challenge_is_active_audited_and_unsaturated(self):
        plan = build_trial_plan(
            load_document(
                CONFIG_ROOT
                / "bilateral_calibration_v3_energy_challenge.json"
            )
        )
        payload = execute_bilateral_simulation(plan["trials"][0])
        metrics = derive_trial_metrics(
            payload.samples,
            load_document(
                CONFIG_ROOT
                / "metrics_bilateral_v3_energy_challenge.json"
            ),
        )

        self.assertTrue(metrics["h4_energy_audit_pass"])
        self.assertTrue(metrics["energy_challenge_gate_pass"])
        self.assertGreater(metrics["h4_shadow_floor_deficit_J"], 0.0)
        self.assertGreater(
            metrics["bilateral_energy_probe_raw_work_J"], 0.0
        )
        self.assertGreaterEqual(
            metrics["bilateral_projection_intervention_fraction"], 0.02
        )
        self.assertLessEqual(
            metrics["bilateral_projection_intervention_fraction"], 0.30
        )
        self.assertEqual(metrics["bilateral_force_limit_hit_fraction"], 0.0)
        self.assertEqual(metrics["bilateral_limit_active_fraction"], 0.0)
        self.assertTrue(metrics["bilateral_stable_contact_gate_pass"])
        self.assertTrue(np.all(payload.samples["h3_eligible"] == 0))
        self.assertFalse(payload.metadata["h3_eligible"])
        self.assertEqual(
            payload.metadata["energy_probe"]["window"],
            "hann_squared_sine",
        )
        with self.assertRaisesRegex(MetricError, "H3 is ineligible"):
            derive_trial_metrics(
                payload.samples,
                {"enabled": ["h3"], "h3": {}},
            )


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