#!/usr/bin/env python3
"""Strict Stage-B network metrics and paired-artifact tests."""

import csv
import json
from pathlib import Path
import sys
import tempfile
import unittest

import numpy as np


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

from analysis.make_paper_artifacts import (  # noqa: E402
    _apply_network_paired_gates,
    generate_paper_source_data,
)
from analysis.metrics import MetricError, derive_trial_metrics  # noqa: E402
from experiments.io import TrialPayload  # noqa: E402
from experiments.plan import build_trial_plan  # noqa: E402
from experiments.runner import run_trial_plan  # noqa: E402


LIMIT_FIELDS = (
    "master_joint_limit_active",
    "slave_joint_limit_active",
    "master_velocity_limit_active",
    "slave_velocity_limit_active",
    "master_acceleration_limit_active",
    "slave_acceleration_limit_active",
    "master_torque_saturation_active",
    "slave_torque_saturation_active",
    "haptic_rate_limit_active",
    "haptic_torque_saturation_active",
)


def network_samples(
    *,
    contact_expected: bool = True,
    contact_force_N: float = 0.2,
    zero_delay_error_rad: float = 0.02,
    reference_lag_error_rad: float = 0.01,
    feedback_lag_error_Nm: float = 0.01,
) -> dict[str, np.ndarray]:
    count = 5
    contact = np.full(count, contact_force_N, dtype=float)
    zeros = np.zeros(count, dtype=np.int8)
    samples: dict[str, np.ndarray] = {
        "energy_before_J": np.ones(count),
        "energy_after_J": np.ones(count),
        "tau_master_candidate": np.zeros((count, 2)),
        "tau_master_applied": np.zeros((count, 2)),
        "qd_master": np.zeros((count, 2)),
        "dt": np.full(count, 0.01),
        "configured_energy_min_J": np.zeros(count),
        "configured_energy_max_J": np.full(count, 2.0),
        "master_tracking_error": np.full(count, 0.01),
        "slave_tracking_error": np.full(count, 0.02),
        "contact_force_norm": contact,
        "rho": np.ones(count),
        "wall_force_raw_N": contact.copy(),
        "wall_force_applied_N": contact.copy(),
        "wall_force_saturation_active": zeros.copy(),
        "configured_wall_force_limit_N": np.full(count, 20.0),
        "energy_probe_raw_work_J": np.zeros(count),
        "forward_packet_state": np.array([4, 2, 1, 2, 1]),
        "forward_packet_active": np.ones(count, dtype=np.int8),
        "forward_packet_fresh": np.array([1, 0, 1, 0, 1]),
        "forward_packet_age": np.array([0.0, 0.01, 0.0, 0.01, 0.0]),
        "forward_packet_seq": np.array([0, 0, 1, 1, 2]),
        "return_packet_state": np.array([0, 4, 2, 1, 2]),
        "return_packet_active": np.array([0, 1, 1, 1, 1]),
        "return_packet_fresh": np.array([0, 1, 0, 1, 0]),
        "return_packet_age": np.array([np.nan, 0.0, 0.01, 0.0, 0.01]),
        "return_packet_seq": np.array([-1, 0, 0, 1, 1]),
        "slave_zero_delay_tracking_error": np.full(
            count, zero_delay_error_rad
        ),
        "slave_reference_lag_error": np.full(
            count, reference_lag_error_rad
        ),
        "return_feedback_lag_error": np.full(
            count, feedback_lag_error_Nm
        ),
        "contact_expected": np.full(
            count, int(contact_expected), dtype=np.int8
        ),
        "configured_forward_delay_s": np.zeros(count),
        "configured_return_delay_s": np.zeros(count),
        "configured_forward_jitter_s": np.zeros(count),
        "configured_return_jitter_s": np.zeros(count),
        "configured_forward_packet_loss": np.zeros(count),
        "configured_return_packet_loss": np.zeros(count),
        "configured_forward_timeout_s": np.full(count, 0.2),
        "configured_return_timeout_s": np.full(count, 0.2),
    }
    for name in LIMIT_FIELDS:
        samples[name] = zeros.copy()
    return samples


def network_metric_configuration(*, paired: bool = False) -> dict:
    configuration = {
        "enabled": ["h4", "bilateral", "network"],
        "h4": {"audit_tolerance_J": 1e-12},
        "bilateral": {
            "contact_force_threshold_N": 1e-6,
            "gates": {
                "minimum_contact_fraction": 0.0,
                "minimum_contact_rms_N": 0.0,
                "maximum_force_limit_hit_fraction": 0.0,
                "minimum_force_headroom_N": 1.0,
                "maximum_master_tracking_rmse_rad": 0.08,
                "maximum_slave_tracking_rmse_rad": 0.08,
                "minimum_projection_intervention_fraction": 0.0,
                "maximum_projection_intervention_fraction": 0.05,
                "maximum_limit_active_fraction": 0.0,
                "minimum_energy_probe_raw_work_J": 0.0,
            },
        },
        "network": {
            "gates": {
                "maximum_slave_zero_delay_tracking_rmse_rad": 0.1,
                "minimum_forward_active_fraction": 0.8,
                "minimum_return_active_fraction": 0.8,
                "minimum_forward_fresh_fraction": 0.5,
                "minimum_return_fresh_fraction": 0.4,
                "maximum_forward_internal_missing_fraction": 0.0,
                "maximum_return_internal_missing_fraction": 0.0,
                "maximum_forward_timeout_fraction": 0.0,
                "maximum_return_timeout_fraction": 0.0,
                "maximum_forward_packet_age_s": 0.2,
                "maximum_return_packet_age_s": 0.2,
                "minimum_contact_fraction": 0.5,
                "minimum_contact_rms_N": 0.1,
                "maximum_free_space_contact_force_N": 1e-6,
            }
        },
    }
    if paired:
        configuration["network"]["paired_gates"] = {
            "nominal_profile_id": "nominal",
            "maximum_tracking_rmse_delta_vs_nominal_rad": 0.01,
            "maximum_abs_contact_rms_relative_change_vs_nominal": 0.2,
        }
    return configuration


def paired_network_executor(trial):
    profile = trial["factors"]["network_profile"]["profile_id"]
    impaired = profile != "nominal"
    return TrialPayload(
        samples=network_samples(
            zero_delay_error_rad=0.025 if impaired else 0.02,
            feedback_lag_error_Nm=0.015 if impaired else 0.01,
            contact_force_N=0.22 if impaired else 0.20,
        ),
        metadata={
            "bilateral_data_group_id": "bilateral-data-test",
            "bilateral_data_seed_record_hash": "bilateral-seed-hash",
            "network_pair_group_id": "network-pair-test",
        },
    )


class NetworkMetricsV3Test(unittest.TestCase):
    def test_strict_packet_evidence_rejects_nan_state_and_sequence_forgery(self):
        configuration = network_metric_configuration()
        cases = []

        active_nan = network_samples()
        active_nan["forward_packet_age"][0] = np.nan
        cases.append((active_nan, "finite and non-negative while active"))

        inactive_number = network_samples()
        inactive_number["return_packet_age"][0] = 0.0
        cases.append((inactive_number, "must be NaN while inactive"))

        forged_state = network_samples()
        forged_state["forward_packet_fresh"][1] = 1
        cases.append((forged_state, "disagrees with forward_packet_state"))

        stale_sequence = network_samples()
        stale_sequence["forward_packet_seq"][2] = 0
        cases.append((stale_sequence, "strictly increase on fresh packets"))

        noninteger_state = network_samples()
        noninteger_state["return_packet_state"] = (
            noninteger_state["return_packet_state"].astype(float)
        )
        noninteger_state["return_packet_state"][2] = 2.5
        cases.append((noninteger_state, "integer states 0..4"))

        for samples, message in cases:
            with self.subTest(message=message):
                with self.assertRaisesRegex(MetricError, message):
                    derive_trial_metrics(samples, configuration)

    def test_freshness_and_internal_missing_gates_are_independent(self):
        configuration = network_metric_configuration()
        baseline = derive_trial_metrics(
            network_samples(),
            configuration,
        )
        for name in (
            "network_forward_fresh_gate_pass",
            "network_return_fresh_gate_pass",
            "network_forward_internal_missing_gate_pass",
            "network_return_internal_missing_gate_pass",
        ):
            self.assertTrue(baseline[name])

        compatibility_configuration = network_metric_configuration()
        for name in (
            "minimum_forward_fresh_fraction",
            "minimum_return_fresh_fraction",
            "maximum_forward_internal_missing_fraction",
            "maximum_return_internal_missing_fraction",
        ):
            compatibility_configuration["network"]["gates"].pop(name)
        compatibility = derive_trial_metrics(
            network_samples(),
            compatibility_configuration,
        )
        self.assertTrue(compatibility["network_local_gate_pass"])

        forward_stale = network_samples()
        forward_stale["forward_packet_state"][-1] = 2
        forward_stale["forward_packet_fresh"][-1] = 0
        forward_stale["forward_packet_seq"][-1] = 1

        return_stale = network_samples()
        return_stale["return_packet_state"][3] = 2
        return_stale["return_packet_fresh"][3] = 0
        return_stale["return_packet_seq"][3:] = 0

        forward_missing = network_samples()
        forward_missing["forward_packet_seq"][2:] = [2, 2, 3]

        return_missing = network_samples()
        return_missing["return_packet_seq"][3:] = 2

        cases = (
            (
                forward_stale,
                "network_forward_fresh_gate_pass",
            ),
            (
                return_stale,
                "network_return_fresh_gate_pass",
            ),
            (
                forward_missing,
                "network_forward_internal_missing_gate_pass",
            ),
            (
                return_missing,
                "network_return_internal_missing_gate_pass",
            ),
        )
        for samples, failed_gate in cases:
            with self.subTest(failed_gate=failed_gate):
                metrics = derive_trial_metrics(samples, configuration)
                self.assertFalse(metrics[failed_gate])
                self.assertFalse(metrics["network_local_gate_pass"])
                self.assertFalse(metrics["network_local_full_gate_pass"])

    def test_contact_and_free_space_gates_are_condition_specific(self):
        configuration = network_metric_configuration()
        contact = derive_trial_metrics(network_samples(), configuration)
        self.assertTrue(contact["network_contact_expected"])
        self.assertTrue(contact["network_contact_condition_gate_pass"])
        self.assertTrue(contact["network_local_full_gate_pass"])

        missing_contact = derive_trial_metrics(
            network_samples(contact_force_N=0.0),
            configuration,
        )
        self.assertFalse(
            missing_contact["network_contact_fraction_gate_pass"]
        )
        self.assertFalse(missing_contact["network_contact_rms_gate_pass"])
        self.assertFalse(missing_contact["network_local_full_gate_pass"])

        free_space = derive_trial_metrics(
            network_samples(
                contact_expected=False,
                contact_force_N=0.0,
            ),
            configuration,
        )
        self.assertFalse(free_space["network_contact_expected"])
        self.assertTrue(free_space["network_free_space_peak_gate_pass"])
        self.assertTrue(free_space["network_local_full_gate_pass"])

        false_contact = derive_trial_metrics(
            network_samples(
                contact_expected=False,
                contact_force_N=0.01,
            ),
            configuration,
        )
        self.assertFalse(
            false_contact["network_free_space_peak_gate_pass"]
        )
        self.assertFalse(false_contact["network_local_full_gate_pass"])

    def test_network_family_requires_h4_and_bilateral_first(self):
        with self.assertRaisesRegex(
            MetricError, "must follow H4 and bilateral"
        ):
            derive_trial_metrics(
                network_samples(),
                {"enabled": ["network"], "network": {}},
            )

    def test_paired_artifacts_preserve_ids_and_compute_nominal_deltas(self):
        plan = build_trial_plan(
            {
                "study_id": "network_paired_source_data",
                "split": "pilot",
                "root_seed": 13,
                "replicates": 1,
                "methods": ["proposed_energy"],
                "trajectories": [
                    {
                        "trajectory_id": "contact_network",
                        "family": "contact_roundtrip",
                    }
                ],
                "factors": {
                    "network_profile": [
                        {"profile_id": "nominal"},
                        {"profile_id": "delay"},
                    ]
                },
            }
        )
        configuration = network_metric_configuration(paired=True)
        with tempfile.TemporaryDirectory() as temporary:
            batch = Path(temporary) / "batch"
            run_trial_plan(plan, batch, paired_network_executor)
            manifest = generate_paper_source_data(batch, configuration)
            self.assertEqual(manifest["family_row_counts"]["network"], 2)
            rows = [
                json.loads(line)
                for line in (
                    batch / "derived" / "trial_metrics.jsonl"
                ).read_text(encoding="utf-8").splitlines()
            ]
            nominal = next(
                row
                for row in rows
                if row["network_profile_id"] == "nominal"
            )
            delayed = next(
                row
                for row in rows
                if row["network_profile_id"] == "delay"
            )
            self.assertEqual(nominal["bilateral_data_group_id"], "bilateral-data-test")
            self.assertEqual(
                nominal["bilateral_data_seed_record_hash"],
                "bilateral-seed-hash",
            )
            self.assertEqual(
                nominal["network_pair_group_id"], "network-pair-test"
            )
            self.assertEqual(
                nominal[
                    "network_zero_delay_tracking_rmse_delta_vs_nominal_rad"
                ],
                0.0,
            )
            self.assertEqual(
                nominal[
                    "network_return_feedback_lag_rmse_delta_vs_nominal_Nm"
                ],
                0.0,
            )
            self.assertAlmostEqual(
                delayed[
                    "network_zero_delay_tracking_rmse_delta_vs_nominal_rad"
                ],
                0.005,
            )
            self.assertAlmostEqual(
                delayed[
                    "network_return_feedback_lag_rmse_delta_vs_nominal_Nm"
                ],
                0.005,
            )
            self.assertAlmostEqual(
                delayed[
                    "network_contact_rms_relative_change_vs_nominal"
                ],
                0.1,
            )
            self.assertTrue(delayed["network_paired_gate_pass"])
            self.assertTrue(delayed["network_full_gate_pass"])

            with (
                batch / "paper" / "source_data" / "network.csv"
            ).open(newline="", encoding="utf-8") as stream:
                table = list(csv.DictReader(stream))
            self.assertEqual(len(table), 2)
            for name in (
                "bilateral_data_group_id",
                "bilateral_data_seed_record_hash",
                "network_pair_group_id",
                "network_profile_id",
            ):
                self.assertIn(name, table[0])

    def test_free_space_paired_contact_metrics_are_not_applicable(self):
        configuration = network_metric_configuration(paired=True)

        def row(profile, tracking_rmse):
            return {
                "trial_id": f"trial-{profile}",
                "network_pair_group_id": "free-space-group",
                "network_profile_id": profile,
                "method_id": "proposed_energy",
                "trajectory_id": "free-space",
                "replicate": 0,
                "network_local_full_gate_pass": True,
                "network_contact_expected": False,
                "network_slave_zero_delay_tracking_rmse_rad": tracking_rmse,
                "network_return_feedback_lag_rmse_Nm": 0.01,
                "network_contact_force_rms_N": 0.0,
            }

        rows = [row("nominal", 0.02), row("delay", 0.04)]
        _apply_network_paired_gates(rows, configuration)

        nominal, delayed = rows
        for result in rows:
            self.assertIsNone(
                result[
                    "network_contact_rms_relative_change_vs_nominal"
                ]
            )
            self.assertIsNone(
                result[
                    "network_contact_force_rms_relative_change_vs_nominal"
                ]
            )
            self.assertIsNone(
                result["network_paired_contact_gate_pass"]
            )
            self.assertEqual(
                result["network_paired_gate_pass"],
                result["network_paired_tracking_gate_pass"],
            )
        self.assertTrue(nominal["network_paired_gate_pass"])
        self.assertFalse(delayed["network_paired_gate_pass"])

    def test_paired_artifacts_reject_missing_or_duplicate_nominal(self):
        configuration = network_metric_configuration(paired=True)

        def row(profile):
            return {
                "trial_id": f"trial-{profile}",
                "network_pair_group_id": "group",
                "network_profile_id": profile,
                "method_id": "proposed_energy",
                "trajectory_id": "contact",
                "replicate": 0,
                "network_local_full_gate_pass": True,
                "network_contact_expected": True,
                "network_slave_zero_delay_tracking_rmse_rad": 0.02,
                "network_return_feedback_lag_rmse_Nm": 0.01,
                "network_contact_force_rms_N": 0.2,
            }

        with self.assertRaisesRegex(ValueError, "missing nominal"):
            _apply_network_paired_gates([row("delay")], configuration)
        with self.assertRaisesRegex(ValueError, "duplicate nominal"):
            _apply_network_paired_gates(
                [row("nominal"), row("nominal")],
                configuration,
            )


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