inspect_rag_ensemble.py

"""
inspect_rag_ensemble.py — Ensemble RAG inspector (no LLM) v4.

Three-signal ensemble:
  1. Dense retrieval   — FAISS + E5 embeddings (HF API)  → finds right topic
  2. Sparse retrieval  — TF-IDF over protocols (ICD-augmented) → keyword match
  3. ICD code re-ranker — TF-IDF over ALL ICD-10 descriptions → matches
     symptoms directly to disease names, resolving category↔subcode ambiguity

Plus:
  - Abbreviation expansion (abbr2term.yaml)
  - Medical synonym expansion (106 patterns)
  - Prefix matching for evaluation
  - Subcode expansion

Run:
    uv run python inspect_rag_ensemble.py --n 10 --seed 42
"""

import argparse
import json
import logging
import random
import re
from pathlib import Path

import numpy as np
import yaml
from sklearn.feature_extraction.text import TfidfVectorizer

from rag_hf import FAISSRetrieverHF, RetrievalResult

logging.basicConfig(level=logging.ERROR)
logging.getLogger("rag_hf").setLevel(logging.INFO)

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
ROOT = Path(__file__).parent
ABBR_PATH   = ROOT / "abbr2term.yaml"
MKB_PATH    = ROOT / "mkb2descr.yaml"
CORPUS_PATH = ROOT / "data" / "corpus" / "corpus" / "protocols_corpus_enriched.jsonl"
TEST_DIR    = ROOT / "data" / "test_set"

DENSE_WEIGHT   = 0.50
SPARSE_WEIGHT  = 0.25
ICD_RE_WEIGHT  = 0.25   # ICD description re-ranker
TOP_K_DENSE    = 15
TOP_K_FINAL    = 5

# ---------------------------------------------------------------------------
# Medical synonym dictionary  (colloquial patient → formal medical terms)
# ---------------------------------------------------------------------------
MEDICAL_SYNONYMS: dict[str, list[str]] = {
    "болит": ["боль", "болевой синдром"],
    "ноет": ["ноющая боль"],
    "стреляет": ["стреляющая боль", "невралгия"],
    "печёт": ["жжение"],
    "жжет": ["жжение"],
    "давит": ["давящая боль", "сдавление"],
    "тянет": ["тянущая боль"],
    "немеет": ["онемение", "парестезия"],
    "головная боль": ["цефалгия"],
    "кружится голова": ["головокружение", "вертиго"],
    "задышка": ["одышка", "диспноэ"],
    "одышка": ["диспноэ", "дыхательная недостаточность"],
    "не хватает воздуха": ["одышка", "диспноэ"],
    "кашель": ["кашлевой синдром"],
    "кашляю": ["кашель", "кашлевой синдром"],
    "мокрота": ["экспекторация"],
    "хрипы": ["стридор"],
    "сердцебиение": ["тахикардия", "пальпитация"],
    "давление": ["гипертензия", "гипотензия"],
    "боль в груди": ["торакалгия", "стенокардия"],
    "тошнота": ["наузея"],
    "тошнит": ["тошнота", "наузея"],
    "рвота": ["эмезис"],
    "вырвало": ["рвота"],
    "понос": ["диарея"],
    "запор": ["констипация"],
    "живот болит": ["абдоминальная боль"],
    "изжога": ["гастроэзофагеальный рефлюкс"],
    "вздутие": ["метеоризм"],
    "аппетит плохой": ["снижение аппетита", "анорексия"],
    "нет аппетита": ["анорексия"],
    "сыпь": ["экзантема", "дерматит"],
    "зуд": ["пруритус"],
    "чешется": ["зуд", "пруритус"],
    "покраснение": ["гиперемия", "эритема"],
    "отек": ["эдема"],
    "опухло": ["отек", "эдема"],
    "синяк": ["гематома"],
    "температура": ["лихорадка", "гипертермия"],
    "жар": ["лихорадка", "гипертермия"],
    "озноб": ["фебрильная реакция"],
    "знобит": ["озноб"],
    "потливость": ["гипергидроз"],
    "горло болит": ["фарингит", "тонзиллит"],
    "ангина": ["тонзиллит", "острый тонзиллит"],
    "насморк": ["ринит", "ринорея"],
    "спина болит": ["дорсалгия", "люмбалгия"],
    "поясница": ["люмбалгия"],
    "сустав": ["артралгия", "артрит"],
    "суставы": ["полиартралгия"],
    "перелом": ["фрактура"],
    "ушиб": ["контузия"],
    "моча темная": ["гематурия"],
    "мало писает": ["олигурия"],
    "больно писать": ["дизурия"],
    "часто писает": ["поллакиурия"],
    "кровь в моче": ["гематурия"],
    "слабость": ["астения", "утомляемость"],
    "похудел": ["потеря массы тела"],
    "похудела": ["потеря массы тела"],
    "обморок": ["синкопе"],
    "судороги": ["конвульсии", "судорожный синдром"],
    "тревога": ["тревожное расстройство"],
    "депрессия": ["депрессивное расстройство"],
    "пьет": ["алкогольная зависимость", "употребление алкоголя"],
    "запой": ["алкогольная зависимость", "абстинентное состояние"],
    "упал": ["травма", "падение"],
    "упала": ["травма", "падение"],
    "удар": ["ушиб", "контузия"],
    "ожог": ["термическое повреждение", "комбустия"],
    "отравление": ["интоксикация", "токсическое действие"],
    "выпил": ["отравление", "интоксикация"],
    "диабет": ["сахарный диабет"],
    "жажда": ["полидипсия"],
    "аллергия": ["аллергическая реакция", "гиперчувствительность"],
    "крапивница": ["уртикария"],
}


def load_abbreviations(path: Path) -> dict[str, str]:
    if not path.exists():
        return {}
    with open(path, "r", encoding="utf-8") as f:
        data = yaml.safe_load(f)
    return {str(k).strip(): str(v).strip() for k, v in data.items()} if data else {}


def load_icd_descriptions(path: Path) -> dict[str, str]:
    if not path.exists():
        return {}
    with open(path, "r", encoding="utf-8") as f:
        data = yaml.safe_load(f)
    return {str(k).strip(): str(v).strip() for k, v in data.items()} if data else {}


def expand_abbreviations(text: str, abbr_map: dict[str, str]) -> str:
    if not abbr_map:
        return text
    words = re.split(r"(\s+)", text)
    expanded = []
    for w in words:
        clean = w.strip(".,;:!?()[]«»\"'")
        if clean.upper() in abbr_map:
            expanded.append(f"{w} ({abbr_map[clean.upper()]})")
        else:
            expanded.append(w)
    return "".join(expanded)


def expand_medical_synonyms(text: str) -> str:
    text_lower = text.lower()
    found: list[str] = []
    for phrase, synonyms in sorted(MEDICAL_SYNONYMS.items(), key=lambda x: -len(x[0])):
        if phrase.lower() in text_lower:
            found.extend(synonyms)
    if found:
        seen: set[str] = set()
        unique = [t for t in found if t.lower() not in seen and not seen.add(t.lower())]
        return text + " " + " ".join(unique)
    return text


# ---------------------------------------------------------------------------
# ICD code utilities
# ---------------------------------------------------------------------------
def build_subcode_map(icd_descr: dict[str, str]) -> dict[str, list[str]]:
    cat_to_sub: dict[str, list[str]] = {}
    for code in icd_descr:
        if "." in code:
            cat = code.split(".")[0]
            cat_to_sub.setdefault(cat, []).append(code)
    return cat_to_sub


def expand_codes_to_subcodes(codes: list[str], subcode_map: dict[str, list[str]]) -> list[str]:
    expanded, seen = [], set()
    for code in codes:
        if code in seen:
            continue
        seen.add(code)
        expanded.append(code)
        if "." not in code and code in subcode_map:
            for sub in subcode_map[code]:
                if sub not in seen:
                    seen.add(sub)
                    expanded.append(sub)
    return expanded


def prefix_match(predicted: str, valid_set: set[str]) -> bool:
    if predicted in valid_set:
        return True
    for v in valid_set:
        if v.startswith(predicted + ".") or predicted.startswith(v + "."):
            return True
    return False


def any_prefix_match(codes: list[str], valid_set: set[str]) -> bool:
    return any(prefix_match(c, valid_set) for c in codes)


# ---------------------------------------------------------------------------
# ICD Description Re-ranker
# ---------------------------------------------------------------------------
class ICDReRanker:
    """
    TF-IDF index over ALL ICD-10 code descriptions.
    Given patient symptoms, finds which ICD codes match best by description.
    This adds a third signal to the ensemble that directly resolves
    the category→subcode problem.
    """

    def __init__(self):
        self.icd_codes: list[str] = []
        self.icd_texts: list[str] = []
        self.vectorizer: TfidfVectorizer | None = None
        self.tfidf_matrix = None

    def load(self, icd_descr: dict[str, str]) -> None:
        print("  [ICD-Reranker] Building index over ICD descriptions …")
        # Only include actual disease codes (skip range headers like "A00-A09")
        for code, desc in icd_descr.items():
            if "-" in code or len(code) < 3:
                continue
            self.icd_codes.append(code)
            self.icd_texts.append(f"{code} {desc}")

        self.vectorizer = TfidfVectorizer(
            max_features=30000,
            sublinear_tf=True,
            ngram_range=(1, 2),
            min_df=1,
            max_df=0.95,
        )
        self.tfidf_matrix = self.vectorizer.fit_transform(self.icd_texts)
        print(
            f"  [ICD-Reranker] Ready: {len(self.icd_codes)} codes, "
            f"vocab={len(self.vectorizer.vocabulary_)}"
        )

    def score_query(self, query: str, top_k: int = 20) -> list[tuple[str, float]]:
        """Score query against all ICD descriptions, return top matching codes."""
        q_vec = self.vectorizer.transform([query])
        scores = (self.tfidf_matrix @ q_vec.T).toarray().flatten()
        top_idx = np.argsort(scores)[::-1][:top_k]
        return [(self.icd_codes[i], float(scores[i])) for i in top_idx if scores[i] > 0]


# ---------------------------------------------------------------------------
# TF-IDF retriever — ICD-augmented
# ---------------------------------------------------------------------------
class TFIDFRetriever:
    def __init__(self, corpus_path: Path = CORPUS_PATH):
        self.corpus_path = corpus_path
        self.protocol_ids: list[str] = []
        self.titles: list[str] = []
        self.icd_codes_list: list[list[str]] = []
        self.texts: list[str] = []
        self.vectorizer: TfidfVectorizer | None = None
        self.tfidf_matrix = None

    def load(self, icd_descr: dict[str, str] | None = None) -> None:
        print("  [TF-IDF] Loading corpus …")
        with open(self.corpus_path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                obj = json.loads(line)
                self.protocol_ids.append(obj.get("protocol_id", ""))
                self.titles.append(obj.get("title", ""))
                icd_codes = obj.get("icd_codes", [])
                self.icd_codes_list.append(icd_codes)

                icd_str = " ".join(icd_codes)
                icd_desc_parts = []
                if icd_descr:
                    for code in icd_codes:
                        desc = icd_descr.get(code, "")
                        if desc:
                            icd_desc_parts.append(desc)
                        if "." not in code:
                            for fc, fd in icd_descr.items():
                                if fc.startswith(code + "."):
                                    icd_desc_parts.append(f"{fc} {fd}")

                combined = (
                    f"{obj.get('title','')} {icd_str} "
                    f"{' '.join(icd_desc_parts)} "
                    f"{obj.get('text','')[:3000]}"
                )
                self.texts.append(combined)

        print(f"  [TF-IDF] Fitting on {len(self.texts)} protocols …")
        self.vectorizer = TfidfVectorizer(
            max_features=50000,
            sublinear_tf=True,
            ngram_range=(1, 2),
            min_df=1,
            max_df=0.95,
        )
        self.tfidf_matrix = self.vectorizer.fit_transform(self.texts)
        print(f"  [TF-IDF] Ready: vocab={len(self.vectorizer.vocabulary_)}")

    def retrieve(self, query: str, top_k: int = 10) -> list[tuple[int, float]]:
        q_vec = self.vectorizer.transform([query])
        scores = (self.tfidf_matrix @ q_vec.T).toarray().flatten()
        top_idx = np.argsort(scores)[::-1][:top_k]
        return [(int(i), float(scores[i])) for i in top_idx if scores[i] > 0]


# ---------------------------------------------------------------------------
# Ensemble fusion
# ---------------------------------------------------------------------------
def ensemble_retrieve(
    query: str,
    retriever_dense: FAISSRetrieverHF,
    retriever_sparse: TFIDFRetriever,
    icd_reranker: ICDReRanker,
    abbr_map: dict[str, str],
    subcode_map: dict[str, list[str]],
    top_k_dense: int = TOP_K_DENSE,
    top_k_final: int = TOP_K_FINAL,
    dense_weight: float = DENSE_WEIGHT,
    sparse_weight: float = SPARSE_WEIGHT,
    icd_re_weight: float = ICD_RE_WEIGHT,
) -> tuple[list[dict], list[tuple[str, float]]]:
    """
    Three-signal ensemble:
      1. FAISS dense retrieval → protocol-level
      2. TF-IDF sparse → protocol-level
      3. ICD re-ranker → code-level (matches symptoms to disease names)
    Returns: (fused protocol results, top ICD code matches).
    """
    expanded_query = expand_abbreviations(query, abbr_map)
    synonym_query = expand_medical_synonyms(expanded_query)

    # 1. Dense retrieval
    dense_results = retriever_dense.retrieve(expanded_query, top_k=top_k_dense)
    dense_scores: dict[str, float] = {}
    dense_info: dict[str, RetrievalResult] = {}
    for r in dense_results:
        dense_scores[r.protocol_id] = r.score
        dense_info[r.protocol_id] = r

    # 2. Sparse retrieval (with synonym expansion)
    sparse_hits = retriever_sparse.retrieve(synonym_query, top_k=top_k_dense)
    sparse_scores: dict[str, float] = {}
    for idx, score in sparse_hits:
        pid = retriever_sparse.protocol_ids[idx]
        sparse_scores[pid] = score

    # 3. ICD re-ranker — find ICD codes whose descriptions match the symptoms
    icd_matches = icd_reranker.score_query(synonym_query, top_k=20)
    # Map ICD matches back to protocols that contain those codes
    icd_code_scores: dict[str, float] = {}
    for code, score in icd_matches:
        icd_code_scores[code] = score

    # Build protocol-level ICD scores: for each protocol, take the max
    # ICD re-ranker score among its codes
    icd_proto_scores: dict[str, float] = {}
    # Check both dense and sparse candidates
    all_pids = set(dense_scores.keys()) | set(sparse_scores.keys())
    for pid in all_pids:
        # Get protocol's ICD codes
        if pid in dense_info:
            p_codes = dense_info[pid].icd_codes
        else:
            idx_list = [i for i, p in enumerate(retriever_sparse.protocol_ids) if p == pid]
            p_codes = retriever_sparse.icd_codes_list[idx_list[0]] if idx_list else []

        # Expand to subcodes and check against ICD re-ranker matches
        expanded = expand_codes_to_subcodes(p_codes, subcode_map)
        max_icd_score = 0.0
        for code in expanded:
            if code in icd_code_scores:
                max_icd_score = max(max_icd_score, icd_code_scores[code])
        if max_icd_score > 0:
            icd_proto_scores[pid] = max_icd_score

    # Also add protocols from sparse retriever that match top ICD codes
    # even if not in dense results
    top_icd_codes = set(code for code, _ in icd_matches[:10])
    for i, pid in enumerate(retriever_sparse.protocol_ids):
        if pid in all_pids:
            continue
        p_codes = retriever_sparse.icd_codes_list[i]
        expanded = expand_codes_to_subcodes(p_codes, subcode_map)
        if top_icd_codes & set(expanded):
            all_pids.add(pid)
            max_score = max((icd_code_scores.get(c, 0) for c in expanded), default=0)
            if max_score > 0:
                icd_proto_scores[pid] = max_score

    # Normalize
    max_dense  = max(dense_scores.values())  if dense_scores  else 1.0
    max_sparse = max(sparse_scores.values()) if sparse_scores else 1.0
    max_icd    = max(icd_proto_scores.values()) if icd_proto_scores else 1.0

    fused = []
    for pid in all_pids:
        d = dense_scores.get(pid, 0.0) / max_dense if max_dense > 0 else 0
        s = sparse_scores.get(pid, 0.0) / max_sparse if max_sparse > 0 else 0
        ic = icd_proto_scores.get(pid, 0.0) / max_icd if max_icd > 0 else 0
        combined = dense_weight * d + sparse_weight * s + icd_re_weight * ic

        if pid in dense_info:
            info = dense_info[pid]
            title, icd_codes, context = info.title, info.icd_codes, info.context
        else:
            try:
                idx_list = [i for i, p in enumerate(retriever_sparse.protocol_ids) if p == pid]
                if idx_list:
                    idx = idx_list[0]
                    title = retriever_sparse.titles[idx]
                    icd_codes = retriever_sparse.icd_codes_list[idx]
                    context = retriever_sparse.texts[idx][:800]
                else:
                    title, icd_codes, context = pid, [], ""
            except Exception:
                title, icd_codes, context = pid, [], ""

        expanded_codes = expand_codes_to_subcodes(icd_codes, subcode_map)

        fused.append({
            "protocol_id": pid,
            "title": title,
            "icd_codes": icd_codes,
            "icd_codes_expanded": expanded_codes,
            "context": context,
            "dense_score": dense_scores.get(pid, 0.0),
            "sparse_score": sparse_scores.get(pid, 0.0),
            "icd_re_score": icd_proto_scores.get(pid, 0.0),
            "fused_score": combined,
        })

    fused.sort(key=lambda x: x["fused_score"], reverse=True)
    return fused[:top_k_final], icd_matches[:10]


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main():
    parser = argparse.ArgumentParser(description="Ensemble RAG Inspector v4")
    parser.add_argument("--n",    type=int, default=10)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--topk", type=int, default=5)
    parser.add_argument("--dense-weight",  type=float, default=DENSE_WEIGHT)
    parser.add_argument("--sparse-weight", type=float, default=SPARSE_WEIGHT)
    parser.add_argument("--icd-weight",    type=float, default=ICD_RE_WEIGHT)
    args = parser.parse_args()

    dw, sw, iw = args.dense_weight, args.sparse_weight, args.icd_weight

    print("=" * 100)
    print("ENSEMBLE RAG INSPECTOR v4  (FAISS + TF-IDF + ICD Re-ranker)")
    print("=" * 100)

    print("\n1. Loading abbreviation dictionary …")
    abbr_map = load_abbreviations(ABBR_PATH)
    print(f"   {len(abbr_map)} abbreviation mappings")

    print("2. Loading ICD-10 descriptions …")
    icd_descr = load_icd_descriptions(MKB_PATH)
    print(f"   {len(icd_descr)} ICD-10 descriptions")

    print("3. Building subcode expansion map …")
    subcode_map = build_subcode_map(icd_descr)
    print(f"   {len(subcode_map)} category codes")

    print(f"4. Medical synonyms: {len(MEDICAL_SYNONYMS)} patterns")

    print("5. Loading FAISS dense retriever …")
    retriever_dense = FAISSRetrieverHF()
    retriever_dense.load()
    print("   FAISS ready")

    print("6. Loading TF-IDF sparse retriever …")
    retriever_sparse = TFIDFRetriever()
    retriever_sparse.load(icd_descr=icd_descr)

    print("7. Loading ICD description re-ranker …")
    icd_reranker = ICDReRanker()
    icd_reranker.load(icd_descr)

    print(f"\n   Weights: dense={dw:.2f}, sparse={sw:.2f}, icd={iw:.2f}")
    print(f"   Top-K: {args.topk}")
    print(f"   Prefix matching: ENABLED | Synonyms: ENABLED\n")

    test_files = list(TEST_DIR.glob("*.json"))
    sample_files = random.Random(args.seed).sample(test_files, min(args.n, len(test_files)))

    print("=" * 100)
    print(f"INSPECTING {len(sample_files)} RANDOM TEST CASES  (seed={args.seed})")
    print("=" * 100)

    exact_acc1 = exact_rec3 = exact_rec5 = exact_rec10 = 0
    pfx_acc1 = pfx_rec3 = pfx_rec5 = pfx_rec10 = 0

    for i, f in enumerate(sample_files, 1):
        data  = json.loads(f.read_text())
        query = data.get("query") or ""
        gt    = data.get("gt", "")
        valid = set(data.get("icd_codes", []))

        if not query:
            print(f"\n[{i}/{len(sample_files)}]  SKIP {f.name} — empty query")
            continue

        expanded = expand_abbreviations(query, abbr_map)
        triggered = []
        text_lower = query.lower()
        for phrase, syns in sorted(MEDICAL_SYNONYMS.items(), key=lambda x: -len(x[0])):
            if phrase.lower() in text_lower:
                triggered.append(f"{phrase}→{','.join(syns[:2])}")

        print(f"\n{'─' * 100}")
        print(f"[{i}/{len(sample_files)}]  FILE: {f.name}")
        print(f"{'─' * 100}")
        print(f"SYMPTOMS  : {query[:500]}")
        if expanded != query:
            diff_parts = []
            for w in set(re.findall(r"\b[A-ZА-ЯЁ]{2,}\b", query)):
                if w.upper() in abbr_map:
                    diff_parts.append(f"{w}→{abbr_map[w.upper()]}")
            if diff_parts:
                print(f"ABBR      : {', '.join(diff_parts[:5])}")
        if triggered:
            print(f"SYNONYMS  : {' | '.join(triggered[:8])}")

        gt_descr = icd_descr.get(gt, "—")
        print(f"GT CODE   : {gt}  →  {gt_descr}")
        print(f"VALID     : {sorted(valid)}")
        for vc in sorted(valid):
            vd = icd_descr.get(vc, "")
            if vd:
                print(f"            {vc}: {vd}")
        print()

        results, icd_matches = ensemble_retrieve(
            query, retriever_dense, retriever_sparse, icd_reranker,
            abbr_map, subcode_map,
            top_k_final=args.topk, dense_weight=dw, sparse_weight=sw, icd_re_weight=iw,
        )

        # Show ICD re-ranker top matches
        print("ICD RE-RANKER TOP MATCHES:")
        for code, score in icd_matches[:5]:
            desc = icd_descr.get(code, "")
            hit = "✅" if code in valid or any(prefix_match(code, valid) for _ in [1]) else "  "
            print(f"  {hit} {code}: {desc[:80]}  (score={score:.4f})")
        print()

        # Collect expanded codes
        ordered_codes: list[str] = []
        seen_codes: set[str] = set()
        for r in results:
            for code in r["icd_codes_expanded"]:
                if code and code not in seen_codes:
                    seen_codes.add(code)
                    ordered_codes.append(code)

        top1 = ordered_codes[0] if ordered_codes else ""
        top3 = ordered_codes[:3]
        top5 = ordered_codes[:5]
        top10 = ordered_codes[:10]

        e_hit1 = 1 if top1 == gt else 0
        e_hit3 = 1 if any(c in valid for c in top3) else 0
        e_hit5 = 1 if any(c in valid for c in top5) else 0
        e_hit10 = 1 if any(c in valid for c in top10) else 0
        exact_acc1 += e_hit1; exact_rec3 += e_hit3; exact_rec5 += e_hit5; exact_rec10 += e_hit10

        p_hit1 = 1 if prefix_match(top1, valid) else 0
        p_hit3 = 1 if any_prefix_match(top3, valid) else 0
        p_hit5 = 1 if any_prefix_match(top5, valid) else 0
        p_hit10 = 1 if any_prefix_match(top10, valid) else 0
        pfx_acc1 += p_hit1; pfx_rec3 += p_hit3; pfx_rec5 += p_hit5; pfx_rec10 += p_hit10

        print(f"TOP {args.topk} ENSEMBLE RESULTS:")
        for j, r in enumerate(results, 1):
            orig = r["icd_codes"][:5]
            exp = r["icd_codes_expanded"][:8]
            codes_str = ", ".join(orig)
            hit_exact = any(c in valid for c in exp)
            hit_pfx = any_prefix_match(exp, valid)
            marker = "✅" if hit_exact else ("🔶" if hit_pfx else "❌")

            descs = [f"{c}: {icd_descr.get(c,'')}" for c in orig[:2] if icd_descr.get(c)]
            new_subs = [c for c in exp if c not in orig][:5]

            print(
                f"  {j}. {marker} [{codes_str}]  "
                f"fused={r['fused_score']:.4f}  "
                f"(d={r['dense_score']:.3f} s={r['sparse_score']:.3f} icd={r.get('icd_re_score',0):.3f})"
            )
            print(f"     {r['title'][:100]}")
            if new_subs:
                print(f"     +subcodes: {', '.join(new_subs)}")
            for d in descs[:2]:
                print(f"     📋 {d}")

        if e_hit1: status = "✅ EXACT"
        elif p_hit1: status = "🔶 PREFIX"
        elif e_hit3: status = "🔶 top-3 exact"
        elif p_hit3: status = "🟡 top-3 prefix"
        elif e_hit5: status = "🟡 top-5 exact"
        elif p_hit5: status = "🟠 top-5 prefix"
        elif e_hit10: status = "🟠 top-10 exact"
        elif p_hit10: status = "🟤 top-10 prefix"
        else: status = "❌ MISS"
        print(f"\n  → {status}  top1={top1}  gt={gt}")
        if not e_hit1 and top1:
            print(f"  → Pred: {top1}: {icd_descr.get(top1,'')}")
            print(f"  → Want: {gt}: {gt_descr}")

    n = len(sample_files)
    print(f"\n{'═' * 100}")
    print(f"ENSEMBLE v4 SUMMARY  (d={dw:.2f} s={sw:.2f} icd={iw:.2f}, topK={args.topk})")
    print(f"{'═' * 100}")
    print(f"  Cases: {n}")
    print(f"  EXACT:  Acc@1={exact_acc1}/{n} ({exact_acc1/n*100:.1f}%)  "
          f"Rec@3={exact_rec3}/{n} ({exact_rec3/n*100:.1f}%)  "
          f"Rec@5={exact_rec5}/{n} ({exact_rec5/n*100:.1f}%)  "
          f"Rec@10={exact_rec10}/{n} ({exact_rec10/n*100:.1f}%)")
    print(f"  PREFIX: Acc@1={pfx_acc1}/{n} ({pfx_acc1/n*100:.1f}%)  "
          f"Rec@3={pfx_rec3}/{n} ({pfx_rec3/n*100:.1f}%)  "
          f"Rec@5={pfx_rec5}/{n} ({pfx_rec5/n*100:.1f}%)  "
          f"Rec@10={pfx_rec10}/{n} ({pfx_rec10/n*100:.1f}%)")
    print(f"{'═' * 100}")


if __name__ == "__main__":
    main()