"""v0.4 D6 — retrieval-attribution audit.

For each predicted node, classify into 4 categories based on (a) whether
the BFS layer that produced it had a high-cosine top retrieved edge
(>= good_threshold, default 0.5) and (b) whether the predicted node
survived to a gold match.

  both_work             matched + good retrieval
  llm_saved_retrieval   matched + poor retrieval (LLM rescued)
  llm_ignored_retrieval not matched + good retrieval (LLM ignored or fabricated)
  both_fail             not matched + poor retrieval
  no_retrieval_data     layer or cosine missing from trace

Aggregate ratios across 6 events answer "is over-prediction the LLM's
fault or retrieval's fault" — directly informs Phase 2 axis selection.

Inputs:
  - data/evaluation/diagnostics/<event>_bfs_full.json
  - data/evaluation/gold/<event>.json (matches, predicted_chain.cascade_events)

Outputs:
  - data/evaluation/diagnostics/v04/retrieval_attribution.{json,md}

BFS trace shape (verified live, v0.4):
  trace = {
    "event_id": ...,
    "trace": [
      {
        "layer": N,
        "produced_ids": ["E1", ...],
        "retrieved_edges": {frontier_id: [{"similarity": float, ...}]},
        ...
      },
      ...
    ]
  }
The 2025-0632-ROU outlier is filtered out (configured in
`evaluation.outlier_event_ids`).
"""

from __future__ import annotations

import argparse
import json
from collections import Counter
from pathlib import Path
from typing import Optional

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

_DEFAULT_GOOD_THRESHOLD = 0.5
_OUTLIER_EVENT_IDS = {"2025-0632-ROU"}
_CATEGORIES = (
    "both_work",
    "llm_saved_retrieval",
    "llm_ignored_retrieval",
    "both_fail",
    "no_retrieval_data",
)


def classify_node(
    node_id: str,
    matched: bool,
    layer_top_cosine: Optional[float],
    good_threshold: float = _DEFAULT_GOOD_THRESHOLD,
) -> str:
    if layer_top_cosine is None:
        return "no_retrieval_data"
    good = layer_top_cosine >= good_threshold
    if matched and good:
        return "both_work"
    if matched and not good:
        return "llm_saved_retrieval"
    if not matched and good:
        return "llm_ignored_retrieval"
    return "both_fail"


def summarize_event(nodes: list[dict], good_threshold: float = _DEFAULT_GOOD_THRESHOLD) -> dict:
    cnt: Counter[str] = Counter()
    for n in nodes:
        cls = classify_node(
            node_id=n.get("node_id", ""),
            matched=bool(n.get("matched", False)),
            layer_top_cosine=n.get("layer_top_cosine"),
            good_threshold=good_threshold,
        )
        cnt[cls] += 1
    return {c: int(cnt[c]) for c in _CATEGORIES}


def _extract_node_layer_cosine(trace: dict) -> dict[str, Optional[float]]:
    """Map predicted node_id -> top retrieved-edge similarity in the layer
    that produced it. Returns None for nodes whose layer has no retrieval.

    Verified BFS trace shape (v0.4):
      trace = {"event_id": ..., "trace": [{"layer": N, "produced_ids": [...],
               "retrieved_edges": {frontier_id: [{"similarity": float, ...}]},
               ...}]}
    """
    out: dict[str, Optional[float]] = {}
    layers = trace.get("trace") or []
    for layer in layers:
        edges_block = layer.get("retrieved_edges") or {}
        if isinstance(edges_block, dict):
            edges_flat = [e for edges in edges_block.values() for e in (edges or [])]
        else:
            edges_flat = edges_block or []
        cosines: list[float] = []
        for e in edges_flat:
            c = e.get("similarity") or e.get("cosine") or e.get("score")
            if isinstance(c, (int, float)):
                cosines.append(float(c))
        top = max(cosines) if cosines else None
        for nid in (layer.get("produced_ids") or []):
            out[nid] = top
    return out


def _load_matched_set(event_id: str) -> set[str]:
    cache = GOLD_CACHE_DIR / f"{event_id}.json"
    if not cache.exists():
        return set()
    d = json.loads(cache.read_text())
    return {m["p_id"] for m in d.get("matches", [])}


def _load_predicted_ids(event_id: str) -> list[str]:
    cache = GOLD_CACHE_DIR / f"{event_id}.json"
    if not cache.exists():
        return []
    d = json.loads(cache.read_text())
    return [n["id"] for n in d.get("predicted_chain", {}).get("cascade_events", [])]


def audit_event(event_id: str, good_threshold: float = _DEFAULT_GOOD_THRESHOLD) -> dict:
    trace_path = TRACE_DIR / f"{event_id}_bfs_full.json"
    if not trace_path.exists():
        return {"event_id": event_id, "error": "no_trace_file"}
    trace = json.loads(trace_path.read_text())
    node_to_cosine = _extract_node_layer_cosine(trace)
    matched = _load_matched_set(event_id)
    predicted_ids = _load_predicted_ids(event_id)
    nodes = [
        {
            "node_id": pid,
            "matched": pid in matched,
            "layer_top_cosine": node_to_cosine.get(pid),
        }
        for pid in predicted_ids
    ]
    return {
        "event_id": event_id,
        "summary": summarize_event(nodes, good_threshold),
        "nodes": nodes,
    }


def render_md(per_event: list[dict], total: dict, good_threshold: float) -> str:
    lines = [f"# v0.4 D6 — Retrieval Attribution Audit (good_threshold={good_threshold})", ""]
    lines += ["## Per-event category counts", "",
              "| event | both_work | llm_saved | llm_ignored | both_fail | no_data |",
              "|---|---:|---:|---:|---:|---:|"]
    for e in per_event:
        if "error" in e:
            lines.append(f"| {e['event_id']} | (error: {e['error']}) | | | | |")
            continue
        s = e["summary"]
        lines.append(
            f"| {e['event_id']} | {s['both_work']} | {s['llm_saved_retrieval']} | "
            f"{s['llm_ignored_retrieval']} | {s['both_fail']} | {s['no_retrieval_data']} |"
        )
    lines += ["", "## Aggregate (sum across events)", ""]
    for c in _CATEGORIES:
        lines.append(f"- {c}: {total[c]}")
    return "\n".join(lines) + "\n"


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

    event_ids = sorted(
        p.stem.replace("_bfs_full", "")
        for p in TRACE_DIR.glob("*_bfs_full.json")
        if p.stem.replace("_bfs_full", "") not in _OUTLIER_EVENT_IDS
    )
    per_event = [audit_event(eid, args.good_threshold) for eid in event_ids]

    total = {c: 0 for c in _CATEGORIES}
    for e in per_event:
        if "summary" in e:
            for c in _CATEGORIES:
                total[c] += e["summary"][c]

    payload = {
        "good_threshold": args.good_threshold,
        "per_event": per_event,
        "total": total,
    }
    (args.out_dir / "retrieval_attribution.json").write_text(json.dumps(payload, indent=2))
    (args.out_dir / "retrieval_attribution.md").write_text(
        render_md(per_event, total, args.good_threshold)
    )
    print(f"Wrote {args.out_dir / 'retrieval_attribution.json'}")
    print(f"Wrote {args.out_dir / 'retrieval_attribution.md'}")


if __name__ == "__main__":
    main()