engine.py

"""
In-process single-note engine for the Clinical Simplifier Web UI.

This is the glue that drives the *unmodified* V7 pipeline modules (copied
verbatim into ./pipeline_core/) for ONE clinical note at a time, in-process,
and returns a structured payload the front-end can render:

  • input with hard-word spans (what the extractor flagged + why)
  • pipeline (RAG) simplification with per-edit spans, each carrying its
    ContextCite source attribution + a natural-language rationale
  • baseline simplification (same model, NO glossary, NO attribution)
  • full metric comparison (everything in evaluate.py EXCEPT the judge panel)

Nothing here re-implements pipeline logic — it imports the V7 stage functions
and calls them directly:

  Novita   → extraction + classification        (extract_common helpers)
  local    → knowledge-base attribution         (attribution.lookup_by_kind)
  Fireworks→ simplification + ContextCite        (simplifier.simplify_one)
  Fireworks→ baseline                            (baseline.simplify_one_baseline)
  Fireworks→ rationalisation                     (rationalize.rationalize_op)
  metrics  → SARI/FKGL/FRE/SMOG/CLI/MEPR/BERT/NLI/citation/bootstrap/Wilcoxon

The V7 folder itself is NEVER imported or written to — we only read its data
files and roberta checkpoints (paths wired through env below).
"""

from __future__ import annotations

import os
import re
import sys
import threading
import time
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional

# --------------------------------------------------------------------------
# Paths. WEB_ROOT/pipeline_core holds the copied V7 modules; V7_DIR holds the
# read-only data + roberta models we reuse (never modified).
# --------------------------------------------------------------------------
WEB_ROOT = Path(__file__).resolve().parent
CORE_DIR = WEB_ROOT / "pipeline_core"
V7_DIR = WEB_ROOT.parent / "v7"
# Data lives in v7/data locally; for deployment we bundle a copy in webapp/data.
DATA_DIR = (V7_DIR / "data") if (V7_DIR / "data").exists() else (WEB_ROOT / "data")


def _read_key(*candidates: Path) -> str:
    for c in candidates:
        p = Path(c)
        if p.exists():
            lines = p.read_text(encoding="utf-8").strip().splitlines()
            if lines and lines[0].strip():
                return lines[0].strip()
    return ""


# Lexical-familiarity metrics: unlike FKGL (syllable-based, blind to jargon),
# these reward replacing rare technical words with common everyday ones.
def _lexical(text: str) -> Dict[str, Any]:
    out: Dict[str, Any] = {"dale_chall": None, "jargon_rate": None, "zipf_mean": None}
    t = (text or "").strip()
    if not t:
        return out
    try:
        import textstat as _ts
        out["dale_chall"] = round(_ts.dale_chall_readability_score(t), 2)
    except Exception:
        pass
    try:
        from wordfreq import zipf_frequency as _zf
        words = re.findall(r"[A-Za-z]+", t.lower())
        if words:
            zs = [_zf(w, "en") for w in words]
            out["zipf_mean"] = round(sum(zs) / len(zs), 2)
            out["jargon_rate"] = round(sum(1 for z in zs if z < 3.0) / len(words), 4)
    except Exception:
        pass
    return out


# --------------------------------------------------------------------------
# Configuration knobs (overridable via environment before first import).
# Mirrors run_local.py's CONFIG block.
# --------------------------------------------------------------------------
CC_NUM_ABLATIONS = os.environ.get("CC_NUM_ABLATIONS", "32")
ENABLE_TORCH_EVALS = os.environ.get("ENABLE_TORCH_EVALS", "1") not in ("0", "false", "False")

# Simplifier model menu. Each carries its provider so that model names which
# might overlap across providers stay unambiguous. All of these return echo
# logprobs, which ContextCite needs. Fireworks ids look like
# "accounts/fireworks/..."; everything else is routed to Together.
FW_BASE = "https://api.fireworks.ai/inference/v1"
TG_BASE = "https://api.together.xyz/v1"
SIMPLIFIER_MODELS = [
    {"key": "fw:deepseek-v4-flash", "id": "accounts/fireworks/models/deepseek-v4-flash",
     "provider": "Firework", "base": FW_BASE, "reasoning": "none",
     "tier": "fast", "label": "DeepSeek V4 Flash", "default": True},
    {"key": "fw:deepseek-v4-pro", "id": "accounts/fireworks/models/deepseek-v4-pro",
     "provider": "Firework", "base": FW_BASE, "reasoning": "none",
     "tier": "best quality", "label": "DeepSeek V4 Pro"},
    {"key": "tg:qwen3.5-9b", "id": "Qwen/Qwen3.5-9B",
     "provider": "Together", "base": TG_BASE, "reasoning": "",
     "tier": "balanced, slower", "label": "Qwen3.5 9B"},
    {"key": "tg:gemma-3n-e4b", "id": "google/gemma-3n-E4B-it",
     "provider": "Together", "base": TG_BASE, "reasoning": "",
     "tier": "small, slower", "label": "Gemma 3n E4B"},
]
MODELS_BY_KEY = {m["key"]: m for m in SIMPLIFIER_MODELS}
DEFAULT_MODEL_KEY = "fw:deepseek-v4-flash"
DEFAULT_FW_MODEL = os.environ.get("FIREWORKS_MODEL", "accounts/fireworks/models/deepseek-v4-flash")

# Classification model menu (Novita).
NOVITA_MODELS_MENU = [
    {"id": "deepseek/deepseek-v4-flash", "label": "DeepSeek V4 Flash", "blurb": "fast, cheap", "default": True},
    {"id": "deepseek/deepseek-v4-pro", "label": "DeepSeek V4 Pro", "blurb": "stronger"},
    {"id": "zai-org/glm-5", "label": "GLM-5", "blurb": "alternate"},
]
# LLM-as-judge panel (Novita) — off by default, costs money.
JUDGE_PANEL = ["deepseek/deepseek-v4-pro", "zai-org/glm-5", "minimax/minimax-m3"]


def _setup_env() -> Dict[str, str]:
    """Populate os.environ exactly the way run_local.set_env() does, so the
    copied stage modules read the right models / data / keys at import time."""
    fireworks_key = os.environ.get("FIREWORKS_API_KEY", "").strip() or _read_key(V7_DIR / "firework.txt")
    novita_key = os.environ.get("NOVITA_API_KEY", "").strip() or _read_key(
        V7_DIR / "novita.txt", V7_DIR / "novita_api")
    together_key = os.environ.get("TOGETHER_API_KEY", "").strip() or _read_key(V7_DIR / "togetherAI_KEY")
    if not fireworks_key:
        raise SystemExit(f"[fatal] no Fireworks key in {V7_DIR / 'firework.txt'}")
    if not novita_key:
        raise SystemExit(f"[fatal] no Novita key in {V7_DIR / 'novita.txt'}")

    env = {
        "PYTHONUTF8": "1",
        "PYTHONIOENCODING": "utf-8",
        # backend
        "LLM_BACKEND": "fireworks",
        # Fireworks (simplifier / baseline / rationalize + ContextCite echo)
        "FIREWORKS_API_KEY": fireworks_key,
        # Together.ai (alternative simplifier provider; only used if picked)
        "TOGETHER_API_KEY": together_key,
        "FIREWORKS_MODEL": os.environ.get(
            "FIREWORKS_MODEL", "accounts/fireworks/models/deepseek-v4-flash"),
        "FIREWORKS_FALLBACKS": ",".join([
            "accounts/fireworks/models/minimax-m3",
            "accounts/fireworks/models/qwen3p7-plus",
        ]),
        "FIREWORKS_NUM_WORKERS": os.environ.get("FIREWORKS_NUM_WORKERS", "8"),
        "FIREWORKS_REASONING_EFFORT": "none",
        # Novita (extraction + classification)
        "NOVITA_API_KEY": novita_key,
        "NOVITA_MODELS": os.environ.get("NOVITA_MODELS", "deepseek/deepseek-v4-flash"),
        "NOVITA_TEMPERATURE": "0.2",
        "NOVITA_MAX_TOKENS": "8192",
        "NOVITA_BASE_URL": os.environ.get("NOVITA_BASE_URL", "https://api.novita.ai/openai"),
        # data + IO paths (READ-ONLY use of v7/data)
        "DATA_DIR": str(DATA_DIR),
        "AOA_FILE": str(DATA_DIR / "en.aoa.csv"),
        "README_FILE": str(DATA_DIR / "readme_exp_good.jsonl"),
        # generation
        "MAX_NEW_TOKENS": "512",
        "SIMPL_MAX_NEW_TOKENS": "512",
        "TEMPERATURE": "0.0",
        "SEED": "1234",
        # simplifier style — single-voice doctor arm, replace tag mode (thesis cfg)
        "SIMPLIFICATION_MODE": "doctor",
        "SIMPLIFIER_TAG_MODE": "replace",
        # ContextCite
        "CC_NUM_ABLATIONS": str(CC_NUM_ABLATIONS),
        "CC_LAMBDA": "0.01",
        "CC_BATCH_SIZE": "8",
        "CC_MAX_SOURCES": "48",
        "CC_PER_OP": "1",
        # no UMLS / scispaCy on the laptop
        "USE_UMLS": "0",
        # heavy (torch) evals — BERTScore + NLI faithfulness + NLI citation
        "USE_BERTSCORE": "1" if ENABLE_TORCH_EVALS else "0",
        "BERTSCORE_MODEL_PATH": str(V7_DIR / "roberta-large") if ENABLE_TORCH_EVALS else "",
        "USE_NLI": "1" if ENABLE_TORCH_EVALS else "0",
        "NLI_MODEL_PATH": str(V7_DIR / "roberta-large-mnli") if ENABLE_TORCH_EVALS else "",
    }
    for k, v in env.items():
        os.environ[k] = v
    return env


# Set env BEFORE importing the stage modules (they read it at import time),
# then put pipeline_core on the path so their bare `import pipeline` etc. work.
_ENV = _setup_env()
if str(CORE_DIR) not in sys.path:
    sys.path.insert(0, str(CORE_DIR))

# --------------------------------------------------------------------------
# Import the copied V7 stage modules (this is the "copy 99%" — same code).
# --------------------------------------------------------------------------
import pipeline as v7_pipeline                       # noqa: E402
import medvocab                                      # noqa: E402
import extract_common                                # noqa: E402
import attribution as v7_attribution                 # noqa: E402
import simplifier as v7_simplifier                   # noqa: E402
import baseline as v7_baseline                       # noqa: E402
import rationalize as v7_rationalize                 # noqa: E402
import evaluate as v7_evaluate                       # noqa: E402
import judge as v7_judge                             # noqa: E402
from llm_utils import get_llm, config_from_env, FireworksConfig, FireworksLLM  # noqa: E402

WORD_RE = medvocab.WORD_RE

# Tag parser for the simplifier's <replace>/<elaborate>/<abbr> output.
_TAG_RE = re.compile(
    r'<(?P<tag>replace|elaborate|abbr)\b(?P<attrs>[^>]*)>(?P<body>.*?)</(?P=tag)>',
    re.IGNORECASE | re.DOTALL)
_ORIG_RE = re.compile(r'orig\s*=\s*"([^"]*)"', re.IGNORECASE)

ProgressCb = Optional[Callable[[str, str, Optional[Dict[str, Any]]], None]]


# ==========================================================================
# Engine singleton — loads glossaries, AoA, and the Fireworks LLM ONCE.
# ==========================================================================
class Engine:
    def __init__(self) -> None:
        self._loaded = False
        self._lock = threading.Lock()
        self._run_lock = threading.Lock()
        self.sources: Dict[str, Any] = {}
        self.aoa_map: Dict[str, Any] = {}
        self.llm = None
        self._llms: Dict[str, Any] = {}   # fireworks model id -> loaded llm
        self.load_status: Dict[str, Any] = {"loaded": False, "detail": "not started"}

    # ---- one-time warmup --------------------------------------------------
    def ensure_loaded(self, progress: ProgressCb = None) -> None:
        if self._loaded:
            return
        with self._lock:
            if self._loaded:
                return
            self._emit(progress, "load", "running", {"detail": "loading glossaries"})
            self._load_sources()
            self._emit(progress, "load", "running", {"detail": "loading age-of-acquisition table"})
            self.aoa_map = v7_pipeline.load_aoa(v7_pipeline.AOA_FILE)
            self._emit(progress, "load", "running", {"detail": "probing Fireworks echo backend"})
            self.llm = get_llm(config_from_env())
            self.llm.load()
            self._llms[DEFAULT_MODEL_KEY] = self.llm
            self._loaded = True
            self.load_status = {"loaded": True, "detail": "ready",
                                "model": getattr(self.llm, "chosen_model", None)}
            self._emit(progress, "load", "done", self.load_status)

    def get_llm_for(self, key: Optional[str], progress: ProgressCb = None):
        """Return a loaded LLM for the requested model KEY (provider-aware),
        caching one per key. Falls back to the default if key is empty/unknown.
        Fireworks and Together are both OpenAI-compatible /completions backends;
        we just swap base_url + api_key + reasoning per the menu entry."""
        m = MODELS_BY_KEY.get(key or "")
        if not m:
            return self.llm
        if m["key"] in self._llms:
            return self._llms[m["key"]]
        if m["provider"] == "Together":
            api_key = os.environ.get("TOGETHER_API_KEY", "").strip()
            if not api_key:
                raise RuntimeError("No Together API key. Set the TOGETHER_API_KEY "
                                   "env var (or add v7/togetherAI_KEY) to use "
                                   f"{m['label']}.")
        else:
            api_key = os.environ.get("FIREWORKS_API_KEY", "").strip()
        self._emit(progress, "load", "running",
                   {"detail": f"probing {m['label']} on {m['provider']}"})
        # Build a dedicated LLM instance (NOT the get_llm singleton, which would
        # hand back the already-loaded default). Fireworks + Together are both
        # OpenAI-compatible /completions with echo, so only base_url/key/reasoning
        # change. No cross-provider fallback.
        _is_tg = m["provider"] == "Together"
        cfg = FireworksConfig(
            model_name=m["id"], fallback_models=[], base_url=m["base"],
            api_key=api_key, reasoning_effort=m["reasoning"],
            max_new_tokens=int(os.environ.get("SIMPL_MAX_NEW_TOKENS", "512")),
            temperature=float(os.environ.get("TEMPERATURE", "0.0")),
            seed=int(os.environ.get("SEED", "1234")),
            # Together's serverless endpoints 503 under parallel load, so use
            # low concurrency + more patient retries there; Fireworks handles 8.
            num_workers=2 if _is_tg else int(os.environ.get("FIREWORKS_NUM_WORKERS", "8")),
            max_retries=7 if _is_tg else 6,
            timeout_s=180.0 if _is_tg else 120.0,
        )
        llm = FireworksLLM(cfg)
        llm.load()
        self._llms[m["key"]] = llm
        return llm

    def _load_sources(self) -> None:
        s: Dict[str, Any] = {}
        s["readme_by_key"], s["readme_rows"] = medvocab.load_readme(
            Path(os.environ["README_FILE"]))
        s["nih"] = medvocab.load_nih(DATA_DIR / "nih.json")
        s["medlane"] = medvocab.load_medlane(
            DATA_DIR / "medlane_abbreviation_dictionary_flat.csv")
        s["dictionary"] = medvocab.load_dictionary(DATA_DIR / "dictonary.csv")
        s["iowa"] = medvocab.load_glossary_json(DATA_DIR / "Iowa.json")
        s["michigan"] = medvocab.load_glossary_json(DATA_DIR / "michigan_plmd.json")
        s["justplainclear"] = medvocab.load_glossary_json(
            DATA_DIR / "justplainclear_en_clean.json")
        s["thesaurus"] = medvocab.load_glossary_json(DATA_DIR / "thesarus.json")
        self.sources = s

    @staticmethod
    def _emit(progress: ProgressCb, stage: str, status: str,
              data: Optional[Dict[str, Any]] = None) -> None:
        if progress:
            try:
                progress(stage, status, data)
            except Exception:
                pass

    # ---- stage 1: extraction + classification (Novita) --------------------
    def extract(self, note_text: str, model: Optional[str] = None) -> List[Dict[str, Any]]:
        from openai import OpenAI
        client = OpenAI(api_key=os.environ["NOVITA_API_KEY"],
                        base_url=os.environ["NOVITA_BASE_URL"], timeout=120)
        user_msg = extract_common.USER_TEMPLATE.format(note=note_text)
        budget = int(os.environ.get("NOVITA_MAX_TOKENS", "8192"))
        novita_model = (model or os.environ["NOVITA_MODELS"].split(",")[0]).strip()
        resp = client.chat.completions.create(
            model=novita_model,
            messages=[{"role": "system", "content": extract_common.SYSTEM_PROMPT},
                      {"role": "user", "content": user_msg}],
            temperature=float(os.environ.get("NOVITA_TEMPERATURE", "0.2")),
            max_tokens=budget,
        )
        raw = resp.choices[0].message.content or ""
        parsed = extract_common.parse_json_strict(extract_common.strip_json_fences(raw))
        if not isinstance(parsed, list):
            # one retry with a bigger budget if it looks truncated
            parsed = []
        parsed = extract_common.dedupe_records(parsed)
        parsed, _audit = extract_common.sanitize_records(parsed, note_text)
        final_items, _errors = extract_common.normalize_records(parsed, self.aoa_map)
        return final_items

    # ---- stage 2: knowledge-base attribution (local glossaries) -----------
    def attribute(self, final_items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        split_on = os.environ.get("ATTRIBUTION_SPLIT_UNCOVERED", "1") not in ("0", "false", "False")
        hard_records: List[Dict[str, Any]] = []
        for it in final_items:
            text = it["text"]
            kind = it.get("kind", "word")
            attribs = v7_attribution.lookup_by_kind(text, kind, self.sources)
            rec = {**it, "attributions": attribs}
            hard_records.append(rec)
            # hierarchical splitting when a multi-word phrase has no whole def
            if (split_on and not attribs and len(WORD_RE.findall(text)) >= 2):
                try:
                    splits = v7_attribution.split_uncovered_phrase(text, self.sources)
                except Exception:
                    splits = []
                for sp in splits:
                    hard_records.append({
                        "text": sp["text"], "kind": sp["kind"],
                        "category": sp["category"], "action": sp["action"],
                        "reason": it.get("reason", ""),
                        "features": it.get("features", {}),
                        "from_chunk": text, "origin": "split",
                        "attributions": sp["attributions"],
                    })
        return hard_records

    # ---- full single-note run --------------------------------------------
    def run(self, note_text: str, progress: ProgressCb = None,
            ablations: Optional[int] = None, fw_model: Optional[str] = None,
            novita_model: Optional[str] = None, reference: Optional[str] = None,
            run_judge: bool = False) -> Dict[str, Any]:
        note_text = (note_text or "").strip()
        if not note_text:
            raise ValueError("empty note")
        note_text = v7_pipeline.strip_mimic_placeholders(note_text)
        reference = (reference or "").strip()

        with self._run_lock:   # the Fireworks echo backend isn't re-entrant per run
            self.ensure_loaded(progress)
            llm = self.get_llm_for(fw_model, progress)
            _m = MODELS_BY_KEY.get(fw_model or "") or MODELS_BY_KEY.get(DEFAULT_MODEL_KEY)
            # Reasoning models (gpt-oss) spend tokens on an analysis channel
            # before the answer; give them room so the harmony "final" channel is
            # reached and the raw thinking doesn't leak into the output.
            _tok = 2048 if "gpt-oss" in _m["id"] else 512
            v7_simplifier.SIMPL_MAX_NEW_TOKENS = _tok
            v7_baseline.SIMPL_MAX_NEW_TOKENS = _tok
            # Together's per-call latency is high, so cap ContextCite ablations to
            # keep a run to a few minutes instead of appearing frozen.
            _abl = int(ablations) if ablations else int(os.environ.get("CC_NUM_ABLATIONS", "32"))
            if _m["provider"] == "Together":
                _abl = min(_abl, 8)
            v7_simplifier.NUM_ABLATIONS = _abl
            os.environ["CC_NUM_ABLATIONS"] = str(_abl)
            t0 = time.time()

            # multi-sentence: process every sentence, render as one note.
            sents = v7_evaluate._split_sentences(note_text) or [note_text]

            all_items: List[Dict[str, Any]] = []
            all_records: List[Dict[str, Any]] = []
            all_sources: List[Dict[str, Any]] = []
            all_ops: List[Dict[str, Any]] = []
            all_rats: List[Dict[str, Any]] = []
            pipe_tagged_parts, base_tagged_parts = [], []
            pipe_parts, base_parts = [], []

            self._emit(progress, "extract", "running",
                       {"detail": f"Novita - reading {len(sents)} sentence(s)"})
            n_terms = 0
            for si, sent in enumerate(sents):
                final_items = self.extract(sent, novita_model)
                hard_records = self.attribute(final_items)
                attrib_note = {"note_idx": si + 1, "note_text": sent,
                               "hard": hard_records, "border": []}
                classify_rows = [{"term": it["text"], "category": it.get("category", ""),
                                  "action": it.get("action", ""), "reason": it.get("reason", "")}
                                 for it in final_items]
                sources = v7_simplifier.build_sources_for_note(attrib_note, classify_rows)
                hard_terms = v7_simplifier.collect_hard_terms_for_note(attrib_note)
                n_terms += len(final_items)
                self._emit(progress, "attribution", "running",
                           {"detail": f"sentence {si+1}/{len(sents)} - lay definitions"})
                simpl = v7_simplifier.simplify_one(llm, sent, sources, hard_terms, verbose=False)
                base = v7_baseline.simplify_one_baseline(llm, sent, hard_terms, verbose=False)
                ops = simpl.get("doctor_operations") or []
                rats = self._rationalize(ops, sources)
                all_items += final_items
                all_records += hard_records
                all_sources += sources
                all_ops += ops
                all_rats += rats
                pipe_tagged_parts.append(simpl.get("doctor_tagged") or simpl.get("doctor_simplified", ""))
                base_tagged_parts.append(base.get("doctor_tagged") or base.get("doctor_simplified", ""))
                pipe_parts.append(simpl.get("doctor_simplified", ""))
                base_parts.append(base.get("doctor_simplified", ""))
            self._emit(progress, "extract", "done", {"count": n_terms,
                       "terms": [it["text"] for it in all_items]})
            n_defs = sum(len(h.get("attributions") or []) for h in all_records)
            self._emit(progress, "attribution", "done", {"definitions": n_defs})
            self._emit(progress, "simplify", "done", {"edits": len(all_ops),
                       "ablations": int(os.environ.get("CC_NUM_ABLATIONS", "32"))})
            self._emit(progress, "baseline", "done", None)
            self._emit(progress, "rationalize", "done", {"count": len(all_rats)})

            simpl = {"doctor_simplified": " ".join(p for p in pipe_parts if p).strip(),
                     "doctor_tagged": " ".join(p for p in pipe_tagged_parts if p).strip()}
            base = {"doctor_simplified": " ".join(p for p in base_parts if p).strip(),
                    "doctor_tagged": " ".join(p for p in base_tagged_parts if p).strip()}
            audit = self._audit_rationales(all_ops, all_rats)

            # evaluation (all metrics, no judge) + optional gold reference
            self._emit(progress, "evaluate", "running",
                       {"detail": "SARI/FKGL/FRE/SMOG/MEPR/NLI/citation/safety"})
            metrics = self._evaluate(note_text, simpl, base, all_records, all_ops, reference)
            self._emit(progress, "evaluate", "done", None)

            judge = None
            if run_judge:
                self._emit(progress, "judge", "running", {"detail": "Novita panel scoring"})
                judge = self._judge(note_text, simpl["doctor_simplified"], base["doctor_simplified"])
                self._emit(progress, "judge", "done", None)

            payload = self._assemble(note_text, all_items, all_records, all_sources,
                                     simpl, base, all_ops, all_rats, audit, metrics)
            payload["benchmark"] = bool(reference)
            payload["reference"] = reference
            payload["judge"] = judge
            payload["elapsed_s"] = round(time.time() - t0, 1)
            _mm = MODELS_BY_KEY.get(fw_model or "") or MODELS_BY_KEY.get(DEFAULT_MODEL_KEY)
            payload["model"] = {
                "fireworks": getattr(llm, "chosen_model", None),
                "provider": _mm["provider"], "label": _mm["label"],
                "novita": (novita_model or os.environ["NOVITA_MODELS"].split(",")[0]).strip(),
                "ablations": int(os.environ.get("CC_NUM_ABLATIONS", "32")),
            }
            self._emit(progress, "complete", "done", None)
            return payload

    # ---- rationalisation ---------------------------------------------------
    def _rationalize(self, doc_ops: List[Dict[str, Any]],
                     sources: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        out: List[Dict[str, Any]] = []
        prev: List[str] = []
        for op in doc_ops:
            try:
                r = v7_rationalize.rationalize_op(self.llm, op, self.aoa_map, sources, prev)
            except Exception as e:
                r = {"op_kind": op.get("type", "replace"),
                     "src_text": op.get("src_text", ""),
                     "tgt_text": op.get("tgt_text", ""),
                     "ground_truth": {}, "rationale_text": f"(rationale unavailable: {e})"}
            out.append(r)
            if r.get("rationale_text"):
                prev.append(r["rationale_text"])
        return out

    # ---- hallucination audit (reuse hallucination_check if importable) -----
    def _audit_rationales(self, doc_ops, rationales) -> Dict[str, Any]:
        try:
            import hallucination_check as hc
        except Exception:
            hc = None
        total = len(rationales)
        if not total:
            return {"n": 0}
        grounded = 0
        for r in rationales:
            gt = r.get("ground_truth") or {}
            txt = (r.get("rationale_text") or "").lower()
            # a rationale is "grounded" if it quotes the frequency score and
            # references the cited source label (cheap structural check that
            # mirrors hallucination_check.py's slot/source coverage)
            sf = gt.get("src_features") or {}
            zipf_ok = (sf.get("zipf") is None) or ("frequency score" in txt or "zipf" in txt)
            src_label = (gt.get("source_short") or "").lower()
            src_ok = (not src_label) or (src_label in txt) or (gt.get("source_label", "").split()[0].lower() in txt if gt.get("source_label") else False)
            if zipf_ok and src_ok:
                grounded += 1
        return {"n": total, "grounded": grounded,
                "grounded_rate": round(grounded / total, 3) if total else None,
                "module": bool(hc)}

    # ---- LLM-as-judge panel (Novita) -- optional, costs money --------------
    def _judge(self, inp: str, pipe: str, base: str) -> Dict[str, Any]:
        from openai import OpenAI
        import random
        client = OpenAI(api_key=os.environ["NOVITA_API_KEY"],
                        base_url=os.environ["NOVITA_BASE_URL"], timeout=180)
        # randomise which arm is A vs B to remove position bias
        pipe_is_a = random.random() < 0.5
        a, b = (pipe, base) if pipe_is_a else (base, pipe)
        user = v7_judge.USER_TEMPLATE.format(inp=inp, a=a, b=b)
        panel, wins = [], {"pipeline": 0, "baseline": 0, "tie": 0}
        for model in JUDGE_PANEL:
            try:
                resp = client.chat.completions.create(
                    model=model, temperature=0.0, max_tokens=2048,
                    messages=[{"role": "system", "content": v7_judge.JUDGE_SYSTEM},
                              {"role": "user", "content": user}])
                parsed = v7_judge._parse_judge_json(resp.choices[0].message.content or "")
                norm, err = v7_judge._validate_judge(parsed) if parsed else (None, "no json")
                if not norm:
                    panel.append({"model": model, "error": err}); continue
                pipe_s = norm["A"] if pipe_is_a else norm["B"]
                base_s = norm["B"] if pipe_is_a else norm["A"]
                w = norm["winner"]
                who = "tie" if w == "tie" else (("pipeline" if w == "a" else "baseline")
                                                if pipe_is_a else ("baseline" if w == "a" else "pipeline"))
                wins[who] += 1
                panel.append({"model": model.split("/")[-1], "pipeline": pipe_s,
                              "baseline": base_s, "winner": who, "reason": norm["reason"]})
            except Exception as e:
                panel.append({"model": model.split("/")[-1], "error": str(e)[:120]})
        def avg(arm, ax):
            xs = [p[arm][ax] for p in panel if isinstance(p.get(arm), dict)]
            return round(sum(xs) / len(xs), 2) if xs else None
        axes = ["faithfulness", "simplicity", "fluency"]
        return {"panel": panel, "wins": wins,
                "pipeline": {ax: avg("pipeline", ax) for ax in axes},
                "baseline": {ax: avg("baseline", ax) for ax in axes}}

    # ---- evaluation (everything in evaluate.py except the judge) -----------
    def _evaluate(self, note_text: str, simpl: Dict[str, Any], base: Dict[str, Any],
                  hard_records: List[Dict[str, Any]],
                  doc_ops: List[Dict[str, Any]],
                  reference: str = "") -> Dict[str, Any]:
        ev = v7_evaluate
        pipe_out = simpl.get("doctor_simplified", "") or ""
        base_out = base.get("doctor_simplified", "") or ""
        refs: List[str] = [reference] if reference and reference.strip() else []

        in_metrics = ev.readability(note_text)

        def arm_metrics(sys_out: str, ops: Optional[List[Dict[str, Any]]]) -> Dict[str, Any]:
            m = ev.metrics_one(sys_out, note_text, refs)
            # readability deltas vs the original
            if m.get("fkgl") is not None and in_metrics.get("fkgl") is not None:
                m["fkgl_drop"] = round(in_metrics["fkgl"] - m["fkgl"], 2)
            if m.get("fre") is not None and in_metrics.get("fre") is not None:
                m["fre_gain"] = round(m["fre"] - in_metrics["fre"], 2)
            if m.get("smog") is not None and in_metrics.get("smog") is not None:
                m["smog_drop"] = round(in_metrics["smog"] - m["smog"], 2)
            if m.get("coleman_liau") is not None and in_metrics.get("coleman_liau") is not None:
                m["coleman_liau_drop"] = round(in_metrics["coleman_liau"] - m["coleman_liau"], 2)
            # MEPR — medical entity preservation (reference-free)
            m["mepr"] = ev.mepr_for_output(sys_out, hard_records)
            m["mepr_breakdown"] = ev.mepr_breakdown(sys_out, hard_records)
            # NLI faithfulness (reference-free)
            try:
                m["nli_faith"] = ev.nli_faithfulness(note_text, sys_out)
            except Exception:
                m["nli_faith"] = None
            # clinical-safety slot preservation (reference-free)
            try:
                cs = ev.clinical_safety_metrics(note_text, sys_out)
                if cs:
                    m.update(cs)
            except Exception:
                pass
            # lexical familiarity (captures jargon removal FKGL is blind to)
            m.update(_lexical(sys_out))
            # reference-based edit/replacement F1 (benchmark mode only)
            if refs:
                try:
                    er = ev.edit_replacement_metrics(note_text, refs[0], sys_out, hard_records)
                    if er:
                        m.update(er)
                except Exception:
                    pass
            # citation precision/recall/F1 (pipeline only; baseline has no sources)
            if ops:
                try:
                    cit = ev.citation_metrics_for_ops(ops)
                except Exception:
                    cit = None
                if cit:
                    m["citation_precision"] = cit["precision"]
                    m["citation_recall"] = cit["recall"]
                    m["citation_f1"] = cit["f1"]
                    m["citation_n_ops"] = cit["n_ops"]
            else:
                m["citation_precision"] = 0.0
                m["citation_recall"] = 0.0
                m["citation_f1"] = 0.0
            # BERTScore only if a reference is present
            if refs:
                try:
                    m["bertscore_f1"] = ev._bertscore_one(sys_out, refs[0])
                except Exception:
                    m["bertscore_f1"] = None
            return m

        pipe_m = arm_metrics(pipe_out, doc_ops)
        base_m = arm_metrics(base_out, None)

        # ContextCite intrinsic diagnostics (aggregate over pipeline edits)
        lds_vals, drop_vals = [], []
        for op in doc_ops:
            cc = op.get("contextcite") or {}
            if isinstance(cc.get("lds"), (int, float)):
                lds_vals.append(float(cc["lds"]))
            lp_o, lp_e = cc.get("original_log_prob"), cc.get("empty_log_prob")
            if isinstance(lp_o, (int, float)) and isinstance(lp_e, (int, float)):
                drop_vals.append(abs(float(lp_o) - float(lp_e)))
        cc_diag = {
            "mean_lds": round(sum(lds_vals) / len(lds_vals), 3) if lds_vals else None,
            "mean_source_gap": round(sum(drop_vals) / len(drop_vals), 3) if drop_vals else None,
            "n_edits": len(doc_ops),
        }

        return {"input": in_metrics, "pipeline": pipe_m, "baseline": base_m,
                "contextcite": cc_diag, "has_reference": bool(refs)}

    # ---- assemble the front-end payload -----------------------------------
    def _assemble(self, note_text, final_items, hard_records, sources,
                  simpl, base, doc_ops, rationales, audit, metrics) -> Dict[str, Any]:
        input_spans = _input_highlight_spans(note_text, final_items)
        rat_by_key = {(_norm(r.get("src_text")), _norm(r.get("tgt_text"))): r
                      for r in rationales}
        op_by_key = {(_norm(op.get("src_text")), _norm(op.get("tgt_text"))): op
                     for op in doc_ops}

        pipe_segments = _segment_tagged(
            simpl.get("doctor_tagged") or simpl.get("doctor_simplified", ""),
            rat_by_key, op_by_key)
        base_segments = _segment_tagged(
            base.get("doctor_tagged") or base.get("doctor_simplified", ""),
            {}, {})

        return {
            "ok": True,
            "input": {"text": note_text, "spans": input_spans},
            "pipeline": {
                "text": simpl.get("doctor_simplified", ""),
                "segments": pipe_segments,
                "n_edits": len(doc_ops),
            },
            "baseline": {
                "text": base.get("doctor_simplified", ""),
                "segments": base_segments,
            },
            "sources_used": [
                {"id": s["id"], "term": s["term"], "source": s["source"],
                 "text": s["text"], "passage": s["passage"]}
                for s in sources
            ],
            "hard_terms": [
                {"text": it["text"], "kind": it.get("kind"),
                 "category": it.get("category"), "action": it.get("action"),
                 "reason": it.get("reason", ""),
                 "features": it.get("features", {}),
                 "n_sources": len([h for h in hard_records
                                   if _norm(h.get("text")) == _norm(it["text"])
                                   for _ in (h.get("attributions") or [])])}
                for it in final_items
            ],
            "rationale_audit": audit,
            "metrics": metrics,
        }


# ==========================================================================
# Rendering helpers
# ==========================================================================
def _norm(s: Optional[str]) -> str:
    return re.sub(r"\s+", " ", (s or "")).strip().lower()


def _input_highlight_spans(note_text: str,
                           final_items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """Find char spans of each flagged hard term in the original note. Longest
    terms first so a phrase wins over its component words; non-overlapping."""
    low = note_text.lower()
    claimed: List[tuple] = []
    spans: List[Dict[str, Any]] = []
    items = sorted(final_items, key=lambda it: -len(it.get("text", "")))
    for it in items:
        term = it.get("text", "")
        if not term:
            continue
        start = 0
        tl = term.lower()
        while True:
            idx = low.find(tl, start)
            if idx == -1:
                break
            end = idx + len(term)
            if not any(idx < ce and end > cs for cs, ce in claimed):
                claimed.append((idx, end))
                f = it.get("features", {}) or {}
                spans.append({
                    "start": idx, "end": end, "text": note_text[idx:end],
                    "category": it.get("category", ""), "action": it.get("action", ""),
                    "reason": it.get("reason", ""),
                    "zipf": f.get("zipf"), "aoa": f.get("aoa"),
                    "syllables": f.get("syllables"),
                })
            start = idx + len(term)
    spans.sort(key=lambda s: s["start"])
    return spans


def _segment_tagged(tagged: str, rat_by_key: Dict, op_by_key: Dict) -> List[Dict[str, Any]]:
    """Split a <replace>/<elaborate>/<abbr>-tagged string into ordered segments
    so the front-end can render plain text + highlighted edits. Each edit
    segment carries its rationale + ContextCite source (when available)."""
    segments: List[Dict[str, Any]] = []
    last = 0
    for m in _TAG_RE.finditer(tagged):
        if m.start() > last:
            segments.append({"type": "text", "text": tagged[last:m.start()]})
        body = m.group("body")
        orig_m = _ORIG_RE.search(m.group("attrs") or "")
        src = (orig_m.group(1).strip() if orig_m else "")
        tgt = re.sub(r"<[^>]+>", "", body).strip()
        seg = {"type": "edit", "text": tgt, "src": src,
               "tag": m.group("tag").lower()}
        key = (_norm(src), _norm(tgt))
        r = rat_by_key.get(key)
        op = op_by_key.get(key)
        if r:
            seg["rationale"] = r.get("rationale_text", "")
            gt = r.get("ground_truth") or {}
            seg["source_label"] = gt.get("source_label")
            seg["source_short"] = gt.get("source_short")
            seg["source_passage"] = gt.get("source_passage")
            seg["attribution_weight"] = gt.get("top_attribution_weight")
            seg["src_features"] = gt.get("src_features")
            seg["tgt_features"] = gt.get("tgt_features")
            # The rationale's source_label is "source (term)" from pick_top_source
            # (the edit's ACTUAL glossary source) — use its term, not ContextCite's
            # noisy top-ranked one, so the tooltip source matches the edit.
            _lbl = gt.get("source_label") or ""
            _mt = re.search(r"\((.*)\)\s*$", _lbl)
            if _mt:
                seg["source_term"] = _mt.group(1).strip()
        if op:
            cc = op.get("contextcite") or {}
            seg["cc_lds"] = cc.get("lds")
            lp_o, lp_e = cc.get("original_log_prob"), cc.get("empty_log_prob")
            if isinstance(lp_o, (int, float)) and isinstance(lp_e, (int, float)):
                seg["cc_source_gap"] = round(abs(float(lp_o) - float(lp_e)), 3)
            ranked = []
            for ts in (op.get("top_sources") or [])[:5]:
                ranked.append({
                    "source": ts.get("source"),
                    "term": ts.get("term_in_source") or ts.get("term"),
                    "weight": ts.get("weight"),
                    "passage": (ts.get("passage") or ts.get("text") or "")[:200],
                })
            seg["ranked_sources"] = ranked
            # Only fall back to ContextCite's source if the rationale gave none.
            if "source_label" not in seg and ranked:
                seg["source_short"] = ranked[0]["source"]
                seg["source_passage"] = ranked[0]["passage"]
                if not seg.get("source_term"):
                    seg["source_term"] = ranked[0]["term"]
        segments.append(seg)
        last = m.end()
    if last < len(tagged):
        segments.append({"type": "text", "text": tagged[last:]})
    if not segments:
        segments.append({"type": "text", "text": tagged})
    return segments


# Module-level singleton
ENGINE = Engine()