#!/usr/bin/env python3
"""End-to-end source-data generation from completed raw trials."""

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 generate_paper_source_data  # 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


def wrench_executor(_trial):
    reference = np.zeros((4, 6), dtype=float)
    estimate = reference.copy()
    estimate[:, 0] = 1.0
    return TrialPayload(
        samples={
            "wrench_estimated": estimate,
            "wrench_reference": reference,
        }
    )


class PaperSourceDataTest(unittest.TestCase):
    def test_source_tables_are_generated_from_raw_npz(self):
        plan = build_trial_plan(
            {
                "study_id": "h2_source_data",
                "split": "pilot",
                "root_seed": 4,
                "replicates": 1,
                "methods": ["dls", "undamped"],
                "trajectories": [
                    {"trajectory_id": "load_001", "family": "static_load"}
                ],
                "factors": {"axis": ["x"]},
            }
        )
        configuration = {"enabled": ["h2"], "h2": {}}
        with tempfile.TemporaryDirectory() as temporary:
            batch = Path(temporary) / "batch"
            run_trial_plan(plan, batch, wrench_executor)
            manifest = generate_paper_source_data(batch, configuration)
            self.assertEqual(manifest["row_count"], 2)
            metric_path = batch / "derived" / "trial_metrics.jsonl"
            rows = [
                json.loads(line)
                for line in metric_path.read_text(encoding="utf-8").splitlines()
            ]
            self.assertEqual(len(rows), 2)
            self.assertTrue(
                all(row["h2_force_rmse_N"] == 1.0 for row in rows)
            )
            source_path = batch / "paper" / "source_data" / "h2.csv"
            with source_path.open(newline="", encoding="utf-8") as stream:
                table = list(csv.DictReader(stream))
            self.assertEqual(len(table), 2)
            self.assertIn("h2_force_rmse_N", table[0])
            self.assertTrue((batch / "paper" / "artifact_manifest.json").is_file())

    def test_metric_family_method_filter_prevents_mixed_supervisors(self):
        plan = build_trial_plan(
            {
                "study_id": "filtered_source_data",
                "split": "pilot",
                "root_seed": 5,
                "replicates": 1,
                "methods": ["dls", "undamped"],
                "trajectories": ["load_001"],
            }
        )
        configuration = {
            "enabled": ["h2"],
            "h2": {"include_methods": ["dls"]},
        }
        with tempfile.TemporaryDirectory() as temporary:
            batch = Path(temporary) / "batch"
            run_trial_plan(plan, batch, wrench_executor)
            manifest = generate_paper_source_data(batch, configuration)
            self.assertEqual(manifest["row_count"], 2)
            self.assertEqual(manifest["family_row_counts"], {"h2": 1})

            metric_rows = [
                json.loads(line)
                for line in (batch / "derived" / "trial_metrics.jsonl")
                .read_text(encoding="utf-8")
                .splitlines()
            ]
            dls_row = next(row for row in metric_rows if row["method_id"] == "dls")
            undamped_row = next(
                row for row in metric_rows if row["method_id"] == "undamped"
            )
            self.assertIn("h2_force_rmse_N", dls_row)
            self.assertNotIn("h2_force_rmse_N", undamped_row)

            with (
                batch / "paper" / "source_data" / "h2.csv"
            ).open(newline="", encoding="utf-8") as stream:
                table = list(csv.DictReader(stream))
            self.assertEqual([row["method_id"] for row in table], ["dls"])


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