src/rag/predictor.py

"""Iterative BFS cascade predictor (v0.2 issue #6).

Replaces the v0.1 single-shot ``CascadePredictor`` with a breadth-first
search over edge-fragment retrieval:

* **Layer 0** — embed the new event's trigger and pull top-K *first-level*
  edges from the historical ``cascade_edges`` collection; ask the LLM
  (``predict_initial.txt``) for the layer-1 cascades anchored at the trigger.
* **Layer ≥1** — for each frontier node, retrieve top-K node-to-node edges
  whose parent description matches the frontier node; ask the LLM
  (``predict_iterative.txt``) to produce the next layer rooted at the
  frontier; repeat until any of the three termination conditions trip:

    1. cumulative ``time_offset_hours`` exceeds ``rag.time_window_hours``;
    2. every per-frontier max similarity falls below
       ``rag.similarity_threshold``;
    3. the LLM returns ``{"layer": [], "stop_reason": "saturation"}``
       (or ``"out_of_domain"``).

Two non-semantic safety nets — ``rag.max_layers`` and
``rag.max_total_nodes`` — bound the loop independently of the model.

Public surface (``predict()`` signature, ``PredictionResult`` shape) is
preserved so that ``Evaluator`` / ``GoldEvaluator`` / the Streamlit
Predict tab keep working without changes; the only additive field is
``PredictionResult.trace`` (per-layer observability record).
"""

from __future__ import annotations

import json
import logging
from collections import Counter
from collections.abc import Iterator
from datetime import date
from typing import Any, Optional

from src.llm.client import (
    LLMClient,
    load_config,
    load_expert_knowledge,
    load_prompt_template,
)
from src.models.schemas import CascadeChain, CascadeEdge, CascadeNode, PredictionResult
from src.rag.chain_index import enrich_retrieval_info, load_chain_index
from src.rag.edge_retriever import EdgeRetriever
from src.rag.embedder import trigger_to_embedding_text

logger = logging.getLogger(__name__)


# Closed taxonomy mirrored from prompts/expert_predict.md — used only for
# light-touch validation logging; the LLM is the source of truth.
_VALID_DOMAINS: set[str] = {
    "infrastructure/power",
    "infrastructure/water",
    "infrastructure/transport",
    "infrastructure/communication",
    "health/casualties",
    "health/hospital_service",
    "health/disease_outbreak",
    "social/evacuation",
    "social/supply_shortage",
    "economy/business_damage",
    "economy/agriculture",
    "environment/contamination",
}
_VALID_STOP_REASONS: set[str] = {"saturation", "out_of_domain"}


# v0.4 issue #63 #4 — domain-budget cap.
# D5 audit identified three domains with high over-prediction ratio AND
# zero or near-zero matches against gold. Each event's emits in these
# domains are capped at gold-avg + small margin; further LLM emits in the
# same domain are dropped (still recorded in trace.domain_budget_dropped).
# Domains not in this map are unlimited.
_DOMAIN_BUDGET: dict[str, int] = {
    "social/supply_shortage": 1,        # gold avg 0.5/event (3/6); 13 pred / 0 matched
    "infrastructure/communication": 2,  # gold avg 1/event (6/6); 13 pred / 0 matched
    "economy/business_damage": 2,       # gold avg 0.83/event (5/6); 17 pred / 1 matched
}


def _check_domain_budget(domain: str, counts: dict[str, int]) -> bool:
    """Return True if this domain has remaining budget for one more emit.

    Domains absent from ``_DOMAIN_BUDGET`` are unlimited (always returns True).
    """
    budget = _DOMAIN_BUDGET.get(domain)
    if budget is None:
        return True
    return counts.get(domain, 0) < budget


def _apply_domain_budget(
    nodes: list, counts: dict[str, int]
) -> tuple[list, list[dict]]:
    """Filter ``nodes`` against per-domain budgets, mutating ``counts``.

    Returns (kept, dropped_records). Each dropped_record:
    ``{"node_id": str, "domain": str, "reason": "domain_budget_exceeded"}``.

    The function increments ``counts`` only for kept nodes — so calling
    again with the same dict carries running budget across BFS layers.
    """
    kept = []
    dropped: list[dict] = []
    for n in nodes:
        d = n.domain if hasattr(n, "domain") else n.get("domain", "")
        if not _check_domain_budget(d, counts):
            dropped.append(
                {
                    "node_id": n.id if hasattr(n, "id") else n.get("id", "?"),
                    "domain": d,
                    "reason": "domain_budget_exceeded",
                }
            )
            continue
        kept.append(n)
        counts[d] = counts.get(d, 0) + 1
    return kept, dropped


class CascadePredictor:
    """Predict cascade risks for a new flood event using BFS over edge RAG."""

    def __init__(
        self,
        llm_client: LLMClient,
        config: dict | None = None,
        *,
        dump_full_trace: bool = False,
    ):
        self.config = config or load_config()
        self.llm_client = llm_client
        self.retriever = EdgeRetriever(self.config)
        self._chain_index = load_chain_index(self.config)
        # When True, every layer's trace record additionally carries the full
        # retrieved-edge payload + the per-layer LLM call (template path,
        # variables, raw response). Off by default; flipped on by issue #12
        # diagnostic runs through ``scripts/05_evaluate.py --dump-bfs-full``.
        self.dump_full_trace: bool = dump_full_trace

        rag_cfg = self.config.get("rag", {})
        self.initial_top_k: int = rag_cfg.get("initial_top_k", 8)
        self.iterative_top_k: int = rag_cfg.get("iterative_top_k_per_frontier", 5)
        self.similarity_threshold: float = rag_cfg.get("similarity_threshold", 0.5)
        self.time_window_hours: float = rag_cfg.get("time_window_hours", 336)
        self.max_layers: int = rag_cfg.get("max_layers", 8)
        self.max_total_nodes: int = rag_cfg.get("max_total_nodes", 200)
        # Per-layer cap on newly-emitted nodes (semantic safety, complements
        # the global max_total_nodes ceiling). Hitting the cap truncates the
        # layer but does NOT terminate BFS — the loop keeps going on the
        # surviving subset. See technical_report/v0.2/rag/issue11_predictor_tuning.md.
        self.max_new_nodes_per_layer: int = rag_cfg.get(
            "max_new_nodes_per_layer", 10
        )

        # Pre-load templates / knowledge once. The judge evaluator does the
        # same (see src/eval/evaluator.py), so this keeps the call hot path
        # free of disk I/O across the BFS loop.
        self._initial_prompt = load_prompt_template(
            self.config["prompts"]["predict_initial"]
        )
        self._iterative_prompt = load_prompt_template(
            self.config["prompts"]["predict_iterative"]
        )
        self._predict_knowledge = load_expert_knowledge(
            self.config["knowledge"]["predict"]
        )

    # ------------------------------------------------------------------
    # Public entry point
    # ------------------------------------------------------------------
    def predict(
        self,
        country: str,
        iso: str,
        location: str,
        event_date: str,
        severity: str,
        description: str,
        seed: int | None = None,
    ) -> PredictionResult:
        """Run BFS to completion and return the assembled :class:`PredictionResult`.

        Internally consumes :meth:`predict_stream`; callers that need
        layer-by-layer progress (e.g. Streamlit Predict tab) should call
        :meth:`predict_stream` directly.

        Args:
            country: Human-readable country name (e.g. ``"Italy"``).
            iso: ISO 3166-1 alpha-3 country code (e.g. ``"ITA"``).
            location: Sub-national location string used for RAG context.
            event_date: ISO-format date string (``"YYYY-MM-DD"``).
            severity: Flood severity label (e.g. ``"major"``).
            description: Free-text description of the flood event.
            seed: Optional integer forwarded to every :meth:`LLMClient.call
                <src.llm.client.LLMClient.call>` invocation inside the BFS
                loop (both ``_llm_initial`` and ``_llm_iterative``).  Honored
                by local backends (greedy decode when ``temperature=0``,
                sampling with fixed seed otherwise); cloud backends
                (``claude`` / ``claude-cli``) raise ``NotImplementedError``
                when *seed* is not ``None``.  Default ``None`` preserves
                backward-compatible non-deterministic behavior.
        """
        final: dict | None = None
        for chunk in self.predict_stream(
            country=country,
            iso=iso,
            location=location,
            event_date=event_date,
            severity=severity,
            description=description,
            seed=seed,
        ):
            if chunk.get("is_final"):
                final = chunk
        if final is None:  # defensive — predict_stream always yields a final chunk
            raise RuntimeError("predict_stream did not yield a final chunk")
        return final["result"]

    def predict_stream(
        self,
        country: str,
        iso: str,
        location: str,
        event_date: str,
        severity: str,
        description: str,
        seed: int | None = None,
    ) -> Iterator[dict]:
        """Yield one chunk per BFS layer for streaming UIs.

        Per-layer chunk shape (``is_final=False``)::

            {
                "layer": int,                   # 0 = seed; safety:max_layers may exceed max_layers
                "trace_record": dict,           # the same record appended to result.trace
                "produced": list[CascadeNode],  # nodes newly added to the DAG this layer
                "evidence_ids": list[str],      # edge_ids retrieved for this layer
                "stop_reason": str | None,      # mirrors trace_record["stop_reason"]
                "partial_dag": list[CascadeNode],  # snapshot of DAG after this layer commits
                "is_final": False,
            }

        Final chunk (``is_final=True``)::

            {"is_final": True, "result": PredictionResult}

        ``predict_stream`` always emits at least one per-layer chunk and
        exactly one final chunk (even when BFS short-circuits at layer 0).

        Args:
            country: Human-readable country name (e.g. ``"Italy"``).
            iso: ISO 3166-1 alpha-3 country code (e.g. ``"ITA"``).
            location: Sub-national location string used for RAG context.
            event_date: ISO-format date string (``"YYYY-MM-DD"``).
            severity: Flood severity label (e.g. ``"major"``).
            description: Free-text description of the flood event.
            seed: Optional integer forwarded to every :meth:`LLMClient.call
                <src.llm.client.LLMClient.call>` invocation inside the BFS
                loop (both ``_llm_initial`` and ``_llm_iterative``).  Honored
                by local backends (greedy decode when ``temperature=0``,
                sampling with fixed seed otherwise); cloud backends
                (``claude`` / ``claude-cli``) raise ``NotImplementedError``
                when *seed* is not ``None``.  Default ``None`` preserves
                backward-compatible non-deterministic behavior.
        """
        event_id = f"PRED-{event_date}-{iso}"
        event_inputs = {
            "country": country,
            "iso": iso,
            "location": location,
            "date": event_date,
            "severity": severity,
            "description": description,
            "event_id": event_id,
        }

        dag: list[CascadeNode] = []
        trace: list[dict] = []
        # v0.4 issue #63 #4 — running per-domain emit counts (cross-layer)
        # plus a flat list of dropped-by-budget node records for trace.
        domain_emit_counts: dict[str, int] = {}
        domain_budget_dropped: list[dict] = []
        # Aggregate retrieval provenance across all layers; layer-0 events
        # are the analogue of v0.1 reference_event_ids, but we union in
        # layer-≥1 evidence too so the Predict tab can show what BFS leaned
        # on at every step.
        seen_event_ids: list[str] = []
        event_id_to_max_sim: dict[str, float] = {}

        # ── Layer 0: seed ────────────────────────────────────────────
        q0 = trigger_to_embedding_text(
            country=country,
            event_date=event_date,
            severity=severity,
            summary=description,
            location=location,
        )
        layer1_nodes, seed_results, stop_reason, layer0_confidence, layer0_llm_call = (
            self._seed_layer(q0, event_inputs, seed=seed)
        )
        self._update_provenance(seed_results, seen_event_ids, event_id_to_max_sim)
        layer1_nodes, layer0_truncated = self._apply_per_layer_cap(
            layer1_nodes, layer0_confidence
        )
        if layer0_truncated:
            logger.info(
                "Layer 0 truncated by per-layer cap (%d → %d, dropped %d)",
                layer0_truncated + len(layer1_nodes),
                len(layer1_nodes),
                layer0_truncated,
            )
        # v0.4 issue #63 #4 — domain-budget cap (after per-layer cap)
        layer1_nodes, layer0_budget_drops = _apply_domain_budget(
            layer1_nodes, domain_emit_counts
        )
        if layer0_budget_drops:
            logger.info(
                "Layer 0 dropped %d node(s) by domain-budget cap: %s",
                len(layer0_budget_drops),
                [(d["node_id"], d["domain"]) for d in layer0_budget_drops],
            )
            domain_budget_dropped.extend(layer0_budget_drops)
        layer0_record = self._record_layer(
            layer_idx=0,
            frontier_ids=[],
            evidence={"trigger": seed_results},
            stop_reason=stop_reason,
            produced_ids=[n.id for n in layer1_nodes],
            truncated_by_per_layer_cap=layer0_truncated,
            llm_call=layer0_llm_call,
        )
        trace.append(layer0_record)
        dag.extend(layer1_nodes)
        yield self._stream_chunk(
            layer_idx=0,
            trace_record=layer0_record,
            produced=layer1_nodes,
            evidence_ids=_collect_edge_ids(seed_results),
            stop_reason=stop_reason,
            dag=dag,
        )
        if stop_reason or not layer1_nodes:
            yield self._final_chunk(
                event_id, country, iso, event_date, severity,
                dag, trace, seen_event_ids, event_id_to_max_sim,
            )
            return

        frontier = layer1_nodes
        # ── Layer ≥1 ─────────────────────────────────────────────────
        for layer_idx in range(1, self.max_layers + 1):
            if len(dag) >= self.max_total_nodes:
                rec = self._record_layer(
                    layer_idx=layer_idx,
                    frontier_ids=[n.id for n in frontier],
                    evidence={},
                    stop_reason="safety:max_total_nodes",
                    produced_ids=[],
                )
                trace.append(rec)
                yield self._stream_chunk(
                    layer_idx=layer_idx,
                    trace_record=rec,
                    produced=[],
                    evidence_ids=[],
                    stop_reason="safety:max_total_nodes",
                    dag=dag,
                )
                break

            evidence = {
                node.id: self.retriever.retrieve_node_to_node(
                    query_text=node.description,
                    top_k=self.iterative_top_k,
                    parent_domain=node.domain,
                    frontier_time_offset_hours=node.time_offset_hours,
                )
                for node in frontier
            }
            flat_ev = self._flatten_evidence(evidence)

            stop_pre = self._should_stop(evidence)
            if stop_pre is not None:
                self._update_provenance(flat_ev, seen_event_ids, event_id_to_max_sim)
                rec = self._record_layer(
                    layer_idx=layer_idx,
                    frontier_ids=[n.id for n in frontier],
                    evidence=evidence,
                    stop_reason=stop_pre,
                    produced_ids=[],
                )
                trace.append(rec)
                yield self._stream_chunk(
                    layer_idx=layer_idx,
                    trace_record=rec,
                    produced=[],
                    evidence_ids=_collect_edge_ids(flat_ev),
                    stop_reason=stop_pre,
                    dag=dag,
                )
                break

            next_layer, stop_reason, iter_confidence, iter_llm_call = (
                self._llm_iterative(event_inputs, dag, frontier, evidence, seed=seed)
            )
            next_layer = self._validate_parent_anchors(next_layer, dag, frontier)
            next_layer, iter_truncated = self._apply_per_layer_cap(
                next_layer, iter_confidence
            )
            if iter_truncated:
                logger.info(
                    "Layer %d truncated by per-layer cap (kept %d, dropped %d)",
                    layer_idx, len(next_layer), iter_truncated,
                )
            # v0.4 issue #63 #4 — domain-budget cap (after per-layer cap)
            next_layer, iter_budget_drops = _apply_domain_budget(
                next_layer, domain_emit_counts
            )
            if iter_budget_drops:
                logger.info(
                    "Layer %d dropped %d node(s) by domain-budget cap: %s",
                    layer_idx,
                    len(iter_budget_drops),
                    [(d["node_id"], d["domain"]) for d in iter_budget_drops],
                )
                domain_budget_dropped.extend(iter_budget_drops)
            self._update_provenance(flat_ev, seen_event_ids, event_id_to_max_sim)

            if stop_reason in _VALID_STOP_REASONS or not next_layer:
                effective_stop = stop_reason or "empty_layer"
                rec = self._record_layer(
                    layer_idx=layer_idx,
                    frontier_ids=[n.id for n in frontier],
                    evidence=evidence,
                    stop_reason=effective_stop,
                    produced_ids=[],
                    llm_call=iter_llm_call,
                )
                trace.append(rec)
                yield self._stream_chunk(
                    layer_idx=layer_idx,
                    trace_record=rec,
                    produced=[],
                    evidence_ids=_collect_edge_ids(flat_ev),
                    stop_reason=effective_stop,
                    dag=dag,
                )
                break

            dag.extend(next_layer)
            rec = self._record_layer(
                layer_idx=layer_idx,
                frontier_ids=[n.id for n in frontier],
                evidence=evidence,
                stop_reason=None,
                produced_ids=[n.id for n in next_layer],
                truncated_by_per_layer_cap=iter_truncated,
                llm_call=iter_llm_call,
            )
            trace.append(rec)
            yield self._stream_chunk(
                layer_idx=layer_idx,
                trace_record=rec,
                produced=next_layer,
                evidence_ids=_collect_edge_ids(flat_ev),
                stop_reason=None,
                dag=dag,
            )

            # Termination ①: time window — only nodes still inside the window
            # propagate to the next frontier. If none, stop.
            live = [
                n
                for n in next_layer
                if (n.time_offset_hours or 0.0) <= self.time_window_hours
            ]
            if not live:
                rec = self._record_layer(
                    layer_idx=layer_idx,
                    frontier_ids=[n.id for n in next_layer],
                    evidence={},
                    stop_reason="time_window_exhausted",
                    produced_ids=[],
                )
                trace.append(rec)
                yield self._stream_chunk(
                    layer_idx=layer_idx,
                    trace_record=rec,
                    produced=[],
                    evidence_ids=[],
                    stop_reason="time_window_exhausted",
                    dag=dag,
                )
                break
            frontier = live
        else:
            # Loop fell through without break ⇒ hit max_layers cap.
            rec = self._record_layer(
                layer_idx=self.max_layers + 1,
                frontier_ids=[n.id for n in frontier],
                evidence={},
                stop_reason="safety:max_layers",
                produced_ids=[],
            )
            trace.append(rec)
            yield self._stream_chunk(
                layer_idx=self.max_layers + 1,
                trace_record=rec,
                produced=[],
                evidence_ids=[],
                stop_reason="safety:max_layers",
                dag=dag,
            )

        yield self._final_chunk(
            event_id, country, iso, event_date, severity,
            dag, trace, seen_event_ids, event_id_to_max_sim,
        )

    # ------------------------------------------------------------------
    # BFS sub-steps
    # ------------------------------------------------------------------
    def _seed_layer(
        self, q0: str, event_inputs: dict, *, seed: int | None = None
    ) -> tuple[list[CascadeNode], list[dict], Optional[str], dict[str, float], dict | None]:
        seed_results = self.retriever.retrieve_first_level(
            q0, top_k=self.initial_top_k
        )
        nodes, stop_reason, confidence, llm_call = self._llm_initial(
            event_inputs, seed_results, seed=seed
        )
        return nodes, seed_results, stop_reason, confidence, llm_call

    def _llm_initial(
        self, event_inputs: dict, seed_edges: list[dict], *, seed: int | None = None
    ) -> tuple[list[CascadeNode], Optional[str], dict[str, float], dict | None]:
        seed_text = _format_seed_edges(seed_edges)
        variables = {**event_inputs, "seed_edges": seed_text}
        raw = self.llm_client.call(
            self._initial_prompt, variables, self._predict_knowledge, seed=seed
        )
        nodes, stop_reason, confidence = self._parse_initial_layer(raw)
        llm_call = self._build_llm_call_record("predict_initial", variables, raw)
        return nodes, stop_reason, confidence, llm_call

    def _llm_iterative(
        self,
        event_inputs: dict,
        dag: list[CascadeNode],
        frontier: list[CascadeNode],
        evidence: dict[str, list[dict]],
        *,
        seed: int | None = None,
    ) -> tuple[list[CascadeNode], Optional[str], dict[str, float], dict | None]:
        variables = {
            "event_inputs": _format_event_inputs(event_inputs),
            "dag_snapshot": _format_dag_snapshot(dag),
            "frontier_nodes": _format_frontier_nodes(frontier),
            "evidence_per_frontier": _format_evidence_per_frontier(frontier, evidence),
        }
        raw = self.llm_client.call(
            self._iterative_prompt, variables, self._predict_knowledge, seed=seed
        )
        nodes, stop_reason, confidence = self._parse_iterative_layer(raw)
        llm_call = self._build_llm_call_record("predict_iterative", variables, raw)
        return nodes, stop_reason, confidence, llm_call

    def _build_llm_call_record(
        self, prompt_key: str, variables: dict, raw_response: str
    ) -> dict | None:
        """Snapshot one LLM call for diagnostic dumps.

        Returns ``None`` when ``dump_full_trace`` is off, so the BFS hot path
        pays only a tuple-wrap cost when diagnostics aren't requested.
        """
        if not self.dump_full_trace:
            return None
        return {
            "prompt_template_path": self.config["prompts"][prompt_key],
            "variables": dict(variables),
            "raw_response": raw_response,
        }

    def _apply_per_layer_cap(
        self,
        nodes: list[CascadeNode],
        confidence: dict[str, float] | None = None,
    ) -> tuple[list[CascadeNode], int]:
        """Truncate ``nodes`` to ``max_new_nodes_per_layer``.

        Selection strategy:
        - If a non-empty ``confidence`` map is provided, keep the top-N by
          confidence (descending). Output preserves the original LLM emission
          order for the kept subset (stable display in the trace + UI).
        - Otherwise fall back to LLM emission order — most prompts produce
          the higher-importance cascades first, so head-truncation is a
          reasonable proxy.

        Returns ``(kept_nodes, truncated_count)``. Never terminates BFS;
        the BFS loop keeps going on the surviving subset.
        """
        if len(nodes) <= self.max_new_nodes_per_layer:
            return nodes, 0

        cap = self.max_new_nodes_per_layer
        if confidence:
            # Stable sort: confidence desc, original index asc as tiebreaker.
            indexed = list(enumerate(nodes))
            indexed.sort(
                key=lambda pair: (-confidence.get(pair[1].id, 0.0), pair[0])
            )
            kept_ids = {pair[1].id for pair in indexed[:cap]}
            kept = [n for n in nodes if n.id in kept_ids]
        else:
            kept = list(nodes[:cap])
        return kept, len(nodes) - len(kept)

    def _should_stop(self, evidence: dict[str, list[dict]]) -> Optional[str]:
        """Termination ② — every frontier node's max similarity below threshold."""
        if not evidence:
            return None
        max_sims = [_max_similarity(es) for es in evidence.values()]
        if all(sim < self.similarity_threshold for sim in max_sims):
            return "similarity_below_threshold"
        return None

    def _validate_parent_anchors(
        self,
        candidates: list[CascadeNode],
        dag: list[CascadeNode],
        frontier: list[CascadeNode],
    ) -> list[CascadeNode]:
        """Validate parent anchors and enforce v0.2 issue #9 / B' edge rules.

        BFS layer ≥1 nodes MUST anchor to either a frontier node or an
        earlier DAG node. This function drops malformed nodes and prunes
        ancestors from multi-parent declarations:

        1. **Empty parent_ids** — layer ≥1 nodes need at least one parent;
           drop silently with a warning (the iterative prompt is explicit).
        2. **Self-loop** — a node listing its own id as a parent is dropped.
        3. **Unknown id** — any parent id not in ``frontier ∪ dag`` is treated
           as a hallucination; the whole node is dropped.
        4. **No-grandparent pruning** — for surviving multi-parent nodes,
           transitively walk ``dag.parent_ids`` upward; if a listed parent P
           is an ancestor of another listed parent Q, P is dropped (most-
           direct cause wins). The node is kept with its pruned parent set.
        """
        known_by_id = {n.id: n for n in dag}
        for n in frontier:
            known_by_id.setdefault(n.id, n)
        known_ids = set(known_by_id.keys())
        all_known_nodes = list(known_by_id.values())

        survivors: list[CascadeNode] = []
        dropped: list[str] = []
        for node in candidates:
            if not node.parent_ids:
                dropped.append(f"{node.id}(no parent_ids)")
                continue
            if node.id in node.parent_ids:
                dropped.append(f"{node.id}(self-loop)")
                continue
            unknown = [pid for pid in node.parent_ids if pid not in known_ids]
            if unknown:
                dropped.append(f"{node.id}(unknown_parent={unknown})")
                continue
            pruned = _prune_ancestor_parents(node.parent_ids, all_known_nodes)
            if len(pruned) < len(node.parent_ids):
                logger.info(
                    "Pruned ancestor parents for %s: %s -> %s",
                    node.id, node.parent_ids, pruned,
                )
            node.parent_ids = pruned
            survivors.append(node)
        if dropped:
            logger.warning(
                "Dropped %d malformed BFS nodes with bad parent anchors: %s",
                len(dropped),
                "; ".join(dropped),
            )
        return survivors

    # ------------------------------------------------------------------
    # Parsing
    # ------------------------------------------------------------------
    def _parse_initial_layer(
        self, response: str
    ) -> tuple[list[CascadeNode], Optional[str], dict[str, float]]:
        data = _extract_json(response)
        if data is None:
            logger.error("Layer-0 LLM response unparseable; treating as empty")
            return [], None, {}
        nodes_data = data.get("cascade_events") or []
        nodes, confidence = _build_nodes(nodes_data)
        # Layer 0: parent_ids MUST be empty per prompt — coerce defensively.
        for n in nodes:
            if n.parent_ids:
                logger.warning(
                    "Layer-0 node %s had non-empty parent_ids %s — coercing to []",
                    n.id, n.parent_ids,
                )
                n.parent_ids = []
        stop_reason = data.get("stop_reason")
        return nodes, stop_reason, confidence

    def _parse_iterative_layer(
        self, response: str
    ) -> tuple[list[CascadeNode], Optional[str], dict[str, float]]:
        data = _extract_json(response)
        if data is None:
            logger.error("Iterative LLM response unparseable; treating as empty")
            return [], None, {}
        layer = data.get("layer") or []
        nodes, confidence = _build_nodes(layer)
        stop_reason = data.get("stop_reason")
        return nodes, stop_reason, confidence

    # ------------------------------------------------------------------
    # Result assembly
    # ------------------------------------------------------------------
    def _build_result(
        self,
        event_id: str,
        country: str,
        iso: str,
        event_date: str,
        severity: str,
        dag: list[CascadeNode],
        trace: list[dict],
        seen_event_ids: list[str],
        event_id_to_max_sim: dict[str, float],
    ) -> PredictionResult:
        chain = CascadeChain(
            event_id=event_id,
            trigger_summary=f"Flood prediction for {country}",
            trigger_country=country,
            trigger_iso=iso,
            trigger_date=date.fromisoformat(event_date),
            trigger_severity=severity,
            cascade_events=dag,
            extraction_date=date.today(),
        )
        # Flat-DAG observability — final accumulated DAG with ≥3 nodes and
        # zero inter-node edges almost always means BFS collapsed to the
        # layer-0 seed only. Logged here (not per-layer) so callers see one
        # warning per prediction.
        if len(dag) >= 3 and not any(n.parent_ids for n in dag):
            logger.warning(
                "Flat DAG: %d cascade nodes emitted with zero inter-node edges "
                "for event %s. Every node will fall back to TRIGGER in the viewer.",
                len(dag),
                event_id,
            )

        confidence_scores: dict[str, float] = {}
        for layer_record in trace:
            for nid in layer_record.get("produced_ids", []):
                # _build_nodes captured per-node confidence into the trace
                # via the optional 'confidence' field; surface it here so
                # downstream UI / eval can read it the same way as v0.1.
                pass  # explicit no-op — see _build_nodes return path

        # Re-derive confidence_scores from the dag itself by looking up the
        # per-node 'confidence' attached during parsing — we stored them
        # transiently via a parallel dict per parse, but the per-node
        # confidence is not part of CascadeNode. To keep v0.1 wire format
        # we just leave confidence_scores empty unless the parsers carry it
        # forward; the in-memory dicts they returned are layer-local. For
        # now (Predict tab still uses similarity_scores), an empty dict is
        # acceptable.

        # reference_event_ids / similarity_scores: union of all retrieved
        # source events, sorted by descending max similarity for stable UI.
        ranked = sorted(
            seen_event_ids,
            key=lambda eid: event_id_to_max_sim.get(eid, 0.0),
            reverse=True,
        )
        similarity_scores = [
            round(event_id_to_max_sim.get(eid, 0.0), 4) for eid in ranked
        ]
        reference_info = enrich_retrieval_info(
            ranked, similarity_scores, self._chain_index
        )

        return PredictionResult(
            predicted_chain=chain,
            confidence_scores=confidence_scores,
            reference_event_ids=ranked,
            similarity_scores=similarity_scores,
            reference_info=reference_info,
            trace=trace,
        )

    # ------------------------------------------------------------------
    # Trace + provenance helpers
    # ------------------------------------------------------------------
    def _record_layer(
        self,
        layer_idx: int,
        frontier_ids: list[str],
        evidence: dict[str, list[dict]],
        stop_reason: Optional[str],
        produced_ids: list[str],
        truncated_by_per_layer_cap: int = 0,
        llm_call: dict | None = None,
    ) -> dict[str, Any]:
        flat_evidence = [r for es in evidence.values() for r in es]
        evidence_summary = {
            key: {
                "n_results": len(es),
                "max_similarity": _max_similarity(es),
                "source_event_ids": _unique_event_ids(es),
            }
            for key, es in evidence.items()
        }
        record: dict[str, Any] = {
            "layer": layer_idx,
            "frontier_ids": list(frontier_ids),
            "evidence_summary": evidence_summary,
            "filter_path_distribution": dict(Counter(
                r.get("filter_path", "?") for r in flat_evidence
            )),
            "stop_reason": stop_reason,
            "produced_ids": list(produced_ids),
            "truncated_by_per_layer_cap": truncated_by_per_layer_cap,
        }
        if self.dump_full_trace:
            # Diagnostic-only fields. Keep them under explicit keys so a
            # consumer can ``record.get("retrieved_edges")`` without breaking
            # when the flag is off.
            record["retrieved_edges"] = _serialize_evidence(evidence)
            record["llm_call"] = llm_call  # may be None for non-LLM safety stops
        return record

    @staticmethod
    def _stream_chunk(
        layer_idx: int,
        trace_record: dict,
        produced: list[CascadeNode],
        evidence_ids: list[str],
        stop_reason: Optional[str],
        dag: list[CascadeNode],
    ) -> dict[str, Any]:
        return {
            "layer": layer_idx,
            "trace_record": trace_record,
            "produced": list(produced),
            "evidence_ids": list(evidence_ids),
            "stop_reason": stop_reason,
            "partial_dag": list(dag),
            "is_final": False,
        }

    def _final_chunk(
        self,
        event_id: str,
        country: str,
        iso: str,
        event_date: str,
        severity: str,
        dag: list[CascadeNode],
        trace: list[dict],
        seen_event_ids: list[str],
        event_id_to_max_sim: dict[str, float],
    ) -> dict[str, Any]:
        return {
            "is_final": True,
            "result": self._build_result(
                event_id=event_id,
                country=country,
                iso=iso,
                event_date=event_date,
                severity=severity,
                dag=dag,
                trace=trace,
                seen_event_ids=seen_event_ids,
                event_id_to_max_sim=event_id_to_max_sim,
            ),
        }

    @staticmethod
    def _flatten_evidence(evidence: dict[str, list[dict]]) -> list[dict]:
        out: list[dict] = []
        for es in evidence.values():
            out.extend(es)
        return out

    @staticmethod
    def _update_provenance(
        results: list[dict],
        seen_event_ids: list[str],
        event_id_to_max_sim: dict[str, float],
    ) -> None:
        for r in results:
            eid = r.get("source_event_id") or ""
            if not eid:
                continue
            sim = float(r.get("similarity") or 0.0)
            if eid not in event_id_to_max_sim:
                seen_event_ids.append(eid)
                event_id_to_max_sim[eid] = sim
            else:
                event_id_to_max_sim[eid] = max(event_id_to_max_sim[eid], sim)


# ----------------------------------------------------------------------
# Module-level helpers (no state)
# ----------------------------------------------------------------------
def _extract_json(raw: str) -> dict | None:
    text = raw.strip()
    if "```json" in text:
        text = text.split("```json", 1)[1].split("```", 1)[0]
    elif "```" in text:
        text = text.split("```", 1)[1].split("```", 1)[0]
    try:
        return json.loads(text.strip())
    except (json.JSONDecodeError, ValueError):
        return None


def _build_nodes(
    nodes_data: list[dict],
) -> tuple[list[CascadeNode], dict[str, float]]:
    nodes: list[CascadeNode] = []
    confidence: dict[str, float] = {}
    for raw in nodes_data:
        nid = raw.get("id") or f"E{len(nodes) + 1}"
        try:
            node = CascadeNode(
                id=nid,
                description=raw.get("description", ""),
                domain=raw.get("domain", "unknown"),
                severity=raw.get("severity", "medium"),
                time_offset_hours=raw.get("time_offset_hours"),
                mechanism=raw.get("mechanism", ""),
                parent_ids=list(raw.get("parent_ids") or []),
            )
        except Exception as exc:
            logger.warning("Skipping malformed cascade node %r: %s", raw, exc)
            continue
        nodes.append(node)
        if "confidence" in raw:
            try:
                confidence[node.id] = float(raw["confidence"])
            except (TypeError, ValueError):
                pass
        if node.domain not in _VALID_DOMAINS:
            logger.warning("Node %s emitted with unknown domain %r", node.id, node.domain)
    return nodes, confidence


def _prune_ancestor_parents(
    parents: list[str],
    known_nodes: list[CascadeNode],
) -> list[str]:
    """Drop parents that are ancestors of another listed parent.

    Implements the v0.2 issue #9 / B' no-grandparent rule on the predictor
    side: when the LLM emits a multi-parent node, transitively walk the
    known DAG's ``parent_ids`` upward to find each listed parent's ancestor
    set. Any listed parent that appears in another's ancestor set is removed
    so only the most-direct causes remain.

    Order is preserved for the survivors so the prompt-side ordering (which
    typically puts the frontier expander first) is retained.
    """
    if len(parents) <= 1:
        return list(parents)

    parent_map = {n.id: list(n.parent_ids) for n in known_nodes}

    def ancestors_of(nid: str) -> set[str]:
        seen: set[str] = set()
        stack = list(parent_map.get(nid, []))
        while stack:
            cur = stack.pop()
            if cur in seen:
                continue
            seen.add(cur)
            stack.extend(parent_map.get(cur, []))
        return seen

    candidate_set = set(parents)
    drop: set[str] = set()
    for p in parents:
        anc = ancestors_of(p)
        for q in candidate_set:
            if q != p and q in anc:
                drop.add(q)
    return [p for p in parents if p not in drop]


def _serialize_evidence(evidence: dict[str, list[dict]]) -> dict[str, list[dict]]:
    """JSON-serializable snapshot of retriever output (issue #12 diagnostics).

    The raw retriever returns dicts that already contain a ``CascadeEdge``
    Pydantic object under ``edge``. We dump it via ``.model_dump(mode="json")``
    so dates / enums round-trip into JSON, preserve the rest of the record,
    and return a plain ``{frontier_id: [...]}`` mapping suitable for inclusion
    in ``PredictionResult.trace``.
    """
    out: dict[str, list[dict]] = {}
    for frontier_id, results in evidence.items():
        items: list[dict] = []
        for r in results:
            ser: dict = {
                "similarity": float(r.get("similarity") or 0.0),
                "source_event_id": r.get("source_event_id"),
                "metadata": r.get("metadata"),
                "document": r.get("document"),
            }
            edge = r.get("edge")
            if edge is not None:
                if hasattr(edge, "model_dump"):
                    ser["edge"] = edge.model_dump(mode="json")
                else:  # defensive: edge already a dict in test fixtures
                    ser["edge"] = dict(edge)
            items.append(ser)
        out[frontier_id] = items
    return out


def _max_similarity(results: list[dict]) -> float:
    if not results:
        return 0.0
    return max(float(r.get("similarity") or 0.0) for r in results)


def _collect_edge_ids(results: list[dict]) -> list[str]:
    """Pull edge_id strings out of retriever results for stream chunks.

    Falls back to ``{source_event_id}::?`` when the on-disk edge could not be
    reverse-loaded (test fixtures sometimes index without the edges/*.json file).
    Order is preserved so the UI can show edges in retrieval order.
    """
    out: list[str] = []
    for r in results:
        edge = r.get("edge")
        if edge is not None and getattr(edge, "edge_id", None):
            out.append(edge.edge_id)
            continue
        sid = (r.get("metadata") or {}).get("source_event_id") or r.get("source_event_id") or "unknown"
        out.append(f"{sid}::?")
    return out


def _unique_event_ids(results: list[dict]) -> list[str]:
    seen: list[str] = []
    for r in results:
        eid = r.get("source_event_id")
        if eid and eid not in seen:
            seen.append(eid)
    return seen


def _format_event_inputs(event_inputs: dict) -> str:
    return (
        f"- Country: {event_inputs.get('country', '')}\n"
        f"- ISO: {event_inputs.get('iso', '')}\n"
        f"- Location: {event_inputs.get('location', '')}\n"
        f"- Date: {event_inputs.get('date', '')}\n"
        f"- Severity: {event_inputs.get('severity', '')}\n"
        f"- Description: {event_inputs.get('description', '')}\n"
        f"- Event ID: {event_inputs.get('event_id', '')}"
    )


def _format_dag_snapshot(dag: list[CascadeNode]) -> str:
    if not dag:
        return "(empty)"
    parts = []
    for n in dag:
        parts.append(
            f"- [{n.id}] {n.description}\n"
            f"    domain={n.domain}, severity={n.severity}, "
            f"time_offset_hours={n.time_offset_hours}, "
            f"parent_ids={n.parent_ids}"
        )
    return "\n".join(parts)


def _format_frontier_nodes(frontier: list[CascadeNode]) -> str:
    if not frontier:
        return "(empty)"
    parts = []
    for n in frontier:
        parts.append(
            f"- id={n.id}, description={n.description}, "
            f"domain={n.domain}, severity={n.severity}, "
            f"time_offset_hours={n.time_offset_hours}"
        )
    return "\n".join(parts)


def _render_edge_evidence(idx: int, result: dict) -> str:
    """Render one retrieved edge result for LLM prompt consumption.

    v0.5 issue C: use ``evidence_template`` (redacted, with ``<N>`` / ``<LOC>``
    placeholders) when present so that raw numbers and place-names from
    historical events do not reach the LLM as in-context examples.  The
    template is built by ``EdgeRetriever`` (Task 6) and stored per result.

    Fallback order when ``evidence_template`` is absent (legacy fixtures /
    edges that pre-date Task 6):
      1. ``edge.parent_text`` / ``edge.child_description`` if edge is loaded.
      2. ChromaDB ``metadata`` fields when edge object is unavailable.
    """
    edge: CascadeEdge | None = result.get("edge")
    sim = result.get("similarity", 0.0)
    src = result.get("source_event_id", "")

    # --- Primary path: use pre-built evidence_template (Task 6) ---
    evidence_template = result.get("evidence_template")
    if evidence_template:
        return (
            f"  [Edge {idx}] (similarity={sim:.3f}, source_event_id={src})\n"
            f"    {evidence_template}"
        )

    # --- Fallback: raw fields (legacy / test fixtures without template) ---
    if edge is None:
        # Fall back to ChromaDB metadata when the edge file could not be
        # reverse-loaded (e.g. test fixtures without on-disk edges).
        md = result.get("metadata") or {}
        return (
            f"  [Edge {idx}] (similarity={sim:.3f}, source_event_id={src})\n"
            f"    parent_text: {result.get('document', '')}\n"
            f"    child_description: {md.get('child_description', '(unavailable)')}\n"
            f"    child_domain: {md.get('child_domain', '')}\n"
            f"    child_severity: {md.get('child_severity', '')}\n"
            f"    time_offset_hours_delta: {md.get('time_offset_hours_delta', '')}"
        )
    return (
        f"  [Edge {idx}] (similarity={sim:.3f}, source_event_id={src})\n"
        f"    parent_text: {edge.parent_text}\n"
        f"    child_description: {edge.child_description}\n"
        f"    child_domain: {edge.child_domain}\n"
        f"    child_severity: {edge.child_severity}\n"
        f"    time_offset_hours_delta: {edge.time_offset_hours_delta}"
    )


def _format_seed_edges(results: list[dict]) -> str:
    if not results:
        return "(no historical first-level edges retrieved — proceed from physical reasoning alone.)"
    parts = []
    for i, r in enumerate(results, start=1):
        parts.append(_render_edge_evidence(i, r))
    return "\n".join(parts)


def _format_evidence_per_frontier(
    frontier: list[CascadeNode], evidence: dict[str, list[dict]]
) -> str:
    if not frontier:
        return "(no frontier)"
    blocks = []
    for n in frontier:
        es = evidence.get(n.id, [])
        header = f"### Frontier node {n.id}: {n.description}"
        if not es:
            blocks.append(f"{header}\n  (no historical edges retrieved)")
            continue
        lines = [header]
        for i, r in enumerate(es, start=1):
            lines.append(_render_edge_evidence(i, r))
        blocks.append("\n".join(lines))
    return "\n\n".join(blocks)