#!/usr/bin/env python3
"""H2 data pairing and conditioning-stratum regression tests."""

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 experiments.executors import execute_h2_synthetic  # noqa: E402
from experiments.plan import build_trial_plan, load_document  # noqa: E402


def h2_pairing_specification():
    return {
        "study_id": "h2_pairing_contract",
        "split": "calibration",
        "root_seed": 37,
        "replicates": 2,
        "methods": ["scaled_dls", "undamped_svd"],
        "trajectories": [
            {
                "trajectory_id": "short_dynamic_wrench",
                "family": "synthetic_dynamic",
                "sample_count": 12,
            }
        ],
        "factors": {
            "characteristic_length_m": [0.2, 0.4],
            "damping": [0.02, 0.05],
            "h2_data_root_seed": [1701],
            "min_scaled_singular": [0.02],
            "model_error_std": [0.02],
            "operational_min_scaled_singular": [0.05],
            "torque_noise_std_Nm": [0.002, 0.02],
            "truth_characteristic_length_m": [0.3],
        },
    }


def select_trial(
    plan,
    *,
    replicate,
    method,
    characteristic_length,
    damping,
    noise,
):
    matches = [
        trial
        for trial in plan["trials"]
        if trial["replicate"] == replicate
        and trial["method"]["method_id"] == method
        and trial["factors"]["characteristic_length_m"]
        == characteristic_length
        and trial["factors"]["damping"] == damping
        and trial["factors"]["torque_noise_std_Nm"] == noise
    ]
    if len(matches) != 1:
        raise AssertionError(f"expected one H2 trial, found {len(matches)}")
    return matches[0]


class H2SyntheticPairingTest(unittest.TestCase):
    def test_algorithm_candidates_reuse_identical_physical_data(self):
        plan = build_trial_plan(h2_pairing_specification())
        first_trial = select_trial(
            plan,
            replicate=0,
            method="scaled_dls",
            characteristic_length=0.2,
            damping=0.02,
            noise=0.002,
        )
        second_trial = select_trial(
            plan,
            replicate=0,
            method="undamped_svd",
            characteristic_length=0.4,
            damping=0.05,
            noise=0.002,
        )
        # The general experiment pair changes with ell/damping.  H2's explicit
        # physical-data group must still bind both candidates to the same data.
        self.assertNotEqual(first_trial["pair_id"], second_trial["pair_id"])
        self.assertNotEqual(first_trial["seeds"], second_trial["seeds"])

        first = execute_h2_synthetic(first_trial)
        second = execute_h2_synthetic(second_trial)
        self.assertEqual(
            first.metadata["h2_data_group_id"],
            second.metadata["h2_data_group_id"],
        )
        self.assertEqual(
            first.metadata["h2_data_seed_record"],
            second.metadata["h2_data_seed_record"],
        )
        self.assertNotEqual(
            first.metadata["calibration_id"],
            second.metadata["calibration_id"],
            "scanned estimator candidates must not share a frozen-looking ID",
        )
        for field in (
            "wrench_reference",
            "qd_slave",
            "jacobian_truth",
            "jacobian_estimator",
            "sensor_noise_Nm",
            "tau_residual_raw",
        ):
            np.testing.assert_array_equal(
                first.samples[field], second.samples[field]
            )
        self.assertFalse(
            np.array_equal(
                first.samples["wrench_estimated"],
                second.samples["wrench_estimated"],
            )
        )

    def test_replicates_and_physical_factor_cells_get_distinct_data(self):
        plan = build_trial_plan(h2_pairing_specification())
        baseline = execute_h2_synthetic(
            select_trial(
                plan,
                replicate=0,
                method="scaled_dls",
                characteristic_length=0.2,
                damping=0.02,
                noise=0.002,
            )
        )
        next_replicate = execute_h2_synthetic(
            select_trial(
                plan,
                replicate=1,
                method="scaled_dls",
                characteristic_length=0.2,
                damping=0.02,
                noise=0.002,
            )
        )
        different_noise = execute_h2_synthetic(
            select_trial(
                plan,
                replicate=0,
                method="scaled_dls",
                characteristic_length=0.2,
                damping=0.02,
                noise=0.02,
            )
        )
        self.assertNotEqual(
            baseline.metadata["h2_data_group_id"],
            next_replicate.metadata["h2_data_group_id"],
        )
        self.assertNotEqual(
            baseline.metadata["h2_data_group_id"],
            different_noise.metadata["h2_data_group_id"],
        )
        self.assertFalse(
            np.array_equal(
                baseline.samples["jacobian_truth"],
                next_replicate.samples["jacobian_truth"],
            )
        )
        self.assertFalse(
            np.array_equal(
                baseline.samples["tau_residual_raw"],
                different_noise.samples["tau_residual_raw"],
            )
        )

    def test_operational_flag_is_separate_from_numerical_rank(self):
        specification = h2_pairing_specification()
        specification["replicates"] = 1
        specification["factors"]["characteristic_length_m"] = [0.3]
        specification["factors"]["damping"] = [0.03]
        specification["factors"]["torque_noise_std_Nm"] = [0.01]
        specification["factors"]["operational_min_scaled_singular"] = [10.0]
        trial = build_trial_plan(specification)["trials"][0]
        payload = execute_h2_synthetic(trial)

        np.testing.assert_array_equal(
            payload.samples["solver_numerical_rank_deficient"],
            np.zeros(12, dtype=np.int8),
        )
        np.testing.assert_array_equal(
            payload.samples["solver_operationally_ill_conditioned"],
            np.ones(12, dtype=np.int8),
        )
        np.testing.assert_array_equal(
            payload.samples["operational_min_scaled_singular_threshold"],
            np.full(12, 10.0),
        )
        definition = payload.metadata[
            "operational_ill_conditioning_definition"
        ]
        self.assertTrue(definition["numerical_rank_is_reported_separately"])

    def test_v2_scan_has_sixteen_physical_groups_and_576_trials(self):
        specification = load_document(
            CODE_ROOT / "config" / "experiments" / "h2_calibration_v2.json"
        )
        plan = build_trial_plan(specification)
        self.assertEqual(plan["pair_count"], 144)
        self.assertEqual(plan["trial_count"], 576)

        group_ids = set()
        for trial in plan["trials"]:
            payload = execute_h2_synthetic(
                {
                    **trial,
                    "trajectory": {
                        **trial["trajectory"],
                        "sample_count": 8,
                    },
                }
            )
            group_ids.add(payload.metadata["h2_data_group_id"])
            if len(group_ids) == 16:
                break
        self.assertEqual(len(group_ids), 16)


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