bundle/v0/tests/test_timing_harness.py

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pytest

from loss_aware_dro_repro.core import CONFIG_ROOT, ContractError
from loss_aware_dro_repro.timing_harness import (
    TIMING_EVIDENCE_SCALE,
    _acceptance_summary,
    _encode_trace_row,
    _run_fixed_iterations,
    _mutate_launch_ledger,
    _solution_iterations,
    _solver_iteration_drift,
    _timing_drift,
    _validate_solver_record,
    _validate_timing_directory,
    load_timing_manifest,
    load_timing_escalation_manifest,
    run_timing_sample,
)


def test_frozen_manifest_declares_exactly_21_supported_local_timing_tasks():
    manifest = load_timing_manifest(CONFIG_ROOT / "cpu_timing_sample_v1.json")
    assert len(manifest["task_selectors"]) == 21
    assert len(set(manifest["task_selectors"])) == 21
    assert manifest["optimizer_identity"]["iterations"] == 100
    assert manifest["evidence_scale"] == TIMING_EVIDENCE_SCALE
    assert not any(manifest["authority"].values())


def test_escalation_manifest_is_bound_to_parent_receipt_and_16_tasks():
    manifest = load_timing_escalation_manifest(
        CONFIG_ROOT / "cpu_timing_escalation_v1.json"
    )
    assert len(manifest["task_selectors"]) == 16
    assert manifest["optimizer_identity"]["iterations"] == 1000
    assert manifest["parent_timing_receipt"]["run_key"].startswith("sha256:")


def test_manifest_rejects_any_convergence_claim(tmp_path):
    manifest = load_timing_manifest(CONFIG_ROOT / "cpu_timing_sample_v1.json")
    manifest["acceptance"]["converged_label_allowed"] = True
    candidate = tmp_path / "bad.json"
    candidate.write_text(json.dumps(manifest), encoding="utf-8")
    with pytest.raises(ContractError, match="converged label"):
        load_timing_manifest(candidate)


def test_trace_byte_binding_is_exact():
    encoded = _encode_trace_row({"record_type": "timing_iteration", "value": [1, 2, 3]})
    row = json.loads(encoded)
    assert encoded.endswith(b"\n")
    assert row["trace_bytes"] == len(encoded)


def test_timing_drift_uses_predeclared_first_and_last_halves():
    stable = _timing_drift([1.0] * 100, 1.25)
    assert stable["acceptance_pass"] is True
    assert stable["last_half_to_first_half_step_time_p95_ratio"] == 1.0

    drifting = _timing_drift([1.0] * 50 + [2.0] * 50, 1.25)
    assert drifting["acceptance_pass"] is False
    assert drifting["last_half_to_first_half_step_time_p95_ratio"] == 2.0

    stable_1000 = _timing_drift([1.0] * 1000, 1.25)
    assert stable_1000["half_size"] == 500
    assert stable_1000["acceptance_pass"] is True


def test_timing_drift_rejects_incomplete_iteration_series():
    with pytest.raises(ContractError, match="100 or 1000"):
        _timing_drift([1.0] * 99, 1.25)


def test_cost_freeze_is_blocked_when_fixed100_never_reaches_released_stop():
    summary = _acceptance_summary(
        [
            {
                "task_id": "portfolio_discrete/d001/r00/n010",
                "timing_drift": {"acceptance_pass": True},
                "solver_iteration_drift": {"acceptance_pass": True},
                "stopping_observation": {"released_stop_trigger_observed": False},
            }
        ]
    )
    assert summary["fixed100_gate_pass"] is False
    assert summary["optimizer_extrapolation_eligible"] is False
    assert summary["portfolio_cost_freeze_eligible"] is False
    assert summary["fixed1000_escalation_selectors"] == [
        "portfolio_discrete/d001/r00/n010"
    ]
    assert "grouped_10m_oos_timing_receipt_required" in summary[
        "portfolio_cost_freeze_blockers"
    ]


def test_solver_iteration_drift_is_field_aware_and_predeclared():
    stable = _solver_iteration_drift(
        [{"solver_iterations": 4}] * 1000, 2.0
    )
    assert stable["acceptance_pass"] is True
    deteriorating = _solver_iteration_drift(
        [{"solver_iterations": 2}] * 500 + [{"solver_iterations": 8}] * 500,
        2.0,
    )
    assert deteriorating["acceptance_pass"] is False


def test_solution_iteration_count_fails_closed_when_missing():
    class Missing:
        pass

    with pytest.raises(ContractError, match="iteration count"):
        _solution_iterations(Missing())


@pytest.mark.parametrize("bad_residual", [float("nan"), float("inf"), -1e-12])
def test_solver_record_rejects_nonfinite_or_negative_residuals(bad_residual):
    with pytest.raises(ContractError, match="finite and nonnegative"):
        _validate_solver_record(
            {
                "solver_status": "optimal",
                "solver_residuals": {"primal": bad_residual},
                "transport_residuals": None,
            },
            {"optimal", "optimal_inaccurate"},
        )


def test_solver_record_rejects_unaccepted_status():
    with pytest.raises(ContractError, match="is not accepted"):
        _validate_solver_record(
            {
                "solver_status": "max_iterations",
                "solver_residuals": {"primal": 0.0},
                "transport_residuals": None,
            },
            {"optimal", "optimal_inaccurate"},
        )


def test_fixed_loop_records_first_trigger_and_still_executes_100(monkeypatch):
    task = {
        "task_id": "synthetic/timing-only",
        "hyperparameters": {
            "gradient_clip": [-10.0, 10.0],
            "learning_rate": 0.01,
            "metric_eigenvalue_clip": [0.1, 10.0],
        },
    }
    calls = 0

    def evaluate(_L):
        nonlocal calls
        calls += 1
        objective = 10.0 if calls == 1 else 11.0
        return {
            "total_objective": objective,
            "total_gradient": np.zeros((1, 1)),
            "lower": {
                "objective": objective,
                "status": "optimal",
                "raw_status": "Solved",
                "residuals": {"primal": 0.0},
            },
        }

    class Captures:
        def __enter__(self):
            return [{"iterations": 3, "raw_status": "Solved"}]

        def __exit__(self, *_args):
            return False

    monkeypatch.setattr(
        "loss_aware_dro_repro.timing_harness._capture_clarabel_solves", Captures
    )
    _, rows, stopping, _, _ = _run_fixed_iterations(
        task=task,
        initial_L=np.eye(1),
        evaluate=evaluate,
        state_payload=lambda state: {
            "solver_residuals": state["lower"]["residuals"]
        },
        iterations=100,
        tolerance=1e-5,
    )
    assert calls == 100
    assert len(rows) == 100
    assert stopping["iterations_executed"] == 100
    assert stopping["first_released_stop_trigger"] == {
        "iteration": 1,
        "signed_relative_improvement": -0.1,
        "stop_trigger_caused_by_worsening": True,
    }
    assert stopping["first_paper_stop_trigger"] == {
        "iteration": 1,
        "signed_relative_improvement": -0.1,
        "stop_trigger_caused_by_worsening": True,
        "denominator_sign_inversion_risk": False,
    }
    final_row = json.loads(rows[-1])
    assert final_row["iteration"] == 99
    assert set(final_row["stopping_diagnostics"]) >= {
        "paper_total_phi_literal",
        "paper_total_phi_abs_denominator",
        "released_lower_objective_abs_denominator",
    }


def test_sample_refuses_dirty_lane_before_any_task(monkeypatch, tmp_path):
    monkeypatch.setattr(
        "loss_aware_dro_repro.timing_harness._repository_state",
        lambda: ("a" * 40, False),
    )
    with pytest.raises(ContractError, match="paper lane is not clean"):
        run_timing_sample(
            CONFIG_ROOT / "cpu_timing_sample_v1.json", tmp_path / "must-not-exist"
        )
    assert not (tmp_path / "must-not-exist").exists()


def test_artifact_validator_rejects_missing_files(tmp_path):
    manifest = load_timing_manifest(CONFIG_ROOT / "cpu_timing_sample_v1.json")
    (tmp_path / "timing-sample.json").write_text(
        json.dumps({"task_count": 21, "tasks": [], "task_artifacts": {}}),
        encoding="utf-8",
    )
    with pytest.raises(ContractError, match="file set"):
        _validate_timing_directory(tmp_path, manifest)


def test_launch_ledger_refuses_same_run_identity_twice(tmp_path):
    ledger = tmp_path / "ledger.json"
    output = tmp_path / "output"
    _mutate_launch_ledger(ledger, "sha256:fixed", output_dir=output, status="reserved")
    with pytest.raises(ContractError, match="already reserved"):
        _mutate_launch_ledger(
            ledger, "sha256:fixed", output_dir=output, status="reserved"
        )
    _mutate_launch_ledger(ledger, "sha256:fixed", output_dir=output, status="success")
    assert json.loads(ledger.read_text(encoding="utf-8"))["runs"]["sha256:fixed"][
        "status"
    ] == "success"