"""v0.4 D4 — severity distribution audit.

Compares predicted vs gold severity distributions per event (and per
domain across events) on 4-rank ordinal scale {low=1, medium=2, high=3,
critical=4}. Reports two diagnostic numbers:

  delta_mean = mean(pred_rank) - mean(gold_rank)  # signed direction
  kl_divergence = D_KL(pred ‖ gold)               # Laplace-smoothed 1e-3

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

Outputs:
  - data/evaluation/diagnostics/v04/severity_dist.json
  - data/evaluation/diagnostics/v04/severity_dist.md

Read-only; no LLM.
"""

from __future__ import annotations

import argparse
import json
import math
from collections import defaultdict
from pathlib import Path
from typing import Optional

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"

_RANK = {"low": 1, "medium": 2, "high": 3, "critical": 4}
_SEV_KEYS = ("low", "medium", "high", "critical")
_KL_SMOOTH = 1e-3


def severity_to_rank(s: str) -> Optional[int]:
    if not isinstance(s, str):
        return None
    return _RANK.get(s.strip().lower())


def rank_mean(severities: list[str]) -> Optional[float]:
    ranks = [r for s in severities if (r := severity_to_rank(s)) is not None]
    if not ranks:
        return None
    return sum(ranks) / len(ranks)


def severity_distribution(severities: list[str]) -> dict[str, float]:
    """Normalized 4-bin distribution; missing/unknown severities are ignored."""
    counts = {k: 0 for k in _SEV_KEYS}
    for s in severities:
        k = (s or "").strip().lower()
        if k in counts:
            counts[k] += 1
    total = sum(counts.values())
    if total == 0:
        return {k: 0.0 for k in _SEV_KEYS}
    return {k: counts[k] / total for k in _SEV_KEYS}


def kl_divergence(pred: dict[str, float], gold: dict[str, float]) -> float:
    """KL(pred ‖ gold) with Laplace smoothing (additive 1e-3 to each bin then renorm)."""
    p = {k: pred.get(k, 0.0) + _KL_SMOOTH for k in _SEV_KEYS}
    g = {k: gold.get(k, 0.0) + _KL_SMOOTH for k in _SEV_KEYS}
    p_sum = sum(p.values()); g_sum = sum(g.values())
    p = {k: v / p_sum for k, v in p.items()}
    g = {k: v / g_sum for k, v in g.items()}
    return sum(p[k] * math.log(p[k] / g[k]) for k in _SEV_KEYS)


def _load_predicted_severities(event_id: str) -> list[str]:
    """Pull predicted severities from gold cache (predicted_chain)."""
    cache = GOLD_CACHE_DIR / f"{event_id}.json"
    if not cache.exists():
        return []
    d = json.loads(cache.read_text())
    nodes = d.get("predicted_chain", {}).get("cascade_events", [])
    return [n.get("severity", "") for n in nodes]


def _load_gold_severities(event_id: str) -> list[str]:
    """Pull gold severities from cascade_chains (per-event gold JSON)."""
    chain = GOLD_CHAINS_DIR / f"{event_id}.json"
    if not chain.exists():
        return []
    d = json.loads(chain.read_text())
    nodes = d.get("cascade_events", [])
    return [n.get("severity", "") for n in nodes]


def _per_domain_severities(event_id: str) -> tuple[dict[str, list[str]], dict[str, list[str]]]:
    """Returns (pred_by_domain, gold_by_domain) for one event."""
    pred = defaultdict(list)
    gold = defaultdict(list)
    cache = GOLD_CACHE_DIR / f"{event_id}.json"
    if cache.exists():
        d = json.loads(cache.read_text())
        for n in d.get("predicted_chain", {}).get("cascade_events", []):
            pred[n.get("domain", "unknown")].append(n.get("severity", ""))
    chain = GOLD_CHAINS_DIR / f"{event_id}.json"
    if chain.exists():
        d = json.loads(chain.read_text())
        for n in d.get("cascade_events", []):
            gold[n.get("domain", "unknown")].append(n.get("severity", ""))
    return dict(pred), dict(gold)


def audit_event(event_id: str) -> dict:
    pred_sev = _load_predicted_severities(event_id)
    gold_sev = _load_gold_severities(event_id)
    pred_dist = severity_distribution(pred_sev)
    gold_dist = severity_distribution(gold_sev)
    return {
        "event_id": event_id,
        "predicted_count": len(pred_sev),
        "gold_count": len(gold_sev),
        "predicted_dist": pred_dist,
        "gold_dist": gold_dist,
        "predicted_mean_rank": rank_mean(pred_sev),
        "gold_mean_rank": rank_mean(gold_sev),
        "delta_mean": (
            rank_mean(pred_sev) - rank_mean(gold_sev)
            if (rank_mean(pred_sev) is not None and rank_mean(gold_sev) is not None)
            else None
        ),
        "kl_divergence": kl_divergence(pred_dist, gold_dist) if pred_sev and gold_sev else None,
    }


def render_md(events: list[dict], domain_summary: dict) -> str:
    lines = ["# v0.4 D4 — Severity Distribution Audit", "",
             "## Per-event severity drift", "",
             "| event_id | pred_n | gold_n | pred_mean_rank | gold_mean_rank | Δ_mean | KL |",
             "|---|---:|---:|---:|---:|---:|---:|"]
    for e in events:
        lines.append(
            f"| {e['event_id']} | {e['predicted_count']} | {e['gold_count']} | "
            f"{e['predicted_mean_rank']:.3f} | {e['gold_mean_rank']:.3f} | "
            f"{e['delta_mean']:+.3f} | {e['kl_divergence']:.4f} |"
            if e['delta_mean'] is not None else
            f"| {e['event_id']} | {e['predicted_count']} | {e['gold_count']} | – | – | – | – |"
        )
    lines += ["", "## Per-domain severity drift (aggregated across events)", "",
              "| domain | pred_n | gold_n | pred_mean_rank | gold_mean_rank | Δ_mean | KL |",
              "|---|---:|---:|---:|---:|---:|---:|"]
    for dom, row in sorted(domain_summary.items()):
        lines.append(
            f"| {dom} | {row['predicted_count']} | {row['gold_count']} | "
            f"{row['predicted_mean_rank']:.3f} | {row['gold_mean_rank']:.3f} | "
            f"{row['delta_mean']:+.3f} | {row['kl_divergence']:.4f} |"
            if row['delta_mean'] is not None else
            f"| {dom} | {row['predicted_count']} | {row['gold_count']} | – | – | – | – |"
        )
    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"))
    events = [audit_event(eid) for eid in event_ids]

    # Per-domain aggregate
    pred_by_dom: dict[str, list[str]] = defaultdict(list)
    gold_by_dom: dict[str, list[str]] = defaultdict(list)
    for eid in event_ids:
        p, g = _per_domain_severities(eid)
        for d, sevs in p.items():
            pred_by_dom[d].extend(sevs)
        for d, sevs in g.items():
            gold_by_dom[d].extend(sevs)

    domains = set(pred_by_dom) | set(gold_by_dom)
    domain_summary: dict[str, dict] = {}
    for d in domains:
        ps, gs = pred_by_dom.get(d, []), gold_by_dom.get(d, [])
        pd_, gd_ = severity_distribution(ps), severity_distribution(gs)
        domain_summary[d] = {
            "predicted_count": len(ps),
            "gold_count": len(gs),
            "predicted_mean_rank": rank_mean(ps),
            "gold_mean_rank": rank_mean(gs),
            "delta_mean": (
                rank_mean(ps) - rank_mean(gs)
                if (rank_mean(ps) is not None and rank_mean(gs) is not None) else None
            ),
            "kl_divergence": kl_divergence(pd_, gd_) if ps and gs else None,
        }

    payload = {"events": events, "domains": domain_summary}
    (args.out_dir / "severity_dist.json").write_text(json.dumps(payload, indent=2))
    (args.out_dir / "severity_dist.md").write_text(render_md(events, domain_summary))
    print(f"Wrote {args.out_dir / 'severity_dist.json'}")
    print(f"Wrote {args.out_dir / 'severity_dist.md'}")


if __name__ == "__main__":
    main()