tutor/visual.py

"""
tutor/visual.py
===============
PIL-based visual renderer for counting and comparison tasks.

Generates child-friendly images that accompany curriculum items.
No heavy ML model required — all rendering is done with Pillow (CPU, instant).

Usage
-----
    from tutor.visual import render_item_image

    img_path = render_item_image("goats_5")   # returns path to saved PNG
    img_path = render_item_image("compare_4_7")
    img_path = render_item_image("beads_2_plus_3")
"""

from __future__ import annotations

import os
import re
import math
from pathlib import Path
from typing import Optional

from PIL import Image, ImageDraw, ImageFont

# ── Output directory ──────────────────────────────────────────────────────────
_ROOT      = Path(__file__).parent.parent
ASSETS_DIR = _ROOT / "assets"
ASSETS_DIR.mkdir(exist_ok=True)

# ── Visual config ─────────────────────────────────────────────────────────────
IMG_W, IMG_H = 480, 320
BG_COLOR     = "#FFF9F0"   # warm off-white (child-friendly)
ACCENT       = "#FF8C42"   # orange
BLUE         = "#4A90D9"
GREEN        = "#5CB85C"
RED          = "#D9534F"
PURPLE       = "#9B59B6"

# Emoji-like object representations (text characters rendered via PIL)
OBJECT_EMOJI: dict[str, str] = {
    "apples":   "🍎",
    "goats":    "🐐",
    "stars":    "⭐",
    "birds":    "🐦",
    "fish":     "🐟",
    "bananas":  "🍌",
    "mangoes":  "🥭",
    "stones":   "🪨",
    "beads":    "●",
    "blocks":   "■",
    "kids":     "👦",
    "cookies":  "🍪",
    "tomatoes": "🍅",
    "drums":    "🥁",
    "beans":    "🫘",
}

# Fallback plain character when emoji rendering is unavailable
FALLBACK_CHAR = "●"

# Try to find a font that supports emoji; fall back to default
def _get_font(size: int) -> ImageFont.ImageFont:
    candidates = [
        "seguiemj.ttf",       # Windows Segoe UI Emoji
        "NotoColorEmoji.ttf", # Linux
        "Apple Color Emoji.ttc",
    ]
    for name in candidates:
        try:
            return ImageFont.truetype(name, size)
        except (IOError, OSError):
            pass
    try:
        return ImageFont.load_default(size=size)
    except TypeError:
        return ImageFont.load_default()


def _get_plain_font(size: int) -> ImageFont.ImageFont:
    candidates = [
        "arialbd.ttf", "Arial Bold.ttf", "DejaVuSans-Bold.ttf",
        "LiberationSans-Bold.ttf", "FreeSansBold.ttf",
    ]
    for name in candidates:
        try:
            return ImageFont.truetype(name, size)
        except (IOError, OSError):
            pass
    try:
        return ImageFont.load_default(size=size)
    except TypeError:
        return ImageFont.load_default()


# ── Core renderer helpers ─────────────────────────────────────────────────────

def _new_canvas() -> tuple[Image.Image, ImageDraw.ImageDraw]:
    img  = Image.new("RGB", (IMG_W, IMG_H), BG_COLOR)
    draw = ImageDraw.Draw(img)
    return img, draw


def _draw_title(draw: ImageDraw.ImageDraw, text: str, color: str = "#333333") -> None:
    font = _get_plain_font(22)
    draw.text((IMG_W // 2, 18), text, fill=color, font=font, anchor="mt")


def _draw_objects(draw: ImageDraw.ImageDraw, count: int,
                  symbol: str, color: str = ACCENT) -> None:
    """
    Draw `count` symbols arranged in a grid in the centre of the canvas.
    Falls back to colored circles if symbol rendering looks broken.
    """
    cols = min(count, 5)
    rows = math.ceil(count / cols)

    cell_w = min(64, (IMG_W - 60) // cols)
    cell_h = min(64, (IMG_H - 100) // max(rows, 1))
    font   = _get_plain_font(int(cell_w * 0.7))

    total_w = cols * cell_w
    total_h = rows * cell_h
    x0 = (IMG_W - total_w) // 2
    y0 = 60 + (IMG_H - 100 - total_h) // 2

    idx = 0
    for r in range(rows):
        for c in range(cols):
            if idx >= count:
                break
            cx = x0 + c * cell_w + cell_w // 2
            cy = y0 + r * cell_h + cell_h // 2
            r_px = cell_w // 3
            draw.ellipse(
                [cx - r_px, cy - r_px, cx + r_px, cy + r_px],
                fill=color, outline="#ffffff", width=2
            )
            idx += 1


def _draw_number(draw: ImageDraw.ImageDraw, number: int,
                 x: int, y: int, color: str = "#333333",
                 size: int = 72) -> None:
    font = _get_plain_font(size)
    draw.text((x, y), str(number), fill=color, font=font, anchor="mm")


# ── Specific renderers ────────────────────────────────────────────────────────

def _render_counting(count: int, obj_name: str) -> Image.Image:
    """Counting image: N colored circles with label."""
    img, draw = _new_canvas()
    label = f"Count the {obj_name}!"
    _draw_title(draw, label)
    _draw_objects(draw, count, obj_name, color=ACCENT)
    # Count label bottom-right (hidden from child — for debug builds only)
    return img


def _render_compare(a: int, b: int) -> Image.Image:
    """Comparison image: two large numbers side by side with VS divider."""
    img, draw = _new_canvas()
    _draw_title(draw, "Which is bigger?")

    mid = IMG_W // 2
    # Left number
    draw.rectangle([30, 80, mid - 20, IMG_H - 40], fill=BLUE, outline="#ffffff", width=3)
    _draw_number(draw, a, mid // 2, IMG_H // 2 + 10, color="#ffffff", size=80)

    # Right number
    draw.rectangle([mid + 20, 80, IMG_W - 30, IMG_H - 40], fill=GREEN, outline="#ffffff", width=3)
    _draw_number(draw, b, mid + (IMG_W - mid) // 2, IMG_H // 2 + 10, color="#ffffff", size=80)

    # VS divider
    font = _get_plain_font(28)
    draw.text((mid, IMG_H // 2 + 10), "VS", fill="#888888", font=font, anchor="mm")
    return img


def _render_addition(a: int, b: int) -> Image.Image:
    """Addition image: two groups of dots with + sign."""
    img, draw = _new_canvas()
    _draw_title(draw, f"{a}  +  {b}  =  ?")

    mid = IMG_W // 2

    # Left group
    cols_a = min(a, 4)
    rows_a = math.ceil(a / cols_a) if a > 0 else 1
    cell   = 44
    x0 = 30
    y0 = 80
    idx = 0
    for r in range(rows_a):
        for c in range(cols_a):
            if idx >= a:
                break
            cx = x0 + c * cell + cell // 2
            cy = y0 + r * cell + cell // 2
            draw.ellipse([cx-16, cy-16, cx+16, cy+16], fill=BLUE, outline="#fff", width=2)
            idx += 1

    # Plus sign
    font = _get_plain_font(48)
    draw.text((mid, IMG_H // 2), "+", fill=ACCENT, font=font, anchor="mm")

    # Right group
    cols_b = min(b, 4)
    rows_b = math.ceil(b / cols_b) if b > 0 else 1
    x0r    = mid + 30
    idx    = 0
    for r in range(rows_b):
        for c in range(cols_b):
            if idx >= b:
                break
            cx = x0r + c * cell + cell // 2
            cy = y0   + r * cell + cell // 2
            draw.ellipse([cx-16, cy-16, cx+16, cy+16], fill=GREEN, outline="#fff", width=2)
            idx += 1

    return img


def _render_subtraction(a: int, b: int) -> Image.Image:
    """Subtraction image: `a` dots, `b` crossed out in red."""
    img, draw = _new_canvas()
    _draw_title(draw, f"{a}  −  {b}  =  ?")

    cols  = min(a, 6)
    rows  = math.ceil(a / cols) if a > 0 else 1
    cell  = 44
    x0    = (IMG_W - cols * cell) // 2
    y0    = 80

    for i in range(a):
        r  = i // cols
        c  = i % cols
        cx = x0 + c * cell + cell // 2
        cy = y0 + r * cell + cell // 2
        color  = RED   if i >= (a - b) else BLUE
        draw.ellipse([cx-16, cy-16, cx+16, cy+16], fill=color, outline="#fff", width=2)
        if i >= (a - b):
            # Cross out
            draw.line([cx-14, cy-14, cx+14, cy+14], fill="#ffffff", width=3)
            draw.line([cx+14, cy-14, cx-14, cy+14], fill="#ffffff", width=3)

    return img


def _render_number_line(lo: int, hi: int) -> Image.Image:
    """Number line with lo and hi visible, middle blank (?)."""
    img, draw = _new_canvas()
    _draw_title(draw, "What number comes between?")

    mid_x  = IMG_W // 2
    mid_y  = IMG_H // 2 + 20
    gap    = 120

    # Draw line
    draw.line([mid_x - gap - 50, mid_y, mid_x + gap + 50, mid_y],
              fill="#888888", width=3)

    # Tick marks and numbers
    font = _get_plain_font(36)
    for val, xpos in [(lo, mid_x - gap), (hi, mid_x + gap)]:
        draw.line([xpos, mid_y - 15, xpos, mid_y + 15], fill="#333", width=3)
        draw.text((xpos, mid_y + 30), str(val), fill="#333333", font=font, anchor="mt")

    # Middle tick with ?
    draw.line([mid_x, mid_y - 15, mid_x, mid_y + 15], fill=ACCENT, width=3)
    draw.text((mid_x, mid_y + 30), "?", fill=ACCENT, font=font, anchor="mt")

    return img


def _render_word_problem(label: str) -> Image.Image:
    """Generic word-problem card with a friendly icon."""
    img, draw = _new_canvas()
    _draw_title(draw, "Word Problem", color=PURPLE)

    # Draw a simple "thinking" icon — 3 dots
    for i, x in enumerate([IMG_W//2 - 40, IMG_W//2, IMG_W//2 + 40]):
        draw.ellipse([x-18, IMG_H//2-18, x+18, IMG_H//2+18],
                     fill=PURPLE, outline="#fff", width=2)

    font = _get_plain_font(16)
    draw.text((IMG_W//2, IMG_H - 40), "Read the question above",
              fill="#888888", font=font, anchor="mm")
    return img


# ── Public API ────────────────────────────────────────────────────────────────

def render_item_image(visual_key: str, force: bool = False) -> str:
    """
    Render a curriculum item image from its `visual` field key.

    Parameters
    ----------
    visual_key : str
        The `visual` field from a CurriculumItem, e.g. "goats_5",
        "compare_4_7", "beads_2_plus_3", "blocks_8_minus_3".
    force : bool
        Re-render even if the file already exists on disk.

    Returns
    -------
    str
        Absolute path to the saved PNG file.
    """
    out_path = ASSETS_DIR / f"{visual_key}.png"
    if out_path.exists() and not force:
        return str(out_path)

    img = _dispatch(visual_key)
    img.save(str(out_path))
    return str(out_path)


def _dispatch(visual_key: str) -> Image.Image:
    """Parse visual_key and call the appropriate renderer."""

    # compare_A_B  → comparison
    m = re.fullmatch(r"compare_(\d+)_(\d+)", visual_key)
    if m:
        return _render_compare(int(m.group(1)), int(m.group(2)))

    # number_line_A_B  → number line
    m = re.fullmatch(r"number_line_(\d+)_(\d+)", visual_key)
    if m:
        return _render_number_line(int(m.group(1)), int(m.group(2)))

    # beads_A_plus_B  → addition dots
    m = re.fullmatch(r"\w+_(\d+)_plus_(\d+)", visual_key)
    if m:
        return _render_addition(int(m.group(1)), int(m.group(2)))

    # blocks_A_minus_B or drums_A_minus_B  → subtraction dots
    m = re.fullmatch(r"\w+_(\d+)_minus_(\d+)", visual_key)
    if m:
        return _render_subtraction(int(m.group(1)), int(m.group(2)))

    # obj_N  → counting (e.g. "goats_5", "apples_3")
    m = re.fullmatch(r"([a-z]+)_(\d+)", visual_key)
    if m:
        return _render_counting(int(m.group(2)), m.group(1))

    # Fallback: word problem card
    return _render_word_problem(visual_key)


def prerender_all(loader) -> dict[str, str]:
    """
    Pre-render images for every curriculum item that has a visual key.
    Returns a dict mapping visual_key → file path.
    """
    paths: dict[str, str] = {}
    items = loader.all_items()
    rendered = skipped = 0
    for item in items:
        if not item.visual:
            continue
        path = render_item_image(item.visual)
        paths[item.visual] = path
        if Path(path).stat().st_size > 0:
            rendered += 1
        else:
            skipped += 1
    print(f"✅  Pre-rendered {rendered} images → assets/  (skipped {skipped})")
    return paths


# ── Smoke-test ────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    test_keys = [
        "goats_5",
        "apples_3",
        "compare_4_7",
        "beads_2_plus_3",
        "blocks_8_minus_3",
        "number_line_47_49",
        "kids_3_cookies_9",
    ]
    print("Rendering test images …")
    for key in test_keys:
        path = render_item_image(key, force=True)
        size = os.path.getsize(path)
        print(f"  {key:<28} → {path}  ({size} bytes)")
    print("\nDone. Check the assets/ folder.")