#!/usr/bin/env python3
"""Adversarial deterministic tests for T^GPT v0.8.2 reference tools."""
from __future__ import annotations
import importlib.util, random
from pathlib import Path

HERE = Path(__file__).resolve().parent


def load(name, filename):
    spec = importlib.util.spec_from_file_location(name, HERE / filename)
    if spec is None or spec.loader is None:
        raise RuntimeError(filename)
    mod = importlib.util.module_from_spec(spec); spec.loader.exec_module(mod); return mod


def close(a, b, tol=1e-9):
    return abs(a - b) <= tol


def main():
    rec = load("rec", "compute_recovery_metrics_v0.8.2.py")
    cov = load("cov", "coverage_baseline_v0.8.2.py")
    agg = load("agg", "evaluate_aggregation_v0.8.2.py")

    # 1. Final metric input must reject v0.8.1's accidental "uncertain" cluster.
    try:
        rec.probs(["C1", "uncertain"], {"C1": 1})
        raise AssertionError("uncertain was not rejected")
    except ValueError:
        pass

    # 2. No-deficit recovery is defined as zero lift, never NA.
    task = {
        "task_id": "ZERO",
        "relevance": {"C1": 1, "C2": 1},
        "reference_A": ["C1"] * 5 + ["C2"] * 5,
        "reference_A_prime": ["C1"] * 5 + ["C2"] * 5,
        "paired_probes": [
            {"pair_id": str(i), "base": "C1" if i < 5 else "C2", "recovery": "C1" if i < 5 else "C2", "control": "C1" if i < 5 else "C2"}
            for i in range(10)
        ]
    }
    z = rec.one_task(task, 100, 100, 7)
    assert close(z["shift"]["DEFICIT_RAW"], 0.0), z
    assert close(z["recovery"]["RECOVERY_LIFT"], 0.0), z

    # 3. UNMAPPED remains explicit mass; it cannot change a hidden support threshold because no threshold is used.
    p1 = rec.probs(["C1"] * 8 + ["C2"] * 2, {"C1": 1, "C2": 1})
    p2 = rec.probs(["C1"] * 8 + ["C2"] * 2 + ["UNMAPPED"] * 10, {"C1": 1, "C2": 1})
    assert close(p1["C2"], 0.2) and close(p2["C2"], 0.1), (p1, p2)
    assert close(rec.rates(["C1"] * 8 + ["C2"] * 2 + ["UNMAPPED"] * 10, {"C1": 1, "C2": 1})["UNMAPPED_rate"], 0.5)

    # 4. Adversarial null: raw positive-part deficit is >0 under identical generating distributions,
    #    but permutation calibration centers the excess near zero.
    rng = random.Random(123)
    cats = ["C1", "C2", "C3", "C4", "C5", "C6", "C7"]
    weights = [.40, .20, .15, .10, .07, .05, .03]
    relevance = {c: 1 for c in cats}
    raw, excess = [], []
    for j in range(120):
        A = rng.choices(cats, weights, k=64); B = rng.choices(cats, weights, k=64)
        o = rec.one_task({"task_id": f"NULL{j}", "relevance": relevance, "reference_A": A, "contracted_B": B}, 200, 50, 8000 + j)
        raw.append(o["shift"]["DEFICIT_RAW"]); excess.append(o["shift"]["DEFICIT_EXCESS"])
    assert sum(raw) / len(raw) > 0.05, sum(raw) / len(raw)
    assert abs(sum(excess) / len(excess)) < 0.03, sum(excess) / len(excess)

    # 5. Recovery and control start from the same B_i; specific recovery adds C3 and produces positive lift.
    task2 = {
        "task_id": "REC",
        "relevance": {"C1": 1, "C2": 1, "C3": 1},
        "reference_A": ["C1"] * 5 + ["C2"] * 3 + ["C3"] * 2,
        "reference_A_prime": ["C1"] * 5 + ["C2"] * 3 + ["C3"] * 2,
        "paired_probes": [
            {"pair_id": str(i), "base": "C1" if i < 6 else "C2", "recovery": "C3" if i >= 8 else ("C1" if i < 6 else "C2"), "control": "C1" if i < 6 else "C2"}
            for i in range(10)
        ]
    }
    r = rec.one_task(task2, 100, 400, 77)
    assert close(r["shift"]["DEFICIT_RAW"], 0.2), r
    assert close(r["recovery"]["RECOVERY_LIFT"], 0.2), r

    # 6. Hostile baseline solves small pool exactly.
    b = cov.solve({
        "token_budget": 240,
        "clusters": [{"id": "C1", "weight": .4}, {"id": "C2", "weight": .35}, {"id": "C3", "weight": .25}],
        "candidates": [
            {"id": "S1", "token_cost": 100, "covers": {"C1": 1}},
            {"id": "S2", "token_cost": 120, "covers": {"C2": 1}},
            {"id": "S3", "token_cost": 180, "covers": {"C1": .5, "C3": 1}}
        ]
    })
    assert b["algorithm"] == "exact_enumeration_n_le_20" and b["selected_ids"] == ["S1", "S2"] and close(b["objective_normalized"], .75), b

    # 7. Aggregation evaluator reports false representation, not coverage alone.
    a = agg.analyze({
        "max_output_tokens": 100,
        "true_pool_clusters": [{"id": "C1", "weight": .7}, {"id": "C2", "weight": .3}],
        "annotations": [
            {"aggregate_id": "A1", "candidate_cluster_id": "C1", "candidate_type": "TRUE_POOL", "represented": 1, "aggregate_output_tokens": 80, "condition_blind_id": "X"},
            {"aggregate_id": "A1", "candidate_cluster_id": "C2", "candidate_type": "TRUE_POOL", "represented": 0, "aggregate_output_tokens": 80, "condition_blind_id": "X"},
            {"aggregate_id": "A1", "candidate_cluster_id": "D1", "candidate_type": "DISTRACTOR", "represented": 1, "aggregate_output_tokens": 80, "condition_blind_id": "X"},
            {"aggregate_id": "A1", "candidate_cluster_id": "D2", "candidate_type": "DISTRACTOR", "represented": 0, "aggregate_output_tokens": 80, "condition_blind_id": "X"}
        ]
    })["aggregates"][0]
    assert close(a["AGG_COVERAGE"], .7) and close(a["REPRESENTATION_FPR"], .5), a

    print("T_GPT_V0_8_2_ADVERSARIAL_SELFTEST_OK")


if __name__ == "__main__":
    main()