"""Choose the arm of the §4.4 screen (pure Python). Run as ``python -m training.screen RUN... --data BUNDLE``.

Arms: B-last, I-last, B-mean, I-mean (checkpoint base/it x pooling last/mean), 1 epoch each. Each run's best val
predictions (``RUN/predictions/val.jsonl``) are scored on val ``eval`` posts:
- primary: mean val log-loss over relevance, direction, commercial (lower is better);
- secondary: direction macro-F1 (gated argmax, uncalibrated).
Rule (classifier-plan §4.4, made exact here): B-last is the default. Another arm beats it if its primary is more
than 2% lower than B-last's, or if its primary is within 2% of B-last's and its secondary is at least 1 point
higher. Among arms that beat B-last the lowest primary wins (then the higher secondary, then arm order).
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

from .calibrate import nll, outputs, pair_up
from .data import load_jsonl, load_split
from .taxonomy import Taxonomy

ARMS = ("B-last", "I-last", "B-mean", "I-mean")
DEFAULT = "B-last"
PRIMARY_FIELDS = ("relevance", "direction", "commercial")
REL_TOL = 0.02
F1_MARGIN = 0.01


def macro_f1(gold: list[str], pred: list[str], classes: list[str]) -> float | None:
    scores = []
    for c in classes:
        tp = sum(1 for g, p in zip(gold, pred) if g == c and p == c)
        support = sum(1 for g in gold if g == c)
        predicted = sum(1 for p in pred if p == c)
        if support + predicted:
            scores.append(2 * tp / (support + predicted))
    return sum(scores) / len(scores) if scores else None


def score(preds: list[dict], rows: list[dict], tax: Taxonomy) -> dict:
    pairs = [(p, r) for p, r in pair_up(preds, rows) if r.get("eval")]
    losses = [nll(pairs, tax, f, 1.0) for f in PRIMARY_FIELDS]
    losses = [v for v in losses if v is not None]
    gold, pred = [], []
    for p, r in pairs:
        if "direction" in r.get("mask", ()) or r["targets"].get("direction") is None:
            continue
        gold.append(r["targets"]["direction"])
        pred.append(outputs(p["logits"], tax)["labels"]["direction"])
    return {"primary": sum(losses) / len(losses) if losses else None,
            "secondary": macro_f1(gold, pred, tax.fields["direction"]), "eval_rows": len(pairs)}


def choose(scores: dict[str, dict]) -> dict:
    """Apply the rule to {arm: {"primary", "secondary"}}; returns the winner and the reasons."""
    if DEFAULT not in scores or scores[DEFAULT]["primary"] is None:
        raise ValueError(f"the {DEFAULT} arm is required (it is the default and the reference)")
    ref = scores[DEFAULT]
    beats = []
    for arm, s in scores.items():
        if arm == DEFAULT or s["primary"] is None:
            continue
        rel = s["primary"] / ref["primary"] - 1.0
        better_primary = rel < -REL_TOL
        better_secondary = abs(rel) <= REL_TOL and s["secondary"] is not None and ref["secondary"] is not None \
            and s["secondary"] >= ref["secondary"] + F1_MARGIN
        if better_primary or better_secondary:
            beats.append(arm)
    order = {a: i for i, a in enumerate(ARMS)}
    if beats:
        beats.sort(key=lambda a: (scores[a]["primary"], -(scores[a]["secondary"] or 0.0), order.get(a, 99)))
        winner = beats[0]
    else:
        winner = DEFAULT
    return {"winner": winner, "beats_default": beats, "scores": scores}


def main(argv=None) -> int:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("runs", nargs="+", help="Run directories (their config.json names the arm)")
    p.add_argument("--data", required=True)
    p.add_argument("--out")
    args = p.parse_args(argv)
    tax = Taxonomy.load(args.data)
    rows = load_split(args.data, "val")
    scores = {}
    for run in map(Path, args.runs):
        cfg = json.loads((run / "config.json").read_text())
        arm = cfg.get("arm") or run.name
        preds = load_jsonl(run / "predictions" / "val.jsonl")
        by_id = {r["id"] for r in preds}
        scores[arm] = score(preds, [r for r in rows if r["id"] in by_id], tax)
        scores[arm]["run"] = str(run)
    result = choose(scores)
    text = json.dumps(result, indent=1)
    if args.out:
        Path(args.out).write_text(text + "\n")
    print(text)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())