"""The probe ladder for alert-triage. Invoked by `envs/probe.py`.

Every rung is a policy from `policies.py`, played over the main-slice seeds
(held-out seeds are never probed). The assertions are the reward's claims,
checked per weighted component rather than as one blended float. Any cheat a
reviewer named is a rung here with an assertion that it loses — a cheat that
is measured is a cheat that cannot come back unnoticed.
"""

from __future__ import annotations

from .generator import is_held_out, world_for_seed
from .policies import POLICIES, run_policy
from .reward import (
    CAUGHT_WEIGHT,
    EVIDENCE_WEIGHT,
    ROLES,
    WEIGHTS,
    escalate_break_even,
    metrics,
    score,
    total,
)

# Main-slice seeds only: fnv1a(seed) % 8 == 7 is never probed. 1024 seeds
# because the margins asserted in rung 8 are population claims, and a
# 200-seed sample can draw 11% hidden instead of 15% and turn one into noise.
SEEDS = [s for s in range(1024) if not is_held_out(s)]
COMPONENTS = ("caught", "hours", "evidence")
GOOD = ("fast", "targeted", "thorough")
CHEATS = ("close_blind", "close_cites_all", "cite_unseen", "escalate_blind", "fast_defensive", "rule_table")

# The one-slider counterfactual the page names: move `hours` here and the
# screen-only analyst wins. The other two weights keep their shipped ratio.
COUNTERFACTUAL_HOURS = 0.45
MARGIN = 0.02


def counterfactual_weights(hours_weight: float) -> dict[str, float]:
    rest = 1.0 - hours_weight
    share = CAUGHT_WEIGHT + EVIDENCE_WEIGHT
    return {"caught": rest * CAUGHT_WEIGHT / share, "hours": hours_weight, "evidence": rest * EVIDENCE_WEIGHT / share}


def run(seeds: list[int] | None = None) -> int:
    seeds = list(seeds or SEEDS)
    worlds = [world_for_seed(s) for s in seeds]
    tiers = {w["tier"] for w in worlds}
    print(f"Probing alert-triage over {len(seeds)} main-slice seeds "
          f"({sum(w['tier'] == 'benign' for w in worlds)} benign, "
          f"{sum(w['tier'] == 'visible' for w in worlds)} visible, "
          f"{sum(w['tier'] == 'hidden' for w in worlds)} hidden)\n")

    results: dict[str, dict[str, float]] = {}
    cf: dict[str, float] = {}
    per_seed: dict[str, list] = {}
    for name, policy in POLICIES.items():
        episodes = [run_policy(policy, w).episode() for w in worlds]
        scores = [score(e) for e in episodes]
        assert all(v is not None for s in scores for v in s.values()), f"{name}: a component was None on a probe seed"
        results[name] = {c: sum(s[c] for s in scores) / len(scores) for c in COMPONENTS}   # type: ignore[misc]
        results[name]["total"] = sum(total(e) for e in episodes) / len(episodes)             # type: ignore[misc]
        cf[name] = sum(total(e, counterfactual_weights(COUNTERFACTUAL_HOURS)) for e in episodes) / len(episodes)  # type: ignore[misc]
        per_seed[name] = episodes

    width = max(len(n) for n in POLICIES)
    header = f"{'policy':<{width}}  " + "  ".join(f"{c:>10}" for c in COMPONENTS) + f"  {'TOTAL':>8}  {'hours@' + str(COUNTERFACTUAL_HOURS):>10}"
    print(header)
    print("-" * len(header))
    for name in POLICIES:
        row = results[name]
        cells = "  ".join(f"{row[c]:>10.4f}" for c in COMPONENTS)
        print(f"{name:<{width}}  {cells}  {row['total']:>8.4f}  {cf[name]:>10.4f}")

    print(f"\nweights: {WEIGHTS}")
    print(f"roles:   {ROLES}")
    print(f"blind-escalation break-even: escalate whenever P(benign) < {escalate_break_even():.2f}")
    mean_hours = {n: round(sum(metrics(e)['hours_spent'] for e in per_seed[n]) / len(seeds), 2) for n in GOOD}
    print(f"hours spent (mean): {mean_hours}")
    hidden = [i for i, w in enumerate(worlds) if w["tier"] == "hidden"]
    print("caught on hidden seeds:", {n: round(sum(score(per_seed[n][i])['caught'] for i in hidden) / max(1, len(hidden)), 3) for n in GOOD})  # type: ignore[arg-type]

    failures: list[str] = []
    r = results

    # 0. The seed set must contain every tier, or the assertions below measure noise.
    if tiers != {"benign", "visible", "hidden"}:
        failures.append(f"probe seeds cover tiers {sorted(tiers)}; need all three")
    if any(w["reference_minutes"] is None for w in worlds):
        failures.append("a probe seed has no reference_minutes — no shipped policy got it right")

    # 1. Doing nothing, and doing it malformed, must measure exactly nothing.
    for name in ("inaction", "malformed"):
        if abs(r[name]["total"]) > 1e-12:
            failures.append(f"{name} scored {r[name]['total']:.6f}, must be 0.000")
        if any(e.rejected != 8 or e.disposition is not None for e in per_seed[name]):
            failures.append(f"{name} must be rejected on all eight turns and abort")

    # 2. The counterweight is attainable: the investigating references catch everything.
    for name in ("targeted", "thorough"):
        if abs(r[name]["caught"] - 1.0) > 1e-12:
            failures.append(f"{name} caught {r[name]['caught']:.4f}, must be exactly 1.0")

    # 3. And it bites the policy that defines the hours denominator on most seeds.
    if r["fast"]["caught"] >= 1.0 - 1e-12:
        failures.append("fast caught 1.0 — the hidden tier leaked into the screen")

    # 4. Paying the counterweight's price is a real price.
    if not (r["fast"]["hours"] > r["targeted"]["hours"] > r["thorough"]["hours"]):
        failures.append("hours must order fast > targeted > thorough")

    # 5. No good policy dominates another on every component.
    for a, b in (("fast", "targeted"), ("fast", "thorough"), ("targeted", "thorough")):
        if all(r[a][c] >= r[b][c] - 1e-12 for c in COMPONENTS):
            failures.append(f"{a} dominates {b} on every component — no real trade-off")
        if all(r[b][c] >= r[a][c] - 1e-12 for c in COMPONENTS):
            failures.append(f"{b} dominates {a} on every component — no real trade-off")

    # 6. Every named cheat loses to every good policy.
    for cheat in CHEATS:
        for good in GOOD:
            if r[cheat]["total"] >= r[good]["total"] - 1e-12:
                failures.append(f"{cheat} ({r[cheat]['total']:.4f}) >= {good} ({r[good]['total']:.4f})")

    # 7. The hour meter binds: a wasted month costs, and nothing else changes.
    if r["thorough_wasteful"]["total"] >= r["thorough"]["total"] - 1e-12:
        failures.append("thorough_wasteful >= thorough — the hour budget does not bind")
    for c in ("caught", "evidence"):
        if abs(r["thorough_wasteful"][c] - r["thorough"][c]) > 1e-12:
            failures.append(f"thorough_wasteful differs from thorough on {c}")

    # 8. The headline claim, measured: targeted wins at shipped weights by a
    #    stated margin, and the one-slider counterfactual hands it to fast.
    if r["targeted"]["total"] - r["fast"]["total"] < MARGIN:
        failures.append(f"targeted beats fast by only {r['targeted']['total'] - r['fast']['total']:.4f} at shipped weights (need {MARGIN})")
    if cf["fast"] - cf["targeted"] < MARGIN:
        failures.append(f"fast beats targeted by only {cf['fast'] - cf['targeted']:.4f} at hours={COUNTERFACTUAL_HOURS} (need {MARGIN})")

    # 9. Citing an unopened document is rejected, every time.
    if any(e.rejected < 1 for e in per_seed["cite_unseen"]):
        failures.append("cite_unseen was not rejected on some seed")

    # 10. F1's precision term works: cite-everything earns less than half of its caught share in evidence.
    if r["close_cites_all"]["evidence"] >= 0.5 * r["close_cites_all"]["caught"]:
        failures.append("close_cites_all earns too much evidence — precision is not biting")
    if r["close_cites_all"]["evidence"] >= r["fast"]["evidence"]:
        failures.append("close_cites_all evidence >= fast evidence")

    # 11. History buys evidence on the visible typologies, and nothing else.
    if r["thorough"]["evidence"] <= r["fast"]["evidence"]:
        failures.append("thorough evidence <= fast evidence — itemised history buys nothing")

    # 12. No weighted component may be flat across the ladder.
    for component in COMPONENTS:
        values = [r[n][component] for n in POLICIES]
        if max(values) - min(values) < 1e-9:
            failures.append(f"component '{component}' is flat across the ladder — it measures nothing")

    print()
    if failures:
        for f in failures:
            print(f"FAIL: {f}")
        return 1
    print("All 13 ladder assertions hold.")
    return 0
