generate_pdf.py

"""
Tri-Netra - Patient Diagnostic Summary PDF Generator
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Generates a clean, official-looking PDF report with structured sections
for patient info, model inference statistics, a Grad-CAM heatmap
placeholder, and a clinician signature block.

Requires: fpdf2  (pip install fpdf2)

Author : Anannya Vyas
Email  : 
"""

from __future__ import annotations

import os
from datetime import datetime
from pathlib import Path

try:
    from fpdf import FPDF
except ImportError:
    raise ImportError(
        "fpdf2 is required.  Install it with:  pip install fpdf2"
    )


# - Brand colours (teal / emerald / gold matching Tri-Netra theme) -
_TEAL = (79, 70, 229)       # primary
_DARK_TEAL = (30, 41, 59)     # header bg
_GOLD = (245, 158, 11)       # accent
_LIGHT_BG = (248, 250, 252)  # section bg
_WHITE = (255, 255, 255)
_BLACK = (30, 30, 30)
_GRAY = (120, 120, 120)


class _ReportPDF(FPDF):
    """Custom FPDF subclass with Tri-Netra branded header / footer."""

    def header(self):
        # - Dark-teal banner -
        self.set_fill_color(*_DARK_TEAL)
        self.rect(0, 0, 210, 28, style="F")

        # Load and place logo if available
        logo_path = Path(__file__).resolve().parent / "Dashboard_Images" / "logo.png"
        text_start_x = 10

        # Brand name
        self.set_font("Helvetica", "B", 18)
        self.set_text_color(*_WHITE)
        self.set_xy(text_start_x, 6)
        self.cell(0, 10, "Tri-Netra", align="L")

        # Subtitle
        self.set_font("Helvetica", "", 9)
        self.set_text_color(200, 230, 225)
        self.set_xy(text_start_x, 16)
        self.cell(0, 6, "Patient Diagnostic Summary", align="L")

        # Right-side tagline
        self.set_font("Helvetica", "I", 8)
        self.set_text_color(200, 230, 225)
        self.set_xy(-70, 10)
        self.cell(60, 6, "AI-Assisted MRI Analysis", align="R")

        self.ln(30)

    def footer(self):
        self.set_y(-18)
        self.set_draw_color(*_TEAL)
        self.line(10, self.get_y(), 200, self.get_y())
        self.ln(2)
        self.set_font("Helvetica", "I", 7)
        self.set_text_color(*_GRAY)
        self.cell(
            0, 5,
            "Generated by Tri-Netra  |  Anannya Vyas",
            align="C",
        )
        self.ln(3)
        self.set_font("Helvetica", "", 7)
        self.cell(0, 4, f"Page {self.page_no()}/{{nb}}", align="C")


# - Helper: section heading -

def _section_heading(pdf: _ReportPDF, title: str) -> None:
    pdf.set_font("Helvetica", "B", 11)
    pdf.set_text_color(*_DARK_TEAL)
    pdf.set_fill_color(*_LIGHT_BG)
    pdf.cell(0, 8, f"  {title}", new_x="LMARGIN", new_y="NEXT", fill=True)
    pdf.ln(2)


def _label_value_row(pdf: _ReportPDF, label: str, value: str) -> None:
    pdf.set_font("Helvetica", "B", 9)
    pdf.set_text_color(*_BLACK)
    pdf.cell(55, 6, label, border=0)
    pdf.set_font("Helvetica", "", 9)
    pdf.set_text_color(60, 60, 60)
    pdf.cell(0, 6, value, new_x="LMARGIN", new_y="NEXT")


# - Public API -

def generate_report(
    output_path: str | Path = "diagnostic_summary.pdf",
    prediction_stats: dict | None = None,
    gradcam_path: str | Path | None = None,
    patient_name: str = "",
    patient_id: str = "",
    diagnostic_note: str = "",
) -> Path:
    """Generate a branded Patient Diagnostic Summary PDF.

    Parameters
    ----------
    output_path : str or Path
        Where to save the PDF.
    prediction_stats : dict or None
        Model inference statistics.  Expected keys (all optional):
        ``model_type``, ``prediction_confidence``, ``tumor_class``,
        ``inference_time_ms``, ``timestamp``.
    gradcam_path : str, Path or None
        Absolute or relative path to a Grad-CAM heatmap image (PNG/JPG).
        If ``None``, a placeholder box is drawn instead.
    patient_name : str
        Leave blank to show a fillable placeholder.
    patient_id : str
        Leave blank to show a fillable placeholder.
    diagnostic_note : str
        Free-text note (e.g. from ``report_assistant`` LLM output).

    Returns
    -------
    Path
        The resolved path to the generated PDF file.
    """

    stats = prediction_stats or {}
    output_path = Path(output_path)
    pdf = _ReportPDF(orientation="P", unit="mm", format="A4")
    pdf.alias_nb_pages()
    pdf.set_auto_page_break(auto=True, margin=15)
    pdf.add_page()

    # - 1. Patient Information -
    _section_heading(pdf, "PATIENT INFORMATION")
    _label_value_row(pdf, "Patient Name:", patient_name or "______________________________")
    _label_value_row(pdf, "Patient ID:", patient_id or "______________________________")
    _label_value_row(pdf, "Report Date:", datetime.now().strftime("%Y-%m-%d  %H:%M"))
    _label_value_row(pdf, "Referring Physician:", "______________________________")
    pdf.ln(4)

    # - 2. Analysis Results -
    _section_heading(pdf, "ANALYSIS RESULTS")
    
    tumor_class = stats.get("tumor_class", "-")
    verdict = "Tumor Detected" if tumor_class.lower() != "no_tumor" else "No Abnormalities Detected"
    _label_value_row(pdf, "Diagnosis:", verdict)
    
    _label_value_row(
        pdf, "Confidence:",
        f"{stats.get('unified_confidence', stats.get('prediction_confidence', '-'))}%"
    )
    
    if "volume_cm3" in stats:
        _label_value_row(pdf, "Estimated Tumor Volume:", f"{stats['volume_cm3']} cm³")
        
    pdf.ln(4)
    
    # - 3. Risk Assessment & Follow Up -
    _section_heading(pdf, "RISK ASSESSMENT & FOLLOW-UP")
    _label_value_row(pdf, "Risk Score (0-100):", str(stats.get('risk_score', 'N/A')))
    _label_value_row(pdf, "Risk Level:", str(stats.get('risk_level', 'N/A')))
    _label_value_row(pdf, "Recommended Follow-Up:", str(stats.get('follow_up', 'Please consult your doctor.')))
    pdf.ln(4)

    # - 4. Grad-CAM Heatmap -
    _section_heading(pdf, "VISUALIZATION")

    if gradcam_path and Path(gradcam_path).exists():
        # Insert actual image, scaled to fit
        img_x = 30
        img_w = 150
        pdf.image(str(gradcam_path), x=img_x, w=img_w)
    else:
        # Draw placeholder box
        box_x, box_y = 30, pdf.get_y()
        box_w, box_h = 150, 60
        pdf.set_draw_color(*_TEAL)
        pdf.set_line_width(0.5)
        pdf.rect(box_x, box_y, box_w, box_h, style="D")

        # Label
        pdf.set_xy(box_x, box_y + 25)
        pdf.set_font("Helvetica", "I", 10)
        pdf.set_text_color(*_GRAY)
        pdf.cell(box_w, 8, "[ No Heatmap Available ]", align="C")

        pdf.set_y(box_y + box_h + 4)
    pdf.ln(4)

    # - 5. Clinical Disclaimer -
    _section_heading(pdf, "CLINICAL DISCLAIMER")
    pdf.set_font("Helvetica", "I", 9)
    pdf.set_text_color(220, 50, 50)
    pdf.multi_cell(0, 5, "DISCLAIMER: This report is generated by an AI research prototype. It is NOT a verified clinical diagnosis. The risk scores and follow-up recommendations are for informational purposes only. Do not make medical decisions based on this report. Please consult a qualified healthcare professional.")
    pdf.ln(6)

    # - 6. Clinician Signature Block -
    _section_heading(pdf, "CLINICIAN REVIEW & SIGNATURE")
    pdf.ln(2)

    col_w = 90
    start_x = pdf.get_x()
    start_y = pdf.get_y()

    # Left column
    pdf.set_xy(start_x, start_y)
    _label_value_row(pdf, "Clinician Name:", "______________________________")
    _label_value_row(pdf, "Designation:", "______________________________")
    _label_value_row(pdf, "Hospital / Institution:", "______________________________")

    # Right column (signature box)
    sig_x = start_x + col_w + 10
    pdf.set_draw_color(*_TEAL)
    pdf.set_line_width(0.4)
    pdf.rect(sig_x, start_y, 80, 25, style="D")
    pdf.set_xy(sig_x, start_y + 8)
    pdf.set_font("Helvetica", "I", 8)
    pdf.set_text_color(*_GRAY)
    pdf.cell(80, 6, "Signature", align="C")

    pdf.set_y(start_y + 30)
    _label_value_row(pdf, "Date of Review:", "______________________________")

    pdf_out = Path(output_path)
    pdf.output(str(pdf_out))
    return pdf_out