"""Contract tests for OID v2 no-clobber, alignment, and mathematics."""
from __future__ import annotations

from datetime import datetime, timezone
from pathlib import Path

import numpy as np
import pytest

from OID_analysis.v2.alignment import exact_time_pair, phase_conditioned_correction, valid_delayed_pairs
from OID_analysis.v2.analysis import fit_pls_type, fit_pod, modal_assurance, physical_pls_modes, principal_angles
from CCD_analysis.utils.run_contract import (
    ContractError,
    IntegrityError,
    array_sha256,
    atomic_publish_npz,
    create_run_directory,
    mark_complete,
)
from OID_analysis.v2.artifacts import ArtifactError, atomic_save_npz, atomic_write_json, validate_field_artifact
from OID_analysis.v2.contracts import (
    OID_FORMAL_SCHEMA_ID,
    BRANCH_SCHEMA_ID,
    exact_branch_gates,
    load_formal_run,
    telemetry_hashes,
    validate_realization_contract,
)
from OID_analysis.v2.layout import RunPaths, make_run_id
from OID_analysis.v2.manifest import validate_execution_environment


def test_run_layout_dry_run_and_no_clobber(tmp_path):
    run_id = make_run_id("karman_re100", "pilot", {"a": 1}, "abcdef1234",
                         now=datetime(2026, 7, 22, tzinfo=timezone.utc))
    dry = RunPaths.create(run_id, runs_root=tmp_path, dry_run=True)
    assert not dry.partial.exists() and not dry.final.exists()
    paths = RunPaths.create(run_id, runs_root=tmp_path)
    assert paths.partial.is_dir()
    with pytest.raises(FileExistsError):
        RunPaths.create(run_id, runs_root=tmp_path)


def test_atomic_npz_refuses_overwrite_and_validates_timeline(tmp_path):
    steps = np.array([3, 6, 9], dtype=np.int64)
    payload = dict(ux=np.zeros((3, 2, 4), np.float32), uy=np.zeros((3, 2, 4), np.float32),
                   lattice_steps=steps, control_indices=(steps - 1) // 8,
                   control_offsets=(steps - 1) % 8 + 1, control_interval=np.asarray(8, np.int64))
    path = atomic_save_npz(tmp_path / "fields.npz", **payload)
    with np.load(path, allow_pickle=False) as data:
        assert validate_field_artifact(data)["timeline"]["count"] == 3
    with pytest.raises(FileExistsError):
        atomic_save_npz(path, **payload)


def test_field_schema_fails_closed_on_missing_timeline():
    with pytest.raises(ArtifactError):
        validate_field_artifact({"ux": np.zeros((2, 2, 2)), "uy": np.zeros((2, 2, 2))})


def test_exact_time_pair_rejects_index_only_pairing():
    q = np.zeros((3, 2))
    with pytest.raises(ValueError):
        exact_time_pair(q, np.array([1, 2, 3]), q, np.array([1, 3, 4]))


def test_valid_delay_has_no_padding_or_cross_segment_pairs():
    steps = np.arange(0, 60, 10)
    q = np.arange(6)[:, None]
    y = (10 + np.arange(6))[:, None]
    pairs = valid_delayed_pairs(q, y, steps, delay_steps=20, segment_ids=[0, 0, 0, 1, 1, 1])
    assert pairs.left[:, 0].tolist() == [0, 3]
    assert pairs.right[:, 0].tolist() == [12, 15]
    assert pairs.dropped == 4


def test_pod_truncated_energy_uses_full_spectrum_denominator():
    data = np.diag([4.0, 3.0, 2.0, 1.0])
    data = np.vstack([data, -data])
    pod = fit_pod(data, rank=2)
    assert pod.energy.sum() < 1.0
    assert np.isclose(pod.energy.sum(), np.sum(pod.singular_values ** 2) / pod.total_energy)


def test_physical_pls_modes_recovers_phi_d_u_and_sign_metrics():
    rng = np.random.default_rng(4)
    phi = np.linalg.qr(rng.normal(size=(30, 3)))[0]
    d = np.array([4.0, 2.0, 0.5])
    u = np.linalg.qr(rng.normal(size=(3, 3)))[0]
    expected = phi @ np.diag(d) @ u
    actual = physical_pls_modes(phi, d, u)
    assert np.allclose(actual, expected)
    assert np.isclose(modal_assurance(actual[:, 0], -actual[:, 0]), 1.0)
    assert np.allclose(principal_angles(actual[:, :2], -actual[:, :2]), 0.0, atol=1e-7)


def test_pls_fitted_only_to_given_development_data():
    rng = np.random.default_rng(8)
    development_a = rng.normal(size=(30, 4))
    development_y = development_a[:, :2] + 0.01 * rng.normal(size=(30, 2))
    first = fit_pls_type(development_a, development_y)
    validation_a = rng.normal(loc=100.0, size=(20, 4))
    validation_y = rng.normal(loc=-100.0, size=(20, 2))
    second = fit_pls_type(development_a, development_y)
    assert np.array_equal(first["a_mean"], second["a_mean"])
    assert not np.allclose(first["a_mean"], np.vstack([development_a, validation_a]).mean(axis=0))


def test_phase_correction_allows_distinct_phase_sequences():
    phase_ctl = np.tile(np.array([0.1, 1.7, 3.2, 4.8]), 3)
    phase_blk = np.tile(np.array([0.2, 1.8, 3.3, 4.9]), 3)
    ctl = np.sin(phase_ctl)[:, None] + 2.0
    blk = np.sin(phase_blk)[:, None]
    delta, meta = phase_conditioned_correction(ctl, phase_ctl, blk, phase_blk, bins=4, min_count=3)
    assert delta.shape == (4, 1)
    assert np.all(meta["controlled_counts"] == 3)


def _realizations():
    result = []
    for index in range(3):
        parent = f"parent-{index}"
        checkpoint = str(index + 1) * 64
        result.append({
            "realization_id": f"r{index}", "parent_id": parent,
            "checkpoint_sha256": checkpoint,
            "initialization": {"seed": index, "independent_spinup": True},
            "phase": {"origin": "checkpoint", "offset": index},
            "sample_range": [index * 100, (index + 1) * 100],
            "branches": {role: {"role": role, "parent_id": parent,
                "checkpoint_sha256": checkpoint}
                for role in ("controlled", "zero_baseline", "baseline_twin")},
        })
    return result


def test_three_realization_contract_is_explicit_and_non_overlapping():
    contract = validate_realization_contract(_realizations())
    assert contract["count"] == 3
    assert set(contract["realizations"][0]["branches"]) == {
        "controlled", "zero_baseline", "baseline_twin"}
    overlapping = _realizations()
    overlapping[2]["sample_range"] = [50, 150]
    with pytest.raises(ContractError, match="overlap"):
        validate_realization_contract(overlapping)
    duplicated = _realizations()
    duplicated[2]["checkpoint_sha256"] = duplicated[0]["checkpoint_sha256"]
    with pytest.raises(ContractError, match="unique"):
        validate_realization_contract(duplicated)


def test_oid_gpu_policy_records_physical_two_and_uses_logical_zero(monkeypatch):
    monkeypatch.setenv("CONDA_DEFAULT_ENV", "pycuda_3_10")
    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "2")
    execution = validate_execution_environment(device=2, require_gpu=True)
    assert execution == {"conda_env": "pycuda_3_10", "physical_gpu": 2, "logical_device": 0}
    with pytest.raises(ContractError):
        validate_execution_environment(device=0, require_gpu=True)


def test_telemetry_hashes_are_field_level_and_timeline_strict():
    telemetry = {
        "sensors": np.ones((3, 2), np.float32), "forces": np.zeros((3, 2), np.float32),
        "lattice_steps": np.asarray([10, 20, 30], np.int64),
        "sample_ids": np.asarray([0, 1, 2], np.int64),
        "times": np.asarray([0.1, 0.2, 0.3], np.float64),
    }
    hashes = telemetry_hashes(telemetry)
    assert hashes["sensors"] == array_sha256(telemetry["sensors"])
    broken = dict(telemetry, lattice_steps=np.asarray([10, 10, 30], np.int64))
    with pytest.raises(ContractError):
        telemetry_hashes(broken)


def _branch_manifests():
    common = {"schema_id": BRANCH_SCHEMA_ID, "parent_id": "parent-r0",
              "checkpoint_sha256": "a" * 64, "timeline_sha256": "b" * 64,
              "sample_ids_sha256": "c" * 64, "times_sha256": "d" * 64}
    return {
        "controlled": {**common, "role": "controlled", "action_kind": "controlled",
                       "fields_sha256": "e" * 64, "telemetry_sha256": "f" * 64},
        "zero_baseline": {**common, "role": "zero_baseline", "action_kind": "zero",
                          "fields_sha256": "1" * 64, "telemetry_sha256": "2" * 64},
        "baseline_twin": {**common, "role": "baseline_twin", "action_kind": "zero",
                          "fields_sha256": "1" * 64, "telemetry_sha256": "2" * 64},
    }


def test_exact_timeline_parent_and_baseline_twin_gates():
    assert exact_branch_gates(_branch_manifests())["parent_id"] == "parent-r0"
    bad_timeline = _branch_manifests()
    bad_timeline["baseline_twin"]["timeline_sha256"] = "9" * 64
    with pytest.raises(ContractError, match="timeline_sha256"):
        exact_branch_gates(bad_timeline)
    bad_twin = _branch_manifests()
    bad_twin["baseline_twin"]["fields_sha256"] = "8" * 64
    twin_report = exact_branch_gates(bad_twin)
    assert twin_report["requires_numerical_twin_gate"] is True
    assert twin_report["twin_reproducibility"]["byte_identical_fields"] is False


def test_numerical_twin_gate_accepts_small_noise_and_rejects_material_noise(tmp_path):
    from OID_analysis.v2.twin import numerical_twin_gate

    zero, twin = tmp_path / "zero", tmp_path / "twin"
    zero.mkdir(); twin.mkdir()
    steps = np.arange(1, 7, dtype=np.int64)
    rng = np.random.default_rng(4)
    ux = rng.normal(size=(6, 3, 2)).astype(np.float32)
    uy = rng.normal(size=(6, 3, 2)).astype(np.float32)
    sensors = rng.normal(size=(6, 6)).astype(np.float32)
    forces = rng.normal(size=(6, 6)).astype(np.float32)
    for path, noise in ((zero, 0.0), (twin, 1e-5)):
        np.savez(path / "fields.npz", ux=ux + noise, uy=uy + noise, lattice_steps=steps)
        np.savez(path / "telemetry.npz", sensors=sensors + noise, forces=forces + noise,
                 lattice_steps=steps)
    assert numerical_twin_gate(zero, twin)["passed"]
    with np.load(twin / "fields.npz") as data:
        np.savez(twin / "fields.bad.npz", **{key: data[key] for key in data.files})
    np.savez(twin / "fields.npz", ux=ux + 1.0, uy=uy + 1.0, lattice_steps=steps)
    with pytest.raises(RuntimeError, match="numerical gate failed"):
        numerical_twin_gate(zero, twin)


def test_strict_formal_loader_binds_artifacts_and_legacy_is_opt_in(tmp_path):
    run = create_run_directory(tmp_path / "runs", "formal-new")
    artifact_name = "analysis/result.npz"
    arrays = {"x": np.asarray([1.0, 2.0], np.float32),
              "lattice_steps": np.asarray([10, 20], np.int64)}
    metadata = atomic_publish_npz(run / artifact_name, arrays)
    contract = validate_realization_contract(_realizations())
    mark_complete(run, {"oid_schema_id": OID_FORMAL_SCHEMA_ID,
        "realization_contract": contract,
        "artifacts": {artifact_name: metadata},
        "timeline_hashes": {artifact_name: array_sha256(arrays["lattice_steps"])},
        "array_hashes": {f"{artifact_name}:x": array_sha256(arrays["x"])}})
    assert load_formal_run(run)["realization_contract"]["count"] == 3
    payload = bytearray((run / artifact_name).read_bytes()); payload[-1] ^= 1
    (run / artifact_name).write_bytes(payload)
    with pytest.raises(IntegrityError):
        load_formal_run(run)

    legacy = tmp_path / "formal-karman-20260722"
    legacy.mkdir()
    atomic_write_json(legacy / "manifest.json", {"schema_version": 1})
    with pytest.raises(Exception):
        load_formal_run(legacy)
    assert load_formal_run(legacy, allow_legacy_realization0=True)["realization_id"] == "0"


def test_dense_field_metric_is_guarded_before_factorization():
    from OID_analysis.v2.analysis import _guard_dense_metric

    oversized = np.lib.stride_tricks.as_strided(
        np.ones(1), shape=(5000, 5000), strides=(0, 0), writeable=False
    )
    with pytest.raises(ValueError, match="one-dimensional vector"):
        _guard_dense_metric(oversized)
