utils/report.py

"""Report and heatmap export utilities."""
import io
from datetime import datetime

import cv2
import numpy as np
from PIL import Image
from reportlab.lib.pagesizes import A4
from reportlab.lib.units import inch
from reportlab.lib.utils import ImageReader
from reportlab.pdfgen import canvas

from config.constants import COLORMAPS


def _draw_pdf_footer(c, margin=72, footer_y=40, include_web_app=False):
    """Draw common footer for S2F PDF reports."""
    c.setFont("Helvetica", 8)
    c.setFillColorRGB(0.4, 0.4, 0.4)
    gen_date = datetime.now().strftime("%Y-%m-%d %H:%M")
    c.drawString(margin, footer_y, f"Generated by Shape2Force (S2F) on {gen_date}")
    c.drawString(margin, footer_y - 12, "Model: https://huggingface.co/Angione-Lab/Shape2Force")
    if include_web_app:
        c.drawString(margin, footer_y - 24, "Web app: https://huggingface.co/spaces/Angione-Lab/Shape2force")
    c.setFillColorRGB(0, 0, 0)


def _pdf_image_layout(page_w_pt, page_h_pt, margin=72, n_images=2):
    """Return layout dict for centered side-by-side images: img_w, img_h, img_gap, img_left, bf_x, hm_x, img_bottom, y_top."""
    img_w = 2.8 * inch
    img_h = 2.8 * inch
    img_gap = 20
    total_img_width = n_images * img_w + (n_images - 1) * img_gap
    img_left = margin + (page_w_pt - 2 * margin - total_img_width) / 2
    bf_x = img_left
    hm_x = img_left + img_w + img_gap
    y_top = page_h_pt - 50
    img_bottom = y_top - 35 - img_h  # header (title + image name) takes 35pt
    return {
        "img_w": img_w,
        "img_h": img_h,
        "img_gap": img_gap,
        "img_left": img_left,
        "bf_x": bf_x,
        "hm_x": hm_x,
        "img_bottom": img_bottom,
        "y_top": y_top,
    }


def heatmap_to_rgb(scaled_heatmap, colormap_name="Jet", zmin=None, zmax=None):
    """Convert scaled heatmap to RGB array using the given colormap.
    If zmin, zmax are provided (e.g. for Range mode), map [zmin,zmax] to 0-1 for coloring."""
    arr = np.asarray(scaled_heatmap, dtype=np.float32)
    if zmin is not None and zmax is not None and zmax > zmin:
        arr = np.clip((arr - zmin) / (zmax - zmin), 0, 1)
    else:
        arr = np.clip(arr, 0, 1)
    heatmap_uint8 = (arr * 255).astype(np.uint8)
    cv2_colormap = COLORMAPS.get(colormap_name, cv2.COLORMAP_JET)
    heatmap_rgb = cv2.cvtColor(cv2.applyColorMap(heatmap_uint8, cv2_colormap), cv2.COLOR_BGR2RGB)
    return heatmap_rgb


def heatmap_to_rgb_with_contour(scaled_heatmap, colormap_name="Jet", cell_mask=None, zmin=None, zmax=None):
    """Convert heatmap to RGB, optionally drawing red cell contour. Mask must match heatmap shape."""
    heatmap_rgb = heatmap_to_rgb(scaled_heatmap, colormap_name, zmin=zmin, zmax=zmax)
    if cell_mask is not None and np.any(cell_mask > 0):
        contours, _ = cv2.findContours(cell_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if contours:
            cv2.drawContours(heatmap_rgb, contours, -1, (255, 0, 0), 3)
    return heatmap_rgb


def heatmap_to_png_bytes(scaled_heatmap, colormap_name="Jet", cell_mask=None, zmin=None, zmax=None):
    """Convert scaled heatmap to PNG bytes buffer. Optionally draw red cell contour.
    If zmin, zmax provided (Range mode), map that range to full colormap."""
    heatmap_rgb = heatmap_to_rgb_with_contour(scaled_heatmap, colormap_name, cell_mask, zmin=zmin, zmax=zmax)
    buf = io.BytesIO()
    Image.fromarray(heatmap_rgb).save(buf, format="PNG")
    buf.seek(0)
    return buf


def create_pdf_report(img, display_heatmap, raw_heatmap, pixel_sum, force, base_name, colormap_name="Jet",
                      cell_mask=None, cell_pixel_sum=None, cell_force=None, cell_mean=None, zmin=None, zmax=None):
    """Create a PDF report with input image, heatmap, and metrics."""
    buf = io.BytesIO()
    c = canvas.Canvas(buf, pagesize=A4)
    c.setTitle("Shape2Force")
    c.setAuthor("Angione-Lab")
    page_w_pt, page_h_pt = A4[0], A4[1]
    margin = 72
    layout = _pdf_image_layout(page_w_pt, page_h_pt, margin)
    img_w = layout["img_w"]
    img_h = layout["img_h"]
    bf_x = layout["bf_x"]
    hm_x = layout["hm_x"]
    img_bottom = layout["img_bottom"]
    y_top = layout["y_top"]

    _draw_pdf_footer(c, margin=margin, include_web_app=True)
    c.setFont("Helvetica-Bold", 16)
    c.drawString(margin, y_top, "Shape2Force (S2F) - Prediction Report")
    c.setFont("Helvetica", 10)
    c.drawString(margin, y_top - 14, f"Image: {base_name}")
    y_top -= 35

    img_pil = Image.fromarray(img) if img.ndim == 2 else Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    img_buf = io.BytesIO()
    img_pil.save(img_buf, format="PNG")
    img_buf.seek(0)
    c.drawImage(ImageReader(img_buf), bf_x, img_bottom, width=img_w, height=img_h, preserveAspectRatio=True)
    c.setFont("Helvetica", 9)
    bf_label_w = c.stringWidth("Bright-field", "Helvetica", 9)
    c.drawString(bf_x + (img_w - bf_label_w) / 2, img_bottom - 14, "Bright-field")

    heatmap_rgb = heatmap_to_rgb_with_contour(display_heatmap, colormap_name, cell_mask, zmin=zmin, zmax=zmax)
    hm_buf = io.BytesIO()
    Image.fromarray(heatmap_rgb).save(hm_buf, format="PNG")
    hm_buf.seek(0)
    c.drawImage(ImageReader(hm_buf), hm_x, img_bottom, width=img_w, height=img_h, preserveAspectRatio=True)
    hm_label = "Force map (red = estimated boundary)" if cell_mask is not None and np.any(cell_mask > 0) else "Force map"
    hm_label_w = c.stringWidth(hm_label, "Helvetica", 9)
    c.drawString(hm_x + (img_w - hm_label_w) / 2, img_bottom - 14, hm_label)

    # Metrics section with spacing
    row_height = 14
    y = img_bottom - 14 - row_height
    c.setFont("Helvetica-Bold", 10)
    c.drawString(margin, y, "Metrics")
    c.setFont("Helvetica", 9)
    y -= 18
    if cell_pixel_sum is not None and cell_force is not None and cell_mean is not None:
        metrics = [
            ("Cell sum (estimated boundary)", f"{cell_pixel_sum:.2f}"),
            ("Cell force (scaled)", f"{cell_force:.2f}"),
            ("Heatmap max", f"{np.max(raw_heatmap):.4f}"),
            ("Heatmap mean (estimated boundary)", f"{cell_mean:.4f}"),
        ]
    else:
        metrics = [
            ("Sum of all pixels", f"{pixel_sum:.2f}"),
            ("Cell force (scaled)", f"{force:.2f}"),
            ("Heatmap max", f"{np.max(raw_heatmap):.4f}"),
            ("Heatmap mean", f"{np.mean(raw_heatmap):.4f}"),
        ]
    for label, val in metrics:
        c.drawString(margin, y, f"{label}: {val}")
        y -= 16

    c.save()
    buf.seek(0)
    return buf.getvalue()


def create_measure_pdf_report(bf_img, heatmap_labeled_rgb, table_rows, base_name):
    """
    Create PDF report for measure tool.
    Contents: bright-field, heatmap with region labels (R1, R2...), table.
    """
    buf = io.BytesIO()
    c = canvas.Canvas(buf, pagesize=A4)
    c.setTitle("Shape2Force - Region Measurement")
    c.setAuthor("Angione-Lab")
    page_w_pt, page_h_pt = A4[0], A4[1]
    margin = 72
    layout = _pdf_image_layout(page_w_pt, page_h_pt, margin)
    img_w = layout["img_w"]
    img_h = layout["img_h"]
    bf_x = layout["bf_x"]
    hm_x = layout["hm_x"]
    img_bottom = layout["img_bottom"]
    y_top = layout["y_top"]

    _draw_pdf_footer(c, margin=margin)
    c.setFont("Helvetica-Bold", 14)
    c.drawString(margin, y_top, "Region Measurement Report")
    c.setFont("Helvetica", 10)
    c.drawString(margin, y_top - 14, f"Image: {base_name}")

    bf_pil = Image.fromarray(bf_img) if bf_img.ndim == 2 else Image.fromarray(
        cv2.cvtColor(bf_img, cv2.COLOR_BGR2RGB)
    )
    bf_buf = io.BytesIO()
    bf_pil.save(bf_buf, format="PNG")
    bf_buf.seek(0)
    c.drawImage(ImageReader(bf_buf), bf_x, img_bottom, width=img_w, height=img_h, preserveAspectRatio=True)
    c.setFont("Helvetica", 9)
    bf_label_w = c.stringWidth("Bright-field", "Helvetica", 9)
    c.drawString(bf_x + (img_w - bf_label_w) / 2, img_bottom - 14, "Bright-field")

    hm_labeled_buf = io.BytesIO()
    Image.fromarray(heatmap_labeled_rgb).save(hm_labeled_buf, format="PNG")
    hm_labeled_buf.seek(0)
    c.drawImage(ImageReader(hm_labeled_buf), hm_x, img_bottom, width=img_w, height=img_h, preserveAspectRatio=True)
    hm_label_w = c.stringWidth("Force map", "Helvetica", 9)
    c.drawString(hm_x + (img_w - hm_label_w) / 2, img_bottom - 14, "Force map")

    # Spacing: one line between images/labels and measurements section
    row_height = 14
    y_table_top = img_bottom - 14 - row_height  # label + one line space

    # Table with borders
    n_cols = len(table_rows[0]) if table_rows else 0
    table_width = page_w_pt - 2 * margin
    col_w = table_width / n_cols if n_cols else 1
    cell_pad = 4

    c.setFont("Helvetica-Bold", 9)
    c.drawString(margin, y_table_top, "Measurements")
    y_table_top -= row_height + 8  # extra space before table

    for ri, row in enumerate(table_rows):
        y_cell_bottom = y_table_top - (ri + 1) * row_height
        for ci, cell in enumerate(row):
            x_left = margin + ci * col_w
            # Draw cell border
            c.rect(x_left, y_cell_bottom, col_w, row_height, stroke=1, fill=0)
            # Draw text (with padding)
            if ri == 0:
                c.setFont("Helvetica-Bold", 8)
            c.drawString(x_left + cell_pad, y_cell_bottom + 4, str(cell)[:22])
            if ri == 0:
                c.setFont("Helvetica", 8)

    c.save()
    buf.seek(0)
    return buf.getvalue()