"""Evaluation report on val and test (pure Python). Run as ``python -m training.evaluate``.

Spec: docs/classifier-plan.md §5.1-5.3. All "accuracy" is agreement with the teacher. Per split:
- **variants**: raw/gated x uncalibrated/calibrated; per field accuracy, macro-F1, kappa, log-loss, ECE
  (gold = ``targets[field]`` unless the field is in ``mask``); the headline macro-F1 leaves out the taxonomy's
  ``macro_f1_exclude`` options (nsb-v2: emotion_fine ``mixed``; ``--macro-f1-exclude`` overrides), which stay in
  ``per_class`` and in ``macro_f1_all``;
- **key classes** (gated, calibrated, eval rows only): precision, recall, class ECE, predicted vs teacher share
  weighted by 1/inclusion_prob;
- **agreement** vs teacher passes r1 and r2 separately, next to r1-vs-r2 (the noise ceiling); the student's gated
  calibrated argmax is compared with the full pass labels (gate-decided absences included);
- **aggregate parity** per stock group x KST day (§7.2 primary index): student shares from calibrated gated
  probabilities, teacher shares from one-hot targets, on the same eval posts. A cell counts when the teacher side has
  >= MIN_AUTHORS contributing authors and the student side at least one. Unweighted (eval posts are random/audit);
- **breakdowns** per code, stock group, session tag, flag (``possible_injection`` separately);
- **intervals**: 95% author-cluster bootstrap of kappa and accuracy for the product fields.
Across splits: crossfit calibration on val (out-of-sample ECE) and drift (val minus test accuracy, points).
"""
from __future__ import annotations

import argparse
import json
from collections import defaultdict
from datetime import datetime
from pathlib import Path
from typing import Iterable, Mapping

from . import metrics as mt
from .calibrate import crossfit, pair_up, probs
from .kst import day_bucket, session_tag
from .taxonomy import AD_OPTIONS_ANY, COMMERCIAL_AD, PRODUCT_FIELDS, RELEVANT, Taxonomy, gate, gate_fill

VERSION = "nsb-eval-v1"
VARIANTS = ("raw_uncalibrated", "raw_calibrated", "gated_uncalibrated", "gated_calibrated")


def key_class_list(ad_options: Iterable[str] = COMMERCIAL_AD) -> list[tuple[str, str]]:
    """The §5.1 key classes; the commercial ones are the taxonomy's ad options."""
    return ([("relevance", RELEVANT), ("direction", "bullish"), ("direction", "bearish")]
            + [("commercial", o) for o in ad_options]
            + [("emotion_valence", "positive"), ("emotion_valence", "negative")])


KEY_CLASSES = key_class_list()  # nsb-v1
FLAGS = ("image_dependent", "truncated_input")
MIN_AUTHORS = 10
MIN_WEIGHT = 0.5
SPLITS = ("val", "test")


# ---------------------------------------------------------------------- helpers
def gold(row: Mapping, field: str) -> str | None:
    return None if field in set(row.get("mask") or ()) else row["targets"].get(field)


def distributions(pred: Mapping, tax: Taxonomy, temps: Mapping[str, float]) -> dict[str, dict]:
    raw_u = probs(pred["logits"], tax, None)
    raw_c = probs(pred["logits"], tax, temps)
    return {"raw_uncalibrated": raw_u, "raw_calibrated": raw_c, "gated_uncalibrated": gate(raw_u),
            "gated_calibrated": gate(raw_c)}


def scores(dist: Mapping[str, Mapping[str, float]]) -> dict[str, float]:
    """§7.2 per-post quantities r, t, b, s, q from one set of distributions (or one-hot labels)."""
    return {"r": dist["relevance"][RELEVANT], "t": dist["target"]["this_stock"], "b": dist["direction"]["bullish"],
            "s": dist["direction"]["bearish"],
            "q": 1.0 - sum(dist["commercial"].get(o, 0.0) for o in AD_OPTIONS_ANY)}


def one_hot(labels: Mapping[str, str | None], tax: Taxonomy) -> dict[str, dict[str, float]] | None:
    """Teacher labels as distributions; None when a field the shares need is unresolved."""
    filled = gate_fill(labels)
    if any(filled.get(f) is None for f in ("relevance", "target", "direction", "commercial")):
        return None
    return {f: {o: (1.0 if filled.get(f) == o else 0.0) for o in tax.fields[f]} for f in tax.names}


def author_balanced(posts: Iterable[tuple[str, Mapping[str, float]]]) -> dict:
    """Author-balanced bull/bear shares: per author B_a = sum(w b)/sum(w), w = r t q; mean over authors with
    sum(w) >= MIN_WEIGHT (one author, one vote)."""
    acc: dict[str, list[float]] = defaultdict(lambda: [0.0, 0.0, 0.0])
    for author, s in posts:
        w = s["r"] * s["t"] * s["q"]
        a = acc[author]
        a[0] += w
        a[1] += w * s["b"]
        a[2] += w * s["s"]
    contrib = [(b / w, s / w) for w, b, s in acc.values() if w >= MIN_WEIGHT]
    if not contrib:
        return {"n_authors": len(acc), "n_contrib": 0, "bull": None, "bear": None, "net": None}
    bull = sum(x for x, _ in contrib) / len(contrib)
    bear = sum(y for _, y in contrib) / len(contrib)
    return {"n_authors": len(acc), "n_contrib": len(contrib), "bull": bull, "bear": bear, "net": bull - bear}


def _written(row: Mapping) -> datetime:
    return datetime.fromisoformat(row["written_at_kst"])


# ---------------------------------------------------------------------- sections
def variant_summaries(pairs, dists, tax: Taxonomy) -> dict:
    out = {}
    for v in VARIANTS:
        out[v] = {f: mt.field_summary([gold(r, f) for _, r in pairs], [d[v][f] for d in dists], tax.fields[f],
                                      tax.macro_f1_exclude.get(f, ()))
                  for f in tax.names}
    return out


def key_classes(pairs, dists, ad_options: Iterable[str] = COMMERCIAL_AD) -> dict:
    keep = [i for i, (_, r) in enumerate(pairs) if r.get("eval")]
    weights = mt.inverse_weights(pairs[i][1].get("inclusion_prob") for i in keep)
    out = {}
    for field, cls in key_class_list(ad_options):
        out[f"{field}={cls}"] = mt.key_class([gold(pairs[i][1], field) for i in keep],
                                             [dists[i]["gated_calibrated"][field] for i in keep], cls, weights)
    return out


def agreement(pairs, dists, tax: Taxonomy) -> dict:
    """Student vs r1, vs r2, and r1 vs r2 (ceiling), per field; rows with both passes only."""
    keep = [i for i, (_, r) in enumerate(pairs) if (r.get("passes") or {}).get("r1") and r["passes"].get("r2")]
    out: dict = {"n": len(keep), "r1": {}, "r2": {}, "ceiling": {}}
    for f in tax.names:
        student = [mt.top(dists[i]["gated_calibrated"][f]) for i in keep]
        r1 = [gate_fill(pairs[i][1]["passes"]["r1"]).get(f) for i in keep]
        r2 = [gate_fill(pairs[i][1]["passes"]["r2"]).get(f) for i in keep]
        out["r1"][f] = {"kappa": mt.kappa(r1, student), "accuracy": mt.accuracy(r1, student)}
        out["r2"][f] = {"kappa": mt.kappa(r2, student), "accuracy": mt.accuracy(r2, student)}
        out["ceiling"][f] = {"kappa": mt.kappa(r1, r2), "accuracy": mt.accuracy(r1, r2)}
    return out


def aggregate_parity(pairs, dists, tax: Taxonomy, groups: Mapping[str, str]) -> dict:
    cells: dict[tuple, dict[str, list]] = defaultdict(lambda: {"student": [], "teacher": []})
    for (_, row), d in zip(pairs, dists):
        if not row.get("eval"):
            continue
        teacher = one_hot(row["targets"], tax)
        if teacher is None:
            continue
        key = (groups.get(row["code"], row["code"]), day_bucket(_written(row)).isoformat())
        cells[key]["student"].append((row["author_key"], scores(d["gated_calibrated"])))
        cells[key]["teacher"].append((row["author_key"], scores(teacher)))
    rows = []
    for (group, day), c in sorted(cells.items()):
        st, te = author_balanced(c["student"]), author_balanced(c["teacher"])
        if te["n_contrib"] < MIN_AUTHORS or not st["n_contrib"]:
            continue
        rows.append({"group": group, "day": day, "n_authors": te["n_authors"], "teacher_contrib": te["n_contrib"],
                     "student_contrib": st["n_contrib"], "student": {k: st[k] for k in ("bull", "bear", "net")},
                     "teacher": {k: te[k] for k in ("bull", "bear", "net")}})

    def series(side, k):
        return [c[side][k] for c in rows]

    bull_mae = mt.mae(series("student", "bull"), series("teacher", "bull"))
    bear_mae = mt.mae(series("student", "bear"), series("teacher", "bear"))
    return {"n_cells": len(rows), "cells_considered": len(cells), "min_authors": MIN_AUTHORS,
            "bull_mae_points": None if bull_mae is None else 100 * bull_mae,
            "bear_mae_points": None if bear_mae is None else 100 * bear_mae,
            "net_r": mt.pearson(series("student", "net"), series("teacher", "net")), "cells": rows}


def _slice(pairs, dists, idx: list[int], tax: Taxonomy) -> dict:
    out: dict = {"n": len(idx), "n_eval": sum(1 for i in idx if pairs[i][1].get("eval")), "accuracy": {}, "kappa": {}}
    for f in tax.names:
        g = [gold(pairs[i][1], f) for i in idx]
        p = [mt.top(dists[i]["gated_calibrated"][f]) for i in idx]
        out["accuracy"][f] = mt.accuracy(g, p)
        if f in PRODUCT_FIELDS:
            out["kappa"][f] = mt.kappa(g, p)
    return out


def breakdowns(pairs, dists, tax: Taxonomy, groups: Mapping[str, str], holidays) -> dict:
    keyed: dict[str, dict[str, list[int]]] = {"code": defaultdict(list), "group": defaultdict(list),
                                              "session": defaultdict(list), "flags": defaultdict(list)}
    injection: list[int] = []
    for i, (_, row) in enumerate(pairs):
        keyed["code"][row["code"]].append(i)
        keyed["group"][groups.get(row["code"], row["code"])].append(i)
        keyed["session"][session_tag(_written(row), holidays)].append(i)
        flags = set(row.get("flags") or ())
        for flag in FLAGS:
            if flag in flags:
                keyed["flags"][flag].append(i)
        if "possible_injection" in flags:
            injection.append(i)
    out = {k: {name: _slice(pairs, dists, idx, tax) for name, idx in sorted(v.items())} for k, v in keyed.items()}
    out["possible_injection"] = _slice(pairs, dists, injection, tax)
    authors = {r["author_key"] for _, r in pairs}
    flagged = {pairs[i][1]["author_key"] for i in injection}
    out["injection_author_share"] = len(flagged) / len(authors) if authors else None
    return out


def intervals(pairs, dists, reps: int, seed: str) -> dict:
    by_id = {r["id"]: (r, d) for (_, r), d in zip(pairs, dists)}
    authors = {i: by_id[i][0]["author_key"] for i in by_id}
    out = {}
    for f in PRODUCT_FIELDS:
        def stat(ids, kind, f=f):
            g = [gold(by_id[i][0], f) for i in ids]
            p = [mt.top(by_id[i][1]["gated_calibrated"][f]) for i in ids]
            return mt.kappa(g, p) if kind == "kappa" else mt.accuracy(g, p)
        out[f] = {k: mt.bootstrap(list(by_id), authors, lambda ids, k=k: stat(ids, k), reps=reps,
                                  seed=f"{seed}|{f}|{k}")["interval"] for k in ("kappa", "accuracy")}
    return out


def split_report(rows, preds, tax, temps, groups, holidays, reps) -> dict:
    pairs = pair_up(preds, rows)
    dists = [distributions(p, tax, temps) for p, _ in pairs]
    split = rows[0].get("split", "?") if rows else "?"
    return {"n": len(pairs), "n_eval": sum(1 for _, r in pairs if r.get("eval")),
            "variants": variant_summaries(pairs, dists, tax),
            "key_classes": key_classes(pairs, dists, tax.commercial_ad),
            "agreement": agreement(pairs, dists, tax), "aggregate_parity": aggregate_parity(pairs, dists, tax, groups),
            "breakdowns": breakdowns(pairs, dists, tax, groups, holidays),
            "intervals": intervals(pairs, dists, reps, f"{mt.SEED}|{split}")}


def crossfit_ece(rows, preds, tax: Taxonomy) -> dict:
    """Out-of-sample calibration check: temperatures fitted on one author half of val, ECE on the other."""
    pairs = pair_up(preds, rows)
    cf = crossfit(pairs, tax)
    out = {"temperatures_by_half": cf["temperatures_by_half"], "sizes": cf["sizes"], "ece": {}, "ece_gated": {}}
    for f in tax.names:
        g = [gold(r, f) for _, r in pairs]
        out["ece"][f] = mt.ece(g, [cf["calibrated"][p["id"]][f] for p, _ in pairs])
        out["ece_gated"][f] = mt.ece(g, [gate(cf["calibrated"][p["id"]])[f] for p, _ in pairs])
    return out


def report(rows: list[dict], preds: list[dict], tax: Taxonomy, temps: Mapping[str, float], *,
           groups: Mapping[str, str] | None = None, holidays: Iterable = (), reps: int = mt.REPS) -> dict:
    """The full §5 report. ``rows`` and ``preds`` may mix splits; each split present is reported."""
    holidays = set(holidays)
    groups = dict(groups or {})
    by_split: dict[str, list[dict]] = defaultdict(list)
    for r in rows:
        by_split[r["split"]].append(r)
    preds_by_id = {p["id"]: p for p in preds}
    out: dict = {"version": VERSION, "metrics_version": mt.VERSION, "taxonomy_version": tax.version,
                 "macro_f1_exclude": tax.macro_f1_exclude, "temperatures": dict(temps), "splits": {}}
    for split in SPLITS:
        if not by_split.get(split):
            continue
        split_preds = [preds_by_id[r["id"]] for r in by_split[split] if r["id"] in preds_by_id]
        out["splits"][split] = split_report(by_split[split], split_preds, tax, temps, groups, holidays, reps)
    if "val" in out["splits"]:
        out["calibration"] = {"crossfit": crossfit_ece(by_split["val"],
                                                       [preds_by_id[r["id"]] for r in by_split["val"]], tax)}
    if "val" in out["splits"] and "test" in out["splits"]:
        va = out["splits"]["val"]["variants"]["gated_calibrated"]
        te = out["splits"]["test"]["variants"]["gated_calibrated"]
        out["drift"] = {f: (None if va[f]["accuracy"] is None or te[f]["accuracy"] is None
                            else 100 * (va[f]["accuracy"] - te[f]["accuracy"])) for f in tax.names}
    return out


def parse_exclude(items: Iterable[str]) -> dict[str, list[str]]:
    """``["emotion_fine=mixed", ...]`` -> ``{"emotion_fine": ["mixed"]}``; ``["none"]`` -> ``{}``."""
    out: dict[str, list[str]] = {}
    for item in items:
        if item == "none":
            continue
        field, sep, option = item.partition("=")
        if not sep or not field or not option:
            raise SystemExit(f"--macro-f1-exclude wants FIELD=OPTION, got {item!r}")
        out.setdefault(field, []).append(option)
    return out


def main(argv=None) -> int:
    from .data import load_jsonl, load_split
    from .kst import load_groups, load_holidays

    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("--data", required=True, help="Bundle directory (val.jsonl, test.jsonl, taxonomy.json)")
    p.add_argument("--run", required=True, help="Run directory (predictions/{val,test}.jsonl, temperatures.json)")
    p.add_argument("--groups", help="config/stocks.yaml (code -> stock group)")
    p.add_argument("--holidays", help="config/krx_holidays.yaml")
    p.add_argument("--reps", type=int, default=mt.REPS)
    p.add_argument("--macro-f1-exclude", action="append", metavar="FIELD=OPTION",
                   help="Leave an option out of the headline macro-F1 (repeatable; 'none' clears; default: the "
                        "taxonomy's macro_f1_exclude)")
    p.add_argument("--out", required=True)
    args = p.parse_args(argv)
    if Path(args.out).exists():
        raise SystemExit("refusing to overwrite " + args.out)
    run = Path(args.run)
    tax = Taxonomy.load(args.data)
    if args.macro_f1_exclude:
        tax = tax.replace(macro_f1_exclude=parse_exclude(args.macro_f1_exclude))
    temps = json.loads((run / "temperatures.json").read_text())
    rows, preds = [], []
    for split in SPLITS:
        path = run / "predictions" / f"{split}.jsonl"
        if path.exists():
            rows += load_split(args.data, split)
            preds += load_jsonl(path)
    groups = load_groups(args.groups, {r["code"] for r in rows})
    rep = report(rows, preds, tax, temps, groups=groups, holidays=load_holidays(args.holidays), reps=args.reps)
    Path(args.out).write_text(json.dumps(rep, indent=1, ensure_ascii=False) + "\n")
    print(json.dumps({s: {"n": v["n"], "n_eval": v["n_eval"]} for s, v in rep["splits"].items()}))
    return 0


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