src/utils/retriever.py

"""
Hierarchical TF-IDF Retriever for ICD-10 candidate generation.

Three-stage retrieval:
  1. Chapter-level scoring  (22 chapters → top-5)
  2. Category-level scoring (within selected chapters → top-20)
  3. Full-code scoring      (within selected categories → top-K)
"""
import numpy as np
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

from utils.config import ICD10_CHAPTERS
from utils.preprocessing import get_icd_chapter, get_icd_category


class HierarchicalTFIDFRetriever:
    """Three-level hierarchical ICD-10 code retriever using TF-IDF."""

    def __init__(self):
        self.icd_metadata_df = None
        self.chapter_tfidf = TfidfVectorizer(max_features=10000, sublinear_tf=True, ngram_range=(1, 2))
        self.category_tfidf = TfidfVectorizer(max_features=20000, sublinear_tf=True, ngram_range=(1, 2))
        self.code_tfidf = TfidfVectorizer(max_features=30000, sublinear_tf=True, ngram_range=(1, 2))
        self.chapter_vectors = None
        self.category_vectors = None
        self.code_vectors = None
        self.chapter_df = None
        self.category_df = None
        self._fitted = False

    def fit(self, bilingual_lookup: dict):
        """Build the retriever from the bilingual lookup dictionary."""
        rows = []
        for code, info in bilingual_lookup.items():
            eng = info.get("english", "")
            chi = info.get("chinese", "")
            chapter_letter = get_icd_chapter(code)
            chapter = ICD10_CHAPTERS.get(chapter_letter, "Unknown")
            category = get_icd_category(code)
            combined_desc = f"{eng} {chi}".strip()
            rows.append({
                "code": code,
                "chapter": chapter,
                "category": category,
                "english_description": eng,
                "chinese_description": chi,
                "combined_description": combined_desc,
            })

        self.icd_metadata_df = pd.DataFrame(rows)

        # ── Stage 1: Chapter-level TF-IDF ──
        self.chapter_df = (
            self.icd_metadata_df
            .groupby("chapter")["combined_description"]
            .apply(lambda x: " ".join(x))
            .reset_index()
        )
        self.chapter_vectors = self.chapter_tfidf.fit_transform(self.chapter_df["combined_description"])

        # ── Stage 2: Category-level TF-IDF ──
        self.category_df = (
            self.icd_metadata_df
            .groupby("category")["combined_description"]
            .apply(lambda x: " ".join(x))
            .reset_index()
        )
        self.category_vectors = self.category_tfidf.fit_transform(self.category_df["combined_description"])

        # ── Stage 3: Code-level TF-IDF ──
        self.code_vectors = self.code_tfidf.fit_transform(self.icd_metadata_df["combined_description"])

        self._fitted = True
        return self

    def retrieve(self, note_text: str, top_k: int = 100,
                 top_chapters: int = 5, top_categories: int = 20) -> list:
        """
        Retrieve top-K ICD-10 candidate codes for a clinical note.

        Returns: list of (code, score) tuples sorted by relevance.
        """
        if not self._fitted:
            raise RuntimeError("Retriever not fitted. Call .fit() first.")

        # ── Stage 1: Score chapters ──
        note_chapter_vec = self.chapter_tfidf.transform([note_text])
        chapter_scores = cosine_similarity(note_chapter_vec, self.chapter_vectors).flatten()
        top_chapter_idx = np.argsort(chapter_scores)[::-1][:top_chapters]
        selected_chapters = set(self.chapter_df.iloc[top_chapter_idx]["chapter"].tolist())

        # ── Stage 2: Score categories within selected chapters ──
        chapter_to_cats = (
            self.icd_metadata_df[self.icd_metadata_df["chapter"].isin(selected_chapters)]
            ["category"].unique()
        )
        cat_mask = self.category_df["category"].isin(chapter_to_cats)
        cat_indices = self.category_df[cat_mask].index.tolist()

        if not cat_indices:
            cat_indices = list(range(len(self.category_df)))

        note_cat_vec = self.category_tfidf.transform([note_text])
        cat_scores = cosine_similarity(note_cat_vec, self.category_vectors[cat_indices]).flatten()
        top_cat_local = np.argsort(cat_scores)[::-1][:top_categories]
        top_cat_idx = [cat_indices[i] for i in top_cat_local]
        selected_categories = set(self.category_df.iloc[top_cat_idx]["category"].tolist())

        # ── Stage 3: Score codes within selected categories ──
        code_mask = self.icd_metadata_df["category"].isin(selected_categories)
        code_indices = self.icd_metadata_df[code_mask].index.tolist()

        if not code_indices:
            code_indices = list(range(len(self.icd_metadata_df)))

        note_code_vec = self.code_tfidf.transform([note_text])
        code_scores = cosine_similarity(note_code_vec, self.code_vectors[code_indices]).flatten()
        top_code_local = np.argsort(code_scores)[::-1][:top_k]

        results = []
        for i in top_code_local:
            global_idx = code_indices[i]
            row = self.icd_metadata_df.iloc[global_idx]
            results.append((row["code"], float(code_scores[i])))

        return results

    @property
    def num_codes(self):
        return len(self.icd_metadata_df) if self.icd_metadata_df is not None else 0