bundle/v0/tests/test_empirical_task_runner.py

from __future__ import annotations

import hashlib
import json
from pathlib import Path

import numpy as np
import pytest

import loss_aware_dro_repro.empirical_task as empirical_task_module
from loss_aware_dro_repro.artifacts import validate_result
from loss_aware_dro_repro.core import ContractError, load_plan, plan_hash
from loss_aware_dro_repro.datasets import generate_dataset
from loss_aware_dro_repro.empirical_task import (
    _portfolio_true_mean,
    _regression_oos_losses,
    normalize_empirical_execution_config,
    require_empirical_execution_snapshot,
    run_empirical_task,
    select_empirical_task,
    verify_empirical_execution_snapshot_unchanged,
)
from loss_aware_dro_repro.matrix import expand_tasks


TEST_EXECUTION = {
    "schema_version": 2,
    "execution_scale": "test",
    "max_outer_iterations": 1,
    "oos_sample_count": 257,
    "oos_chunk_size": 31,
}

STOPPING_STORAGE_EXECUTION = {
    "schema_version": 2,
    "execution_scale": "stopping_storage_sample",
    "max_outer_iterations": 5000,
}


def _load(path: Path):
    return json.loads(path.read_text(encoding="utf-8"))


def _sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


@pytest.mark.parametrize(
    "task_id,suite",
    [
        ("portfolio_discrete/d001/r00/n010", "portfolio_discrete"),
        ("portfolio_gmm/d001/r00/n010", "portfolio_gmm"),
        ("regression_absolute_main/d001/r00/n010", "regression_absolute_main"),
        ("regression_squared/d001/r00/n010", "regression_squared"),
    ],
)
def test_selects_each_supported_empirical_suite(task_id, suite):
    assert select_empirical_task(task_id)["suite"] == suite


def test_rejects_gaussian_suite():
    with pytest.raises(ContractError, match="not supported"):
        select_empirical_task("portfolio_gaussian_main/d001/r00/n010")


def test_execution_contract_freezes_paper_scale_and_bounds_tests():
    plan = load_plan()
    exact = normalize_empirical_execution_config(
        plan, suite="regression_absolute_main"
    )
    squared = normalize_empirical_execution_config(
        plan, suite="regression_squared"
    )
    assert exact["max_outer_iterations"] == 1_000_000
    assert exact["optimizer"] == "paper_algorithm_plain_gradient_descent"
    assert exact["oos_sample_count"] == 10_000_000
    assert exact["oos_sample_count_source"] == "main_or_exemplar_primary"
    assert squared["oos_sample_count"] == 1_000_000
    assert squared["oos_sample_count_source"] == "appendix_f3_multi_experiment"
    with pytest.raises(ContractError, match="may not override"):
        normalize_empirical_execution_config(
            plan,
            {**exact, "oos_sample_count": 100},
            suite="regression_absolute_main",
        )
    with pytest.raises(ContractError, match=r"\[1, 5\]"):
        normalize_empirical_execution_config(
            plan, {**TEST_EXECUTION, "max_outer_iterations": 6}
        )
    changed_optimizer = {
        **plan,
        "paper_hyperparameters": {
            **plan["paper_hyperparameters"],
            "optimizer": "adam",
        },
    }
    with pytest.raises(ContractError, match="frozen paper optimizer"):
        normalize_empirical_execution_config(changed_optimizer)


@pytest.mark.parametrize(
    "suite,oos_sample_count,oos_sample_count_source",
    [
        ("regression_absolute_main", 10_000_000, "main_or_exemplar_primary"),
        ("regression_squared", 1_000_000, "appendix_f3_multi_experiment"),
    ],
)
def test_stopping_storage_sample_preserves_suite_specific_paper_oos_contract(
    suite, oos_sample_count, oos_sample_count_source
):
    normalized = normalize_empirical_execution_config(
        load_plan(),
        STOPPING_STORAGE_EXECUTION,
        suite=suite,
    )
    assert normalized["execution_scale"] == "stopping_storage_sample"
    assert normalized["max_outer_iterations"] == 5000
    assert normalized["store_every"] == 100
    assert normalized["relative_objective_improvement_tolerance"] == 1e-6
    assert normalized["oos_sample_count"] == oos_sample_count
    assert normalized["oos_sample_count_source"] == oos_sample_count_source
    assert normalized["oos_chunk_size"] == 100_000
    assert normalized["claim_eligible"] is False


@pytest.mark.parametrize(
    "override",
    [
        {"max_outer_iterations": 4999},
        {"max_outer_iterations": 5000.0},
        {"schema_version": 2.0},
        {"store_every": 1},
        {"optimizer": "adam"},
        {"oos_sample_count": 257},
        {"oos_chunk_size": 31},
    ],
)
def test_empirical_stopping_storage_sample_rejects_scientific_overrides(override):
    with pytest.raises(ContractError, match="exact frozen 5000-step request"):
        normalize_empirical_execution_config(
            load_plan(),
            {**STOPPING_STORAGE_EXECUTION, **override},
            suite="regression_absolute_main",
        )


def test_empirical_paper_scale_requires_clean_snapshot_but_test_scale_does_not():
    paper = normalize_empirical_execution_config(load_plan())
    with pytest.raises(ContractError, match="paper lane is not clean"):
        require_empirical_execution_snapshot(paper, False)
    sample = normalize_empirical_execution_config(
        load_plan(), STOPPING_STORAGE_EXECUTION
    )
    with pytest.raises(ContractError, match="paper lane is not clean"):
        require_empirical_execution_snapshot(sample, False)
    require_empirical_execution_snapshot(TEST_EXECUTION, False)


def test_empirical_stopping_storage_sample_rejects_changed_snapshot(monkeypatch):
    sample = normalize_empirical_execution_config(
        load_plan(), STOPPING_STORAGE_EXECUTION
    )
    monkeypatch.setattr(
        empirical_task_module,
        "_repository_state",
        lambda: ("b" * 40, True),
    )
    with pytest.raises(ContractError, match="source snapshot changed"):
        verify_empirical_execution_snapshot_unchanged(sample, "a" * 40, True)


@pytest.mark.parametrize(
    "task_id",
    [
        "portfolio_discrete/d001/r00/n010",
        "portfolio_gmm/d001/r00/n010",
        "regression_absolute_main/d001/r00/n010",
        "regression_squared/d001/r00/n010",
    ],
)
def test_executes_empirical_task_and_binds_receipts(tmp_path, task_id):
    output = tmp_path / task_id.split("/")[0]
    result = run_empirical_task(
        task_id, output, requested_execution_config=TEST_EXECUTION
    )
    expected = next(task for task in expand_tasks() if task["task_id"] == task_id)
    validate_result(
        result,
        expected,
        plan_hash(load_plan()),
        artifact_root=output,
        solver_residual_max=load_plan()["numerics"]["solver_residual_max"],
    )
    assert set(path.name for path in output.iterdir()) == {
        "lineage_receipt.json",
        "oos_receipt.json",
        "result.json",
        "solver_receipt.json",
        "stopping_receipt.json",
        "iteration_trace.jsonl.gz",
        "checkpoint_states.jsonl.gz",
    }
    lineage = _load(output / "lineage_receipt.json")
    assert lineage["task_hash"] == expected["task_hash"]
    assert lineage["execution_config_hash"].startswith("sha256:")
    assert lineage["source_tree_hash"].startswith("sha256:")
    assert lineage["claim_eligible"] is False
    assert set(lineage["scientific_verdicts"].values()) == {"HOLD"}
    assert result["artifact_hashes"]["iteration_trace"] == _sha256(
        output / "iteration_trace.jsonl.gz"
    )
    assert result["artifact_hashes"]["checkpoint_states"] == _sha256(
        output / "checkpoint_states.jsonl.gz"
    )
    assert result["optimization"]["iteration_trace_records"] == 1
    assert result["optimization"]["checkpoint_state_records"] == 2
    solver = _load(output / "solver_receipt.json")
    stopping = _load(output / "stopping_receipt.json")
    for receipt in (solver, stopping):
        assert receipt["task_hash"] == expected["task_hash"]
        assert receipt["execution_config_hash"] == lineage["execution_config_hash"]
        assert receipt["source_tree_hash"] == lineage["source_tree_hash"]
    assert solver["all_pass"] is True
    assert stopping["iterations"] == 1
    assert stopping["stop_reason"] == "maximum_outer_iterations_reached"
    assert stopping["primary_stopping_rule"] == "paper_total_phi_literal"
    assert stopping["safety_diagnostic"] == "paper_total_phi_abs_denominator"
    assert stopping["released_source_diagnostic"] == "released_lower_objective_abs_denominator"
    assert stopping["terminal_diagnostics"] == result["optimization"]["stopping"]
    oos = _load(output / "oos_receipt.json")
    if expected["problem"] == "portfolio":
        assert oos["evaluation"] == "exact_generating_distribution_expectation"
    else:
        assert oos["evaluation"] == "deterministic_chunked_monte_carlo"
        assert oos["sample_count"] == TEST_EXECUTION["oos_sample_count"]
        assert oos["sample_count_source"] == "test_override"
        assert oos["chunk_size"] == TEST_EXECUTION["oos_chunk_size"]


@pytest.mark.parametrize(
    "task_id",
    [
        "portfolio_discrete/d001/r00/n010",
        "portfolio_gmm/d001/r00/n010",
    ],
)
def test_portfolio_oos_mean_is_exact_for_generating_distribution(task_id):
    task = select_empirical_task(task_id)
    _, metadata = generate_dataset(task)
    mean = _portfolio_true_mean(task, metadata)
    parameters = metadata["parameters"]
    if task["suite"] == "portfolio_discrete":
        expected = np.asarray(parameters["weights"]) @ np.asarray(parameters["support"])
    else:
        expected = np.asarray(parameters["weights"]) @ np.asarray(parameters["means"])
    np.testing.assert_allclose(mean, expected, rtol=0.0, atol=0.0)


@pytest.mark.parametrize(
    "task_id",
    [
        "regression_absolute_main/d001/r00/n010",
        "regression_squared/d001/r00/n010",
    ],
)
def test_regression_oos_is_deterministic_and_stable_across_chunk_sizes(task_id):
    task = select_empirical_task(task_id)
    _, metadata = generate_dataset(task)
    initial = np.asarray([0.25])
    final = np.asarray([0.75])
    first = _regression_oos_losses(
        task,
        metadata,
        initial,
        final,
        sample_count=101,
        chunk_size=7,
    )
    second = _regression_oos_losses(
        task,
        metadata,
        initial,
        final,
        sample_count=101,
        chunk_size=19,
    )
    repeated = _regression_oos_losses(
        task,
        metadata,
        initial,
        final,
        sample_count=101,
        chunk_size=7,
    )
    assert first == repeated
    np.testing.assert_allclose(first[:2], second[:2], rtol=0.0, atol=1e-12)


def test_empirical_output_is_immutable(tmp_path):
    task_id = "portfolio_discrete/d001/r00/n010"
    output = tmp_path / "immutable"
    run_empirical_task(task_id, output, requested_execution_config=TEST_EXECUTION)
    with pytest.raises(ContractError, match="already exists"):
        run_empirical_task(task_id, output, requested_execution_config=TEST_EXECUTION)


def test_failed_empirical_run_removes_unpublished_streams(tmp_path, monkeypatch):
    import loss_aware_dro_repro.empirical_task as module

    output = tmp_path / "failed"
    monkeypatch.setattr(
        module,
        "_outer_state",
        lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("injected")),
    )
    with pytest.raises(RuntimeError, match="injected"):
        run_empirical_task(
            "portfolio_discrete/d001/r00/n010",
            output,
            requested_execution_config=TEST_EXECUTION,
        )
    assert not output.exists()
    assert not list(tmp_path.glob(".empirical-task-*"))