"""v0.4 D5 — per-domain attribution audit.

For each cascade domain across the 6 evaluable test events, sum:
  - predicted_total: predicted nodes with this domain
  - gold_total: gold nodes with this domain
  - matched_total: matches where the predicted node has this domain
  → P_domain = matched / predicted, R_domain = matched / gold
  → over_prediction_ratio = predicted / max(gold, 1)

Domains with over_prediction_ratio >> 1 are candidates for domain-budget
cap intervention in Phase 2.

Inputs:
  - data/evaluation/gold/<event>.json (predicted_chain + matches)
  - data/processed/cascade_chains/<event>.json (gold chain)

Outputs:
  - data/evaluation/diagnostics/v04/domain_attribution.json
  - data/evaluation/diagnostics/v04/domain_attribution.md
"""

from __future__ import annotations

import argparse
import json
from collections import defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
GOLD_CACHE_DIR = ROOT / "data/evaluation/gold"
GOLD_CHAINS_DIR = ROOT / "data/processed/cascade_chains"
OUT_DIR = ROOT / "data/evaluation/diagnostics/v04"


def aggregate_domain_counts(
    events: list[tuple[list[dict], list[dict], list[dict]]],
) -> dict[str, dict[str, int]]:
    """Aggregate (pred, gold, matches) tuples into per-domain counts.

    Each tuple = (pred_nodes, gold_nodes, matches) for ONE event.
    pred_nodes / gold_nodes: list of dicts with 'id' + 'domain'.
    matches: list of dicts with 'p_id' + 'g_id'.

    Match attribution rule: a match contributes to the PREDICTED node's
    domain (not the gold's). This means a cross-domain match (rare —
    metrics.match_nodes only allows same-domain) would still attribute
    to the predicted side.
    """
    counts: dict[str, dict[str, int]] = defaultdict(
        lambda: {"predicted": 0, "gold": 0, "matched": 0}
    )
    for pred_nodes, gold_nodes, matches in events:
        pred_dom = {n["id"]: n.get("domain", "unknown") for n in pred_nodes}
        for n in pred_nodes:
            counts[n.get("domain", "unknown")]["predicted"] += 1
        for n in gold_nodes:
            counts[n.get("domain", "unknown")]["gold"] += 1
        for m in matches:
            d = pred_dom.get(m["p_id"], "unknown")
            counts[d]["matched"] += 1
    return dict(counts)


def compute_domain_metrics(counts: dict[str, dict[str, int]]) -> dict[str, dict[str, float]]:
    out: dict[str, dict[str, float]] = {}
    for dom, c in counts.items():
        p = c["predicted"]; g = c["gold"]; m = c["matched"]
        out[dom] = {
            "predicted": p,
            "gold": g,
            "matched": m,
            "precision": (m / p) if p else 0.0,
            "recall": (m / g) if g else 0.0,
            "over_prediction_ratio": p / max(g, 1),
        }
    return out


def _load_event_triple(event_id: str) -> tuple[list[dict], list[dict], list[dict]]:
    pred_nodes: list[dict] = []
    matches: list[dict] = []
    cache = GOLD_CACHE_DIR / f"{event_id}.json"
    if cache.exists():
        d = json.loads(cache.read_text())
        pred_nodes = d.get("predicted_chain", {}).get("cascade_events", [])
        matches = d.get("matches", [])
    chain = GOLD_CHAINS_DIR / f"{event_id}.json"
    gold_nodes: list[dict] = []
    if chain.exists():
        d = json.loads(chain.read_text())
        gold_nodes = d.get("cascade_events", [])
    return pred_nodes, gold_nodes, matches


def render_md(metrics: dict[str, dict[str, float]]) -> str:
    rows = sorted(
        metrics.items(), key=lambda kv: -kv[1]["over_prediction_ratio"]
    )
    lines = ["# v0.4 D5 — Per-domain Attribution Audit", "",
             "Sorted by over_prediction_ratio (desc). Ratios >> 1 are over-predicted; << 1 are under-predicted.",
             "",
             "| domain | predicted | gold | matched | P | R | over_pred_ratio |",
             "|---|---:|---:|---:|---:|---:|---:|"]
    for d, m in rows:
        lines.append(
            f"| {d} | {m['predicted']} | {m['gold']} | {m['matched']} | "
            f"{m['precision']:.3f} | {m['recall']:.3f} | {m['over_prediction_ratio']:.2f} |"
        )
    return "\n".join(lines) + "\n"


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--out-dir", type=Path, default=OUT_DIR)
    args = ap.parse_args()
    args.out_dir.mkdir(parents=True, exist_ok=True)

    event_ids = sorted(
        p.stem for p in GOLD_CACHE_DIR.glob("*.json")
        if not p.name.endswith(".diag.json")
    )
    triples = [_load_event_triple(eid) for eid in event_ids]
    counts = aggregate_domain_counts(triples)
    metrics = compute_domain_metrics(counts)

    payload = {"per_domain": metrics, "events_included": event_ids}
    (args.out_dir / "domain_attribution.json").write_text(json.dumps(payload, indent=2))
    (args.out_dir / "domain_attribution.md").write_text(render_md(metrics))
    print(f"Wrote {args.out_dir / 'domain_attribution.json'}")
    print(f"Wrote {args.out_dir / 'domain_attribution.md'}")


if __name__ == "__main__":
    main()