src/eval/evaluator.py

"""Orchestrate predict → load news → judge → cache for one test event."""

from __future__ import annotations

import hashlib
import json
import logging
from datetime import date
from pathlib import Path
from typing import Any, Protocol

from pydantic import ValidationError

from src.data.cascade_extractor import _infer_severity
from src.eval.news_loader import load_articles
from src.llm.client import (
    LLMClient,
    load_expert_knowledge,
    load_prompt_template,
)
from src.models.schemas import (
    EventEvaluation,
    FloodEvent,
    MissedCascade,
    NodeEvaluation,
    PredictionResult,
)
from src.rag.chain_index import enrich_retrieval_info, load_chain_index

logger = logging.getLogger(__name__)


class SupportsPredict(Protocol):
    def predict(self, **kwargs: Any) -> PredictionResult: ...


# Bumped together with the v0.2 BFS predictor (issue #6). The judge prompt /
# knowledge files themselves did not change, but the *predicted chain shape*
# the judge is grading did (BFS-produced layered DAG vs v0.1 single-shot),
# so prior judge caches must be regenerated.
JUDGE_PROMPT_VERSION = "v0.2"


def judge_fingerprint(
    prompt_text: str,
    knowledge_text: str,
    prompt_version: str = JUDGE_PROMPT_VERSION,
) -> str:
    """Stable fingerprint of the judge's prompt+knowledge, for cache invalidation.

    ``prompt_version`` is mixed in so a release that changes the *predictor*
    (and therefore the JSON shape the judge sees) can invalidate caches even
    when the prompt/knowledge files themselves are byte-identical.
    """
    h = hashlib.sha256()
    h.update(prompt_text.encode("utf-8"))
    h.update(b"\x00")
    h.update(knowledge_text.encode("utf-8"))
    h.update(b"\x00")
    h.update(prompt_version.encode("utf-8"))
    return h.hexdigest()[:16]


def _format_news_for_prompt(articles: list[dict]) -> str:
    if not articles:
        return "(no articles)"
    parts = []
    for i, a in enumerate(articles, start=1):
        parts.append(
            f"[Article {i}] {a.get('title', '')}\n"
            f"URL: {a.get('url', '')}\n"
            f"Date: {a.get('date', '')}\n"
            f"Body: {a.get('text', '')}"
        )
    return "\n\n".join(parts)


def _format_event_metadata(event: FloodEvent) -> str:
    lines = [
        f"- event_id: {event.event_id}",
        f"- country: {event.country}",
        f"- location: {event.location or event.country}",
        f"- date: {event.start_date}",
    ]
    if event.origin:
        lines.append(f"- origin: {event.origin}")
    if event.total_affected is not None:
        lines.append(f"- total_affected: {event.total_affected}")
    if event.total_deaths is not None:
        lines.append(f"- total_deaths: {event.total_deaths}")
    return "\n".join(lines)


def _build_description(event: FloodEvent) -> str:
    """Compose a predictor-facing description richer than a one-liner.

    Pulls magnitude / damage / casualties from EM-DAT when present so the
    predictor can pick severity-aware analogues; falls back gracefully.
    """
    parts = [
        f"Flood in {event.location or event.country} ({event.iso})"
        f" starting {event.start_date}."
    ]
    if event.origin:
        parts.append(f"Origin: {event.origin}.")
    if event.magnitude is not None and event.magnitude_scale:
        parts.append(f"Magnitude: {event.magnitude} {event.magnitude_scale}.")
    if event.total_affected is not None:
        parts.append(f"Affected: {event.total_affected} people.")
    if event.total_deaths is not None:
        parts.append(f"Deaths: {event.total_deaths}.")
    if event.total_damage_k_usd is not None:
        parts.append(f"Reported damage: {event.total_damage_k_usd:.0f} k USD.")
    return " ".join(parts)


def _parse_judge_response(raw: str) -> dict | None:
    """Extract JSON from a judge response, tolerating markdown code fences."""
    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:
        return None


class Evaluator:
    """Run predict + judge for test events."""

    def __init__(
        self,
        llm_client: LLMClient,
        predictor: SupportsPredict,
        articles_dir: Path | str,
        config: dict,
    ):
        self.llm_client = llm_client
        self.predictor = predictor
        self.articles_dir = Path(articles_dir)
        self.config = config
        eval_cfg = config["evaluation"]
        self.max_articles = eval_cfg["max_articles_per_event"]
        self.max_chars = eval_cfg["max_chars_per_article"]
        self.output_dir = Path(eval_cfg["output_dir"])

        # Pre-load judge prompt + knowledge once so `fingerprint` stays stable
        # across multi-event runs and cache lookups don't re-read files.
        self._judge_prompt = load_prompt_template(
            config["prompts"]["judge_cascade"]
        )
        self._judge_knowledge = load_expert_knowledge(
            config["knowledge"]["judge"]
        )
        self._judge_fingerprint = judge_fingerprint(
            self._judge_prompt, self._judge_knowledge
        )

        # cascade_chains_index.json enriches retrieval_info with country /
        # date / node count so the Streamlit panel and the report can show
        # meaningful provenance rather than bare event ids.
        self._chain_index: dict[str, dict] = self._load_chain_index()

        # Per-run counters, reset each time `run_many` is invoked via the
        # script; kept separate so downstream callers can inspect.
        self.last_cache_status: str | None = None  # "hit" | "rejudge" | "skip_no_news"

    @property
    def judge_fingerprint_value(self) -> str:
        return self._judge_fingerprint

    def _load_chain_index(self) -> dict[str, dict]:
        return load_chain_index(self.config)

    def _enrich_retrieval_info(
        self, event_ids: list[str], similarities: list[float]
    ) -> list[dict]:
        return enrich_retrieval_info(event_ids, similarities, self._chain_index)

    def _load_cached(self, event_id: str) -> EventEvaluation | None:
        path = self.output_dir / f"{event_id}.json"
        if not path.exists():
            return None
        try:
            return EventEvaluation.model_validate_json(path.read_text())
        except (ValidationError, json.JSONDecodeError, OSError) as exc:
            logger.warning(
                "Discarding unreadable cache for %s: %s", event_id, exc
            )
            return None

    def evaluate_event(
        self,
        event: FloodEvent,
        write_cache: bool = False,
        force_rejudge: bool = False,
    ) -> EventEvaluation:
        # 0. Cache check — respect fingerprint unless forced
        if not force_rejudge:
            cached = self._load_cached(event.event_id)
            if (
                cached is not None
                and cached.judge_fingerprint == self._judge_fingerprint
                and cached.judge_fingerprint != ""
            ):
                logger.info(
                    "cache hit (fingerprint match) for %s — skipping predict+judge",
                    event.event_id,
                )
                self.last_cache_status = "hit"
                return cached

        # 1. Predict (with v2-aligned severity and richer description)
        pred = self.predictor.predict(
            country=event.country,
            iso=event.iso,
            location=event.location or event.country,
            event_date=str(event.start_date),
            severity=_infer_severity(event),
            description=_build_description(event),
        )
        chain = pred.predicted_chain

        # 2. Load news
        articles = load_articles(
            event.event_id, self.articles_dir,
            max_articles=self.max_articles,
            max_chars_per_article=self.max_chars,
        )

        retrieval_info = self._enrich_retrieval_info(
            pred.reference_event_ids, pred.similarity_scores
        )

        if not articles:
            self.last_cache_status = "skip_no_news"
            result = EventEvaluation(
                event_id=event.event_id,
                node_evaluations=[],
                missed_cascades=[],
                summary="No news articles available for this event — judge not invoked.",
                news_sources_used=[],
                retrieval_info=retrieval_info,
                predicted_chain=chain,
                evaluated_at=date.today(),
                judge_fingerprint="",
            )
            if write_cache:
                self._write_cache(result)
            return result

        # 3. Judge
        variables = {
            "event_metadata": _format_event_metadata(event),
            "predicted_chain": chain.model_dump_json(indent=2),
            "news_articles": _format_news_for_prompt(articles),
        }
        logger.info(
            "Calling judge LLM for %s (fingerprint=%s)",
            event.event_id, self._judge_fingerprint,
        )
        raw = self.llm_client.call(
            self._judge_prompt, variables, self._judge_knowledge
        )
        parsed = _parse_judge_response(raw)
        self.last_cache_status = "rejudge"

        if parsed is None:
            logger.warning("Judge returned malformed JSON for %s", event.event_id)
            result = EventEvaluation(
                event_id=event.event_id,
                node_evaluations=[],
                missed_cascades=[],
                summary=f"Judge response malformed; raw prefix: {raw[:200]}",
                news_sources_used=[a["url"] for a in articles],
                retrieval_info=retrieval_info,
                predicted_chain=chain,
                evaluated_at=date.today(),
                judge_fingerprint=self._judge_fingerprint,
            )
            if write_cache:
                self._write_cache(result)
            return result

        # 4. Build EventEvaluation — guard against pydantic ValidationError from
        #    judge JSON that is syntactically valid but semantically off-spec
        #    (e.g. an unknown evidence_level string).
        try:
            node_evals = [
                NodeEvaluation(**n) for n in parsed.get("node_evaluations", [])
            ]
            missed = [
                MissedCascade(**m) for m in parsed.get("missed_cascades", [])
            ]
        except ValidationError as exc:
            logger.warning(
                "Judge returned semantically invalid fields for %s: %s",
                event.event_id, exc,
            )
            result = EventEvaluation(
                event_id=event.event_id,
                node_evaluations=[],
                missed_cascades=[],
                summary=f"Judge response malformed (validation error): {exc}"[:500],
                news_sources_used=[a["url"] for a in articles],
                retrieval_info=retrieval_info,
                predicted_chain=chain,
                evaluated_at=date.today(),
                judge_fingerprint=self._judge_fingerprint,
            )
            if write_cache:
                self._write_cache(result)
            return result

        result = EventEvaluation(
            event_id=event.event_id,
            node_evaluations=node_evals,
            missed_cascades=missed,
            summary=parsed.get("summary", ""),
            news_sources_used=[a["url"] for a in articles],
            retrieval_info=retrieval_info,
            predicted_chain=chain,
            evaluated_at=date.today(),
            judge_fingerprint=self._judge_fingerprint,
        )

        if write_cache:
            self._write_cache(result)
        return result

    def _write_cache(self, evaluation: EventEvaluation) -> None:
        self.output_dir.mkdir(parents=True, exist_ok=True)
        path = self.output_dir / f"{evaluation.event_id}.json"
        path.write_text(evaluation.model_dump_json(indent=2), encoding="utf-8")