tests/test_dominance_selector.py

from __future__ import annotations

import json
from types import SimpleNamespace

import numpy as np
import pytest

from scripts.eval_dominance_selector import (
    _choose_tau,
    _conformal_quantile,
    _evaluate_case,
    _pairwise_calibration_summary as _score_pairwise_calibration_summary,
)
from scripts.eval_learned_dominance_selector import (
    _chart_compat_feature,
    _candidate_feature,
    _choose_thresholds,
    _evaluate_dataset,
    _evaluate_predictions,
    _feature_names,
    _fit_select_ridge,
    _pairwise_calibration_summary,
    _score_shape_matrix,
    _source_evidence_feature,
    _target_value,
    _write_report_artifact,
)
from scripts.eval_nonlinear_dominance_selector import _split_rows, _subset_dataset
from scripts.eval_nonlinear_dominance_selector import _write_provenance as _write_nonlinear_provenance


def test_conformal_quantile_and_fallback_selection() -> None:
    quantile = _conformal_quantile([0.0, 0.2, 0.4, 0.6], alpha=0.25)
    assert quantile == 0.6

    case = {
        "chart_id": "c0",
        "task_id": "pick",
        "seed": "0",
        "train_seed": "0",
        "predicted_margin": 0.5,
        "measured_margin": 0.4,
        "base_utility": 0.3,
        "base_success": 0.0,
        "top_generated_utility": 1.2,
        "top_generated_success": 1.0,
        "proposal_oracle_utility": 1.2,
        "proposal_oracle_success": 1.0,
        "hidden_chart_oracle_utility": 1.5,
        "hidden_chart_oracle_success": 1.0,
        "outcome_ptr": 1.0,
    }
    executed = _evaluate_case(case, residual_quantile=0.1, tau=0.0)
    assert executed["execute_generated"] == 1.0
    assert executed["selected_success"] == 1.0

    fallback = _evaluate_case(case, residual_quantile=0.6, tau=0.0)
    assert fallback["execute_generated"] == 0.0
    assert fallback["fallback_to_base"] == 1.0
    assert fallback["selected_utility"] == 0.3


def test_score_dominance_logs_pairwise_calibration_and_unsafe_execution() -> None:
    case = {
        "chart_id": "c0",
        "task_id": "pick",
        "seed": "0",
        "train_seed": "0",
        "predicted_scores": [3.0, 1.0, -1.0],
        "predicted_margin": 1.0,
        "measured_margin": 0.7,
        "base_utility": 0.2,
        "base_success": 0.0,
        "base_unsafe_known": 0.0,
        "top_generated_utility": 0.9,
        "top_generated_success": 1.0,
        "top_candidate_unsafe_known": 1.0,
        "candidate_safety_label_coverage": 1.0,
        "candidate_unsafe_rate_known": 1.0 / 3.0,
        "generated_utilities": [0.9, 0.4, 0.1],
        "proposal_oracle_utility": 0.9,
        "proposal_oracle_success": 1.0,
        "hidden_chart_oracle_utility": 1.0,
        "hidden_chart_oracle_success": 1.0,
        "outcome_ptr": 1.0,
    }
    calibration = _score_pairwise_calibration_summary([case])
    executed = _evaluate_case(
        case,
        residual_quantile=0.0,
        tau=0.0,
        pairwise_calibration=calibration["rows"][0],
    )
    fallback = _evaluate_case(
        case,
        residual_quantile=2.0,
        tau=0.0,
        pairwise_calibration=calibration["rows"][0],
    )

    assert calibration["num_pairs"] == 3
    assert executed["pairwise_causal_calibration_pairs"] == 3.0
    assert executed["pairwise_causal_calibration_accuracy"] == 1.0
    assert executed["unsafe_execution_label_known"] == 1.0
    assert executed["unsafe_execution_known"] == 1.0
    assert fallback["fallback_to_base"] == 1.0
    assert fallback["unsafe_execution_label_known"] == 1.0
    assert fallback["unsafe_execution_known"] == 0.0


def test_choose_tau_prefers_success_then_utility_then_coverage() -> None:
    cases = [
        {
            "predicted_margin": 2.0,
            "base_utility": 0.5,
            "base_success": 0.0,
            "top_generated_utility": 1.0,
            "top_generated_success": 1.0,
            "proposal_oracle_utility": 1.0,
            "proposal_oracle_success": 1.0,
            "hidden_chart_oracle_utility": 1.0,
            "hidden_chart_oracle_success": 1.0,
            "outcome_ptr": 1.0,
        },
        {
            "predicted_margin": -1.0,
            "base_utility": 0.8,
            "base_success": 1.0,
            "top_generated_utility": 0.1,
            "top_generated_success": 0.0,
            "proposal_oracle_utility": 0.1,
            "proposal_oracle_success": 0.0,
            "hidden_chart_oracle_utility": 1.0,
            "hidden_chart_oracle_success": 1.0,
            "outcome_ptr": 0.0,
        },
    ]
    tau = _choose_tau(cases, residual_quantile=0.0)
    evaluated = [_evaluate_case(case, residual_quantile=0.0, tau=tau) for case in cases]
    assert sum(row["selected_success"] for row in evaluated) == 2.0


def test_learned_dominance_ridge_prefers_positive_candidate() -> None:
    def sample(row: int, candidate: int, score: float, margin: float, success: float) -> dict:
        feature = np.zeros(10, dtype=float)
        feature[0] = 1.0
        feature[1] = score
        feature[2] = margin
        feature[4] = float(candidate)
        return {
            "row_index": row,
            "candidate_index": candidate,
            "feature": feature,
            "target_margin": success - 0.2,
            "candidate_utility": success,
            "candidate_success": success,
            "base_utility": 0.2,
            "base_success": 0.0,
            "proposal_oracle_utility": success,
            "proposal_oracle_success": success,
            "hidden_chart_oracle_utility": 1.0,
            "hidden_chart_oracle_success": 1.0,
            "outcome_ptr": success,
            "chart_id": f"c{row}",
            "task_id": "pick",
            "seed": "0",
            "train_seed": "0",
        }

    dataset = {
        "samples": [
            sample(0, 0, 2.0, 2.0, 1.0),
            sample(0, 1, -1.0, -1.0, 0.0),
            sample(1, 0, 2.0, 2.0, 1.0),
            sample(1, 1, -1.0, -1.0, 0.0),
        ],
        "by_row": {0: [0, 1], 1: [2, 3]},
        "num_rows": 2,
    }
    fit = _fit_select_ridge(dataset, lambdas=[0.1])
    rows = _evaluate_dataset(dataset, fit["weights"], fit["mean"], fit["std"], tau=fit["tau"])
    assert [row["selected_success"] for row in rows] == [1.0, 1.0]


def test_learned_dominance_pairwise_objective_prefers_within_chart_winner() -> None:
    def sample(row: int, candidate: int, feature_value: float, success: float) -> dict:
        feature = np.zeros(10, dtype=float)
        feature[0] = 1.0
        feature[1] = feature_value
        return {
            "row_index": row,
            "candidate_index": candidate,
            "feature": feature,
            "target_margin": success,
            "candidate_utility": success,
            "candidate_success": success,
            "base_utility": 0.0,
            "base_success": 0.0,
            "proposal_oracle_utility": 1.0,
            "proposal_oracle_success": 1.0,
            "hidden_chart_oracle_utility": 1.0,
            "hidden_chart_oracle_success": 1.0,
            "outcome_ptr": 1.0,
            "chart_id": f"c{row}",
            "task_id": "pick",
            "seed": "0",
            "train_seed": "0",
        }

    dataset = {
        "samples": [
            sample(0, 0, 2.0, 1.0),
            sample(0, 1, -1.0, 0.0),
            sample(0, 2, 0.0, 0.0),
            sample(1, 0, 3.0, 1.0),
            sample(1, 1, -2.0, 0.0),
            sample(1, 2, 0.5, 0.0),
        ],
        "by_row": {0: [0, 1, 2], 1: [3, 4, 5]},
        "num_rows": 2,
    }
    fit = _fit_select_ridge(dataset, lambdas=[0.1], fit_objective="pairwise")
    rows = _evaluate_dataset(dataset, fit["weights"], fit["mean"], fit["std"], tau=fit["tau"])
    assert fit["fit_design"]["num_candidate_rows"] == 6
    assert fit["fit_design"]["num_pairwise_rows"] == 8
    assert fit["fit_design"]["num_fit_rows"] == 8
    assert [row["selected_candidate_index"] for row in rows] == [0, 0]
    assert [row["selected_success"] for row in rows] == [1.0, 1.0]


def test_task_scoped_threshold_can_use_visible_task_bucket() -> None:
    dataset = {
        "samples": [
            {
                "row_index": 0,
                "candidate_index": 0,
                "feature": np.ones(2, dtype=float),
                "target_margin": 1.0,
                "candidate_utility": 1.0,
                "candidate_success": 1.0,
                "base_utility": 0.0,
                "base_success": 0.0,
                "proposal_oracle_utility": 1.0,
                "proposal_oracle_success": 1.0,
                "hidden_chart_oracle_utility": 1.0,
                "hidden_chart_oracle_success": 1.0,
                "outcome_ptr": 1.0,
                "chart_id": "a",
                "task_id": "task_execute",
                "seed": "0",
                "train_seed": "0",
            },
            {
                "row_index": 1,
                "candidate_index": 0,
                "feature": np.ones(2, dtype=float),
                "target_margin": -1.0,
                "candidate_utility": 0.0,
                "candidate_success": 0.0,
                "base_utility": 1.0,
                "base_success": 1.0,
                "proposal_oracle_utility": 0.0,
                "proposal_oracle_success": 0.0,
                "hidden_chart_oracle_utility": 1.0,
                "hidden_chart_oracle_success": 1.0,
                "outcome_ptr": 0.0,
                "chart_id": "b",
                "task_id": "task_fallback",
                "seed": "0",
                "train_seed": "0",
            },
        ],
        "by_row": {0: [0], 1: [1]},
        "num_rows": 2,
    }
    predictions = np.asarray([0.6, 0.8], dtype=float)
    global_tau, _global_summary = _choose_thresholds(
        dataset,
        predictions,
        threshold_scope="global",
    )
    task_tau, task_summary = _choose_thresholds(
        dataset,
        predictions,
        threshold_scope="task",
    )
    global_rows = _evaluate_predictions(dataset, predictions, tau=global_tau)
    task_rows = _evaluate_predictions(dataset, predictions, tau=task_tau)
    assert np.mean([row["selected_success"] for row in global_rows]) == 0.5
    assert task_summary["selected_success"] == 1.0
    assert [row["selected_success"] for row in task_rows] == [1.0, 1.0]


def test_pairwise_calibration_ece_is_logged_per_selector_row() -> None:
    def sample(row: int, candidate: int, utility: float) -> dict:
        return {
            "row_index": row,
            "candidate_index": candidate,
            "feature": np.asarray([1.0, float(candidate)]),
            "target_margin": utility,
            "candidate_utility": utility,
            "candidate_success": float(utility > 0.5),
            "base_utility": 0.0,
            "base_success": 0.0,
            "proposal_oracle_utility": 1.0,
            "proposal_oracle_success": 1.0,
            "hidden_chart_oracle_utility": 1.0,
            "hidden_chart_oracle_success": 1.0,
            "outcome_ptr": 1.0,
            "chart_id": f"c{row}",
            "task_id": "pick",
            "seed": "0",
            "train_seed": "0",
        }

    dataset = {
        "samples": [
            sample(0, 0, 1.0),
            sample(0, 1, 0.0),
            sample(0, 2, 0.5),
            sample(1, 0, 1.0),
            sample(1, 1, 0.0),
            sample(1, 2, 0.5),
        ],
        "by_row": {0: [0, 1, 2], 1: [3, 4, 5]},
        "num_rows": 2,
    }
    good_predictions = np.asarray([2.0, -2.0, 0.0, 2.0, -2.0, 0.0], dtype=float)
    bad_predictions = -good_predictions

    good = _pairwise_calibration_summary(dataset, good_predictions)
    bad = _pairwise_calibration_summary(dataset, bad_predictions)
    rows = _evaluate_predictions(
        dataset,
        good_predictions,
        tau=-10.0,
        include_pairwise_calibration=True,
        pairwise_calibration=good,
    )

    assert good["num_pairs"] == 6
    assert good["accuracy"] == 1.0
    assert bad["accuracy"] == 0.0
    assert rows[0]["pairwise_causal_calibration_pairs"] == 3.0
    assert rows[0]["pairwise_causal_calibration_ece"] == pytest.approx(
        good["rows"][0]["ece"]
    )


def test_learned_dominance_tangent_features_and_success_weighted_target() -> None:
    tangent = np.arange(21, dtype=float)
    basic = _candidate_feature(
        score=1.0,
        base_score=0.25,
        score_mean=0.5,
        score_std=0.5,
        candidate_index=1,
        candidate_type="ctt_transport_rank3",
        tangent=tangent,
        num_candidates=8,
        feature_set="basic",
    )
    expanded = _candidate_feature(
        score=1.0,
        base_score=0.25,
        score_mean=0.5,
        score_std=0.5,
        candidate_index=1,
        candidate_type="ctt_transport_rank3",
        tangent=tangent,
        num_candidates=8,
        feature_set="tangent",
    )
    context = _candidate_feature(
        score=1.0,
        base_score=0.25,
        score_mean=0.5,
        score_std=0.5,
        candidate_index=1,
        candidate_type="ctt_transport_rank3",
        tangent=tangent,
        num_candidates=8,
        feature_set="context",
        context={
            "target_task_id": "PickCube-v1",
            "source_task_id": "PickCube-v1",
            "instruction": "Pick up the cube.",
        },
    )
    context_tangent = _candidate_feature(
        score=1.0,
        base_score=0.25,
        score_mean=0.5,
        score_std=0.5,
        candidate_index=1,
        candidate_type="ctt_transport_rank3",
        tangent=tangent,
        num_candidates=8,
        feature_set="context_tangent",
        context={
            "target_task_id": "PickCube-v1",
            "source_task_id": "PickCube-v1",
            "instruction": "Pick up the cube.",
        },
    )
    source_stats = {
        "positive_tangents": [np.zeros(21, dtype=float), np.ones(21, dtype=float)],
        "negative_tangents": [np.full(21, 4.0, dtype=float)],
        "num_nonbase": 4,
        "positive_delta_utilities": [0.4, 0.8],
        "positive_utilities": [1.2, 1.6],
        "positive_success": [1.0, 0.0],
        "positive_progress": [0.8, 0.9],
        "positive_safety": [0.0, 0.0],
    }
    source_evidence = _source_evidence_feature(source_stats, tangent=np.zeros(21, dtype=float))
    context_tangent_source = _candidate_feature(
        score=1.0,
        base_score=0.25,
        score_mean=0.5,
        score_std=0.5,
        candidate_index=1,
        candidate_type="ctt_transport_rank3",
        tangent=tangent,
        num_candidates=8,
        feature_set="context_tangent_source_evidence",
        context={
            "target_task_id": "PickCube-v1",
            "source_task_id": "PickCube-v1",
            "instruction": "Pick up the cube.",
        },
        source_evidence=source_evidence,
    )
    assert len(basic) == len(_feature_names("basic")) == 10
    assert len(expanded) == len(_feature_names("tangent")) == 52
    assert len(context) == len(_feature_names("context")) == 37
    assert len(context_tangent) == len(_feature_names("context_tangent")) == 79
    assert len(context_tangent_source) == len(_feature_names("context_tangent_source_evidence"))
    assert len(context_tangent_source) == 97
    assert basic[2] == 0.75
    assert basic[6] == 3.0
    assert context[-3] == 1.0
    assert source_evidence[0] == 1.0
    assert source_evidence[1] == pytest.approx(np.log1p(2))
    assert source_evidence[-1] == 1.0
    assert _target_value("success_weighted_margin", utility_margin=0.2, candidate_success=1.0) == 1.2
    assert _target_value(
        "success_weighted_margin",
        utility_margin=0.2,
        candidate_success=1.0,
        success_bonus=2.5,
    ) == 2.7


def test_score_shape_feature_set_and_markdown_suppression(tmp_path) -> None:
    score_shape = _score_shape_matrix([3.0, 1.0, 2.0])
    feature = _candidate_feature(
        score=3.0,
        base_score=0.25,
        score_mean=2.0,
        score_std=1.0,
        candidate_index=0,
        candidate_type="ctt_transport_rank0",
        tangent=np.zeros(21, dtype=float),
        num_candidates=3,
        feature_set="score_context_chart_compat",
        context={
            "target_task_id": "PickCube-v1",
            "source_task_id": "PickCube-v1",
            "instruction": "Pick up the cube.",
        },
        chart_compat=np.ones(22, dtype=float),
        score_shape=score_shape[0],
    )

    assert score_shape.shape == (3, 8)
    assert score_shape[0, 0] == 0.0
    assert score_shape[0, 2] == 0.0
    assert score_shape[0, 3] == 1.0
    assert score_shape[1, 0] == 1.0
    assert len(feature) == len(_feature_names("score_context_chart_compat"))
    assert len(feature) == 10 + 8 + 27 + 22

    metrics = {
        "num_calibration_rows": 1,
        "num_eval_rows": 1,
        "selected_lambda": 0.0,
        "tau": 0.0,
        "threshold_scope": "global",
        "fit_objective": "pointwise",
        "feature_set": "basic",
        "target": "utility_margin",
        "eval_summary": {},
        "calibration_summary": {},
    }
    report_path = tmp_path / "report.md"
    report_path.write_text("old\n")
    _write_report_artifact(tmp_path, metrics, no_markdown_report=True)
    assert not report_path.exists()
    _write_report_artifact(tmp_path, metrics, no_markdown_report=False)
    assert report_path.exists()


def test_chart_compat_feature_is_deployment_visible_and_fixed_width() -> None:
    base = np.arange(14, dtype=float).reshape(2, 7) / 10.0
    context = np.linspace(-0.5, 0.5, 18)
    obs = np.linspace(0.0, 1.0, 32)
    obj = np.linspace(1.0, 2.0, 64)
    target = SimpleNamespace(
        base_action=base,
        feature=np.concatenate([base.reshape(-1), context, obs, obj]),
    )
    source = SimpleNamespace(
        base_action=base + 0.1,
        feature=np.concatenate([base.reshape(-1) + 0.1, context, obs + 0.2, obj + 0.3]),
    )

    compat = _chart_compat_feature(
        target,
        source,
        chart_feature_mode="base_context_obs_obj",
    )
    feature = _candidate_feature(
        score=1.0,
        base_score=0.25,
        score_mean=0.5,
        score_std=0.5,
        candidate_index=1,
        candidate_type="ctt_transport_rank3",
        tangent=np.zeros(21, dtype=float),
        num_candidates=8,
        feature_set="chart_compat",
        chart_compat=compat,
    )

    assert compat.shape == (22,)
    assert len(feature) == len(_feature_names("chart_compat")) == 32
    assert compat[0] == 1.0
    assert compat[12] == 1.0
    assert compat[17] == 1.0
    assert np.all(np.isfinite(compat))
    assert np.all(_chart_compat_feature(target, None, chart_feature_mode="base_context_obs_obj") == 0.0)


def test_nonlinear_dominance_split_is_row_disjoint() -> None:
    samples = []
    by_row = {}
    for row in range(10):
        by_row[row] = []
        for candidate in range(2):
            by_row[row].append(len(samples))
            samples.append(
                {
                    "row_index": row,
                    "candidate_index": candidate,
                    "feature": np.asarray([1.0, float(candidate)]),
                    "target_margin": float(candidate),
                    "measured_utility_margin": float(candidate),
                    "candidate_success": float(candidate),
                    "candidate_utility": float(candidate),
                    "base_utility": 0.0,
                    "base_success": 0.0,
                    "proposal_oracle_utility": 1.0,
                    "proposal_oracle_success": 1.0,
                    "hidden_chart_oracle_utility": 1.0,
                    "hidden_chart_oracle_success": 1.0,
                    "outcome_ptr": 1.0,
                    "chart_id": f"c{row}",
                    "task_id": "pick",
                    "seed": "0",
                    "train_seed": "0",
                }
            )
    dataset = {"samples": samples, "by_row": by_row, "num_rows": len(by_row)}

    fit_rows, select_rows = _split_rows(dataset, selection_frac=0.3, seed=7)
    assert set(fit_rows).isdisjoint(select_rows)
    assert set(fit_rows) | set(select_rows) == set(by_row)

    subset = _subset_dataset(dataset, select_rows)
    assert subset["num_rows"] == len(select_rows)
    assert len(subset["samples"]) == 2 * len(select_rows)


def test_nonlinear_chart_compat_provenance_records_default_source_index(tmp_path) -> None:
    calibration_input = tmp_path / "calibration.json"
    eval_input = tmp_path / "eval.json"
    train_index = tmp_path / "train_index.json"
    test_index = tmp_path / "test_index.json"
    calibration_input.write_text("{}\n")
    eval_input.write_text("{}\n")
    train_index.write_text(
        '{"split":"train","content_hash":"train-content","split_hash":"train-split","retrieval_index_allowed":true}\n'
    )
    test_index.write_text(
        '{"split":"test","content_hash":"test-content","split_hash":"test-split","retrieval_index_allowed":false}\n'
    )
    args = SimpleNamespace(
        calibration_input=calibration_input,
        calibration_target_index=train_index,
        eval_input=eval_input,
        eval_target_index=test_index,
        source_index=None,
        checkpoint_template="runs/ctt_residual_base_context_obs_seed{seed}/model.pt",
        out_dir=tmp_path / "out",
        k=16,
        feature_set="chart_compat",
        selector_chart_feature_mode="base_context_obs",
        target="utility_margin",
        model_types="hgb_regressor",
        selection_frac=0.35,
        seed=0,
        bootstrap_samples=10,
    )

    args.out_dir.mkdir()
    _write_nonlinear_provenance(args.out_dir, args)
    data_hash = json.loads((args.out_dir / "data_hash.txt").read_text())
    split_hash = json.loads((args.out_dir / "split_hash.txt").read_text())

    assert "source_index" in data_hash
    assert split_hash["source_index"]["split"] == "train"
    assert split_hash["source_index"]["retrieval_index_allowed"] is True