"""Synthetic CPU-only end-to-end contracts for canonical Phase 3--5 OID."""
from __future__ import annotations

import hashlib
import json
from pathlib import Path

import numpy as np
import pytest

from OID_analysis.v2.oid import (
    _admissible_shift_record,
    _dominant_period,
    _scores,
    _shift,
    build_canonical_paired_product,
    load_canonical_product,
    run_phase_3_5,
)


def _realization(root: Path, realization: int, *, missing_telemetry: str | None = None) -> Path:
    rng = np.random.default_rng(100 + realization)
    n, ny, nx, interval = 72, 3, 4, 5
    steps = np.arange(interval, interval * (n + 1), interval, dtype=np.int64)
    times = steps.astype(float) * 0.02
    sample_ids = np.arange(realization * 1000, realization * 1000 + n, dtype=np.int64)
    latent = rng.normal(size=(n, 3))
    phase = 2 * np.pi * np.arange(n) / 24
    latent[:, 0] += 5.0 * np.sin(phase)
    latent[:, 1] += 4.0 * np.cos(phase)
    spatial = np.linalg.qr(rng.normal(size=(2 * ny * nx, 3)))[0]
    # Keep the physical map shared enough that held-out prediction is meaningful.
    shared = np.random.default_rng(9).normal(size=(2 * ny * nx, 3))
    spatial = np.linalg.qr(shared)[0]
    dq = latent @ spatial.T + 0.01 * rng.normal(size=(n, 2 * ny * nx))
    baseline = 0.05 * rng.normal(size=(n, 2 * ny * nx))
    controlled = baseline + dq
    timeline = {"lattice_steps": steps, "control_indices": (steps - 1) // interval,
                "control_offsets": (steps - 1) % interval + 1,
                "control_interval": np.asarray(interval, np.int64)}
    directory = root / f"raw_{realization}"; directory.mkdir()
    for name, values in (("ctl", controlled), ("blk", baseline)):
        np.savez(directory / f"{name}_fields.npz",
                 ux=values[:, :ny * nx].reshape(n, ny, nx),
                 uy=values[:, ny * nx:].reshape(n, ny, nx), **timeline)
    force_map = np.array([[1.2, -0.4], [0.3, 0.9], [-0.7, 0.5]])
    sensor_map = np.array([[0.4, -0.3, 0.8], [1.0, 0.2, -0.2], [-0.5, 0.7, 0.1]])
    delta_forces = latent @ force_map + 0.02 * rng.normal(size=(n, 2))
    delta_sensors = latent @ sensor_map + 0.02 * rng.normal(size=(n, 3))
    base_forces, base_sensors = rng.normal(size=(n, 2)), rng.normal(size=(n, 3))
    common = {"lattice_steps": steps, "sample_ids": sample_ids, "times": times}
    controlled_t = {**common, "forces": base_forces + delta_forces, "sensors": base_sensors + delta_sensors}
    baseline_t = {**common, "forces": base_forces, "sensors": base_sensors}
    if missing_telemetry:
        controlled_t.pop(missing_telemetry)
    np.savez(directory / "ctl_telemetry.npz", **controlled_t)
    np.savez(directory / "blk_telemetry.npz", **baseline_t)
    product = root / f"product_{realization}"
    build_canonical_paired_product(
        directory / "ctl_fields.npz", directory / "blk_fields.npz",
        directory / "ctl_telemetry.npz", directory / "blk_telemetry.npz", product,
        realization_id=f"r{realization}", fluid_mask=np.ones((ny, nx), bool),
        cell_metric=np.full((ny, nx), 0.25), metric_metadata={"quadrature": "uniform", "cell_area": 0.25},
    )
    return product


def test_phase_3_5_synthetic_end_to_end_is_cpu_only(tmp_path):
    products = [_realization(tmp_path, index) for index in range(3)]
    output = tmp_path / "analysis"
    report = run_phase_3_5(products, output, ranks=(2, 3), lags=(0, 1), surrogate_count=7, seed=4)
    assert report["fold_count"] == 3
    assert len(report["folds"]) == 3
    assert all(fold["held_out_realization"] not in fold["training_realizations"] for fold in report["folds"])
    assert all(0.0 < fold["selection_adjusted_periodic_excluded_empirical_p_value"] <= 1.0
               for fold in report["folds"])
    assert all(len(fold["circular_time_shift_surrogate_scores"]) == 7 for fold in report["folds"])
    assert len(report["fold_mode_comparisons"]) == 3
    assert report["metric_contract"]["sensitivity_kind"] == "global metric scaling is invariant"
    assert report["state_pod_fit_count"] <= report["state_pod_fit_bound"]
    assert report["lag_metadata"]["candidate_sample_lags"] == [0, 1]
    assert report["lag_metadata"]["candidate_lattice_step_lags"] == [0.0, 5.0]
    assert report["lag_metadata"]["candidate_physical_time_lags"] == [0.0, 0.1]
    null_contract = report["circular_time_shift_null_contract"]
    assert null_contract["exclusion_width_samples"] == 2
    assert null_contract["tested_statistic"] == "metric_weighted_r2"
    assert null_contract["periodic_equivalent_exclusions_applied"]
    assert null_contract["admissible_offsets_fixed_before_null_scoring"]
    assert report["metrics"]["primary"] == "metric_weighted_r2"
    assert report["metric_contract"]["observable_metric"] == [1.0, 1.0]
    assert report["metric_contract"]["observable_metric_declaration"] == "H=I"
    assert len(report["metric_contract"]["observable_metric_sha256"]) == 64
    assert (output / report["model_artifact"]["path"]).is_file()
    assert all(len(fold["per_channel_r2"]) == 2 for fold in report["folds"])
    assert all("aggregate_r2" not in fold and fold["tested_statistic"] == "metric_weighted_r2"
               for fold in report["folds"])
    assert all(all(item["tested_statistic"] == "metric_weighted_r2"
                   and "train_only_metric_weighted_r2" in item
                   and "train_only_score" not in item
                   for item in fold["rank_lag_sensitivity"])
               for fold in report["folds"])
    assert all(len(fold["circular_time_shift_surrogate_selections"]) == 7 for fold in report["folds"])
    for fold in report["folds"]:
        null = np.asarray(fold["circular_time_shift_surrogate_scores"])
        expected_p = (1 + np.sum(null >= fold["metric_weighted_r2"])) / 8
        assert fold["selection_adjusted_periodic_excluded_empirical_p_value"] == pytest.approx(expected_p)
        assert "empirical_p_value" not in fold
        assert len(fold["circular_time_shift_admissibility"]) == 2
        control = fold["realization_derangement_negative_control_metric_weighted_r2"]
        assert fold["negative_control_pass"] == (
            fold["metric_weighted_r2"] > control + fold["negative_control_margin"])
        assert fold["realization_derangement_selected_rank"] in (2, 3)
        assert fold["realization_derangement_selected_sample_lag"] in (0, 1)
        assert fold["realization_derangement_rank_lag_sensitivity"]
        assert fold["numerical_status"] == "truncated"
        assert fold["regularization"] is None
        assert fold["rcond"] == 1e-12
    assert all("generated_rank" in fold and "propagator_condition_number" in fold for fold in report["folds"])
    persisted = json.loads((output / "phase_3_5_report.json").read_text())
    assert persisted["schema_id"] == "oid-phase3-5-loro/v2"
    with np.load(output / report["model_artifact"]["path"], allow_pickle=False) as models:
        assert all(np.array_equal(models[f"fold_{index}_observable_metric"], np.ones(2))
                   for index in range(3))


def test_identity_metric_primary_score_is_pooled_not_mean_channel_r2():
    actual = np.array([[-10.0, -1.0], [0.0, 0.0], [10.0, 1.0]])
    predicted = np.array([[-10.0, 0.0], [0.0, 0.0], [10.0, 0.0]])
    score = _scores(actual, predicted, np.zeros(2), np.ones(2))
    assert score["per_channel_r2"] == [1.0, 0.0]
    assert score["unweighted_mean_channel_r2"] == pytest.approx(0.5)
    assert score["metric_weighted_r2"] == pytest.approx(1.0 - 2.0 / 202.0)
    assert score["metric_weighted_r2"] != pytest.approx(score["unweighted_mean_channel_r2"])
    assert "aggregate_r2" not in score


def test_rank_lag_selection_uses_metric_weighted_primary(monkeypatch):
    import OID_analysis.v2.oid as oid

    def conflicting_scores(*args, rank, **kwargs):
        return ({
            "metric_weighted_r2": float(rank),
            "unweighted_mean_channel_r2": float(100 - rank),
        }, {})

    monkeypatch.setattr(oid, "_fit_fold", conflicting_scores)
    rank, lag, table = oid._choose([{}, {}], (1, 2), (0,), "delta_forces", np.ones(2), None)
    assert (rank, lag) == (2, 0)
    assert all(item["tested_statistic"] == "metric_weighted_r2" for item in table)


def test_phase_3_5_fails_closed_below_three_realizations(tmp_path):
    products = [_realization(tmp_path, index) for index in range(2)]
    with pytest.raises(ValueError, match="exactly three canonical realizations"):
        run_phase_3_5(products, tmp_path / "analysis", surrogate_count=2)


def test_canonical_product_rejects_control_end_only_telemetry(tmp_path):
    with pytest.raises(ValueError, match="field-level telemetry missing"):
        _realization(tmp_path, 0, missing_telemetry="times")


def test_phase_3_5_is_deterministic_under_input_and_candidate_reordering(tmp_path):
    products = [_realization(tmp_path, index) for index in range(3)]
    a = run_phase_3_5(products, tmp_path / "a", ranks=(3, 2), lags=(1, 0), surrogate_count=4, seed=17)
    b = run_phase_3_5(list(reversed(products)), tmp_path / "b", ranks=(2, 3), lags=(0, 1), surrogate_count=4, seed=17)
    assert a["folds"] == b["folds"]


def test_surrogate_repeats_selection_without_refitting_state_pod(tmp_path, monkeypatch):
    products = [_realization(tmp_path, index) for index in range(3)]
    import OID_analysis.v2.oid as oid
    calls = 0
    original = oid.fit_pod
    def counted(*args, **kwargs):
        nonlocal calls
        calls += 1
        return original(*args, **kwargs)
    monkeypatch.setattr(oid, "fit_pod", counted)
    report = run_phase_3_5(products, tmp_path / "analysis", ranks=(2, 3), lags=(0, 1), surrogate_count=5, seed=8)
    assert calls == report["state_pod_fit_count"] <= 12
    assert all({x["rank"] for x in fold["circular_time_shift_surrogate_selections"]} <= {2, 3} for fold in report["folds"])
    for fold in report["folds"]:
        allowed = [set(record["admissible_offsets"])
                   for record in fold["circular_time_shift_admissibility"]]
        assert all(all(offset in allowed[index] for index, offset in enumerate(item["offsets"]))
                   for item in fold["circular_time_shift_surrogate_selections"])


def test_surrogate_count_does_not_increase_field_transforms_or_mode_materializations(tmp_path, monkeypatch):
    products = [_realization(tmp_path, index) for index in range(3)]
    import OID_analysis.v2.oid as oid

    physical_mode_calls = 0
    original = oid.fit_schlegel_oid

    def counted(*args, **kwargs):
        nonlocal physical_mode_calls
        physical_mode_calls += kwargs.get("state_modes") is not None
        return original(*args, **kwargs)

    monkeypatch.setattr(oid, "fit_schlegel_oid", counted)
    small = run_phase_3_5(
        products, tmp_path / "small", ranks=(2, 3), lags=(0, 1),
        surrogate_count=1, seed=8,
    )
    small_physical_calls = physical_mode_calls
    large = run_phase_3_5(
        products, tmp_path / "large", ranks=(2, 3), lags=(0, 1),
        surrogate_count=17, seed=8,
    )
    large_physical_calls = physical_mode_calls - small_physical_calls

    assert small["state_field_transform_count"] == large["state_field_transform_count"]
    assert large["state_field_transform_count"] <= large["state_field_transform_bound"]
    assert large["state_field_transform_bound"] == large["state_pod_fit_count"] * 2
    assert small_physical_calls == large_physical_calls == 3
    assert small["physical_mode_materialization_count"] == 3
    assert large["physical_mode_materialization_count"] == 3
    assert large["physical_mode_materialization_count"] <= large["physical_mode_materialization_bound"]


def test_cached_selection_preserves_deterministic_finite_smoke_scores(tmp_path):
    products = [_realization(tmp_path, index) for index in range(3)]
    report = run_phase_3_5(
        products, tmp_path / "analysis", ranks=(2, 3), lags=(0, 1),
        surrogate_count=7, seed=4,
    )
    numeric_scores = [
        (
            fold["metric_weighted_r2"],
            fold["realization_derangement_negative_control_metric_weighted_r2"],
            fold["circular_time_shift_surrogate_scores"],
        )
        for fold in report["folds"]
    ]
    assert len(numeric_scores) == 3
    assert all(np.isfinite([observed, control, *null]).all()
               for observed, control, null in numeric_scores)


def test_negative_control_gate_can_fail_declared_margin(tmp_path):
    products = [_realization(tmp_path, index) for index in range(3)]
    report = run_phase_3_5(products, tmp_path / "analysis", ranks=(2,), lags=(0,), surrogate_count=2,
                           negative_control_margin=1e6)
    assert not report["negative_control_gate"]["all_folds_pass"]
    assert not any(fold["negative_control_pass"] for fold in report["folds"])


def test_physical_lag_truncates_samples(tmp_path):
    products = [_realization(tmp_path, index) for index in range(3)]
    report = run_phase_3_5(products, tmp_path / "analysis", ranks=(2,), lags=(2,), surrogate_count=2)
    assert all(fold["sample_lag"] == 2 and fold["lattice_step_lag"] == 10 and fold["physical_time_lag"] == pytest.approx(0.2) for fold in report["folds"])
    assert all(fold["test_samples"] == 70 for fold in report["folds"])
    assert all(fold["train_discarded_pair_count"] == 4 and fold["test_discarded_pair_count"] == 2
               and fold["total_discarded_pair_count"] == 6 for fold in report["folds"])
    contract = report["circular_time_shift_null_contract"]
    assert contract["exclusion_width_samples"] == 3
    assert contract["exclusion_width_rule"] == "max(2, max(candidate_sample_lags) + 1)"


def test_canonical_loader_and_no_clobber(tmp_path):
    product = _realization(tmp_path, 0)
    loaded = load_canonical_product(product)
    assert loaded["metadata"]["layout_contract"]["component_order"] == ["ux", "uy"]
    with pytest.raises(FileExistsError, match="refusing to overwrite"):
        build_canonical_paired_product(
            tmp_path / "raw_0/ctl_fields.npz", tmp_path / "raw_0/blk_fields.npz",
            tmp_path / "raw_0/ctl_telemetry.npz", tmp_path / "raw_0/blk_telemetry.npz", product,
            realization_id="r0", fluid_mask=np.ones((3, 4), bool),
            cell_metric=np.full((3, 4), 0.25),
        )


def _online_field_hash(path):
    digest = hashlib.sha256()
    with np.load(path, allow_pickle=False) as data:
        for key in ("ux", "uy"):
            value = np.ascontiguousarray(data[key])
            digest.update(str(value.dtype).encode())
            digest.update(np.asarray(value.shape, np.int64).tobytes())
            digest.update(value.tobytes())
    return digest.hexdigest()


def test_source_provenance_validates_online_field_content_semantics(tmp_path):
    raw = tmp_path / "raw"
    raw.mkdir()
    steps = np.array([1, 2], dtype=np.int64)
    timeline = {"lattice_steps": steps, "control_indices": np.array([0, 0]),
                "control_offsets": steps, "control_interval": np.array(2, dtype=np.int64)}
    fields = {}
    for role, offset in (("controlled", 1.0), ("zero_baseline", 0.0)):
        path = raw / f"{role}_fields.npz"
        np.savez(path, ux=np.full((2, 2, 2), offset), uy=np.full((2, 2, 2), offset), **timeline)
        fields[role] = {
            "path": str(path), "declared_content_sha256": _online_field_hash(path),
            "content_hash_semantics": "online.py fields_sha256: SHA-256 over ux then uy; each contributes str(dtype), int64 shape bytes, C-order data bytes",
        }
    common = {"lattice_steps": steps, "sample_ids": steps, "times": steps.astype(float)}
    for role in ("controlled", "zero_baseline"):
        np.savez(raw / f"{role}_telemetry.npz", forces=np.zeros((2, 1)), sensors=np.zeros((2, 1)), **common)
    metadata = build_canonical_paired_product(
        fields["controlled"]["path"], fields["zero_baseline"]["path"],
        raw / "controlled_telemetry.npz", raw / "zero_baseline_telemetry.npz", tmp_path / "out",
        realization_id="r1", fluid_mask=np.ones((2, 2), bool), source_artifacts=fields,
    )
    source = metadata["canonical_pair_sources"]["zero_baseline"]
    assert source["content_hash_validated"] and source["file_sha256"] != source["declared_content_sha256"]
    fields["controlled"]["declared_content_sha256"] = "0" * 64
    with pytest.raises(ValueError, match="declared content hash mismatch"):
        build_canonical_paired_product(
            fields["controlled"]["path"], fields["zero_baseline"]["path"],
            raw / "controlled_telemetry.npz", raw / "zero_baseline_telemetry.npz", tmp_path / "bad",
            realization_id="r1", fluid_mask=np.ones((2, 2), bool), source_artifacts=fields,
        )


def test_rebuild_cli_selects_zero_baseline_and_separates_twin_evidence():
    script = (Path(__file__).parents[1] / "scripts/rebuild_phase3_canonical.py").read_text()
    assert 'sources["zero_baseline_fields"]["path"]' in script
    assert 'sources["zero_baseline_telemetry"]["path"]' in script
    assert '"paired_collection"' in script and '"twin_gate"' in script
    assert '"baseline_twin_fields"' not in script.split("sources =", 1)[1].split("evidence =", 1)[0]


def test_phase_3_5_records_execution_inputs_and_verifies_model(tmp_path):
    products = [_realization(tmp_path, index) for index in range(3)]
    report = run_phase_3_5(
        products, tmp_path / "analysis", ranks=(2,), lags=(0,), surrogate_count=2,
        execution_metadata={"thread_environment": {"OPENBLAS_NUM_THREADS": "2"}},
        convective_scale={"D_over_U0_lbm_steps": 2000.0, "stored_times_units": "LBM lattice-step units"},
    )
    assert report["execution_metadata"]["analysis_wall_time_seconds"] > 0
    assert report["execution_metadata"]["peak_rss_kib"] > 0
    assert report["convective_lag_contract"]["D_over_U0_lbm_steps"] == 2000.0
    assert len(report["input_products"]) == 3
    assert all(len(item["product_sha256"]) == 64 for item in report["input_products"])
    assert report["model_artifact"]["verified_after_write"]


def test_formal_cli_contract_is_fixed_and_no_clobber():
    script = (Path(__file__).parents[1] / "scripts/run_phase3_5_formal.py").read_text()
    assert 'RANKS = (4, 8, 12)' in script and 'LAGS = (0, 1, 2)' in script
    assert 'SURROGATES = 999' in script and 'D, U0 = 20.0, 0.01' in script
    assert 'RCOND = 1e-12' in script and 'rcond=RCOND' in script
    assert 'len(running) < 3' in script and '"OPENBLAS_NUM_THREADS": "2"' in script
    assert 'if destination.exists()' in script and 'refusing to clobber analysis destination' in script
    assert 'selection_adjusted": True' in script and 'gpu_or_cfd_execution": False' in script
    assert '"tested_statistic": "metric_weighted_r2"' in script
    assert 'observable_metric=np.ones(OBSERVABLE_CHANNELS[observable])' in script
    assert 'observable_coordinate_units=OBSERVABLE_COORDINATE_UNITS[observable]' in script
    assert 'ROOT / "src/OID_analysis/v2/analysis.py"' in script


def test_period_estimator_and_periodic_exclusions_are_deterministic():
    n = 60
    time = np.arange(n)
    observable = np.column_stack((
        np.sin(2 * np.pi * time / 10),
        3.0 * np.cos(2 * np.pi * time / 10 + 0.2),
    ))
    first = _dominant_period(observable)
    second = _dominant_period(observable.copy())
    assert first == second
    assert first["dominant_nonzero_bin"] == 6
    assert first["period_samples"] == pytest.approx(10.0)
    assert first["peak_concentration"] == pytest.approx(1.0)

    item = {"metadata": {"realization_id": "r1"}, "delta_forces": observable}
    contract = _admissible_shift_record(item, "delta_forces", 2, 0.25)
    assert contract["period_physical_time"] == pytest.approx(2.5)
    excluded = set(contract["excluded_offsets"])
    for multiple in range(0, 6):
        location = multiple * 10
        for offset in range(max(0, location - 2), min(n, location + 3)):
            assert offset in excluded
    assert {58, 59} <= excluded

    train = [item]
    shifted_a, offsets_a = _shift(train, "delta_forces", np.random.default_rng(7), [contract])
    shifted_b, offsets_b = _shift(train, "delta_forces", np.random.default_rng(7), [contract])
    assert offsets_a == offsets_b
    assert np.array_equal(shifted_a[0]["delta_forces"], shifted_b[0]["delta_forces"])
    assert offsets_a[0] in contract["admissible_offsets"]


def test_period_estimator_fails_closed_for_zero_energy_channel():
    time = np.arange(40)
    observable = np.column_stack((np.sin(2 * np.pi * time / 8), np.ones(40)))
    with pytest.raises(ValueError, match="zero-energy observable channel"):
        _dominant_period(observable)
