math_homework.py

#!/usr/bin/env python3
"""Generate a first-grade math worksheet (Markdown + PDF).

Examples:
    # 12 addition problems, missing slot randomized, numbers 0-20
    python math_homework.py --add 12 --max 20 --output addition_to_20

    # Mixed worksheet
    python math_homework.py --add 6 --sub 6 --max 20 \
        --title "Addition & Subtraction to 20" --output mixed

    # Reproducible (same seed -> same problems)
    python math_homework.py --add 12 --seed 42 --output today
"""

from __future__ import annotations

import argparse
import random
import sys
from dataclasses import dataclass
from datetime import date
from pathlib import Path
from typing import Callable

from reportlab.lib import colors
from reportlab.lib.pagesizes import letter
from reportlab.lib.units import inch
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.platypus import (
    Flowable,
    Paragraph,
    SimpleDocTemplate,
    Table,
    TableStyle,
)
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.enums import TA_CENTER


# Register Segoe UI (Windows system font) for a modern look. Falls back to
# Helvetica if not present (e.g. on macOS/Linux).
FONT_REGULAR = "Helvetica"
FONT_BOLD = "Helvetica-Bold"
_WIN_FONTS = Path("C:/Windows/Fonts")
if (_WIN_FONTS / "segoeui.ttf").exists():
    pdfmetrics.registerFont(TTFont("SegoeUI", str(_WIN_FONTS / "segoeui.ttf")))
    pdfmetrics.registerFont(TTFont("SegoeUI-Bold", str(_WIN_FONTS / "segoeuib.ttf")))
    FONT_REGULAR = "SegoeUI"
    FONT_BOLD = "SegoeUI-Bold"


OPERATORS = {
    "add": "+",
    "sub": "−",  # U+2212 minus sign (typographically nicer than '-')
    "mul": "×",  # U+00D7 multiplication sign
    "div": "÷",  # U+00F7 division sign
}

OP_NAMES = {
    "add": "Addition",
    "sub": "Subtraction",
    "mul": "Multiplication",
    "div": "Division",
}


@dataclass
class Problem:
    a: int
    b: int
    c: int
    op: str  # key into OPERATORS
    missing: str  # 'a', 'b', or 'c'

    def operator_symbol(self) -> str:
        return OPERATORS[self.op]


# ---------------------------------------------------------------------------
# Problem generation
# ---------------------------------------------------------------------------
#
# Each generator returns a Problem (a OP b = c) with all values in [min_v, max_v]
# (with the additional rule that division never has b=0 and never has c=0, so the
# missing-slot puzzle always has a single correct answer).
#
# Adding a new operation:
#   1. Add an entry to OPERATORS and OP_NAMES.
#   2. Write a generator(min_v, max_v, rng) -> Problem.
#   3. Register it in GENERATORS and add a CLI flag in build_parser().


def generate_addition(min_v: int, max_v: int, rng: random.Random) -> Problem:
    a = rng.randint(min_v, max_v)
    b_hi = max(min_v, max_v - a)
    b = rng.randint(min_v, b_hi)
    c = a + b
    if c > max_v:  # only possible if min_v > 0 and the range is too tight
        c = max_v
        b = c - a
    return Problem(a, b, c, "add", rng.choice(("a", "b", "c")))


def generate_subtraction(min_v: int, max_v: int, rng: random.Random) -> Problem:
    # Pick c (difference) and b (subtrahend); a = b + c is then in range.
    c = rng.randint(min_v, max_v)
    b_hi = max(min_v, max_v - c)
    b = rng.randint(min_v, b_hi)
    a = b + c
    return Problem(a, b, c, "sub", rng.choice(("a", "b", "c")))


def generate_multiplication(min_v: int, max_v: int, rng: random.Random) -> Problem:
    # Avoid 0-factor problems when missing='a' or 'b' would be ambiguous
    # (e.g. "0 × □ = 0" has infinitely many answers). Easiest fix: require
    # both factors >= max(1, min_v).
    lo = max(1, min_v)
    a = rng.randint(lo, max(lo, max_v))
    b_hi = max(lo, max_v // a)
    b = rng.randint(lo, b_hi)
    c = a * b
    return Problem(a, b, c, "mul", rng.choice(("a", "b", "c")))


def generate_division(min_v: int, max_v: int, rng: random.Random) -> Problem:
    # b >= 1 (no divide-by-zero), c >= 1 (no ambiguous "0 ÷ □ = 0")
    lo = max(1, min_v)
    b = rng.randint(lo, max(lo, max_v))
    c_hi = max(lo, max_v // b)
    c = rng.randint(lo, c_hi)
    a = b * c
    return Problem(a, b, c, "div", rng.choice(("a", "b", "c")))


GENERATORS: dict[str, Callable[[int, int, random.Random], Problem]] = {
    "add": generate_addition,
    "sub": generate_subtraction,
    "mul": generate_multiplication,
    "div": generate_division,
}


def build_problems(
    counts: dict[str, int], min_v: int, max_v: int, rng: random.Random
) -> list[Problem]:
    problems: list[Problem] = []
    for op, n in counts.items():
        for _ in range(n):
            problems.append(GENERATORS[op](min_v, max_v, rng))
    rng.shuffle(problems)
    return problems


# ---------------------------------------------------------------------------
# Markdown rendering
# ---------------------------------------------------------------------------

CIRCLE = "◯"  # U+25EF LARGE CIRCLE


def problem_md(p: Problem) -> str:
    parts = [
        (str(p.a), "a"),
        (p.operator_symbol(), None),
        (str(p.b), "b"),
        ("=", None),
        (str(p.c), "c"),
    ]
    return " ".join(CIRCLE if slot == p.missing else text for text, slot in parts)


def render_markdown(
    problems: list[Problem], title: str, columns: int, instructions: str
) -> str:
    lines: list[str] = []
    lines.append(f"# {title}")
    lines.append("")
    lines.append(instructions)
    lines.append("")
    lines.append("|" + " |" * columns)
    lines.append("|" + "---|" * columns)
    for i in range(0, len(problems), columns):
        row = problems[i : i + columns]
        cells = [f" {i + idx + 1}. {problem_md(p)} " for idx, p in enumerate(row)]
        while len(cells) < columns:
            cells.append(" ")
        lines.append("|" + "|".join(cells) + "|")
    lines.append("")
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# PDF rendering (ReportLab)
# ---------------------------------------------------------------------------


class ProblemFlowable(Flowable):
    """Renders a single equation with an empty circle in the missing slot."""

    def __init__(
        self,
        problem: Problem,
        font_name: str = FONT_BOLD,
        font_size: int = 24,
        circle_radius: float = 20,
    ):
        super().__init__()
        self.problem = problem
        self.font_name = font_name
        self.font_size = font_size
        self.circle_radius = circle_radius
        self._tokens: list[tuple] = []
        self._content_width: float = 0
        self._width: float = 0
        self._height: float = 0

    def _build_tokens(self) -> list[tuple]:
        p = self.problem
        parts = [
            (str(p.a), "a"),
            (p.operator_symbol(), None),
            (str(p.b), "b"),
            ("=", None),
            (str(p.c), "c"),
        ]
        tokens = []
        for text, slot in parts:
            if slot == p.missing:
                tokens.append(("circle", self.circle_radius * 2))
            else:
                w = pdfmetrics.stringWidth(text, self.font_name, self.font_size)
                tokens.append(("text:" + text, w))
        return tokens

    def wrap(self, available_width: float, available_height: float):
        self._tokens = self._build_tokens()
        spacing = self.font_size * 0.35
        self._content_width = sum(w for _, w in self._tokens) + spacing * (
            len(self._tokens) - 1
        )
        self._width = available_width
        self._height = max(self.font_size * 1.4, self.circle_radius * 2.2) + 6
        return (self._width, self._height)

    def draw(self):
        c = self.canv
        c.setFont(self.font_name, self.font_size)
        spacing = self.font_size * 0.35
        x = (self._width - self._content_width) / 2
        # Vertical center: text baseline is roughly font_size * 0.3 above
        # the visual midpoint of the digits.
        text_y = self._height / 2 - self.font_size * 0.33
        circle_cy = self._height / 2
        for tok, w in self._tokens:
            if tok == "circle":
                c.setLineWidth(1.5)
                c.circle(x + w / 2, circle_cy, self.circle_radius, stroke=1, fill=0)
            else:
                _, _, text = tok.partition(":")
                c.drawString(x, text_y, text)
            x += w + spacing


def render_pdf(
    problems: list[Problem],
    title: str,
    columns: int,
    instructions: str,
    output_path: Path,
) -> None:
    doc = SimpleDocTemplate(
        str(output_path),
        pagesize=letter,
        leftMargin=0.6 * inch,
        rightMargin=0.6 * inch,
        topMargin=0.5 * inch,
        bottomMargin=0.5 * inch,
        title=title,
    )

    styles = getSampleStyleSheet()
    title_style = ParagraphStyle(
        "WorksheetTitle",
        parent=styles["Heading1"],
        fontName=FONT_BOLD,
        alignment=TA_CENTER,
        fontSize=22,
        leading=26,
        spaceAfter=4,
    )
    instructions_style = ParagraphStyle(
        "Instructions",
        parent=styles["BodyText"],
        fontName=FONT_REGULAR,
        alignment=TA_CENTER,
        fontSize=12,
        leading=14,
        spaceAfter=10,
    )

    story: list = [
        Paragraph(title, title_style),
        Paragraph(instructions, instructions_style),
    ]

    rows = []
    for i in range(0, len(problems), columns):
        row_problems = problems[i : i + columns]
        row = [ProblemFlowable(p) for p in row_problems]
        while len(row) < columns:
            row.append("")
        rows.append(row)

    page_width = letter[0] - 1.2 * inch
    col_width = page_width / columns
    grid = Table(rows, colWidths=[col_width] * columns)
    grid.setStyle(
        TableStyle(
            [
                ("GRID", (0, 0), (-1, -1), 1, colors.black),
                ("VALIGN", (0, 0), (-1, -1), "MIDDLE"),
                ("ALIGN", (0, 0), (-1, -1), "CENTER"),
                ("TOPPADDING", (0, 0), (-1, -1), 14),
                ("BOTTOMPADDING", (0, 0), (-1, -1), 14),
            ]
        )
    )
    story.append(grid)
    doc.build(story)


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="Generate a first-grade math worksheet (Markdown + PDF).",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    p.add_argument("--add", type=int, default=0, help="Number of addition problems")
    p.add_argument("--sub", type=int, default=0, help="Number of subtraction problems")
    p.add_argument(
        "--mul", type=int, default=0, help="Number of multiplication problems"
    )
    p.add_argument("--div", type=int, default=0, help="Number of division problems")
    p.add_argument(
        "--min", dest="min_v", type=int, default=0, help="Minimum value (default: 0)"
    )
    p.add_argument(
        "--max", dest="max_v", type=int, default=20, help="Maximum value (default: 20)"
    )
    p.add_argument(
        "--columns", type=int, default=2, help="Number of problem columns (default: 2)"
    )
    p.add_argument(
        "--title",
        default=None,
        help="Worksheet title. If omitted, generated from the operations and range.",
    )
    p.add_argument(
        "--instructions",
        default="Write the missing number in each equation.",
    )
    p.add_argument(
        "--output",
        "-o",
        default=None,
        help=(
            "Output path prefix (no extension). Writes <prefix>.md and <prefix>.pdf. "
            "Default: 'worksheet_<YYYY-MM-DD>'."
        ),
    )
    p.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Random seed for reproducible worksheets.",
    )
    return p


def auto_title(counts: dict[str, int], max_v: int) -> str:
    active = [op for op, n in counts.items() if n > 0]
    if len(active) == 1:
        return f"{OP_NAMES[active[0]]} to {max_v}"
    return " & ".join(OP_NAMES[op] for op in active) + f" to {max_v}"


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)

    counts = {"add": args.add, "sub": args.sub, "mul": args.mul, "div": args.div}
    counts = {op: n for op, n in counts.items() if n > 0}
    if not counts:
        print(
            "Error: specify at least one of --add / --sub / --mul / --div with N > 0.",
            file=sys.stderr,
        )
        return 2
    if args.min_v < 0 or args.max_v < args.min_v:
        print("Error: require 0 <= --min <= --max.", file=sys.stderr)
        return 2
    if args.columns < 1:
        print("Error: --columns must be >= 1.", file=sys.stderr)
        return 2

    rng = random.Random(args.seed)
    problems = build_problems(counts, args.min_v, args.max_v, rng)
    title = args.title or auto_title(counts, args.max_v)

    output = Path(args.output) if args.output else Path(f"worksheet_{date.today()}")
    output.parent.mkdir(parents=True, exist_ok=True)
    md_path = output.with_suffix(".md")
    pdf_path = output.with_suffix(".pdf")

    md_path.write_text(
        render_markdown(problems, title, args.columns, args.instructions),
        encoding="utf-8",
    )
    render_pdf(problems, title, args.columns, args.instructions, pdf_path)

    print(f"Wrote {md_path}")
    print(f"Wrote {pdf_path}")
    return 0


if __name__ == "__main__":
    sys.exit(main())