explainability/llm_explain.py

"""
explainability/llm_explain.py
──────────────────────────────
Local LLM-powered natural language explanation of model predictions.
Built from scratch using HuggingFace Transformers — no external API,
no API key, runs entirely on your machine.

Architecture
────────────
  Tokenizer  : AutoTokenizer  (google/flan-t5-base)
  Model      : T5ForConditionalGeneration
  Inference  : beam search, greedy decode, or sampling
  Fallback   : deterministic template engine (works with no model at all)

FLAN-T5 was chosen because:
  - Instruction-tuned  → responds well to structured prompts
  - No API key needed  → fully self-contained
  - Reasonable size    → flan-t5-base is ~250 MB, runs on CPU
  - Medical text       → handles clinical terminology cleanly
  - You've used it before in RialoLens AI Explainer Engine

Model variants (pass as model_name)
────────────────────────────────────
  "google/flan-t5-small"   ~80 MB   fastest, lowest quality
  "google/flan-t5-base"    ~250 MB  ← default, good balance
  "google/flan-t5-large"   ~780 MB  better quality, needs more RAM
  "google/flan-t5-xl"      ~3 GB    best quality, needs GPU

How it connects to the existing architecture
────────────────────────────────────────────
  model/inference.py        → consumes its prediction dict output
  explainability/gradcam.py → optionally consumes Grad-CAM result dict
  No existing files are modified.

Usage
─────
    from model.inference    import BreastCancerInferencePipeline
    from explainability     import GradCAM, LLMExplainer

    pipeline = BreastCancerInferencePipeline("model/weights.pth")
    result   = pipeline.predict("slide.png")

    # Basic explanation
    llm    = LLMExplainer()                        # downloads model once
    report = llm.explain(result, audience="clinician")
    print(report["summary"])

    # With Grad-CAM spatial context
    cam        = GradCAM(pipeline.model)
    cam_result = cam.explain("slide.png")
    report     = llm.explain_with_gradcam(cam_result, audience="patient")
    print(report["detail"])

Install
───────
  pip install transformers sentencepiece accelerate
"""

from __future__ import annotations

import sys
import textwrap
from pathlib import Path
from typing import Literal, Optional

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

# Audience type
Audience = Literal["clinician", "researcher", "patient"]


# ╔══════════════════════════════════════════════════════════════════════════════╗
# ║  FLAN-T5 BACKBONE                                                           ║
# ╚══════════════════════════════════════════════════════════════════════════════╝

class FlanT5Engine:
    """
    Thin wrapper around FLAN-T5 for text generation.

    Downloads the model once on first use and caches it to
    ~/.cache/huggingface/hub (HuggingFace default).

    Parameters
    ----------
    model_name : str
        HuggingFace model identifier. Default: google/flan-t5-base
    device : str
        "cuda", "mps", or "cpu". Auto-detected if None.
    max_new_tokens : int
        Maximum tokens to generate per explanation. Default 256.
    """

    def __init__(
        self,
        model_name:     str           = "google/flan-t5-large",
        device:         Optional[str] = None,
        max_new_tokens: int           = 256,
    ) -> None:
        self.model_name     = model_name
        self.max_new_tokens = max_new_tokens
        self.device         = self._resolve_device(device)

        self._tokenizer = None
        self._model     = None

    # ── Lazy loading — model loads on first generate() call ──────────────────

    def _load(self) -> None:
        """Download / load tokenizer and model into memory."""
        if self._model is not None:
            return   # already loaded

        try:
            from transformers import AutoTokenizer, T5ForConditionalGeneration
        except ImportError:
            raise ImportError(
                "transformers not installed.\n"
                "Run: pip install transformers sentencepiece"
            )

        import torch

        print(f"[LLMExplainer] Loading {self.model_name} …")
        print(f"[LLMExplainer] Device: {self.device}")
        print("[LLMExplainer] First run downloads ~250 MB, cached afterwards.")

        self._tokenizer = AutoTokenizer.from_pretrained(self.model_name)

        self._model = T5ForConditionalGeneration.from_pretrained(
            self.model_name,
            torch_dtype=torch.float16 if self.device != "cpu" else torch.float32,
        ).to(self.device)

        self._model.eval()
        print(f"[LLMExplainer] Model ready.")

    def generate(self, prompt: str) -> str:
        """
        Generate text from a prompt using FLAN-T5.

        Parameters
        ----------
        prompt : str
            Instruction-style prompt in the T5 format.

        Returns
        -------
        str — generated text, decoded and stripped.
        """
        import torch

        self._load()

        inputs = self._tokenizer(
            prompt,
            return_tensors="pt",
            truncation=True,
            max_length=512,           # T5 input limit
        ).to(self.device)

        with torch.inference_mode():
            output_ids = self._model.generate(
                **inputs,
                max_new_tokens  = self.max_new_tokens,
                num_beams       = 4,          # beam search — better quality
                early_stopping  = True,
                no_repeat_ngram_size = 3,     # reduce repetition
                length_penalty  = 1.2,        # encourage longer outputs
            )

        decoded = self._tokenizer.decode(
            output_ids[0],
            skip_special_tokens=True,
        )
        return decoded.strip()

    def generate_chat(self, prompt: str) -> str:
        """
        Generate a conversational chat response using FLAN-T5.

        Uses temperature sampling instead of beam search to produce
        natural, varied, human-sounding responses rather than mechanical outputs.

        Parameters
        ----------
        prompt : str
            Conversational prompt with full context and instruction.

        Returns
        -------
        str — generated text, decoded and stripped.
        """
        import torch

        self._load()

        inputs = self._tokenizer(
            prompt,
            return_tensors = "pt",
            truncation     = True,
            max_length     = 512,
        ).to(self.device)

        with torch.inference_mode():
            output_ids = self._model.generate(
                **inputs,
                max_new_tokens       = 400,
                do_sample            = True,    # sampling for natural variety
                temperature          = 0.8,     # slightly creative but still coherent
                top_p                = 0.92,    # nucleus sampling
                no_repeat_ngram_size = 3,
                repetition_penalty   = 1.3,
            )

        decoded = self._tokenizer.decode(
            output_ids[0],
            skip_special_tokens = True,
        )
        return decoded.strip()

    @staticmethod
    def _resolve_device(device: Optional[str]) -> str:
        import torch
        if device is not None:
            return device
        if torch.cuda.is_available():
            return "cuda"
        if torch.backends.mps.is_available():
            return "mps"
        return "cpu"


# ╔══════════════════════════════════════════════════════════════════════════════╗
# ║  TEMPLATE ENGINE — deterministic fallback                                   ║
# ╚══════════════════════════════════════════════════════════════════════════════╝

class TemplateEngine:
    """
    Audience-aware deterministic explanation engine.
    Produces genuinely different content for clinician, researcher, and patient.
    """

    TIERS = [
        (0.95, "very high",  "very high confidence"),
        (0.85, "high",       "high confidence"),
        (0.70, "moderate",   "moderate confidence"),
        (0.55, "borderline", "borderline confidence — treat with caution"),
        (0.00, "low",        "low confidence — result unreliable without further review"),
    ]

    def confidence_tier(self, conf: float) -> tuple[str, str]:
        for threshold, label, desc in self.TIERS:
            if conf >= threshold:
                return label, desc
        return "low", "low confidence"

    # Modality-specific facts so explanations name the right model, dataset,
    # reviewer, and sample type. Histopathology = DenseNet/PCam/pathologist;
    # mammogram = EfficientNet-B4 ensemble/RSNA/radiologist.
    MODALITY = {
        "histopathology": {
            "model":        "DenseNet-121",
            "sample":       "patch",
            "sample_plain": "tissue sample",
            "metrics":      "87.5% sensitivity and 88.0% accuracy on the held-out PCam test set",
            "reviewer":     "pathologist",
            "correlation":  "histomorphological correlation",
            "pos_action":   "tissue biopsy and full clinical assessment",
            "neg_action":   "Routine 12-month follow-up is appropriate if clinically indicated.",
            "imaging_desc": "checks tissue under a microscope",
            "abnormal":     "abnormal cells",
            "model_full":   "DenseNet-121 (7.22M params) fine-tuned on 220,025 "
                            "deduplicated PCam patches",
            "train_detail": "OneCycleLR (max_lr=3e-3), Mixup (α=0.4), StainJitter "
                            "(HED, strength=0.05), label smoothing (0.1), "
                            "CrossEntropyLoss (pos_weight=1.469)",
            "metrics_short":"Best test sensitivity: 87.5%, accuracy: 88.0%",
        },
        "mammogram": {
            "model":        "the EfficientNet-B4 ensemble",
            "sample":       "mammogram",
            "sample_plain": "mammogram",
            "metrics":      "0.84 AUC (70.1% sensitivity, 82.4% specificity) on the "
                            "RSNA validation set",
            "reviewer":     "radiologist",
            "correlation":  "radiological correlation",
            "pos_action":   "a diagnostic mammographic work-up and possible biopsy",
            "neg_action":   "Routine screening follow-up is appropriate per BI-RADS guidance.",
            "imaging_desc": "reviews the breast X-ray image",
            "abnormal":     "an abnormal area",
            "model_full":   "a 3-model EfficientNet-B4 ensemble trained on the RSNA 2022 "
                            "mammography dataset (54,706 images)",
            "train_detail": "3× EfficientNet-B4 (seeds 42/123/999), WeightedRandomSampler "
                            "for class balance, AMP mixed precision, CrossEntropyLoss "
                            "(weight=[1,5]), predictions averaged across members",
            "metrics_short":"Ensemble AUC: 0.84 (sensitivity 70.1%, specificity 82.4%)",
        },
    }

    def build(
        self,
        prediction:      str,
        confidence:      float,
        benign_logit:    float,
        malignant_logit: float,
        audience:        Audience,
        gradcam_context: Optional[str] = None,
        modality:        str           = "histopathology",
        birads:          Optional[str] = None,
    ) -> tuple[str, str]:
        tier_label, tier_desc = self.confidence_tier(confidence)
        pct        = f"{confidence:.1%}"
        margin     = abs(malignant_logit - benign_logit)
        is_mal     = prediction == "malignant"
        cam        = gradcam_context or ""
        f          = self.MODALITY.get(modality, self.MODALITY["histopathology"])

        # ── CLINICIAN ─────────────────────────────────────────────────────────
        if audience == "clinician":
            # Use the caller-supplied BI-RADS (from the ensemble) when available,
            # otherwise fall back to a confidence-derived suggestion.
            birad = birads or (
                "BI-RADS 4B — Suspicious" if is_mal and confidence >= 0.75 else
                "BI-RADS 4A — Low suspicion" if is_mal else
                "BI-RADS 2 — Benign finding"
            )
            boundary = (
                "The decision margin of {:.3f} indicates a clear separation "
                "from the decision boundary, supporting diagnostic reliability.".format(margin)
                if margin > 1.5 else
                "The decision margin of {:.3f} places this case near the "
                "classification boundary — a borderline result requiring careful "
                "{}.".format(margin, f["correlation"])
            )
            cam_line = (
                f" Grad-CAM spatial analysis: {cam}"
                if cam else
                " Grad-CAM spatial attention maps are available for region-level review."
            )
            summary = (
                f"{f['model']} classified this {f['sample']} as "
                f"{'MALIGNANT' if is_mal else 'BENIGN'} at {pct} ({tier_desc}). "
                f"Raw logits: benign = {benign_logit:.4f}, malignant = {malignant_logit:.4f} "
                f"(margin: {margin:.4f}). "
                f"Suggested {birad}."
            )
            detail = (
                f"{boundary}"
                f"{cam_line} "
                f"Underlying model performance: {f['metrics']}. "
                f"{('A positive result warrants ' + f['pos_action'] + '.') if is_mal else f['neg_action']} "
                f"This output is AI-assisted and must not replace {f['reviewer']} review."
            )

        # ── RESEARCHER ────────────────────────────────────────────────────────
        elif audience == "researcher":
            softmax_b = round(1 / (1 + 2.718 ** (malignant_logit - benign_logit)), 4)
            softmax_m = round(1 - softmax_b, 4)
            cam_line  = f" Grad-CAM activation summary: {cam}" if cam else ""
            summary = (
                f"Classification output: {prediction.upper()} "
                f"[softmax({benign_logit:.4f}, {malignant_logit:.4f}) = "
                f"({softmax_b:.4f}, {softmax_m:.4f})]. "
                f"Argmax class = {'1 (malignant)' if is_mal else '0 (benign)'}. "
                f"Decision margin |Δlogit| = {margin:.4f} "
                f"({'above' if margin > 1.5 else 'below'} the 1.5 heuristic threshold "
                f"for high-confidence separation)."
            )
            detail = (
                f"Model: {f['model_full']}. Training: {f['train_detail']}. "
                f"{f['metrics_short']}. "
                f"Softmax probabilities are uncalibrated — no temperature scaling applied."
                f"{cam_line}"
            )

        # ── PATIENT ───────────────────────────────────────────────────────────
        else:
            sample = f["sample_plain"]
            if is_mal:
                if confidence >= 0.85:
                    summary = (
                        f"The AI system flagged an area in this {sample} that looks "
                        f"unusual, and it is fairly confident about this ({pct}). "
                        f"This is a signal that a doctor should take a closer look — "
                        f"it does not mean you definitely have cancer."
                    )
                    detail = (
                        f"Think of this AI like a second set of eyes that {f['imaging_desc']}. "
                        f"It spotted a pattern it associates with {f['abnormal']} in this {sample}. "
                        f"{'The AI was also looking at the ' + cam.split('.')[0].lower() + ' area most closely.' if cam else ''} "
                        f"Your doctor will review this result and decide the right next step — "
                        f"this might be a follow-up scan or a biopsy. Please do not worry "
                        f"until you have spoken with your healthcare provider. "
                        f"This AI tool is for screening only, not a final diagnosis."
                    )
                else:
                    summary = (
                        f"The AI system found something in this {sample} that it "
                        f"wasn't entirely sure about ({pct} confidence). "
                        f"The result is uncertain and will need your doctor's review "
                        f"before any conclusions are drawn."
                    )
                    detail = (
                        f"This means the AI could not clearly decide whether the {sample} "
                        f"looks normal or abnormal — it is on the borderline. "
                        f"This happens sometimes with difficult cases. "
                        f"Your doctor is the right person to interpret this alongside "
                        f"your full medical history and any other tests. "
                        f"This AI tool is a screening aid only, not a diagnosis."
                    )
            else:
                summary = (
                    f"The AI system found no signs of an abnormality in this {sample} "
                    f"({pct} confidence). This is a reassuring result."
                )
                detail = (
                    f"The AI {f['imaging_desc']} and did not find features it associates "
                    f"with cancer. "
                    f"This is a good sign, but all AI results should be confirmed by "
                    f"your doctor as part of your complete care. "
                    f"{'The AI was paying attention to ' + cam.split('.')[0].lower() + '.' if cam else ''} "
                    f"Please keep any follow-up appointments your doctor recommends. "
                    f"This AI tool is for screening only, not a final diagnosis."
                )

        detail = detail.strip()
        return summary, detail


# ╔══════════════════════════════════════════════════════════════════════════════╗
# ║  PROMPT BUILDER                                                             ║
# ╚══════════════════════════════════════════════════════════════════════════════╝

class PromptBuilder:
    """
    Builds FLAN-T5 instruction prompts from structured prediction data.

    FLAN-T5 responds best to explicit task instructions in the format:
      "Task description: [context]. Answer:"
    """

    AUDIENCE_CONTEXT = {
        "clinician": (
            "a clinical pathologist who needs precise technical details, "
            "logit scores, confidence calibration, and clinical caveats"
        ),
        "researcher": (
            "an ML researcher who wants to understand the model's decision "
            "in terms of logit scores, softmax probabilities, and feature analysis"
        ),
        "patient": (
            "a patient with no medical background who needs a clear, "
            "compassionate explanation without jargon"
        ),
    }

    def build(
        self,
        prediction:      str,
        confidence:      float,
        benign_logit:    float,
        malignant_logit: float,
        audience:        Audience,
        gradcam_context: Optional[str] = None,
    ) -> str:
        """Construct the instruction prompt for FLAN-T5."""

        tier = (
            "high" if confidence >= 0.85 else
            "moderate" if confidence >= 0.70 else
            "low"
        )

        cam_section = (
            f" The Grad-CAM spatial analysis shows: {gradcam_context}"
            if gradcam_context else ""
        )

        prompt = textwrap.dedent(f"""
            Explain a breast cancer AI classifier result to {self.AUDIENCE_CONTEXT[audience]}.

            Model result:
            - Prediction: {prediction.upper()}
            - Confidence: {confidence:.1%} ({tier} confidence)
            - Benign logit score: {benign_logit:.3f}
            - Malignant logit score: {malignant_logit:.3f}{cam_section}

            Write a clear 3-sentence explanation of what this result means,
            what the confidence level implies, and remind the reader this is
            a research tool that requires clinical confirmation.

            Explanation:
        """).strip()

        return prompt


# ╔══════════════════════════════════════════════════════════════════════════════╗
# ║  MAIN EXPLAINER CLASS                                                       ║
# ╚══════════════════════════════════════════════════════════════════════════════╝

class LLMExplainer:
    """
    Local LLM-powered natural language explainer for breast cancer predictions.

    Uses FLAN-T5 running entirely on your machine — no API key, no internet
    connection required after the initial model download.

    Falls back to a deterministic template engine if:
      - use_llm=False is passed
      - transformers is not installed
      - The model fails to generate meaningful output

    Parameters
    ----------
    model_name : str
        HuggingFace FLAN-T5 variant. Default: google/flan-t5-base (~250 MB).
    device : str | None
        "cuda", "mps", or "cpu". Auto-detected if None.
    max_new_tokens : int
        Max tokens per generated explanation. Default 256.
    use_llm : bool
        Set False to skip FLAN-T5 and use the template engine directly.
        Useful for fast testing or environments without GPU/internet.
    """

    DISCLAIMER = (
        "Research and educational use only. "
        "Not a standalone diagnostic tool. "
        "Clinical confirmation by a qualified pathologist is required."
    )

    def __init__(
        self,
        model_name:     str           = "google/flan-t5-large",
        device:         Optional[str] = None,
        max_new_tokens: int           = 256,
        use_llm:        bool          = True,
    ) -> None:
        self.use_llm         = use_llm
        self._prompt_builder = PromptBuilder()
        self._template       = TemplateEngine()

        if use_llm:
            self._llm = FlanT5Engine(
                model_name     = model_name,
                device         = device,
                max_new_tokens = max_new_tokens,
            )
        else:
            self._llm = None
            print("[LLMExplainer] Running in template-only mode (use_llm=False).")

    # ── Public API ────────────────────────────────────────────────────────────

    def explain(
        self,
        prediction: dict,
        audience:   Audience = "clinician",
        modality:   str      = "histopathology",
    ) -> dict:
        """
        Generate a natural language explanation from an inference.py output dict.

        Parameters
        ----------
        prediction : dict
            Output from a predict() call:
            {"prediction": str, "confidence": float, "logits": Tensor[1,2],
             "birads": str (optional)}
        audience : "clinician" | "researcher" | "patient"
        modality : "histopathology" | "mammogram"

        Returns
        -------
        dict
            {
                "summary"    : str   — plain-language summary
                "detail"     : str   — deeper explanation with confidence context
                "disclaimer" : str   — standard research disclaimer
                "audience"   : str   — target audience
                "engine"     : str   — "flan-t5" | "template"
            }
        """
        pred, conf, b_logit, m_logit = self._unpack(prediction)
        birads = prediction.get("birads") if isinstance(prediction, dict) else None
        return self._generate(pred, conf, b_logit, m_logit, audience,
                              gradcam_context=None, modality=modality, birads=birads)

    def explain_with_gradcam(
        self,
        gradcam_result: dict,
        audience:       Audience = "clinician",
        modality:       str      = "histopathology",
    ) -> dict:
        """
        Generate an explanation that incorporates Grad-CAM spatial findings.

        Parameters
        ----------
        gradcam_result : dict
            Output from a GradCAM.explain() call:
            {"prediction", "confidence", "logits", "heatmap", "birads"(optional), ...}
        audience : "clinician" | "researcher" | "patient"
        modality : "histopathology" | "mammogram"

        Returns
        -------
        Same schema as explain() — adds spatial activation context.
        """
        pred, conf, b_logit, m_logit = self._unpack(gradcam_result)
        gradcam_context = self._summarise_heatmap(gradcam_result["heatmap"])
        birads = gradcam_result.get("birads") if isinstance(gradcam_result, dict) else None
        return self._generate(pred, conf, b_logit, m_logit, audience,
                              gradcam_context=gradcam_context, modality=modality,
                              birads=birads)

    # ── Internal generation flow ──────────────────────────────────────────────

    def _generate(
        self,
        prediction:      str,
        confidence:      float,
        benign_logit:    float,
        malignant_logit: float,
        audience:        Audience,
        gradcam_context: Optional[str],
        modality:        str           = "histopathology",
        birads:          Optional[str] = None,
    ) -> dict:
        """
        Core generation method. Tries FLAN-T5 first, falls back to template.
        """
        summary, detail = self._template.build(
            prediction      = prediction,
            confidence      = confidence,
            benign_logit    = benign_logit,
            malignant_logit = malignant_logit,
            audience        = audience,
            gradcam_context = gradcam_context,
            modality        = modality,
            birads          = birads,
        )
        engine_used = "template"

        # ── Optional FLAN-T5 enhancement ─────────────────────────────────────
        if self.use_llm and self._llm is not None:
            try:
                prompt = self._prompt_builder.build(
                    prediction      = prediction,
                    confidence      = confidence,
                    benign_logit    = benign_logit,
                    malignant_logit = malignant_logit,
                    audience        = audience,
                    gradcam_context = gradcam_context,
                )
                generated = self._llm.generate(prompt)
                if len(generated.split()) >= 20:
                    # Append FLAN-T5 output to template detail — don't replace it
                    detail = detail + " " + generated.strip()
                    engine_used = "flan-t5"
            except Exception as e:
                print(f"[LLMExplainer] FLAN-T5 enhancement failed ({e}) — template only.")

        return {
            "summary":     summary,
            "detail":      detail,
            "disclaimer":  self.DISCLAIMER,
            "audience":    audience,
            "engine":      engine_used,
        }

    # ── Helpers ───────────────────────────────────────────────────────────────

    @staticmethod
    def _unpack(prediction: dict) -> tuple[str, float, float, float]:
        """Extract prediction, confidence, and logit values from output dict."""
        import torch

        pred       = prediction["prediction"]
        confidence = float(prediction["confidence"])
        logits     = prediction["logits"]

        # Handle Tensor or list
        if isinstance(logits, torch.Tensor):
            vals = logits.squeeze().tolist()
        else:
            vals = list(logits)

        if isinstance(vals, float):
            vals = [vals, vals]

        return pred, confidence, round(vals[0], 4), round(vals[1], 4)

    @staticmethod
    def _split_generated(text: str, gradcam_context: Optional[str]) -> tuple[str, str]:
        """
        Split FLAN-T5 output into summary and detail.
        Uses sentence boundary: first 2 sentences → summary, rest → detail.
        """
        import re
        sentences = re.split(r'(?<=[.!?])\s+', text.strip())
        sentences = [s.strip() for s in sentences if s.strip()]

        if len(sentences) >= 3:
            summary = " ".join(sentences[:2])
            detail  = " ".join(sentences[2:])
        elif len(sentences) == 2:
            summary = sentences[0]
            detail  = sentences[1]
        else:
            summary = text
            detail  = ""

        if gradcam_context and gradcam_context not in detail:
            detail += f" Spatial analysis: {gradcam_context}"

        return summary, detail

    @staticmethod
    def _summarise_heatmap(heatmap) -> str:
        """
        Convert a (224, 224) Grad-CAM heatmap into a spatial text description.
        Divides into 3×3 grid, reports regions with activation above threshold.
        """
        import numpy as np

        h, w      = heatmap.shape
        grid_h    = h // 3
        grid_w    = w // 3
        threshold = 0.6

        region_names = [
            ["top-left",    "top-centre",    "top-right"],
            ["middle-left", "centre",        "middle-right"],
            ["bottom-left", "bottom-centre", "bottom-right"],
        ]

        hot_regions = []
        for row in range(3):
            for col in range(3):
                patch = heatmap[
                    row * grid_h : (row + 1) * grid_h,
                    col * grid_w : (col + 1) * grid_w,
                ]
                if float(patch.mean()) > threshold:
                    hot_regions.append(
                        f"{region_names[row][col]} ({patch.mean():.2f})"
                    )

        overall_mean = float(heatmap.mean())
        overall_max  = float(heatmap.max())

        if hot_regions:
            return (
                f"High activation in: {', '.join(hot_regions)}. "
                f"Mean={overall_mean:.3f}, peak={overall_max:.3f}. "
                f"These regions most influenced the prediction."
            )
        return (
            f"No dominant activation region detected. "
            f"Mean={overall_mean:.3f}, peak={overall_max:.3f}. "
            f"Prediction driven by diffuse low-level features."
        )



# ╔══════════════════════════════════════════════════════════════════════════════╗
# ║  CHAT ENGINE — human-like conversational responses via FLAN-T5              ║
# ╚══════════════════════════════════════════════════════════════════════════════╝

class ChatEngine:
    """
    Human-like conversational response engine powered by FLAN-T5-large.

    Uses carefully designed prompts to make FLAN-T5 respond naturally —
    like a knowledgeable colleague — rather than producing stiff templated text.

    Each response incorporates:
    - Patient name, age, and medical history (if provided)
    - Specific scan numbers (confidence, logits, decision margin)
    - Grad-CAM spatial findings
    - Audience-appropriate tone and vocabulary
    - Conversation history for multi-turn context

    Falls back to rich deterministic responses if FLAN-T5 is not loaded.
    """

    def __init__(self, llm: "FlanT5Engine | None" = None) -> None:
        self.llm = llm

    def respond(
        self,
        message:         str,
        audience:        str,
        prediction:      str,
        confidence:      float,
        benign_logit:    float,
        malignant_logit: float,
        spatial_summary: str            = "",
        history:         list           = None,
        patient:         dict           = None,
    ) -> str:
        """
        Generate a human-like conversational response.

        Parameters
        ----------
        message         : the user's question
        audience        : clinician | researcher | patient
        prediction      : benign | malignant
        confidence      : float [0, 1]
        benign_logit    : raw logit for benign class
        malignant_logit : raw logit for malignant class
        spatial_summary : Grad-CAM text description
        history         : list of prior {"role", "content"} dicts
        patient         : dict with name, age, sex, medical_history, symptoms, previous_scans
        """
        patient   = patient or {}
        history   = history or []
        is_mal    = prediction == "malignant"
        pct       = f"{confidence:.1%}"
        margin    = abs(malignant_logit - benign_logit)
        name      = patient.get("name", "")
        age       = patient.get("age", 0)
        p_history = patient.get("medical_history", "")
        symptoms  = patient.get("symptoms", "")

        # Try FLAN-T5 first
        if self.llm is not None:
            try:
                prompt = self._build_prompt(
                    message, audience, prediction, confidence,
                    benign_logit, malignant_logit, spatial_summary,
                    history, patient, is_mal, pct, margin
                )
                response = self.llm.generate_chat(prompt)
                if len(response.split()) >= 15:
                    return response
            except Exception as e:
                print(f"[ChatEngine] FLAN-T5 failed ({e}) — using fallback.")

        # Fallback: rich deterministic responses
        return self._fallback(
            message, audience, prediction, confidence,
            benign_logit, malignant_logit, spatial_summary,
            is_mal, pct, margin, name, age, p_history, symptoms
        )

    # ── Prompt builder ────────────────────────────────────────────────────────

    def _build_prompt(
        self, message, audience, prediction, confidence,
        b_logit, m_logit, cam, history, patient,
        is_mal, pct, margin
    ) -> str:
        """Build a FLAN-T5 conversational prompt with full context."""

        name    = patient.get("name", "")
        age     = patient.get("age", 0)
        sex     = patient.get("sex", "")
        hist    = patient.get("medical_history", "")
        symp    = patient.get("symptoms", "")
        scans   = patient.get("previous_scans", "")

        patient_ctx = ""
        if name or age or hist or symp:
            patient_ctx = f"""
Patient: {name or 'Anonymous'}{', age ' + str(age) if age else ''}{', ' + sex if sex else ''}.
{('Medical history: ' + hist) if hist else ''}
{('Symptoms: ' + symp) if symp else ''}
{('Previous scans: ' + scans) if scans else ''}
""".strip()

        audience_style = {
            "clinician": (
                "a consultant radiologist. Use clinical terminology. "
                "Be precise and collegial. Reference BI-RADS, logit margins, "
                "and clinical decision context naturally."
            ),
            "researcher": (
                "an ML researcher. Be technical. Reference softmax probabilities, "
                "logit values, model architecture (DenseNet-121), training methodology "
                "(OneCycleLR, Mixup, StainJitter), and calibration naturally."
            ),
            "patient": (
                "a patient with no medical background. Be warm, empathetic, and clear. "
                "Use plain English. No jargon. Acknowledge their feelings. "
                "Be reassuring but honest. Address them by name if you know it."
            ),
        }.get(audience, "a medical professional")

        history_ctx = ""
        if history:
            recent = history[-4:]
            lines = [('User' if h.get('role')=='user' else 'Assistant')+': '+h.get('content','') for h in recent]
            history_ctx = 'Previous exchanges: ' + ' | '.join(lines)

        prompt = f"""You are a knowledgeable and empathetic AI medical assistant for the MedAI platform.
You are speaking to {audience_style}

Scan result:
- Classification: {prediction.upper()} ({pct} confidence)
- Logit scores: benign={b_logit:.4f}, malignant={m_logit:.4f} (margin={margin:.4f})
- Grad-CAM: {cam or 'Not available'}
- Model accuracy: 88.0%, sensitivity: 87.5%
{patient_ctx}
{history_ctx}

The person asks: "{message}"

Respond naturally and warmly in 2-4 sentences. Be specific — reference the actual numbers. Sound like a knowledgeable colleague, not a robot. End with a brief reminder that clinical confirmation is required.

Response:"""

        return prompt.strip()

    # ── Fallback: rich deterministic responses ────────────────────────────────

    def _fallback(
        self, message, audience, prediction, confidence,
        b_logit, m_logit, cam, is_mal, pct, margin,
        name, age, p_history, symptoms
    ) -> str:
        """Rich, human-sounding deterministic responses as fallback."""
        msg  = message.lower()
        addr = f"{name.split()[0]}, " if name else ""
        scan = f"{'malignant' if is_mal else 'benign'} at {pct} confidence"

        if audience == "patient":
            return self._patient_fallback(msg, is_mal, pct, margin, cam, addr, name, confidence, symptoms)
        elif audience == "researcher":
            return self._researcher_fallback(msg, is_mal, pct, b_logit, m_logit, margin, cam)
        else:
            return self._clinician_fallback(msg, is_mal, pct, b_logit, m_logit, margin, cam, confidence, p_history)

    def _patient_fallback(self, msg, is_mal, pct, margin, cam, addr, name, conf, symptoms) -> str:
        if any(k in msg for k in ["worry","worried","scared","serious","cancer","bad"]):
            if is_mal:
                return (
                    f"{addr}I completely understand why you might feel anxious right now — "
                    f"that's a very natural reaction. What I can tell you is that this AI "
                    f"flagged something that needs a closer look, and your doctor is the right "
                    f"person to interpret this alongside your full clinical picture. "
                    f"Many findings like this turn out to be benign on further investigation. "
                    f"Please don't make any decisions until you've spoken with your healthcare provider."
                )
            else:
                return (
                    f"{addr}I can hear that this has been worrying for you. "
                    f"The good news is that the AI found no signs of abnormal tissue in this sample — "
                    f"that's a reassuring result. Of course, your doctor will want to confirm this "
                    f"as part of your overall care, but this is genuinely positive."
                )
        if any(k in msg for k in ["mean","understand","explain","what is","tell me"]):
            if is_mal:
                return (
                    f"{addr}think of the AI like a very experienced set of eyes that has studied "
                    f"thousands of tissue samples. It noticed patterns in this image — at {pct} confidence — "
                    f"that it has learned to associate with abnormal cells. "
                    f"That said, this is a screening tool, not a diagnosis. "
                    f"Your doctor will look at this result together with everything else they know about you."
                )
            else:
                return (
                    f"{addr}the AI examined the patterns in this tissue sample and found that they "
                    f"look consistent with normal, healthy tissue — it's {pct} confident in that assessment. "
                    f"That's a really good sign. Your doctor will confirm this at your next appointment."
                )
        if any(k in msg for k in ["next","step","do","happen","biopsy","test"]):
            if is_mal:
                return (
                    f"{addr}the most important next step is to have a conversation with your doctor "
                    f"about this result as soon as possible. They may recommend additional imaging "
                    f"or a biopsy — which is a small, simple procedure to collect a tiny tissue sample "
                    f"for a laboratory to examine more closely. "
                    f"Please don't let anxiety about what might happen stop you from making that appointment."
                )
            else:
                return (
                    f"{addr}with a reassuring result like this, your doctor will likely recommend "
                    f"continuing with your routine screening schedule. Do mention this result at your "
                    f"next appointment so it becomes part of your medical record. "
                    f"Is there anything else you'd like to understand about what this means?"
                )
        if any(k in msg for k in ["accurate","right","trust","sure","certain","reliable"]):
            return (
                f"{addr}that's a really important question to ask. The AI was correct on {pct} "
                f"of test cases it hadn't seen before — which is good, but not perfect. "
                f"No AI system is 100% accurate, which is exactly why your doctor always "
                f"reviews the result before any clinical decision is made. "
                f"Think of it as a very thorough first opinion."
            )
        if any(k in msg for k in ["heatmap","colour","color","red","highlighted","image","overlay"]):
            return (
                f"{addr}the coloured image you're seeing is called a Grad-CAM heatmap. "
                f"The red and orange areas show where the AI was paying the most attention "
                f"when it made its decision — those are the parts of the tissue it found "
                f"most significant. Blue areas were largely ignored. "
                + (f"In your scan, the AI was particularly focused on {cam.split('.')[0].lower()}." if cam else
                   "Your doctor can use this to understand exactly what the AI was looking at.")
            )
        # Generic fallback
        return (
            f"{addr}I'm here to help you make sense of this result. "
            f"The AI classified this sample as {'potentially abnormal' if is_mal else 'normal-looking'} "
            f"at {pct} confidence. "
            f"You can ask me things like 'what does this mean?', 'should I be worried?', "
            f"or 'what happens next?' — and I'll do my best to explain clearly and honestly."
        )

    def _clinician_fallback(self, msg, is_mal, pct, b_logit, m_logit, margin, cam, conf, p_history) -> str:
        birad = ("BI-RADS 4B (Suspicious)" if conf >= 0.75 else "BI-RADS 4A (Low suspicion)") if is_mal else "BI-RADS 2 (Benign)"
        boundary = ("clear decision boundary" if margin > 1.5 else "near the decision boundary — borderline case")

        if any(k in msg for k in ["why","reason","basis","how","drove"]):
            return (
                f"The classifier scored this patch {'malignant' if is_mal else 'benign'} based on "
                f"DenseNet-121 feature activations — logits benign={b_logit:.4f}, malignant={m_logit:.4f}, "
                f"margin={margin:.4f} ({boundary}). "
                + (f"Grad-CAM identifies high activation in {cam.split('.')[0].lower()}, suggesting those "
                   f"spatial regions drove the classification." if cam else
                   f"Grad-CAM overlay is available in the heatmap tab for region-level review.") +
                f" Clinical correlation with morphological features is warranted."
            )
        if any(k in msg for k in ["birad","bi-rad","category","score","stage"]):
            return (
                f"Based on an AI confidence of {pct} and a logit margin of {margin:.3f}, "
                f"a suggested starting point is {birad}. "
                f"This is an AI-assisted recommendation only — final BI-RADS assignment "
                f"requires full clinical, imaging, and patient history correlation by the "
                f"responsible radiologist. "
                + (f"Relevant history: {p_history[:100]}..." if p_history else "")
            )
        if any(k in msg for k in ["biopsy","next","action","recommend","management"]):
            if is_mal:
                return (
                    f"With a {pct} confidence malignant classification and a logit margin of {margin:.3f}, "
                    f"{'tissue biopsy for histological confirmation is recommended' if conf >= 0.70 else 'short-interval follow-up imaging (6 months) may be appropriate given the borderline confidence'}. "
                    f"Full clinical workup — including prior imaging comparison and patient history — "
                    f"should precede any intervention decision."
                )
            else:
                return (
                    f"With a {pct} confidence benign result and margin {margin:.3f}, "
                    f"routine follow-up per standard screening protocol is appropriate. "
                    f"If clinical suspicion remains high despite the AI result, "
                    f"conventional workup should proceed independently of this output."
                )
        if any(k in msg for k in ["confident","confidence","reliable","calibrat"]):
            return (
                f"The softmax confidence of {pct} is uncalibrated — no temperature scaling "
                f"or isotonic regression was applied post-hoc. The underlying model achieved "
                f"87.5% sensitivity and 88.0% accuracy on 32,768 held-out PCam patches. "
                f"{'A margin of ' + str(round(margin,3)) + ' above 1.5 indicates strong separation from the decision boundary.' if margin > 1.5 else 'A margin of ' + str(round(margin,3)) + ' below 1.5 suggests caution — this is a borderline result.'}"
            )
        if any(k in msg for k in ["gradcam","grad-cam","heatmap","attention","region"]):
            return (
                f"Grad-CAM computed ∂score_{('malignant' if is_mal else 'benign')}/∂A_k across "
                f"the norm5 feature layer (1024×7×7 spatial maps), globally average-pooled "
                f"to derive channel importance weights. "
                + (f"High-activation regions: {cam} — these locations contributed most to the "
                   f"classification. The overlay is viewable in the Grad-CAM tab." if cam else
                   "No dominant activation region detected — prediction driven by diffuse features.")
            )
        return (
            f"The model returned {'malignant' if is_mal else 'benign'} at {pct} confidence "
            f"(logits: b={b_logit:.4f}, m={m_logit:.4f}, margin={margin:.4f}). "
            f"I can elaborate on BI-RADS scoring, biopsy guidance, Grad-CAM interpretation, "
            f"confidence calibration, or model performance — what would be most useful?"
        )

    def _researcher_fallback(self, msg, is_mal, pct, b_logit, m_logit, margin, cam) -> str:
        softmax_b = round(1 / (1 + 2.718 ** (m_logit - b_logit)), 4)
        softmax_m = round(1 - softmax_b, 4)

        if any(k in msg for k in ["logit","score","raw","softmax","probability","output"]):
            return (
                f"Raw logits: [benign={b_logit:.6f}, malignant={m_logit:.6f}]. "
                f"After softmax: [P(benign)={softmax_b:.4f}, P(malignant)={softmax_m:.4f}]. "
                f"|Δlogit| = {margin:.6f} — "
                f"{'above the empirical 1.5 threshold for high-confidence separation' if margin > 1.5 else 'below 1.5, suggesting proximity to the decision boundary'}. "
                f"No temperature scaling or calibration applied post-hoc."
            )
        if any(k in msg for k in ["gradcam","grad-cam","gradient","activation","feature","saliency"]):
            return (
                f"Grad-CAM implementation: forward hook on model.features.norm5 (B×1024×7×7). "
                f"Backward pass computes ∂score_{{'malignant' if is_mal else 'benign'}}/∂A_k for each channel k. "
                f"Global average pooling of gradients gives weights α_k. "
                f"CAM = ReLU(Σ_k α_k · A_k), bilinearly upsampled 7×7 → 224×224. "
                + (f"Spatial summary: {cam}." if cam else "No dominant activation detected.")
            )
        if any(k in msg for k in ["train","architecture","model","densenet","weight","epoch"]):
            return (
                f"DenseNet-121 (7,219,330 params) fine-tuned on 220,025 deduplicated PCam patches. "
                f"Training: OneCycleLR(max_lr=3e-3, pct_start=0.3), Mixup(α=0.4), "
                f"StainJitter(HED, strength=0.05), LabelSmoothing(0.1), "
                f"CrossEntropyLoss(pos_weight=1.469). "
                f"Checkpoint selection by val_sensitivity — best epoch 13/20, val_sens=0.903, "
                f"test_sens=0.875, test_acc=0.880."
            )
        if any(k in msg for k in ["dataset","pcam","camelyon","dedup","duplicate","balance"]):
            return (
                f"PatchCamelyon: 262,144 raw training patches → 220,025 after MD5 deduplication "
                f"(42,119 removed, 83.9% retention). "
                f"The original dataset was artificially balanced by duplicating malignant patches — "
                f"post-dedup true distribution: benign=130,908, malignant=89,117 (pos_weight=1.469). "
                f"Deduplication was the single largest contributor to the +6.8pp sensitivity improvement."
            )
        if any(k in msg for k in ["calibrat","uncertain","temperature","reliability","ece"]):
            return (
                f"P({('malignant' if is_mal else 'benign')})={pct} is an uncalibrated softmax output. "
                f"No temperature scaling, Platt scaling, or isotonic regression was applied. "
                f"ECE was not computed on this checkpoint. "
                f"For reliable probability estimates, recommend fitting calibration on a held-out set "
                f"using sklearn.calibration.CalibratedClassifierCV or manual temperature scaling on logits."
            )
        return (
            f"Output: {('malignant' if is_mal else 'benign').upper()} | "
            f"logits=[{b_logit:.4f}, {m_logit:.4f}] | softmax=[{softmax_b:.4f}, {softmax_m:.4f}] | "
            f"|Δlogit|={margin:.4f}. "
            f"I can go deeper on logit analysis, Grad-CAM implementation, model architecture, "
            f"dataset statistics, or calibration. What would you like to explore?"
        )