report_generator.py

"""
report_generator.py v6 — BioMistral-7B powered
------------------------------------------------
Complete file with all methods intact.
"""

import logging
import re
import time
from dataclasses import dataclass

import torch

logger = logging.getLogger(__name__)


@dataclass
class PatientInfo:
    age: int
    gender: str
    symptoms: str
    clinical_history: str


@dataclass
class ImagingInfo:
    modality: str
    body_region: str
    view: str = ""


@dataclass
class RadiologyReport:
    examination_details: str = ""
    clinical_information: str = ""
    findings: str = ""
    impression: str = ""
    recommendations: str = ""
    confidence_score: float = 0.0
    generation_time_seconds: float = 0.0
    raw_text: str = ""
    image_caption: str = ""

    def to_dict(self) -> dict:
        return {
            "examination_details":     self.examination_details,
            "clinical_information":    self.clinical_information,
            "findings":                self.findings,
            "impression":              self.impression,
            "recommendations":         self.recommendations,
            "confidence_score":        self.confidence_score,
            "generation_time_seconds": self.generation_time_seconds,
            "image_caption":           self.image_caption,
        }

    def to_text(self) -> str:
        return "\n".join([
            "RADIOLOGY REPORT",
            "=" * 60,
            "",
            "EXAMINATION DETAILS:",
            self.examination_details,
            "",
            "CLINICAL INFORMATION:",
            self.clinical_information,
            "",
            "FINDINGS:",
            self.findings,
            "",
            "IMPRESSION:",
            self.impression,
            "",
            "RECOMMENDATIONS:",
            self.recommendations,
            "",
            f"CONFIDENCE SCORE: {self.confidence_score:.2f}",
            "",
            "=" * 60,
            "END OF REPORT",
        ])


REGION_FOCUS = {
    "chest":     "lung fields, cardiac silhouette, mediastinum, pleural spaces, bony thorax",
    "brain":     "brain parenchyma, ventricles, sulci, white matter, midline structures, posterior fossa",
    "abdomen":   "liver, spleen, kidneys, pancreas, bowel loops, peritoneal spaces",
    "pelvis":    "pelvic bones, hip joints, bladder, iliac crests, femoral heads, any fixation hardware",
    "spine":     "vertebral bodies, disc spaces, spinal canal, neural foramina, alignment",
    "knee":      "femoral condyles, tibial plateau, patella, joint space, menisci",
    "shoulder":  "glenohumeral joint, rotator cuff, humeral head, acromion",
    "neck":      "cervical spine, soft tissues, thyroid, lymph nodes",
    "breast":    "fibroglandular tissue, lesions, microcalcifications, axillary nodes",
    "cardiac":   "cardiac chambers, great vessels, pericardium",
    "liver":     "hepatic parenchyma, biliary system, portal veins",
    "thyroid":   "thyroid lobes, echogenicity, nodules, vascularity",
    "obstetric": "fetal anatomy, placenta, amniotic fluid, cervix",
}


def build_meditron_prompt(patient: PatientInfo, imaging: ImagingInfo,
                           image_caption: str) -> str:
    view_str    = f", {imaging.view} view" if imaging.view else ""
    focus       = REGION_FOCUS.get(imaging.body_region.lower(),
                                    f"{imaging.body_region} anatomy and pathology")
    vision_line = ""
    if image_caption and len(image_caption) > 20:
        vision_line = f"\nIMAGE DESCRIPTION: {image_caption}"

    prompt = (
        f"### Instruction:\n"
        f"You are a consultant radiologist. Write a structured radiology report "
        f"using standard clinical terminology. Be concise and specific. "
        f"Only report what is visible. Do not speculate beyond the findings.\n\n"
        f"STUDY: {imaging.modality} {imaging.body_region}{view_str}\n"
        f"PATIENT: {patient.age}-year-old {patient.gender}\n"
        f"CLINICAL INDICATION: {patient.symptoms}\n"
        f"HISTORY: {patient.clinical_history}\n"
        f"KEY STRUCTURES: {focus}{vision_line}\n\n"
        f"Write exactly three sections:\n"
        f"FINDINGS: (2-4 sentences describing what is seen)\n"
        f"IMPRESSION: (1-2 sentences — the diagnosis or conclusion)\n"
        f"RECOMMENDATIONS: (1-2 sentences — next clinical steps)\n\n"
        f"### Response:\n"
        f"FINDINGS: The {imaging.body_region} demonstrates"
    )
    return prompt


class ReportGenerator:
    MAX_NEW_TOKENS     = 300
    TEMPERATURE        = 0.1
    TOP_P              = 0.85
    REPETITION_PENALTY = 1.2

    def __init__(self, slm_model=None, slm_tokenizer=None):
        self.model     = slm_model
        self.tokenizer = slm_tokenizer
        self._device   = "cuda" if torch.cuda.is_available() else "cpu"

    def generate(self, patient: PatientInfo, imaging: ImagingInfo,
                 image_caption: str = "", visual_features=None) -> RadiologyReport:
        prompt  = build_meditron_prompt(patient, imaging, image_caption)
        start   = time.time()
        raw     = self._run_inference(prompt)
        elapsed = time.time() - start
        logger.info(f"Raw (250c): {raw[:250]}")
        report  = self._parse(raw, patient, imaging)
        self._validate_consistency(report, imaging)
        report.raw_text                = raw
        report.image_caption           = image_caption
        report.generation_time_seconds = round(elapsed, 2)
        report.confidence_score        = self._score(report, imaging)
        return report

    # ─────────────────────────────────────────────────────────────────
    # Inference
    # ─────────────────────────────────────────────────────────────────

    def _run_inference(self, prompt: str) -> str:
        inputs = self.tokenizer(
            prompt,
            return_tensors="pt",
            truncation=True,
            max_length=800,
            padding=False,
        ).to(self._device)

        with torch.no_grad():
            out = self.model.generate(
                **inputs,
                max_new_tokens=self.MAX_NEW_TOKENS,
                temperature=self.TEMPERATURE,
                top_p=self.TOP_P,
                repetition_penalty=self.REPETITION_PENALTY,
                do_sample=True,
                pad_token_id=self.tokenizer.eos_token_id,
                eos_token_id=self.tokenizer.eos_token_id,
            )
        new = out[0][inputs["input_ids"].shape[1]:]
        return self.tokenizer.decode(new, skip_special_tokens=True)

    # ─────────────────────────────────────────────────────────────────
    # Text cleaning
    # ─────────────────────────────────────────────────────────────────

    def _clean(self, text: str) -> str:
        text = re.sub(r"```[\s\S]*?```", " ", text)
        text = re.sub(r"^###.*$", "", text, flags=re.MULTILINE)
        text = re.sub(r"\n{3,}", "\n\n", text)
        text = re.sub(r"[ \t]{2,}", " ", text)
        return text.strip()

    def _cut_at_sentences(self, text: str, max_sent: int = 4) -> str:
        sentences = re.split(r'(?<=[.!?])\s+', text.strip())
        good = []
        for s in sentences:
            s = s.strip()
            if len(s) < 8:
                continue
            good.append(s)
            if len(good) >= max_sent:
                break
        return " ".join(good)

    def _extract(self, text: str, section: str, max_sent: int = 3) -> str:
        aliases = {
            "findings":        ["FINDINGS", "Findings"],
            "impression":      ["IMPRESSION", "Impression", "CONCLUSION",
                                "Conclusion", "DIAGNOSIS"],
            "recommendations": ["RECOMMENDATIONS", "Recommendations",
                                 "RECOMMENDATION", "PLAN", "Plan"],
        }
        all_hdrs = [h for v in aliases.values() for h in v]
        headers  = aliases.get(section, [section.upper()])
        stopper  = "|".join(re.escape(h) for h in all_hdrs if h not in headers)

        for hdr in headers:
            pat = (rf"(?:^|\n){re.escape(hdr)}\s*:?\s*(.*?)"
                   rf"(?=\n(?:{stopper})\s*:?|\Z)")
            m = re.search(pat, text, re.DOTALL | re.IGNORECASE)
            if m:
                val = self._clean(m.group(1))
                val = self._cut_at_sentences(val, max_sent)
                val = re.sub(r"\s{2,}", " ", val).strip("- \n")
                if len(val) > 15:
                    return val
        return ""

    # ─────────────────────────────────────────────────────────────────
    # Parsing
    # ─────────────────────────────────────────────────────────────────

    def _parse(self, raw: str, patient: PatientInfo,
                imaging: ImagingInfo) -> RadiologyReport:
        report   = RadiologyReport()
        view_str = f", {imaging.view} view" if imaging.view else ""

        report.examination_details = (
            f"{imaging.modality} examination of the {imaging.body_region}"
            f"{view_str} was performed."
        )
        # Fix double periods (e.g. "infection.." -> "infection.")
        symptoms_clean = patient.symptoms.rstrip(". ")
        history_clean  = patient.clinical_history.rstrip(". ")
        report.clinical_information = (
            f"{patient.age}-year-old {patient.gender} presenting with "
            f"{symptoms_clean}. History: {history_clean}."
        )

        full = self._clean(
            f"FINDINGS: The {imaging.body_region} demonstrates " + raw
        )

        report.findings        = self._extract(full, "findings",        3)
        report.impression      = self._extract(full, "impression",      2)
        report.recommendations = self._extract(full, "recommendations", 2)

        # Findings fallback
        if len(report.findings) < 15:
            lines = [l.strip() for l in full.split("\n") if len(l.strip()) > 30]
            if lines:
                candidate = re.sub(r"^FINDINGS\s*:\s*", "", lines[0], flags=re.I)
                report.findings = self._cut_at_sentences(candidate, 3)

        # Impression fallback
        if len(report.impression) < 15:
            f_lower = report.findings.lower()
            abnormal_keywords = [
                "opacity", "opacities", "ground glass", "consolidation",
                "lesion", "mass", "nodule", "effusion", "infiltrate",
                "atelectasis", "fibrosis", "infection",
            ]
            normal_keywords = [
                "no acute", "no evidence", "unremarkable",
                "within normal", "normal", "intact",
            ]
            if any(w in f_lower for w in abnormal_keywords):
                report.impression = (
                    f"Abnormal findings in the {imaging.body_region.lower()} as described. "
                    f"Clinical correlation is recommended."
                )
            elif any(w in f_lower for w in normal_keywords):
                report.impression = (
                    f"No acute {imaging.body_region.lower()} abnormality identified. "
                    f"Clinical correlation is recommended."
                )
            else:
                report.impression = (
                    f"Findings in the {imaging.body_region.lower()} as described. "
                    f"Clinical correlation is recommended."
                )

        # Recommendations fallback
        if len(report.recommendations) < 15:
            report.recommendations = (
                "Clinical correlation with the patient's history and examination "
                "findings is recommended. Follow-up imaging as clinically indicated."
            )

        return report

    # ─────────────────────────────────────────────────────────────────
    # Consistency check
    # ─────────────────────────────────────────────────────────────────

    def _validate_consistency(self, report: RadiologyReport,
                               imaging: ImagingInfo):
        f   = report.findings.lower()
        imp = report.impression.lower()
        abnormal_keywords = [
            "opacity", "opacities", "ground glass", "consolidation",
            "lesion", "mass", "nodule", "effusion", "infiltrate",
        ]
        if any(w in f for w in abnormal_keywords):
            if any(x in imp for x in ["no acute", "normal", "unremarkable"]):
                report.impression = self._build_specific_impression(
                    report, imaging
                )

    def _build_specific_impression(self, report: RadiologyReport,
                                    imaging: ImagingInfo) -> str:
        """
        Builds a specific, clinically meaningful impression by:
        1. Extracting the key finding from FINDINGS
        2. Incorporating clinical history context
        3. Suggesting a differential diagnosis
        """
        f    = report.findings.lower()
        hist = report.clinical_information.lower()
        region = imaging.body_region.lower()

        # ── Identify key finding ───────────────────────────────────
        if "ground glass" in f or "ggo" in f:
            key_finding = "Ground-glass opacities"
            if "bilateral" in f or "both lung" in f:
                key_finding = "Bilateral ground-glass opacities"
        elif "consolidation" in f:
            key_finding = "Airspace consolidation"
            if "bilateral" in f:
                key_finding = "Bilateral consolidation"
            elif "right" in f:
                key_finding = "Right-sided consolidation"
            elif "left" in f:
                key_finding = "Left-sided consolidation"
        elif "opacity" in f or "opacities" in f:
            key_finding = "Pulmonary opacities"
            if "bilateral" in f:
                key_finding = "Bilateral pulmonary opacities"
        elif "nodule" in f:
            key_finding = "Pulmonary nodule(s)"
        elif "effusion" in f:
            key_finding = "Pleural effusion"
            if "bilateral" in f:
                key_finding = "Bilateral pleural effusions"
        elif "mass" in f or "lesion" in f:
            key_finding = f"Focal {region} lesion"
        elif "infiltrate" in f:
            key_finding = "Pulmonary infiltrates"
        else:
            key_finding = f"Abnormal {region} findings"

        # ── Clinical context from history ──────────────────────────
        context = ""
        if any(w in hist for w in ["immunosuppress", "immunocompromis",
                                    "mycophenol", "steroid", "transplant",
                                    "hiv", "aids", "chemotherapy"]):
            context = (
                f" In the context of immunosuppression, opportunistic "
                f"infection (e.g. PCP, fungal, CMV pneumonitis) should "
                f"be considered."
            )
        elif any(w in hist for w in ["infection", "pneumonia", "fever", "cough"]):
            context = (
                f" Findings are consistent with an infective or "
                f"inflammatory aetiology."
            )
        elif any(w in hist for w in ["malignancy", "cancer", "tumour",
                                      "tumor", "oncology", "lymphoma"]):
            context = (
                f" In the context of known malignancy, metastatic "
                f"disease should be considered."
            )
        elif any(w in hist for w in ["autoimmune", "lupus", "rheumatoid",
                                      "sarcoid", "vasculitis"]):
            context = (
                f" In the context of autoimmune disease, an inflammatory "
                f"or drug-related aetiology should be considered."
            )

        # ── Build impression ───────────────────────────────────────
        impression = (
            f"{key_finding} as described above, likely representing "
            f"an infective or inflammatory process.{context} "
            f"Clinical correlation is recommended."
        )
        return impression

    # ─────────────────────────────────────────────────────────────────
    # Confidence scoring
    # ─────────────────────────────────────────────────────────────────

    def _score(self, report: RadiologyReport, imaging: ImagingInfo) -> float:
        score = 0.45
        f     = report.findings

        if len(f) > 40:  score += 0.10
        if len(f) > 100: score += 0.05

        anatomy = {
            "brain":   ["parenchyma", "ventricle", "sulci", "cortex",
                        "signal", "midline", "fossa", "white matter"],
            "chest":   ["lung", "cardiac", "pleural", "mediastinum",
                        "consolidation", "silhouette", "costophrenic"],
            "pelvis":  ["femoral", "acetabul", "iliac", "hip", "sacrum",
                        "hardware", "screw", "fixation", "prosthes"],
            "abdomen": ["liver", "spleen", "kidney", "bowel",
                        "pancreas", "peritoneal", "aorta"],
            "spine":   ["vertebr", "disc", "canal", "foramina",
                        "alignment", "cord", "facet"],
            "knee":    ["condyle", "meniscus", "patella", "tibial",
                        "ligament", "joint space"],
        }
        region_key = imaging.body_region.lower()
        terms      = anatomy.get(region_key, [])
        matched    = sum(1 for t in terms if t in f.lower())
        score     += min(matched * 0.03, 0.15)

        imp  = report.impression
        if len(imp) > 30: score += 0.08
        good = ["consistent with", "no evidence", "no acute", "identified",
                "findings", "normal", "unremarkable", "in keeping with"]
        if any(p in imp.lower() for p in good):
            score += 0.08
        if "clinical presentation of a" not in imp:
            score += 0.05

        if len(report.recommendations) > 20:
            score += 0.04

        return round(max(0.45, min(0.88, score)), 2)