app/tiles.py

"""
File: tiles.py
Author: Dr. Gordon Wright
Description: Composite-tile + wireframe rendering for the six-emotion
             replication challenge.

             Each tile is a single PIL image (one per slot in the 2x3
             grid) and is also what the single-page PDF artefact
             stamps in. A tile shows:

                 - the intended emotion (header bar)
                 - the face crop, OR a landmark wireframe of the same
                   face (Format B — face-free)
                 - a small horizontal bar chart of the classifier's
                   full 7-emotion probability vector
                 - a one-line verdict: classifier agreed / disagreed

             The wireframe is rendered from the 478 MediaPipe Face
             Landmarker points using a hand-coded subset of the
             FACEMESH_* feature connections (face oval, eyes, brows,
             nose bridge, lips). No identifying detail is preserved.

License: MIT License
"""

from __future__ import annotations

import io
import tempfile
from typing import Dict, Optional, List, Tuple

import numpy as np
from PIL import Image, ImageDraw, ImageFont
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages

from app.config import DICT_EMO, COLORS
from app.session import (
    BASIC_EMOTIONS,
    EMOTION_TO_CLASSIFIER,
    Capture,
    empty_session,
)


# ---------- Face-mesh feature connections (subset) ---------------------
#
# Hand-curated edge lists taken from MediaPipe's FACEMESH_* constants —
# the face landmarker exposes 478 indexed points, and these connection
# pairs are the canonical feature outlines.

FACEMESH_FACE_OVAL = [
    (10, 338), (338, 297), (297, 332), (332, 284), (284, 251), (251, 389),
    (389, 356), (356, 454), (454, 323), (323, 361), (361, 288), (288, 397),
    (397, 365), (365, 379), (379, 378), (378, 400), (400, 377), (377, 152),
    (152, 148), (148, 176), (176, 149), (149, 150), (150, 136), (136, 172),
    (172, 58), (58, 132), (132, 93), (93, 234), (234, 127), (127, 162),
    (162, 21), (21, 54), (54, 103), (103, 67), (67, 109), (109, 10),
]

FACEMESH_LIPS = [
    # outer
    (61, 146), (146, 91), (91, 181), (181, 84), (84, 17), (17, 314),
    (314, 405), (405, 321), (321, 375), (375, 291),
    (61, 185), (185, 40), (40, 39), (39, 37), (37, 0), (0, 267),
    (267, 269), (269, 270), (270, 409), (409, 291),
    # inner
    (78, 95), (95, 88), (88, 178), (178, 87), (87, 14), (14, 317),
    (317, 402), (402, 318), (318, 324), (324, 308),
    (78, 191), (191, 80), (80, 81), (81, 82), (82, 13), (13, 312),
    (312, 311), (311, 310), (310, 415), (415, 308),
]

FACEMESH_LEFT_EYE = [
    (263, 249), (249, 390), (390, 373), (373, 374), (374, 380), (380, 381),
    (381, 382), (382, 362),
    (263, 466), (466, 388), (388, 387), (387, 386), (386, 385), (385, 384),
    (384, 398), (398, 362),
]

FACEMESH_RIGHT_EYE = [
    (33, 7), (7, 163), (163, 144), (144, 145), (145, 153), (153, 154),
    (154, 155), (155, 133),
    (33, 246), (246, 161), (161, 160), (160, 159), (159, 158), (158, 157),
    (157, 173), (173, 133),
]

FACEMESH_LEFT_EYEBROW = [
    (276, 283), (283, 282), (282, 295), (295, 285),
    (300, 293), (293, 334), (334, 296), (296, 336),
]

FACEMESH_RIGHT_EYEBROW = [
    (46, 53), (53, 52), (52, 65), (65, 55),
    (70, 63), (63, 105), (105, 66), (66, 107),
]

FACEMESH_NOSE = [
    (168, 6), (6, 197), (197, 195), (195, 5), (5, 4),
    (4, 1), (1, 19), (19, 94), (94, 2),
]

ALL_FEATURE_EDGES = (
    FACEMESH_FACE_OVAL
    + FACEMESH_LIPS
    + FACEMESH_LEFT_EYE
    + FACEMESH_RIGHT_EYE
    + FACEMESH_LEFT_EYEBROW
    + FACEMESH_RIGHT_EYEBROW
    + FACEMESH_NOSE
)


# ---------- Wireframe rendering ----------------------------------------

WIRE_BG = (245, 245, 248)
WIRE_LINE = (40, 40, 60)
WIRE_DOT = (110, 110, 140)


def render_wireframe(
    landmarks: List[Tuple[float, float]],
    image_size: Tuple[int, int],
    bbox: Tuple[int, int, int, int],
    output_size: Tuple[int, int] = (260, 260),
) -> Image.Image:
    """Draw an anonymised landmark mesh into an output image sized to
    match the face crop's aspect ratio.

    `landmarks` is the list of (nx, ny) normalised in [0, 1] over the
    original image (which had `image_size` = (W, H)). `bbox` is the face
    crop in *original* image pixels.
    """
    img_w, img_h = image_size
    sx, sy, ex, ey = bbox
    crop_w = max(1, ex - sx)
    crop_h = max(1, ey - sy)

    out_w, out_h = output_size
    img = Image.new("RGB", (out_w, out_h), WIRE_BG)
    draw = ImageDraw.Draw(img)

    # Map a normalised landmark (nx, ny in [0,1] of full image) into
    # the output canvas via the face bbox.
    def project(nx: float, ny: float) -> Tuple[float, float]:
        px = nx * img_w
        py = ny * img_h
        rel_x = (px - sx) / crop_w
        rel_y = (py - sy) / crop_h
        return rel_x * out_w, rel_y * out_h

    pts = [project(nx, ny) for (nx, ny) in landmarks]

    # Feature outlines
    for a, b in ALL_FEATURE_EDGES:
        if a < len(pts) and b < len(pts):
            draw.line([pts[a], pts[b]], fill=WIRE_LINE, width=1)

    # Dots for every landmark — light texture, no facial detail
    for x, y in pts:
        draw.ellipse([x - 0.6, y - 0.6, x + 0.6, y + 0.6], fill=WIRE_DOT)

    return img


# ---------- Bar chart strip --------------------------------------------

EMOTION_ORDER_CLASSIFIER = [DICT_EMO[i] for i in range(7)]
EMOTION_COLOURS = [COLORS[i] for i in range(7)]


def _bar_strip(
    capture: Capture,
    size: Tuple[int, int] = (260, 80),
) -> Image.Image:
    """Small horizontal bar chart of the classifier's 7-emotion vector.
    Bar matching the intended emotion is outlined; bar that the
    classifier picked as winner is annotated."""
    fig, ax = plt.subplots(figsize=(size[0] / 100, size[1] / 100), dpi=100)
    heights = [capture.emotion_probs.get(emo, 0.0) for emo in EMOTION_ORDER_CLASSIFIER]
    bars = ax.bar(EMOTION_ORDER_CLASSIFIER, heights, color=EMOTION_COLOURS, edgecolor="none")
    intended_classifier = EMOTION_TO_CLASSIFIER.get(capture.intended)
    for bar, emo in zip(bars, EMOTION_ORDER_CLASSIFIER):
        if emo == intended_classifier:
            bar.set_edgecolor("#222")
            bar.set_linewidth(1.5)
    ax.set_ylim(0, 1)
    ax.set_yticks([])
    ax.set_xticks(range(7))
    ax.set_xticklabels([e[:3] for e in EMOTION_ORDER_CLASSIFIER], fontsize=7)
    for spine in ("top", "right", "left"):
        ax.spines[spine].set_visible(False)
    ax.tick_params(axis="x", length=0, pad=1)
    fig.subplots_adjust(left=0.02, right=0.98, top=0.95, bottom=0.25)

    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=100)
    plt.close(fig)
    buf.seek(0)
    return Image.open(buf).convert("RGB")


# ---------- Tile composition -------------------------------------------

TILE_W = 280
TILE_H = 396
HEADER_H = 36
FACE_H = 220
BAR_H = 72
VERDICT_H = 44
PAD = 6

EMPTY_FACE_BG = (250, 250, 252)
EMPTY_FACE_FG = (200, 200, 210)
TILE_BG = (255, 255, 255)
TILE_BORDER = (220, 220, 228)
HEADER_BG = (30, 30, 40)
HEADER_FG = (255, 255, 255)
AGREE_BAND = (34, 139, 84)
DISAGREE_BAND = (200, 90, 40)
VERDICT_FG = (255, 255, 255)
EMPTY_BAND = (240, 240, 244)
EMPTY_BAND_FG = (140, 140, 150)


def _load_font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont:
    """Pillow has no reliable bold-by-default; fall back to default if
    DejaVuSans is missing."""
    candidates = (
        ("/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold
         else "/System/Library/Fonts/Supplemental/Arial.ttf"),
        ("/System/Library/Fonts/Helvetica.ttc"),
        ("DejaVuSans-Bold.ttf" if bold else "DejaVuSans.ttf"),
    )
    for path in candidates:
        try:
            return ImageFont.truetype(path, size)
        except (OSError, IOError):
            continue
    return ImageFont.load_default()


def _tile_canvas() -> Tuple[Image.Image, ImageDraw.ImageDraw]:
    img = Image.new("RGB", (TILE_W, TILE_H), TILE_BG)
    draw = ImageDraw.Draw(img)
    draw.rectangle([0, 0, TILE_W - 1, TILE_H - 1], outline=TILE_BORDER, width=1)
    return img, draw


def _draw_header(draw: ImageDraw.ImageDraw, intended: str, slot_num: int) -> None:
    draw.rectangle([0, 0, TILE_W, HEADER_H], fill=HEADER_BG)
    title = f"{slot_num}.  {intended.upper()}"
    font = _load_font(15, bold=True)
    draw.text((PAD * 2, 9), title, fill=HEADER_FG, font=font)


def _paste_face_region(
    tile: Image.Image,
    face_image: Image.Image,
) -> None:
    region_top = HEADER_H + PAD
    region = (PAD, region_top, TILE_W - PAD, region_top + FACE_H)
    target_w = region[2] - region[0]
    target_h = region[3] - region[1]
    scaled = face_image.resize((target_w, target_h), Image.Resampling.LANCZOS)
    tile.paste(scaled, (region[0], region[1]))


def render_empty_tile(intended: str, slot_num: int) -> Image.Image:
    img, draw = _tile_canvas()
    _draw_header(draw, intended, slot_num)
    region_top = HEADER_H + PAD
    draw.rectangle(
        [PAD, region_top, TILE_W - PAD, region_top + FACE_H],
        fill=EMPTY_FACE_BG,
        outline=EMPTY_FACE_FG,
    )
    font = _load_font(11)
    msg = "not yet attempted"
    bbox = draw.textbbox((0, 0), msg, font=font)
    tw = bbox[2] - bbox[0]
    th = bbox[3] - bbox[1]
    draw.text(
        ((TILE_W - tw) / 2, region_top + (FACE_H - th) / 2),
        msg,
        fill=EMPTY_FACE_FG,
        font=font,
    )
    bar_top = region_top + FACE_H + PAD
    draw.rectangle(
        [PAD, bar_top, TILE_W - PAD, bar_top + BAR_H],
        fill=EMPTY_FACE_BG,
        outline=EMPTY_FACE_FG,
    )
    # Empty verdict band — grey, mirrors the filled band so the grid
    # alignment stays consistent.
    verdict_top = bar_top + BAR_H + PAD
    draw.rectangle(
        [0, verdict_top, TILE_W, verdict_top + VERDICT_H],
        fill=EMPTY_BAND,
    )
    font = _load_font(13, bold=True)
    msg = "awaiting your attempt"
    bbox = draw.textbbox((0, 0), msg, font=font)
    tw = bbox[2] - bbox[0]
    th = bbox[3] - bbox[1]
    draw.text(
        ((TILE_W - tw) / 2, verdict_top + (VERDICT_H - th) / 2 - 2),
        msg,
        fill=EMPTY_BAND_FG,
        font=font,
    )
    return img


def render_filled_tile(
    capture: Capture,
    slot_num: int,
    wireframe: bool = False,
    image_size: Optional[Tuple[int, int]] = None,
) -> Image.Image:
    """Compose a filled tile.

    `image_size` is the (W, H) of the original image the landmarks came
    from — only needed in wireframe mode. Falls back to (1, 1) for
    legacy captures that have no landmarks.
    """
    img, draw = _tile_canvas()
    _draw_header(draw, capture.intended, slot_num)

    # Face / wireframe
    region_top = HEADER_H + PAD
    target_w = TILE_W - 2 * PAD
    if wireframe and capture.landmarks is not None and capture.bbox is not None:
        face_pil = render_wireframe(
            capture.landmarks,
            image_size=capture.image_size or image_size or (1000, 1000),
            bbox=capture.bbox,
            output_size=(target_w, FACE_H),
        )
    else:
        face_pil = Image.fromarray(capture.face)
    _paste_face_region(img, face_pil)

    # Bars
    bar_top = region_top + FACE_H + PAD
    bar_img = _bar_strip(capture, size=(target_w, BAR_H))
    img.paste(bar_img.resize((target_w, BAR_H), Image.Resampling.LANCZOS),
              (PAD, bar_top))

    # Verdict — full-width coloured band, white bold text. Green for
    # classifier agreement, amber-red for disagreement. Two-line layout
    # keeps the type readable across a 280-wide tile and survives in
    # the A4 PDF export at the same scale.
    top_emo, top_p = capture.top_emotion()
    intended_classifier = EMOTION_TO_CLASSIFIER.get(capture.intended)
    agree = top_emo == intended_classifier
    verdict_top = bar_top + BAR_H + PAD
    band_fill = AGREE_BAND if agree else DISAGREE_BAND
    draw.rectangle(
        [0, verdict_top, TILE_W, verdict_top + VERDICT_H],
        fill=band_fill,
    )

    title_font = _load_font(15, bold=True)
    detail_font = _load_font(11, bold=False)
    if agree:
        title = "AGREES"
        detail = f"{top_emo}  ({top_p:.2f})"
    else:
        title = "DISAGREES"
        detail = f"sees {top_emo} ({top_p:.2f}), not {intended_classifier}"

    # Centred two-line stack
    t_bbox = draw.textbbox((0, 0), title, font=title_font)
    d_bbox = draw.textbbox((0, 0), detail, font=detail_font)
    t_w = t_bbox[2] - t_bbox[0]
    d_w = d_bbox[2] - d_bbox[0]
    # Trim detail if it overflows
    if d_w > TILE_W - 10:
        while d_w > TILE_W - 10 and len(detail) > 4:
            detail = detail[:-2]
            d_bbox = draw.textbbox((0, 0), detail + "…", font=detail_font)
            d_w = d_bbox[2] - d_bbox[0]
        detail = detail + "…"

    draw.text(
        ((TILE_W - t_w) / 2, verdict_top + 3),
        title, fill=VERDICT_FG, font=title_font,
    )
    draw.text(
        ((TILE_W - d_w) / 2, verdict_top + 24),
        detail, fill=VERDICT_FG, font=detail_font,
    )

    return img


def render_tile(
    intended: str,
    capture: Optional[Capture],
    slot_num: int,
    wireframe: bool = False,
    image_size: Optional[Tuple[int, int]] = None,
) -> Image.Image:
    if capture is None:
        return render_empty_tile(intended, slot_num)
    return render_filled_tile(capture, slot_num, wireframe=wireframe,
                              image_size=image_size)


def grid_tiles(
    state: Dict[str, Optional[Capture]],
    wireframe: bool = False,
    image_size: Optional[Tuple[int, int]] = None,
) -> List[Image.Image]:
    """Six tiles in `BASIC_EMOTIONS` order, ready for the gr.Gallery."""
    return [
        render_tile(emo, state.get(emo), idx + 1,
                    wireframe=wireframe, image_size=image_size)
        for idx, emo in enumerate(BASIC_EMOTIONS)
    ]


# ---------- Single-page A4 PDF artefact --------------------------------

def export_single_page_pdf(
    state: Dict[str, Optional[Capture]],
    student_name: str = "",
    wireframe: bool = False,
    image_size: Optional[Tuple[int, int]] = None,
) -> str:
    """Render the six-tile session as a single-page A4 portrait PDF.

    The page carries: a one-line title row with the student name and
    chosen format, a 2x3 grid of tiles, and a tight footer with the
    privacy reminder. No second page — the artefact is deliberately
    one-glance.
    """
    suffix = "-wireframe.pdf" if wireframe else "-face.pdf"
    tmp = tempfile.NamedTemporaryFile(
        prefix="lesson4-emotionmap", suffix=suffix, delete=False
    )
    tmp.close()

    A4_W, A4_H = 8.27, 11.69  # inches
    fig = plt.figure(figsize=(A4_W, A4_H))
    fig.patch.set_facecolor("white")

    # Title row
    fmt_name = "Format B — wireframe (face-free)" if wireframe else "Format A — face"
    title = "Lesson 4 EmotionMap — six-emotion replication challenge"
    sub = f"{student_name.strip() or '[your name]'}    ·    {fmt_name}"
    fig.text(0.5, 0.965, title, ha="center", fontsize=13, weight="bold")
    fig.text(0.5, 0.945, sub, ha="center", fontsize=10, color="#555")

    # 2 rows × 3 columns of tile images
    grid = grid_tiles(state, wireframe=wireframe, image_size=image_size)
    rows, cols = 2, 3
    gs = fig.add_gridspec(
        rows, cols,
        left=0.04, right=0.96,
        bottom=0.07, top=0.92,
        hspace=0.08, wspace=0.06,
    )
    for idx, tile in enumerate(grid):
        ax = fig.add_subplot(gs[idx // cols, idx % cols])
        ax.imshow(tile)
        ax.axis("off")

    # Footer
    filled = sum(1 for c in state.values() if c is not None)
    agreed = sum(
        1 for c in state.values()
        if c is not None and c.classifier_agrees()
    )
    footer = (
        f"{filled} / 6 emotions captured · classifier agreed on {agreed} / {filled if filled else 6}"
        if filled
        else "0 / 6 emotions captured"
    )
    fig.text(0.5, 0.055, footer, ha="center", fontsize=9, color="#333")
    fig.text(
        0.5, 0.034,
        "The toy app uploads nothing beyond the model inference call and stores nothing on a server.",
        ha="center", fontsize=7, color="#777",
    )
    fig.text(
        0.5, 0.018,
        "Created by Dr. Gordon Wright — A LittleMonkeyLab caper. "
        "Part of the Goldsmiths MSc in Psychology, Week 3 Part 4.",
        ha="center", fontsize=7, color="#888", style="italic",
    )

    with PdfPages(tmp.name) as pdf:
        pdf.savefig(fig)
    plt.close(fig)
    return tmp.name