#!/usr/bin/env python3
"""Contracts for deterministic paired plans and random substreams."""

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.plan import build_trial_plan, validate_trial_plan  # noqa: E402
from experiments.rng import generator_from_record, named_seed_record  # noqa: E402
from experiments.schema import ContractError  # noqa: E402


def example_specification():
    return {
        "study_id": "g0c_bilateral_contract",
        "split": "calibration",
        "root_seed": 20260727,
        "replicates": 2,
        "methods": ["proposed", "direct"],
        "trajectories": [
            {
                "trajectory_id": "contact_probe_001",
                "family": "approach_contact_return",
                "parameters": {"amplitude": 0.1},
            }
        ],
        "factors": {"delay_ms": [0, 80], "stiffness_N_per_m": [500]},
    }


class ExperimentPlanContractTest(unittest.TestCase):
    def test_plan_is_deterministic_and_paired(self):
        first = build_trial_plan(example_specification())
        second = build_trial_plan(example_specification())
        self.assertEqual(first, second)
        self.assertEqual(first["pair_count"], 4)
        self.assertEqual(first["trial_count"], 8)

        by_pair = {}
        for trial in first["trials"]:
            by_pair.setdefault(trial["pair_id"], []).append(trial)
        for pair in by_pair.values():
            self.assertEqual(len(pair), 2)
            self.assertEqual(pair[0]["seeds"], pair[1]["seeds"])
            self.assertEqual(pair[0]["trajectory"], pair[1]["trajectory"])
            self.assertEqual(pair[0]["factors"], pair[1]["factors"])
            self.assertNotEqual(
                pair[0]["method"]["method_id"],
                pair[1]["method"]["method_id"],
            )

    def test_named_streams_are_order_independent_and_distinct(self):
        forward = named_seed_record(
            9,
            {"pair": 3},
            ["trajectory", "sensor", "network"],
        )
        reverse = named_seed_record(
            9,
            {"pair": 3},
            ["network", "sensor", "trajectory"],
        )
        self.assertEqual(forward, reverse)
        self.assertEqual(len({tuple(value) for value in forward.values()}), 3)

        generator_a = generator_from_record(forward, "trajectory")
        generator_b = generator_from_record(reverse, "trajectory")
        np.testing.assert_array_equal(
            generator_a.integers(0, 2**31, size=16),
            generator_b.integers(0, 2**31, size=16),
        )

    def test_plan_hash_detects_mutation(self):
        plan = build_trial_plan(example_specification())
        plan["trials"][0]["factors"]["delay_ms"] = 999
        with self.assertRaises(ContractError):
            validate_trial_plan(plan)


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