"""v0.5 issue B+C — D6 hallucination re-audit (issue #70 Task 10).

Adapted from scripts/v04_retrieval_audit.py. Measures whether the v0.5
predictor (with structural filter from Task 2/3, templating disabled per
Task 9 B-only decision) hallucinates specific numbers / place names that
aren't grounded in either:

  (a) the new event's description / location / country (legitimate context),
  (b) any retrieved evidence's parent_text or child_description as recorded
      in the per-event BFS trace (legitimate historical anchor).

A predicted node is flagged "hallucination" if its description contains at
least one specific number or known place name that fails BOTH (a) and (b)
groundedness checks.

Note: the BFS trace is not seed-specific (temperature=0 → deterministic
prediction, all three seeds produce identical node lists and trace paths).
The per-seed GoldEvaluation caches are processed individually to confirm
this, but hallucination counts should be identical across seeds.

Reads:
  - data/evaluation/gold/{event_id}_seed{seed}.json  (GoldEvaluation caches)
  - data/evaluation/diagnostics/{event_id}_bfs_full.json  (BFS trace)
  - data/splits/test_events.json  (FloodEvent descriptions for check (a))
  - knowledge/place_names.txt  (vocabulary for LOC detection)

Writes:
  - /tmp/v05_d6_audit.csv  (per-node detail)
  - Stdout summary (pipe to tee /tmp/v05_d6_audit.log)

Usage:
    PYTHONPATH=. python scripts/v05_d6_audit.py | tee /tmp/v05_d6_audit.log
"""

from __future__ import annotations

import csv
import json
import re
from collections import defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
EVAL_DIR = ROOT / "data/evaluation/gold"
TRACE_DIR = ROOT / "data/evaluation/diagnostics"
TEST_EVENTS_PATH = ROOT / "data/splits/test_events.json"
OUT_CSV = Path("/tmp/v05_d6_audit.csv")

_OUTLIER_EVENT_IDS = {"2025-0632-ROU"}

# Same number regex as evidence_templating._NUMBER_RE — kept inline so the
# audit does NOT accidentally import a changed definition from the module.
_NUMBER_RE = re.compile(r"\b\d[\d,.]*[KkMm]?(?=[^0-9KkMm.,]|$)")


# ---------------------------------------------------------------------------
# Place-name loading (mirrors evidence_templating.load_place_names)
# ---------------------------------------------------------------------------

def load_place_names(path: str | Path = "knowledge/place_names.txt") -> set[str]:
    """Read lowercase place-name set; tokens < 3 chars excluded."""
    p = Path(path)
    if not p.exists():
        return set()
    return {
        line.strip().lower()
        for line in p.read_text(encoding="utf-8").splitlines()
        if line.strip() and len(line.strip()) >= 3
    }


# ---------------------------------------------------------------------------
# Token extraction helpers
# ---------------------------------------------------------------------------

def _extract_numbers(text: str) -> set[str]:
    """Return all distinct number tokens in text."""
    if not text:
        return set()
    return {m.group(0) for m in _NUMBER_RE.finditer(text)}


def _extract_places(text: str, place_names: set[str]) -> set[str]:
    """Return all place-name tokens (lowercase) found in text."""
    if not text:
        return set()
    found: set[str] = set()
    text_lower = text.lower()
    for p in place_names:
        if re.search(r"\b" + re.escape(p) + r"\b", text_lower):
            found.add(p)
    return found


# ---------------------------------------------------------------------------
# Hallucination predicate
# ---------------------------------------------------------------------------

def _is_hallucination(
    pred_description: str,
    grounded_text: str,
    place_names: set[str],
) -> tuple[bool, int, int, int]:
    """Return (is_halluc, ungrounded_num_count, ungrounded_place_count, grounded_match_count).

    A node is a hallucination if it contains at least one number or place name
    that does not appear in the grounded_text corpus.
    """
    pred_nums = _extract_numbers(pred_description)
    pred_places = _extract_places(pred_description, place_names)

    grounded_nums = _extract_numbers(grounded_text)
    grounded_places = _extract_places(grounded_text, place_names)

    ungrounded_nums = pred_nums - grounded_nums
    ungrounded_places = pred_places - grounded_places

    grounded_match = len(pred_nums & grounded_nums) + len(pred_places & grounded_places)

    is_halluc = bool(ungrounded_nums) or bool(ungrounded_places)
    return is_halluc, len(ungrounded_nums), len(ungrounded_places), grounded_match


# ---------------------------------------------------------------------------
# BFS trace retrieval-text extraction
# ---------------------------------------------------------------------------

def _trace_path_for(event_id: str) -> Path | None:
    """Resolve the BFS trace path for an event, preferring multi-seed dumps.

    Multi-seed eval writes one trace per seed (filename suffix
    `_seed{seed}_bfs_full.json`); single-seed legacy eval writes a single
    seedless trace (`_bfs_full.json`). Since BFS is deterministic at
    temperature=0, any of the per-seed traces is fine — pick the first
    sorted seed for stability. Return None if no trace exists.
    """
    seed_traces = sorted(TRACE_DIR.glob(f"{event_id}_seed*_bfs_full.json"))
    if seed_traces:
        return seed_traces[0]
    legacy = TRACE_DIR / f"{event_id}_bfs_full.json"
    return legacy if legacy.exists() else None


def _build_retrieval_corpus(event_id: str) -> str:
    """Collect all parent_text + child_description strings from the BFS trace
    for an event. Returns a single concatenated string for substring matching.

    Falls back to empty string if the trace file is missing.
    """
    trace_path = _trace_path_for(event_id)
    if trace_path is None:
        return ""
    try:
        trace_data = json.loads(trace_path.read_text())
    except Exception:
        return ""

    parts: list[str] = []
    for layer in trace_data.get("trace", []):
        edges_block = layer.get("retrieved_edges", {})
        # edges_block is dict[frontier_id, list[edge_record]]
        if isinstance(edges_block, dict):
            edge_records = [e for edges in edges_block.values() for e in (edges or [])]
        else:
            edge_records = list(edges_block or [])

        for rec in edge_records:
            # The document field contains the BFS query text (parent-side)
            if rec.get("document"):
                parts.append(rec["document"])
            # The edge sub-object holds the raw parent_text + child_description
            edge_obj = rec.get("edge") or {}
            if edge_obj.get("parent_text"):
                parts.append(edge_obj["parent_text"])
            if edge_obj.get("child_description"):
                parts.append(edge_obj["child_description"])

    return " ".join(parts)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> None:
    place_names = load_place_names(ROOT / "knowledge/place_names.txt")
    print(f"Loaded {len(place_names)} place names from knowledge/place_names.txt")

    # Build event_id → FloodEvent fields map for check (a)
    test_events_raw: list[dict] = json.loads(TEST_EVENTS_PATH.read_text())
    test_events: dict[str, dict] = {e["event_id"]: e for e in test_events_raw}

    # Walk eval cache files
    files = sorted(EVAL_DIR.glob("*_seed*.json"))
    print(f"Found {len(files)} eval cache files")

    # Staleness guard: trace older than ANY of its event's caches means the
    # last eval run forgot --dump-bfs-full. Abort rather than silently audit
    # against stale evidence (the grounding corpus would not match the
    # prediction the audit is judging).
    stale: list[tuple[str, float, float]] = []
    missing: list[str] = []
    seen_events: set[str] = set()
    for f in files:
        stem = f.stem
        if "_seed" not in stem:
            continue
        event_id = stem.rsplit("_seed", 1)[0]
        if event_id in _OUTLIER_EVENT_IDS or event_id in seen_events:
            continue
        seen_events.add(event_id)
        trace_path = _trace_path_for(event_id)
        if trace_path is None:
            missing.append(event_id)
            continue
        cache_mtime = f.stat().st_mtime
        trace_mtime = trace_path.stat().st_mtime
        if trace_mtime < cache_mtime:
            stale.append((event_id, trace_mtime, cache_mtime))
    if stale or missing:
        print("ERROR: D6 audit cannot run — BFS trace inputs are stale or missing.")
        for event_id, tm, cm in stale:
            print(f"  STALE {event_id}: trace {tm:.0f} < cache {cm:.0f}")
        for event_id in missing:
            print(f"  MISSING {event_id}: no trace file")
        print()
        print("Rerun the eval with --dump-bfs-full to refresh traces:")
        print("  PYTHONPATH=. python scripts/05_evaluate.py --force --dump-bfs-full")
        raise SystemExit(1)

    # retrieval corpus is per-event (not per-seed — deterministic prediction)
    retrieval_corpus_cache: dict[str, str] = {}

    # per_event[event_id] = list of (is_halluc, num_count, place_count, seed)
    per_event: dict[str, list[tuple[bool, int, int, int]]] = defaultdict(list)

    csv_rows: list[dict] = []

    for f in files:
        stem = f.stem
        if "_seed" not in stem:
            continue
        event_id, seed_part = stem.rsplit("_seed", 1)
        if event_id in _OUTLIER_EVENT_IDS:
            continue
        try:
            seed = int(seed_part)
        except ValueError:
            continue

        try:
            gold_data = json.loads(f.read_text())
        except Exception as exc:
            print(f"  skip {f.name}: {exc}")
            continue

        event = test_events.get(event_id)
        if event is None:
            print(f"  skip {f.name}: event not in test_events.json")
            continue

        # ---- Grounding corpus (a): new event context ----
        event_context = " ".join(filter(None, [
            event.get("description") or "",
            event.get("location") or "",
            event.get("country") or "",
        ]))

        # ---- Grounding corpus (b): retrieved evidence texts ----
        if event_id not in retrieval_corpus_cache:
            retrieval_corpus_cache[event_id] = _build_retrieval_corpus(event_id)
        retrieval_corpus = retrieval_corpus_cache[event_id]

        # Combined grounded text
        grounded_text = event_context + " " + retrieval_corpus

        # ---- Walk predicted nodes ----
        predicted_chain = gold_data.get("predicted_chain", {})
        cascade_events = predicted_chain.get("cascade_events", [])

        for node in cascade_events:
            desc = node.get("description") or ""
            is_halluc, n_nums, n_places, n_grounded = _is_hallucination(
                desc, grounded_text, place_names,
            )

            per_event[event_id].append((is_halluc, n_nums, n_places, seed))

            csv_rows.append({
                "event_id": event_id,
                "seed": seed,
                "node_id": node.get("id", ""),
                "domain": node.get("domain", ""),
                "is_hallucination": int(is_halluc),
                "ungrounded_nums": n_nums,
                "ungrounded_places": n_places,
                "grounded_match": n_grounded,
                "description": desc[:200],
            })

    # ---- Write CSV ----
    if csv_rows:
        with OUT_CSV.open("w", newline="", encoding="utf-8") as fh:
            writer = csv.DictWriter(fh, fieldnames=list(csv_rows[0].keys()))
            writer.writeheader()
            writer.writerows(csv_rows)
        print(f"\nWrote {len(csv_rows)} per-node rows to {OUT_CSV}")
    else:
        print("\nNo CSV rows to write.")

    # ---- Per-event aggregate ----
    print("\n=== Per-event hallucination rate ===")
    header = f"{'event_id':<22} {'pred_nodes (×seeds)':>20} {'halluc_nodes':>13} {'rate':>7}"
    print(header)
    print("-" * len(header))

    rates: list[float] = []
    for eid in sorted(per_event):
        rows = per_event[eid]
        total = len(rows)
        halluc = sum(1 for r in rows if r[0])
        rate = halluc / total if total else 0.0
        rates.append(rate)
        print(f"{eid:<22} {total:>20} {halluc:>13} {rate:>7.3f}")

    print("\n=== Macro aggregate ===")
    if rates:
        macro_rate = sum(rates) / len(rates)
        print(f"Macro hallucination rate:  {macro_rate:.3f}  (across {len(rates)} events)")
        print(f"v0.4 D6 baseline:          0.795")
        delta = macro_rate - 0.795
        delta_str = f"{delta:+.3f}"
        print(f"Delta vs v0.4:             {delta_str}")
        print()
        if macro_rate < 0.50:
            print("Target < 0.50:  MET")
        else:
            print("Target < 0.50:  NOT MET (templating is OFF; filter-only improvement modest)")
    else:
        print("No events to aggregate.")


if __name__ == "__main__":
    main()