export.py

import io
import tempfile
from datetime import datetime

from fpdf import FPDF
from PIL import Image as PILImage

PAGE_W   = 210
PAGE_H   = 297
MARGIN   = 15
USABLE_W = PAGE_W - 2 * MARGIN   # 180 mm

N_COLS  = 3
COL_GAP = 3
IMG_W   = (USABLE_W - (N_COLS - 1) * COL_GAP) / N_COLS   # ~58 mm
IMG_H   = IMG_W
CAP_H   = 5

INPUT_W = 80

FS_TITLE = 16
FS_HEAD  = 11
FS_BODY  = 9
FS_SMALL = 7
FS_MONO  = 8

DARK_GRAY  = (50,  50,  50)
MID_GRAY   = (110, 110, 110)
LIGHT_GRAY = (200, 200, 200)
BLACK      = (0,   0,   0)
RED_TEXT   = (180, 30,  30)
GREEN_TEXT = (30,  120, 30)
TBL_HDR    = (230, 238, 248)
TBL_ALT    = (245, 249, 253)
TBL_BOR    = (185, 200, 218)

LABELS = {
    "YOLO conf":    "YOLO Confidence",
    "Seg failed":   "Segmentation Failed",
    "vCDR":         "Vertical CDR",
    "hCDR":         "Horizontal CDR",
    "Inferior rim": "Inferior Rim Ratio",
    "Superior rim": "Superior Rim Ratio",
    "SI compliant": "Inferior > Superior Rim Rule",
    "NRR area":     "NRR Area Ratio",
    "p holistic":   "EfficientNet Probability",
    "p fused":      "Fused Glaucoma Score",
    "Threshold":    "Decision Threshold",
}


def _safe(text: str) -> str:
    text = (text
        .replace('—', '-').replace('–', '-')
        .replace('‘', "'").replace('’', "'")
        .replace('“', '"').replace('”', '"')
        .replace('…', '...').replace('•', '*')
        .replace('°', 'deg').replace('±', '+/-')
    )
    return text.encode('latin-1', errors='replace').decode('latin-1')


def _png(img: PILImage.Image) -> io.BytesIO:
    buf = io.BytesIO()
    img.save(buf, format='PNG')
    buf.seek(0)
    return buf


def _fmt(v) -> str:
    if v is True:  return "Yes"
    if v is False: return "No"
    if isinstance(v, float): return f"{v:.4f}"
    return str(v)


def _hline(pdf: FPDF):
    y = pdf.get_y()
    pdf.set_draw_color(*LIGHT_GRAY)
    pdf.set_line_width(0.25)
    pdf.line(MARGIN, y, PAGE_W - MARGIN, y)
    pdf.set_draw_color(0, 0, 0)
    pdf.set_line_width(0.2)


def _section(pdf: FPDF, title: str, gap_before: float = 7):
    pdf.ln(gap_before)
    pdf.set_font("Helvetica", style='B', size=FS_HEAD)
    pdf.set_text_color(*DARK_GRAY)
    pdf.cell(0, 6, title, new_x='LMARGIN', new_y='NEXT')
    pdf.ln(2)
    pdf.set_text_color(*BLACK)


def _image_grid(pdf: FPDF, items: list):
    """
    3-column grid of (caption, PIL Image).
    Uses set_xy + image(w, h) [no explicit y] per image.
    Caption placed with pdf.text() — zero cursor impact.
    row_y is a local variable, never derived from get_y() after image placement.
    """
    row_y = None

    for i, (caption, img) in enumerate(items):
        col = i % N_COLS

        if col == 0:
            row_needed = IMG_H + CAP_H + 4
            if pdf.get_y() + row_needed > PAGE_H - MARGIN:
                pdf.add_page()
            row_y = pdf.get_y()

        x = MARGIN + col * (IMG_W + COL_GAP)

        # Place image at (x, row_y) without specifying y explicitly
        pdf.set_xy(x, row_y)
        pdf.image(_png(img), w=IMG_W, h=IMG_H)
        # fpdf2 advances cursor to (x + IMG_W, row_y + IMG_H) — we ignore it

        # Caption: pdf.text() places text at absolute coords, does NOT move cursor
        pdf.set_font("Helvetica", style='', size=FS_SMALL)
        pdf.set_text_color(*MID_GRAY)
        cap = _safe(caption)
        tw  = pdf.get_string_width(cap)
        pdf.text(x + (IMG_W - tw) / 2, row_y + IMG_H + 4, cap)
        pdf.set_text_color(*BLACK)

        # After last column (or last item), advance cursor past the row
        if col == N_COLS - 1 or i == len(items) - 1:
            pdf.set_xy(MARGIN, row_y + IMG_H + CAP_H + 4)


def _table(pdf: FPDF, metrics: dict):
    """
    2-column table.  Uses new_y='TOP' for first cell (y stays),
    new_y='NEXT' for last cell (y advances one row).
    No TMARGIN — that resets y to page top and breaks layout.
    """
    lw = USABLE_W * 0.68
    vw = USABLE_W * 0.32
    rh = 7

    pdf.set_draw_color(*BLACK)

    # Header row
    pdf.set_font("Helvetica", style='B', size=FS_SMALL)
    pdf.set_text_color(*BLACK)
    pdf.cell(lw, rh, "BIOMARKER / SCORE", border=1,
             new_x='RIGHT', new_y='TOP')
    pdf.cell(vw, rh, "VALUE", border=1,
             new_x='LMARGIN', new_y='NEXT')

    # Data rows
    for k, v in metrics.items():
        pdf.set_font("Helvetica", style='', size=FS_BODY)
        pdf.set_text_color(*BLACK)
        pdf.cell(lw, rh, LABELS.get(k, k), border=1,
                 new_x='RIGHT', new_y='TOP')
        pdf.set_font("Courier", style='', size=FS_MONO)
        pdf.cell(vw, rh, _fmt(v), border=1,
                 new_x='LMARGIN', new_y='NEXT')


def generate_pdf_report(data: dict) -> str:
    pdf = FPDF(orientation='P', unit='mm', format='A4')
    pdf.set_margins(MARGIN, MARGIN, MARGIN)
    pdf.set_auto_page_break(auto=True, margin=MARGIN)
    pdf.add_page()

    # ── TITLE ────────────────────────────────────────────────────
    pdf.set_font("Helvetica", style='B', size=FS_TITLE)
    pdf.set_text_color(*DARK_GRAY)
    pdf.cell(0, 9, "GlauNET - Glaucoma Screening Report",
             align='C', new_x='LMARGIN', new_y='NEXT')

    pdf.set_font("Helvetica", style='', size=FS_SMALL)
    pdf.set_text_color(*MID_GRAY)
    pdf.cell(0, 5, f"Generated: {datetime.now().strftime('%Y-%m-%d  %H:%M:%S')}",
             align='C', new_x='LMARGIN', new_y='NEXT')

    pdf.set_font("Helvetica", style='I', size=FS_SMALL - 1)
    pdf.cell(0, 4,
             "Research prototype - not a substitute for diagnosis by a qualified ophthalmologist.",
             align='C', new_x='LMARGIN', new_y='NEXT')

    pdf.ln(3)
    _hline(pdf)
    pdf.set_text_color(*BLACK)

    # ── DIAGNOSIS ────────────────────────────────────────────────
    _section(pdf, "Automated Diagnosis", gap_before=5)
    decision    = data["decision"]
    is_positive = "SUSPECTED" in decision
    pdf.set_font("Helvetica", style='B', size=14)
    pdf.set_text_color(*BLACK)
    pdf.cell(0, 9, _safe(decision), align='C', new_x='LMARGIN', new_y='NEXT')

    # ── INPUT FUNDUS IMAGE ───────────────────────────────────────
    _section(pdf, "Input Fundus Image")

    if pdf.get_y() + INPUT_W + 8 > PAGE_H - MARGIN:
        pdf.add_page()

    x_img   = MARGIN + (USABLE_W - INPUT_W) / 2
    img_top = pdf.get_y()                           # save y BEFORE image
    pdf.set_xy(x_img, img_top)
    pdf.image(_png(data["input_image"]), w=INPUT_W, h=INPUT_W)
    # fpdf2 cursor is now at img_top + INPUT_W — we IGNORE it and use img_top

    pdf.set_font("Helvetica", style='I', size=FS_SMALL)
    pdf.set_text_color(*MID_GRAY)
    cap = "Original fundus photograph"
    pdf.text(MARGIN + (USABLE_W - pdf.get_string_width(cap)) / 2,
             img_top + INPUT_W + 4, cap)
    pdf.set_text_color(*BLACK)

    # Advance cursor manually: img_top + image height + caption gap
    pdf.set_xy(MARGIN, img_top + INPUT_W + 8)

    # ── SEGMENTATION PIPELINE (first 5 images) ───────────────────
    all_images = list(data["images"].items())
    _section(pdf, "Segmentation Pipeline")
    _image_grid(pdf, all_images[:5])

    # ── GRADCAM ANALYSIS (last 4 images) ─────────────────────────
    _section(pdf, "GradCAM Analysis")
    _image_grid(pdf, all_images[5:])

    # ── CLINICAL BIOMARKERS ──────────────────────────────────────
    _section(pdf, "Clinical Biomarkers")
    _table(pdf, data["metrics"])
    pdf.ln(3)

    # ── MEDGEMMA CLINICAL REPORT ─────────────────────────────────
    _section(pdf, "MedGemma Clinical Report")
    pdf.set_font("Courier", style='', size=FS_MONO)
    pdf.set_text_color(*DARK_GRAY)
    report = _safe(data.get("report", ""))
    if report:
        pdf.multi_cell(USABLE_W, 5, report)
    else:
        pdf.set_text_color(*MID_GRAY)
        pdf.cell(0, 5, "(No report generated)", new_x='LMARGIN', new_y='NEXT')
    pdf.set_text_color(*BLACK)


    tmp = tempfile.NamedTemporaryFile(
        suffix=".pdf", prefix="glaunet_report_", delete=False
    )
    tmp.close()
    pdf.output(tmp.name)
    return tmp.name