diffu_page/cli.py

"""``diffu-page`` — synthesize full document pages from generated (or placeholder) lines.

    # placeholder ink (no model/GPU) — proves the layout + compositor:
    diffu-page --template multi_column --n 3

    # the whole run from a YAML config (page geometry + renderer + every scan knob):
    diffu-page --config run.yaml

    # real handwriting via a Diffu checkpoint (batched), overriding a couple of fields:
    diffu-page --renderer diffu --ckpt run/final/model.safetensors --style ref_line.png --bed-prob 0.5

The run is a :class:`~diffu_page.config.GenConfig`; ``--config`` loads it from YAML and the flags below
override common fields (CLI > YAML > default). Each page writes ``page_XXXX.png`` + one label file per
format (json / page / alto).
"""

import random
from collections.abc import Callable
from pathlib import Path

import cv2
import numpy as np
import typer
from rich.console import Console

from .compose import compose_page, compose_page_from_inks
from .config import GenConfig, Renderer, Template, load_gen_config
from .dataset import CropRecord, PageRecord, assign_split, export_line_crops, write_manifest
from .labels import EXTENSIONS, LineLabel, PageLabels, export
from .layout import Layout, PageSpec, build_layout, load_xml_layout
from .render import (
    BatchedDiffuRenderer,
    DiffuRenderer,
    ErukuRenderer,
    LineRenderer,
    PlaceholderRenderer,
    render_layouts,
)

console = Console()

# A little public-domain-ish Swedish filler so placeholder pages look like text. Swap via --text-file.
_SAMPLE = [
    "Göteborgs poliskammare den fjärde maj artonhundra åttiotvå.",
    "Anhållen för fylleri och förargelseväckande beteende å allmän gata.",
    "Mannen uppgav sig vara arbetare och boende i Haga församling.",
    "Förd till häktet och förhörd av vakthavande konstapel om natten.",
    "Smörgåsbord, kärlek och ärlighet varar längst sade den gamle mannen.",
]


def _paragraphs(text_file: Path | None) -> list[str]:
    if text_file:
        return [ln for ln in text_file.read_text(encoding="utf-8").splitlines() if ln.strip()]
    return _SAMPLE * 12  # repeat so columns/tables fill


# Plausible archival furniture, auto-generated per page when GenConfig.furniture_prob fires. The exact
# wording is irrelevant to the labels — what matters is that the page carries a header / marginalia /
# page number, like a real protocol, instead of being a bare body.
_HEADERS = [
    ["Göteborgs Poliskammare.", "Protokoll d. {d} Maj 1882."],
    ["Rådhusrätten i Göteborg.", "d. {d} Mars 1884."],
    ["Kongl. Maj:ts Befallningshafvande."],
    ["Till Magistraten.", "Inkommet d. {d} Juni."],
]
_MARGINALIA = [
    "Fylleri.",
    "Erkännande.",
    "Vittnen.",
    "Till domstol.",
    "Anhållan.",
    "Utslag.",
    "Käranden.",
    "Svaranden.",
    "Afslag.",
    "Bifall.",
    "Böter.",
    "Frikänd.",
]


def _furniture(rng: random.Random) -> tuple[list[str], list[str], str]:
    """Auto-generate a header block, a few marginalia notes and a page number for one page."""
    header = [h.replace("{d}", str(rng.randint(1, 28))) for h in rng.choice(_HEADERS)]
    marginalia = rng.sample(_MARGINALIA, k=rng.randint(2, 4))
    return header, marginalia, f"{rng.randint(1, 400)}."


def _with_overrides(cfg: GenConfig, **flags: object) -> GenConfig:
    """Apply the CLI flags that were actually set (non-None) over ``cfg`` — CLI > YAML > default."""
    aug = {k: flags.pop(k) for k in ("bed_prob", "clasp_prob") if flags.get(k) is not None}
    top = {k: v for k, v in flags.items() if v is not None}
    if aug:
        top["augment"] = cfg.augment.model_copy(update=aug)
    return cfg.model_copy(update=top)


def _build_renderer(cfg: GenConfig) -> LineRenderer | BatchedDiffuRenderer:
    """The renderer for ``cfg.renderer``: a batched Diffu renderer (one GPU pass for the whole run), Eruku,
    or the no-model placeholder."""
    if cfg.renderer is Renderer.diffu:
        if not (cfg.ckpt and (cfg.style or cfg.styles)):
            raise typer.BadParameter(
                "renderer=diffu requires ckpt and a style/styles (in the config or via --ckpt/--style)"
            )
        if cfg.fast:  # batched: faster for big runs, but a touch softer (wider canvas spreads strokes)
            return BatchedDiffuRenderer(str(cfg.ckpt), num_steps=cfg.steps, cfg_scale=cfg.cfg_scale)
        style = cfg.style or cfg.styles[0]  # per-line path is single-style; the batched path mixes hands
        return DiffuRenderer(str(cfg.ckpt), str(style), num_steps=cfg.steps, cfg_scale=cfg.cfg_scale)
    if cfg.renderer is Renderer.eruku:
        if not cfg.style:
            raise typer.BadParameter("renderer=eruku requires style")
        return ErukuRenderer(str(cfg.style), cfg_scale=cfg.eruku_cfg)
    return PlaceholderRenderer()


def _page_box_styles(layouts: list[Layout], cfg: GenConfig, *, seed: int) -> list[list[str]]:
    """Per-box writer-style path for each page. A page samples one MAIN hand from the style pool; when the
    pool has more than one, marginalia get a second hand so a page can mix writers (main scribe + annotator).
    Pure + deterministic given ``seed`` so a corpus is reproducible. Empty for the no-style renderers."""
    pool = [str(p) for p in cfg.styles] or ([str(cfg.style)] if cfg.style else [])
    if not pool:
        return [["" for _ in layout.boxes] for layout in layouts]
    out: list[list[str]] = []
    for i, layout in enumerate(layouts):
        rng = random.Random(seed + i)
        main = rng.choice(pool)
        alt = rng.choice(pool) if len(pool) > 1 else main
        out.append([alt if box.region == "marginalia" else main for box in layout.boxes])
    return out


def _load_specs(path: Path) -> list[PageSpec]:
    """Load the FULL control surface from a JSON file — one ``PageSpec`` (object) or many (array), so an
    agent or user can drive every knob (margin_keys, drop_cap, catchword, table columns, ...) without code.
    Pydantic validates each; an out-of-range or unknown field fails fast with a clear error."""
    import json

    data = json.loads(path.read_text(encoding="utf-8"))
    items = data if isinstance(data, list) else [data]
    return [PageSpec.model_validate(item) for item in items]


def _build_layouts(
    cfg: GenConfig,
    *,
    n: int,
    seed: int,
    paragraphs: list[str],
    xml: Path | None,
    measure: Callable[[str, int], int] | None,
    specs: list[PageSpec] | None = None,
) -> list[Layout]:
    """Lay out all ``n`` pages up-front (so the Diffu path can batch every line of every page at once).

    With ``specs`` (from ``--spec``), each page uses a full caller-supplied :class:`PageSpec` verbatim
    (cycled if fewer than ``n``); otherwise a spec is built from ``cfg`` and a ``furniture_prob`` fraction of
    pages get an auto-generated header + marginalia + page number.
    """
    layouts: list[Layout] = []
    for i in range(n):
        if specs is not None:  # full user/agent-authored spec — respected verbatim, no cfg furniture override
            spec = specs[i % len(specs)]
            layouts.append(
                load_xml_layout(str(xml), spec)
                if xml
                else build_layout(spec, paragraphs, seed=seed + i, measure=measure)
            )
            continue
        page_rng = random.Random(seed + i)
        header: list[str] = []
        marginalia: list[str] = []
        page_number: str | None = None
        if page_rng.random() < cfg.furniture_prob:
            header, marginalia, page_number = _furniture(page_rng)
        spec = PageSpec(
            width=cfg.width,
            height=cfg.height,
            template=cfg.template.value,
            n_columns=cfg.n_columns,
            line_height=cfg.line_height,
            line_gap=cfg.line_gap,
            para_gap=cfg.para_gap,
            # mix bound/unbound across the corpus; a spread has its own central gutter, not a single spine
            binding=cfg.template is not Template.spread and page_rng.random() < cfg.binding_prob,
            recto=page_rng.random() < 0.5,  # random page side -> spine + page-no on the matching edge
            binding_ratio=round(page_rng.uniform(0.025, 0.075), 3),  # vary how close text runs to the gutter
            header_lines=header,
            marginalia=marginalia,
            page_number=page_number,
            insertion_prob=cfg.insertion_prob,
            deletion_prob=cfg.deletion_prob,
        )
        layouts.append(
            load_xml_layout(str(xml), spec)
            if xml
            else build_layout(spec, paragraphs, seed=seed + i, measure=measure)
        )
    return layouts


def _render_pages(
    renderer: LineRenderer | BatchedDiffuRenderer,
    layouts: list[Layout],
    cfg: GenConfig,
    *,
    seed: int,
    box_styles: list[list[str]],
) -> list[tuple[np.ndarray, list[LineLabel]]]:
    """Render + compose every page. The batched Diffu renderer goes through one pass over all lines (mixing
    hands per ``box_styles``); the placeholder / Eruku renderers compose page-by-page."""
    paper = str(cfg.paper_dir) if cfg.paper_dir else None
    if isinstance(renderer, BatchedDiffuRenderer):
        inks_per_page = render_layouts(renderer, layouts, box_styles)
        return [
            compose_page_from_inks(lay, inks, paper_dir=paper, seed=seed + i, aug=cfg.augment)
            for i, (lay, inks) in enumerate(zip(layouts, inks_per_page, strict=True))
        ]
    return [
        compose_page(lay, renderer, paper_dir=paper, seed=seed + i, aug=cfg.augment)
        for i, lay in enumerate(layouts)
    ]


def generate(
    config: Path | None = typer.Option(
        None, "--config", help="YAML GenConfig (full control); flags override"
    ),
    out: Path = typer.Option(Path("out/pages"), help="output directory"),
    n: int = typer.Option(4, help="number of pages"),
    seed: int = typer.Option(0),
    text_file: Path | None = typer.Option(None, "--text-file", help="one paragraph per line"),
    xml: Path | None = typer.Option(None, help="replay layout from an ALTO/PAGE XML instead of a template"),
    spec: Path | None = typer.Option(
        None, "--spec", help="full PageSpec JSON (one object or an array, cycled) — the whole control surface"
    ),
    renderer: Renderer | None = typer.Option(None, help="override: line generator"),
    ckpt: Path | None = typer.Option(None, help="override: Diffu checkpoint"),
    style: Path | None = typer.Option(None, help="override: writer style reference line"),
    paper_dir: Path | None = typer.Option(None, "--paper-dir", help="override: blank paper scans"),
    binding_prob: float | None = typer.Option(None, "--binding-prob", min=0.0, max=1.0, help="override"),
    bed_prob: float | None = typer.Option(
        None, "--bed-prob", min=0.0, max=1.0, help="override: scanner-bed fraction"
    ),
    clasp_prob: float | None = typer.Option(
        None, "--clasp-prob", min=0.0, max=1.0, help="override: clasp fraction"
    ),
    fast: bool | None = typer.Option(
        None, "--fast/--no-fast", help="override: batched diffu rendering (faster, slightly softer strokes)"
    ),
    furniture_prob: float | None = typer.Option(
        None, "--furniture-prob", min=0.0, max=1.0, help="override: pages with header + marginalia + page no."
    ),
    crops: bool | None = typer.Option(
        None, "--crops/--no-crops", help="override: also export per-line image+text crops + lines.jsonl"
    ),
    width: int | None = typer.Option(None, help="override: page width in px"),
    height: int | None = typer.Option(None, help="override: page height in px"),
    template: Template | None = typer.Option(
        None, help="override: layout (single_column / multi_column / semi_table / columns / key_value)"
    ),
) -> None:
    """Synthesize full handwritten document pages + ground-truth labels.

    Pass ``--config run.yaml`` for full control over geometry, renderer and every scan knob; the flags
    above override common fields (CLI > YAML > default).
    """
    cfg = _with_overrides(
        load_gen_config(config),
        renderer=renderer,
        ckpt=ckpt,
        style=style,
        paper_dir=paper_dir,
        binding_prob=binding_prob,
        bed_prob=bed_prob,
        clasp_prob=clasp_prob,
        fast=fast,
        furniture_prob=furniture_prob,
        crop_lines=crops,
        width=width,
        height=height,
        template=template,
    )
    out.mkdir(parents=True, exist_ok=True)
    line_renderer = _build_renderer(cfg)
    measure = getattr(line_renderer, "measure_width", None)  # wrap by real width when available
    specs = _load_specs(spec) if spec else None  # --spec: the full control surface, verbatim per page
    layouts = _build_layouts(
        cfg, n=n, seed=seed, paragraphs=_paragraphs(text_file), xml=xml, measure=measure, specs=specs
    )
    box_styles = _page_box_styles(layouts, cfg, seed=seed)
    pages: list[PageRecord] = []
    crops: list[CropRecord] = []
    rendered = _render_pages(line_renderer, layouts, cfg, seed=seed, box_styles=box_styles)
    for i, (page, lines) in enumerate(rendered):
        name = f"page_{i:04d}"
        split = assign_split(i, val_frac=cfg.val_frac, test_frac=cfg.test_frac)
        ph, pw = page.shape[:2]  # real shape (xml replay / bed mount can differ from the requested size)
        cv2.imwrite(str(out / f"{name}.png"), page)
        labels = PageLabels(width=pw, height=ph, image=f"{name}.png", lines=lines)
        for f in cfg.formats:
            (out / f"{name}{EXTENSIONS[f.value]}").write_text(export(labels, f.value), encoding="utf-8")
        if cfg.crop_lines:
            crops.extend(export_line_crops(page, lines, out_dir=out, page_name=name, split=split))
        page_spec = layouts[i].spec
        pages.append(
            PageRecord(
                image=f"{name}.png",
                width=pw,
                height=ph,
                template=page_spec.template,  # the page's real template (may come from --spec, not cfg)
                seed=seed + i,
                binding=page_spec.binding,
                recto=page_spec.recto,
                n_lines=len(lines),
                split=split,
                styles=sorted({s for s in box_styles[i] if s}),
            )
        )
        console.print(
            f"[green]✓[/green] {name}.png  ({len(lines)} lines, {split})  "
            f"+ {', '.join(f.value for f in cfg.formats)}"
        )
    write_manifest(out, pages, crops)
    tail = f" + lines.jsonl ({len(crops)} crops)" if crops else ""
    console.print(f"[green]✓[/green] manifest.jsonl ({len(pages)} pages){tail}")


def main() -> None:
    typer.run(generate)


if __name__ == "__main__":
    main()