tinydoc-5m / tinydoc /synth /documents.py
tinydoc-5m v4: OmniDocBench EN text edit 0.117 (v3 0.127), inline maths as text, invoice summary tables
e24bf15 verified 5 days ago
56.5 kB
"""Synthetic document layouts (receipts, invoices, general pages, newspapers, magazines, academic papers)
with exact boxes, labels and relations.
A Layout can be used without rendering (extraction-model training only needs boxes and text) or
rendered to an RGB page image for the detector, recognizer and orientation models.
"""

from dataclasses import dataclass, field

import numpy as np
from PIL import Image, ImageDraw

from ..charset import normalize
from ..geometry import rect_quad
from . import text as T
from .fonts import FontBank


@dataclass
class Seg:
    text: str
    label: str
    box: tuple            # ink-consistent box x0, y0, x1, y1 (line-height based vertically)
    origin: tuple         # baseline origin for drawing
    font: object
    color: tuple
    row: int | None = None
    block: int | None = None   # layout block (paragraph, list item, table, caption...)


@dataclass
class Layout:
    width: int
    height: int
    doc_type: str
    segs: list = field(default_factory=list)
    kv: list = field(default_factory=list)          # (key_idx, value_idx)
    rules: list = field(default_factory=list)       # (x0, y0, x1, y1, width, color) lines
    fills: list = field(default_factory=list)       # (x0, y0, x1, y1, color) rectangles drawn under text
    barcodes: list = field(default_factory=list)    # (x0, y0, x1, y1)
    fields: dict = field(default_factory=dict)      # ground truth values for end-to-end evaluation
    draws: list = field(default_factory=list)       # (origin, text, font, color) drawn without a segment of their own
    blocks_known: bool = False                      # segment block ids are meaningful (structured pages)

    def segments(self):
        return [{"quad": rect_quad(*s.box), "text": s.text, "label": s.label, "row": s.row, "ignore": False, "block": s.block}
                for s in self.segs]


class Writer:
    """Places text on a Layout, measuring with the real font so boxes are exact."""

    def __init__(self, layout: Layout, bank: FontBank, rng, margin):
        self.L, self.bank, self.rng, self.m = layout, bank, rng, margin
        self.block = 0

    def new_block(self):
        self.block += 1

    def put(self, text, x, baseline, font_spec, label, align="l", row=None, color=(20, 20, 20), max_w=None):
        face, bold_axis, size = font_spec
        font = self.bank.font(face, size, bold_axis)
        text = FontBank.supported(face, normalize(text)).strip()
        if not text:
            return None
        if max_w is not None:
            while len(text) > 1 and font.getlength(text) > max_w:
                text = text[:-1].rstrip()
        l, t, r, b = font.getbbox(text, anchor="ls")
        if r - l < 1:
            return None
        _, rt, _, rb = font.getbbox("Hg", anchor="ls")
        ox = {"l": x - l, "r": x - r, "c": x - (l + r) / 2}[align]
        box = (ox + l, baseline + min(rt, t), ox + r, baseline + max(rb, b))
        if box[0] < 0 or box[2] > self.L.width:
            return None
        self.L.segs.append(Seg(text, label, box, (ox, baseline), font, color, row, self.block))
        return len(self.L.segs) - 1

    def kv(self, key, value, label, baseline, fspec, key_x, val_x, val_align="r", style=None, row=None, vspec=None):
        """Key/value on one line. style: split (two segments + relation), joined (one segment), value."""
        style = style or self.rng.choice(["split", "joined", "value"], p=[0.75, 0.15, 0.10])
        if style == "joined":
            return self.put(f"{key} {value}", key_x, baseline, vspec or fspec, label)
        v = self.put(value, val_x, baseline, vspec or fspec, label, align=val_align, row=row)
        if style == "split" and v is not None:
            k = self.put(key, key_x, baseline, fspec, "key", row=row)
            if k is not None:
                if self.L.segs[k].box[2] + 0.8 * (self.L.segs[k].box[3] - self.L.segs[k].box[1]) > self.L.segs[v].box[0]:
                    self.L.segs.pop()  # key would collide with the value: drop it
                else:
                    self.L.kv.append((k, v))
        return v


def _spec(bank, rng, mono=None, bold=False, size=20):
    face, bold_axis = bank.pick(rng, mono=mono, bold=bold)
    return face, bold_axis, int(size)


def _ink(rng):
    r = rng.random()
    if r < 0.75:
        g = int(rng.integers(0, 60))
        return (g, g, g)
    if r < 0.9:
        return (int(rng.integers(10, 40)), int(rng.integers(20, 60)), int(rng.integers(90, 160)))  # blue ink
    return (int(rng.integers(60, 110)),) * 3                                                         # faded


# ---------------------------------------------------------------------------------------------- receipt
def build_receipt(rng, bank: FontBank) -> Layout:
    W = int(rng.integers(420, 760))
    mono = rng.random() < 0.6
    cpl = rng.uniform(28, 46)
    size = int(np.clip(W / (cpl * (0.6 if mono else 0.5)), 14, 40))
    body = _spec(bank, rng, mono=mono, size=size)
    bold = _spec(bank, rng, mono=mono, bold=True, size=size)
    big = (bold[0], bold[1], int(size * rng.uniform(1.15, 1.9)))
    ink = _ink(rng)
    L = Layout(W, 10_000, "receipt")
    m = W * rng.uniform(0.03, 0.08)
    w = Writer(L, bank, rng, m)
    lh = size * rng.uniform(1.15, 1.5)
    y = m + size * 1.2
    mf = T.MoneyFormat(rng)
    center = rng.random() < 0.7
    hx, halign = (W / 2, "c") if center else (m, "l")

    def nl(k=1.0):
        nonlocal y
        y += lh * k

    merchant = T.merchant_name(rng)
    w.put(merchant, hx, y + big[2] - size, big, "merchant_name", halign, color=ink)
    nl(big[2] / size)
    if rng.random() < 0.85:
        for line in T.address_lines(rng)[: int(rng.integers(1, 3))]:
            w.put(line, hx, y, body, "merchant_address", halign, color=ink)
            nl()
    if rng.random() < 0.7:
        tel = T.phone(rng)
        if rng.random() < 0.5:
            w.put(f"{T.pick(rng, T.KEYS['merchant_phone'])} {tel}", hx, y, body, "merchant_phone", halign, color=ink)
        else:
            w.put(tel, hx, y, body, "merchant_phone", halign, color=ink)
        nl()
    if rng.random() < 0.3:
        w.put(f"{T.pick(rng, T.KEYS['merchant_tax_id'])} {T.tax_id(rng)}", hx, y, body, "merchant_tax_id", halign, color=ink)
        nl()
    nl(rng.uniform(0.3, 1.2))

    y_, mo, d = T.date_value(rng)
    date_s, time_s = T.fmt_date(rng, y_, mo, d), T.fmt_time(rng)
    if rng.random() < 0.5:   # date and time on one line
        w.kv(T.pick(rng, T.KEYS["date"]), date_s, "date", y, body, m, W / 2 - size, val_align="r")
        if rng.random() < 0.8:
            w.put(time_s, W - m, y, body, "time", "r", color=ink)
        nl()
    else:
        w.kv(T.pick(rng, T.KEYS["date"]), date_s, "date", y, body, m, W - m)
        nl()
        if rng.random() < 0.6:
            w.kv(T.pick(rng, T.KEYS["time"]), time_s, "time", y, body, m, W - m)
            nl()
    inv = T.invoice_no(rng)
    if rng.random() < 0.75:
        w.kv(T.pick(rng, T.KEYS["invoice_number"]), inv, "invoice_number", y, body, m, W - m)
        nl()
    else:
        inv = None
    if rng.random() < 0.4:
        w.put(f"{T.pick(rng, ['Cashier', 'Server', 'Table', 'Register', 'Clerk', 'POS'])}: {T.pick(rng, T.FIRST + [str(rng.integers(1, 40))])}", m, y, body, "other", color=ink)
        nl()

    def separator():
        if rng.random() < 0.6:
            ch = T.pick(rng, ["-", "=", "*", "_", "."])
            n = int((W - 2 * m) / max(w.bank.font(body[0], size, body[1]).getlength(ch), 1))
            w.put(ch * max(n - 1, 3), m, y - size * 0.3, body, "other", color=ink)
        else:
            L.rules.append((m, y - size * 0.6, W - m, y - size * 0.6, 1 + int(rng.integers(0, 2)), ink))
        nl(0.8)

    separator()
    style = rng.choice(["desc_amt", "qty_desc_amt", "two_line", "columns"])
    if style == "columns" and rng.random() < 0.6:
        w.put(T.pick(rng, T.ITEM_HEADERS["item_desc"]), m, y, body, "other", color=ink)
        w.put(T.pick(rng, T.ITEM_HEADERS["item_qty"]), W * 0.62, y, body, "other", "r", color=ink)
        w.put(T.pick(rng, T.ITEM_HEADERS["item_amount"]), W - m, y, body, "other", "r", color=ink)
        nl()
    n_items = int(rng.integers(1, 13))
    items, subtotal = [], 0.0
    for i in range(n_items):
        qty = int(rng.choice([1, 1, 1, 1, 2, 2, 3, 4, 5, 6, 10])) if style != "desc_amt" else 1
        unit = T.price_value(rng)
        amt = round(qty * unit, 2)
        subtotal += amt
        desc = T.item_name(rng)
        row = i
        amt_s, unit_s = mf.fmt(amt), mf.fmt(unit)
        amt_w = w.bank.font(body[0], size, body[1]).getlength(amt_s)
        desc_max = W - 2 * m - amt_w - size * 2
        if style == "desc_amt":
            w.put(desc, m, y, body, "item_desc", row=row, color=ink, max_w=desc_max)
            w.put(amt_s, W - m, y, body, "item_amount", "r", row=row, color=ink)
            nl()
        elif style == "qty_desc_amt":
            q = w.put(f"{qty}" + T.pick(rng, ["", "x", " x"]), m, y, body, "item_qty", row=row, color=ink)
            qx = L.segs[q].box[2] + size * 1.5 if q is not None else m
            w.put(desc, qx, y, body, "item_desc", row=row, color=ink, max_w=W - m - qx - amt_w - size * 2)
            w.put(amt_s, W - m, y, body, "item_amount", "r", row=row, color=ink)
            nl()
        elif style == "two_line":
            w.put(desc, m, y, body, "item_desc", row=row, color=ink, max_w=W - 2 * m)
            nl()
            ind = m + size * rng.uniform(0.5, 3)
            if rng.random() < 0.5:
                w.put(f"{qty} {T.pick(rng, ['x', '@', 'X'])} {unit_s}", ind, y, body, "item_qty", row=row, color=ink)
            else:
                w.put(f"{qty}", ind, y, body, "item_qty", row=row, color=ink)
                w.put(unit_s, W * 0.55, y, body, "item_unit_price", "r", row=row, color=ink)
            w.put(amt_s, W - m, y, body, "item_amount", "r", row=row, color=ink)
            nl()
        else:
            w.put(desc, m, y, body, "item_desc", row=row, color=ink, max_w=W * 0.62 - m - size * 3)
            w.put(str(qty), W * 0.62, y, body, "item_qty", "r", row=row, color=ink)
            w.put(amt_s, W - m, y, body, "item_amount", "r", row=row, color=ink)
            nl()
        items.append({"description": desc, "quantity": str(qty), "unit_price": unit_s, "amount": amt_s})
    separator()

    subtotal = round(subtotal, 2)
    discount = round(subtotal * rng.uniform(0.05, 0.2), 2) if rng.random() < 0.15 else 0.0
    rate = float(rng.choice([5, 7.5, 8, 10, 12, 13, 15, 18, 20]))
    tax = round((subtotal - discount) * rate / 100, 2) if rng.random() < 0.75 else 0.0
    total = round(subtotal - discount + tax, mf.decimals)
    key_x = m if rng.random() < 0.6 else W * rng.uniform(0.3, 0.5)
    fields = {"merchant_name": merchant, "date": date_s, "time": time_s, "invoice_number": inv, "items": items}
    if rng.random() < 0.8 or discount or tax:
        w.kv(T.pick(rng, T.KEYS["subtotal"]), mf.fmt(subtotal), "subtotal", y, body, key_x, W - m)
        fields["subtotal"] = mf.fmt(subtotal)
        nl()
    if discount:
        w.kv(T.pick(rng, T.KEYS["discount"]), mf.fmt(-discount), "discount", y, body, key_x, W - m)
        fields["discount"] = mf.fmt(-discount)
        nl()
    if tax:
        key = T.pick(rng, T.KEYS["tax"]).replace("{r}", f"{rate:g}")
        w.kv(key, mf.fmt(tax), "tax", y, body, key_x, W - m)
        fields["tax"] = mf.fmt(tax)
        nl()
    tspec = bold if rng.random() < 0.6 else body
    if rng.random() < 0.3:
        tspec = (bold[0], bold[1], int(size * rng.uniform(1.1, 1.5)))
        nl(0.3)
    w.kv(T.pick(rng, T.KEYS["total"]), mf.fmt(total, with_symbol=True), "total", y + tspec[2] - size, tspec, key_x, W - m, style=rng.choice(["split", "joined"], p=[0.9, 0.1]))
    fields["total"] = mf.fmt(total, with_symbol=True)
    nl(tspec[2] / size)
    if rng.random() < 0.6:
        paid = round(total + (rng.choice([0, 0, 1, 5, 10, 20]) if rng.random() < 0.6 else 0), mf.decimals)
        w.kv(T.pick(rng, T.KEYS["amount_paid"]), mf.fmt(paid), "amount_paid", y, body, key_x, W - m)
        fields["amount_paid"] = mf.fmt(paid)
        nl()
        if paid > total and rng.random() < 0.8:
            w.kv(T.pick(rng, T.KEYS["change"]), mf.fmt(paid - total), "change", y, body, key_x, W - m)
            fields["change"] = mf.fmt(paid - total)
            nl()
    nl(rng.uniform(0.3, 1.5))
    for _ in range(int(rng.integers(0, 4))):
        w.put(T.pick(rng, T.FOOTERS), hx, y, body, "other", halign, color=ink)
        nl()
    if rng.random() < 0.3:
        bw = W * rng.uniform(0.4, 0.7)
        L.barcodes.append((W / 2 - bw / 2, y - size, W / 2 + bw / 2, y + size * 1.5))
        nl(2.5)
    L.height = int(y + m)
    fields["currency"] = mf.code
    L.fields = fields
    return L


# ---------------------------------------------------------------------------------------------- invoice
def _item_table(w, rng, mf, body, bold, ink, m, W, H, y, size, lh):
    """Description / qty / unit price / amount table. Returns (y, items, subtotal, x of the amount column)."""
    L = w.L
    cols = {"desc": m + size * 0.5, "qty": W * rng.uniform(0.55, 0.62), "unit": W * rng.uniform(0.72, 0.78), "amt": W - m - size * 0.5}
    has_unit = rng.random() < 0.85
    grid = rng.random() < 0.5
    if rng.random() < 0.4:
        L.fills.append((m, y - size * 1.3, W - m, y + size * 0.6, tuple(int(v) for v in rng.integers(200, 240, 3))))
    head_y = y
    w.put(T.pick(rng, T.ITEM_HEADERS["item_desc"]), cols["desc"], y, bold, "other", color=ink)
    w.put(T.pick(rng, T.ITEM_HEADERS["item_qty"]), cols["qty"], y, bold, "other", "r", color=ink)
    if has_unit:
        w.put(T.pick(rng, T.ITEM_HEADERS["item_unit_price"]), cols["unit"], y, bold, "other", "r", color=ink)
    w.put(T.pick(rng, T.ITEM_HEADERS["item_amount"]), cols["amt"], y, bold, "other", "r", color=ink)
    L.rules.append((m, y + size * 0.7, W - m, y + size * 0.7, 2, ink))
    y += lh * 1.2
    items, subtotal = [], 0.0
    n_items = int(rng.integers(1, 14))
    for i in range(n_items):
        if y > H * 0.72:
            break
        qty = int(rng.choice([1, 1, 2, 3, 4, 5, 8, 10, 12, 20, 24, 50]))
        unit = T.price_value(rng, scale=rng.choice([1, 5, 20]))
        amt = round(qty * unit, 2)
        subtotal += amt
        desc = " ".join([T.item_name(rng)] + ([T.word(rng)] if rng.random() < 0.4 else []))
        w.put(desc, cols["desc"], y, body, "item_desc", row=i, color=ink, max_w=cols["qty"] - cols["desc"] - size * 4)
        w.put(str(qty), cols["qty"], y, body, "item_qty", "r", row=i, color=ink)
        if has_unit:
            w.put(mf.fmt(unit), cols["unit"], y, body, "item_unit_price", "r", row=i, color=ink)
        w.put(mf.fmt(amt), cols["amt"], y, body, "item_amount", "r", row=i, color=ink)
        items.append({"description": desc, "quantity": str(qty), "unit_price": mf.fmt(unit) if has_unit else None, "amount": mf.fmt(amt)})
        if rng.random() < 0.1:
            y += lh
            w.put(T.sentence(rng, 2, 5), cols["desc"], y, body, "item_desc", row=i, color=ink, max_w=cols["qty"] - cols["desc"] - size * 4)
        if grid:
            L.rules.append((m, y + size * 0.6, W - m, y + size * 0.6, 1, ink))
        y += lh
    if grid:
        for x in (m, cols["qty"] - size * 4, cols["unit"] - size * 5 if has_unit else None, W - m):
            if x is not None:
                L.rules.append((x, head_y - size * 1.3, x, y - lh + size * 0.6, 1, ink))
    y += lh * 0.5
    return y, items, round(subtotal, 2), cols["amt"]


TAX_RATES = [5, 7, 8, 10, 12, 18, 19, 20, 21, 23]
UNITS_OF_MEASURE = ["each", "pcs", "pc", "unit", "kg", "h", "m", "set", "box"]
SUMMARY_HEADS = {
    "rate": ["VAT [%]", "VAT rate", "Tax rate", "Rate", "VAT %", "Tax %", "GST %"],
    "net": ["Net worth", "Net amount", "Net", "Taxable amount", "Taxable value", "Amount excl. VAT", "Net value"],
    "tax": ["VAT", "VAT amount", "Tax", "Tax amount", "GST", "Sales tax"],
    "gross": ["Gross worth", "Gross amount", "Gross", "Amount incl. VAT", "Total incl. tax", "Gross value", "Total"],
}
SUMMARY_TITLES = ["SUMMARY", "Summary", "VAT summary", "Tax summary", "TOTALS", "VAT breakdown", "Tax details"]
SUMMARY_TOTAL_KEYS = ["Total", "TOTAL", "Totals", "Sum", "Grand total", "Total due", "Invoice total"]


def _pct(rng, rate, mf):
    s = f"{rate:g}"
    return (s.replace(".", ",") if mf.euro_style else s) + T.pick(rng, ["%", "%", " %"])


def _header(w, rng, text, x, y, spec, align, max_w, lh, ink, row=None):
    """A column header, wrapped onto two lines when it does not fit its column. Returns (first seg, lines used)."""
    font = w.bank.font(spec[0], spec[2], spec[1])
    words = text.split()
    if len(words) > 1 and (font.getlength(text) > max_w or rng.random() < 0.15):
        cut = (len(words) + 1) // 2
        k = w.put(" ".join(words[:cut]), x, y, spec, "other", align, row=row, color=ink, max_w=max_w)
        w.put(" ".join(words[cut:]), x, y + lh * 0.85, spec, "other", align, color=ink, max_w=max_w)
        return k, 2
    return w.put(text, x, y, spec, "other", align, row=row, color=ink, max_w=max_w), 1


def _vat_item_table(w, rng, mf, body, bold, ink, m, W, H, y, size, lh, rates):
    """Item table with net price, net worth, VAT rate and gross worth columns. Returns (y, items, net by rate)."""
    L = w.L
    jit = lambda f: W * (f + rng.uniform(-0.012, 0.012))
    has_um = rng.random() < 0.7
    cols = {"no": m + size * 0.5, "desc": m + size * rng.uniform(2.4, 3.2), "qty": jit(0.47), "um": jit(0.50),
            "unit": jit(0.62), "net": jit(0.73), "vat": jit(0.83), "gross": W - m - size * 0.5}
    heads = [("no", T.pick(rng, ["No.", "#", "No", "Pos."]), "l"), ("desc", T.pick(rng, T.ITEM_HEADERS["item_desc"]), "l"),
             ("qty", T.pick(rng, T.ITEM_HEADERS["item_qty"]), "r"), ("um", T.pick(rng, ["UM", "Unit", "UoM"]), "l"),
             ("unit", T.pick(rng, ["Net price", "Unit price", "Price"]), "r"),
             ("net", T.pick(rng, ["Net worth", "Net amount", "Net value", "Net"]), "r"),
             ("vat", T.pick(rng, ["VAT [%]", "VAT %", "Tax %", "VAT"]), "r"),
             ("gross", T.pick(rng, ["Gross worth", "Gross amount", "Gross", "Amount", "Total"]), "r")]
    head_y, lines = y, 1
    if rng.random() < 0.4:
        L.fills.append((m, y - size * 1.3, W - m, y + size * 0.6 + lh * 0.85, tuple(int(v) for v in rng.integers(205, 240, 3))))
    for key, text, align in heads:
        if key == "um" and not has_um:
            continue
        lines = max(lines, _header(w, rng, text, cols[key], y, bold, align, W * 0.095 if key != "desc" else W * 0.3, lh, ink)[1])
    y += lh * 0.85 * (lines - 1)
    L.rules.append((m, y + size * 0.7, W - m, y + size * 0.7, 1 + int(rng.random() < 0.5), ink))
    y += lh * 1.3
    items, by_rate, shade = [], {}, rng.random() < 0.4
    for i in range(int(rng.integers(1, 10))):
        if y > H * 0.6:
            break
        qty = int(rng.choice([1, 1, 2, 3, 4, 5, 8, 10, 12]))
        unit = T.price_value(rng, scale=rng.choice([1, 5, 20]))
        net = round(qty * unit, 2)
        rate = float(T.pick(rng, rates))
        by_rate[rate] = by_rate.get(rate, 0.0) + net
        desc = " ".join([T.item_name(rng)] + ([T.word(rng)] if rng.random() < 0.4 else []))
        qty_s = f"{qty:.2f}".replace(".", "," if mf.euro_style else ".") if rng.random() < 0.5 else str(qty)
        if shade and i % 2 == 0:
            L.fills.append((m, y - size * 1.2, W - m, y + size * 0.6, tuple(int(v) for v in rng.integers(220, 240, 3))))
        w.put(f"{i + 1}." if rng.random() < 0.7 else str(i + 1), cols["no"], y, body, "other", row=i, color=ink)
        w.put(desc, cols["desc"], y, body, "item_desc", row=i, color=ink, max_w=cols["qty"] - cols["desc"] - size * 4)
        w.put(qty_s, cols["qty"], y, body, "item_qty", "r", row=i, color=ink)
        if has_um:
            w.put(T.pick(rng, UNITS_OF_MEASURE), cols["um"], y, body, "other", row=i, color=ink)
        w.put(mf.fmt(unit), cols["unit"], y, body, "item_unit_price", "r", row=i, color=ink)
        w.put(mf.fmt(net), cols["net"], y, body, "item_amount", "r", row=i, color=ink)
        w.put(_pct(rng, rate, mf), cols["vat"], y, body, "other", "r", row=i, color=ink)
        w.put(mf.fmt(round(net * (1 + rate / 100), 2)), cols["gross"], y, body, "other", "r", row=i, color=ink)
        items.append({"description": desc, "quantity": qty_s, "unit_price": mf.fmt(unit), "amount": mf.fmt(net)})
        if rng.random() < 0.2:   # description continues on a second line
            y += lh * 0.85
            w.put(T.sentence(rng, 2, 4), cols["desc"], y, body, "item_desc", row=i, color=ink, max_w=cols["qty"] - cols["desc"] - size * 4)
        y += lh * 1.3
        if not shade:
            L.rules.append((m, y - lh * 0.75, W - m, y - lh * 0.75, 1, ink))
    if rng.random() < 0.6:
        for x in (m, W - m):
            L.rules.append((x, head_y - size * 1.3, x, y - lh * 0.75, 1, ink))
    return y + lh * 0.5, items, by_rate


def _vat_summary(w, rng, mf, body, bold, ink, accent, m, W, y, size, lh, by_rate):
    """Tax summary table: a row per VAT rate (net, VAT, gross) and a totals row holding subtotal, tax and total.
    Returns (y, fields)."""
    L = w.L
    if rng.random() < 0.6:
        w.new_block()
        w.put(T.pick(rng, SUMMARY_TITLES), m, y + lh * 0.3, (bold[0], bold[1], int(size * rng.uniform(1.1, 1.6))), "other", color=accent)
        y += lh * 1.9
    w.new_block()
    x0, x1 = (m if rng.random() < 0.5 else W * rng.uniform(0.35, 0.5)), W - m
    order = ["rate", "net", "tax", "gross"] if rng.random() < 0.85 else ["net", "tax", "gross"]
    span0 = x0 + (x1 - x0) * rng.uniform(0.25, 0.4)
    col = dict(zip(order, np.linspace(span0 + (x1 - span0) / len(order), x1 - size * 0.5, len(order))))
    col_w = (x1 - span0) / len(order) * 0.9
    head_y, heads, lines = y, {}, 1
    if rng.random() < 0.4:
        L.fills.append((x0, y - size * 1.3, x1, y + size * 0.6, tuple(int(v) for v in rng.integers(205, 240, 3))))
    for k in order:
        heads[k], n = _header(w, rng, T.pick(rng, SUMMARY_HEADS[k]), col[k], y, bold, "r", col_w, lh, ink, row=100)
        lines = max(lines, n)
    y += lh * 0.85 * (lines - 1) + lh * 1.4
    shade = rng.random() < 0.5
    show_rates = len(by_rate) > 1 or rng.random() < 0.8
    net_t = tax_t = gross_t = 0.0
    for r_i, (rate, net) in enumerate(sorted(by_rate.items())):
        net = round(net, 2)
        tax = round(net * rate / 100, 2)
        net_t, tax_t, gross_t = net_t + net, tax_t + tax, gross_t + net + tax
        if not show_rates:
            continue
        row = 101 + r_i
        if shade:
            L.fills.append((x0, y - size * 1.2, x1, y + size * 0.6, tuple(int(v) for v in rng.integers(220, 240, 3))))
        if "rate" in col:
            w.put(_pct(rng, rate, mf), col["rate"], y, body, "other", "r", row=row, color=ink)
        else:
            w.put(f"VAT {_pct(rng, rate, mf)}", x0 + size * 0.5, y, body, "other", row=row, color=ink)
        for k, v in (("net", net), ("tax", tax), ("gross", net + tax)):
            w.put(mf.fmt(v), col[k], y, body, "other", "r", row=row, color=ink)
        y += lh * 1.4
    if rng.random() < 0.7:
        L.rules.append((x0, y - lh * 0.85, x1, y - lh * 0.85, 1, ink))
    spec = bold if rng.random() < 0.7 else body
    sym = rng.random() < 0.6
    if rng.random() < 0.5:
        key = w.put(T.pick(rng, SUMMARY_TOTAL_KEYS), x0 + size * 0.5, y, spec, "key", row=120, color=ink)
    else:
        key = w.put(T.pick(rng, SUMMARY_TOTAL_KEYS), span0, y, spec, "key", "r", row=120, color=ink)
    fields, vals = {}, {}
    for k, v, label in (("net", net_t, "subtotal"), ("tax", tax_t, "tax"), ("gross", gross_t, "total")):
        v = round(v, mf.decimals if label == "total" else 2)
        vals[k] = w.put(mf.fmt(v, with_symbol=sym), col[k], y, spec, label, "r", row=120, color=ink)
        fields[label] = mf.fmt(v, with_symbol=sym)
        if heads[k] is not None and vals[k] is not None:
            L.kv.append((heads[k], vals[k]))      # the column header names the value in the totals row
    if key is not None and vals["gross"] is not None:
        L.kv.append((key, vals["gross"]))
    if rng.random() < 0.6:
        for x in (x0, x1):
            L.rules.append((x, head_y - size * 1.3, x, y + size * 0.6, 1, ink))
        L.rules.append((x0, head_y - size * 1.3, x1, head_y - size * 1.3, 1, ink))
        L.rules.append((x0, y + size * 0.6, x1, y + size * 0.6, 1, ink))
    y += lh * 1.2
    if rng.random() < 0.25:   # restated below the table
        w.new_block()
        w.kv(T.pick(rng, ["To pay:", "Amount due:", "Total due:", "Balance due:"]), mf.fmt(round(gross_t, mf.decimals), with_symbol=True),
             "total", y + lh * 0.5, bold, W * rng.uniform(0.5, 0.62), x1 - size * 0.5, style="split")
        y += lh * 1.5
    return y, fields


def build_invoice(rng, bank: FontBank) -> Layout:
    W = int(rng.integers(1000, 1500))
    H = int(W * rng.choice([1.414, 1.294]))
    size = int(W * rng.uniform(0.013, 0.019))
    body = _spec(bank, rng, mono=False if rng.random() < 0.9 else True, size=size)
    bold = _spec(bank, rng, bold=True, size=size)
    ink = _ink(rng)
    accent = ink if rng.random() < 0.6 else (int(rng.integers(0, 120)), int(rng.integers(0, 120)), int(rng.integers(60, 200)))
    L = Layout(W, H, "invoice")
    m = W * rng.uniform(0.05, 0.09)
    w = Writer(L, bank, rng, m)
    lh = size * rng.uniform(1.3, 1.7)
    mf = T.MoneyFormat(rng)
    y = m + size

    merchant = T.merchant_name(rng)
    tsize = int(size * rng.uniform(1.4, 2.4))
    w.put(merchant, m, y + tsize - size, (bold[0], bold[1], tsize), "merchant_name", color=accent)
    y_left = y + tsize * 1.3
    for line in T.address_lines(rng):
        w.put(line, m, y_left, body, "merchant_address", color=ink)
        y_left += lh
    if rng.random() < 0.8:
        w.kv(T.pick(rng, T.KEYS["merchant_phone"]), T.phone(rng), "merchant_phone", y_left, body, m, m + size * 6, val_align="l", style=rng.choice(["split", "joined"]))
        y_left += lh
    if rng.random() < 0.6:
        w.put(f"{T.word(rng).lower()}@{T.word(rng).lower()}.com", m, y_left, body, "other", color=ink)
        y_left += lh
    if rng.random() < 0.5:
        w.kv(T.pick(rng, T.KEYS["merchant_tax_id"]), T.tax_id(rng), "merchant_tax_id", y_left, body, m, m + size * 6, val_align="l", style=rng.choice(["split", "joined"]))
        y_left += lh

    title = T.pick(rng, ["INVOICE", "Invoice", "TAX INVOICE", "BILL", "Tax Invoice", "SALES INVOICE"])
    w.put(title, W - m, y + tsize - size, (bold[0], bold[1], int(size * rng.uniform(1.6, 2.6))), "other", "r", color=accent)
    y_right = y + tsize * 1.6
    kx, vx = W * rng.uniform(0.55, 0.65), W - m
    inv = T.invoice_no(rng)
    yy, mo, d = T.date_value(rng)
    date_s = T.fmt_date(rng, yy, mo, d)
    due_s = T.fmt_date(rng, *T.date_value(rng)) if rng.random() < 0.6 else None
    w.kv(T.pick(rng, T.KEYS["invoice_number"]), inv, "invoice_number", y_right, body, kx, vx, style="split")
    y_right += lh
    w.kv(T.pick(rng, T.KEYS["date"]), date_s, "date", y_right, body, kx, vx, style="split")
    y_right += lh
    if due_s:
        w.kv(T.pick(rng, T.KEYS["due_date"]), due_s, "due_date", y_right, body, kx, vx, style="split")
        y_right += lh

    y = max(y_left, y_right) + lh
    k = w.put(T.pick(rng, ["Bill To:", "BILL TO", "Billed To", "Customer:", "To:", "Invoice To:"]), m, y, bold, "key", color=accent)
    y += lh
    customer = T.person_name(rng) if rng.random() < 0.5 else T.merchant_name(rng)
    c = w.put(customer, m, y, body, "customer", color=ink)
    if k is not None and c is not None:
        L.kv.append((k, c))
    y += lh
    for line in T.address_lines(rng):
        w.put(line, m, y, body, "customer", color=ink)
        y += lh
    y += lh

    vat_items = rng.random() < 0.35      # European style: net price, net worth, VAT rate and gross worth per line
    if vat_items:
        rates = sorted({float(r) for r in rng.choice(TAX_RATES, int(rng.choice([1, 1, 1, 2, 3])))})
        y, items, by_rate = _vat_item_table(w, rng, mf, body, bold, ink, m, W, H, y, size, lh, rates)
        subtotal = round(sum(by_rate.values()), 2)
    else:
        y, items, subtotal, amt_x = _item_table(w, rng, mf, body, bold, ink, m, W, H, y, size, lh)
    fields = {"merchant_name": merchant, "invoice_number": inv, "date": date_s, "due_date": due_s, "customer": customer, "items": items}
    if vat_items or rng.random() < 0.1:
        # totals as a tax summary table (a row per VAT rate, then a totals row) instead of a key/value stack
        y, totals = _vat_summary(w, rng, mf, body, bold, ink, accent, m, W, y, size, lh,
                                 by_rate if vat_items else {float(rng.choice(TAX_RATES)): subtotal})
        fields.update(totals)
        y += lh
    else:
        subtotal = round(subtotal, 2)
        discount = round(subtotal * rng.uniform(0.03, 0.15), 2) if rng.random() < 0.2 else 0.0
        rate = float(rng.choice([5, 8, 10, 12, 18, 19, 20, 21]))
        tax = round((subtotal - discount) * rate / 100, 2)
        total = round(subtotal - discount + tax, mf.decimals)
        kx = W * rng.uniform(0.55, 0.65)
        w.kv(T.pick(rng, T.KEYS["subtotal"]), mf.fmt(subtotal), "subtotal", y, body, kx, amt_x, style="split")
        fields["subtotal"] = mf.fmt(subtotal)
        y += lh
        if discount:
            w.kv(T.pick(rng, T.KEYS["discount"]), mf.fmt(-discount), "discount", y, body, kx, amt_x, style="split")
            fields["discount"] = mf.fmt(-discount)
            y += lh
        w.kv(T.pick(rng, T.KEYS["tax"]).replace("{r}", f"{rate:g}"), mf.fmt(tax), "tax", y, body, kx, amt_x, style="split")
        fields["tax"] = mf.fmt(tax)
        y += lh
        L.rules.append((kx, y - size * 0.9, W - m, y - size * 0.9, 1, ink))
        w.kv(T.pick(rng, T.KEYS["total"]), mf.fmt(total, with_symbol=True), "total", y + size * 0.2, bold, kx, amt_x, style="split")
        fields["total"] = mf.fmt(total, with_symbol=True)
        y += lh * 2
    if rng.random() < 0.7:
        for _ in range(int(rng.integers(1, 4))):
            if y > H - m - lh:
                break
            w.put(T.sentence(rng, 4, 12), m, y, body, "other", color=ink, max_w=W - 2 * m)
            y += lh
    if rng.random() < 0.6 and y < H - m - lh:
        w.put(T.pick(rng, ["Thank you for your business!", "Payment due within 30 days.", "Please make checks payable to " + merchant]), m, H - m - lh, body, "other", color=ink, max_w=W - 2 * m)
    fields["currency"] = mf.code
    L.fields = fields
    return L


# ------------------------------------------------------------------------------------------ general page
def _wrap(words, font, max_w):
    lines, cur = [], ""
    for wd in words:
        cand = f"{cur} {wd}".strip()
        if font.getlength(cand) <= max_w or not cur:
            cur = cand
        else:
            lines.append(cur)
            cur = wd
    if cur:
        lines.append(cur)
    return lines


def build_general(rng, bank: FontBank) -> Layout:
    W = int(rng.integers(900, 1500))
    H = int(W * rng.choice([1.414, 1.294]))
    notes = rng.random() < 0.25
    size = int(W * rng.uniform(0.012, 0.02))
    body = _spec(bank, rng, mono=True if notes else (False if rng.random() < 0.9 else None), size=size)
    hface, hax = bank.pick(rng, bold=True)
    ink = _ink(rng)
    hcolor = ink if rng.random() < 0.7 else (int(rng.integers(0, 100)), int(rng.integers(0, 100)), int(rng.integers(50, 180)))
    L = Layout(W, H, "general", blocks_known=True)
    m = W * rng.uniform(0.06, 0.12)
    w = Writer(L, bank, rng, m)
    lh = size * rng.uniform(1.25, 1.7)
    font = bank.font(body[0], body[2], body[1])
    n_cols = 2 if (rng.random() < 0.2 and not notes) else 1
    gutter = W * 0.04
    col_w = (W - 2 * m - (n_cols - 1) * gutter) / n_cols
    y = m + size
    if rng.random() < 0.3:
        w.put(T.sentence(rng, 2, 6).rstrip("."), m, y, (body[0], body[1], int(size * 0.8)), "other", color=ink, max_w=W - 2 * m)
        y += lh * 1.5
    if rng.random() < 0.8:
        ts = int(size * rng.uniform(1.5, 2.4))
        w.new_block()
        for line in _wrap(T.sentence(rng, 2, 8).rstrip(".!?:").split(), bank.font(hface, ts, hax), W - 2 * m)[:2]:
            w.put(line, m if rng.random() < 0.7 else W / 2, y + ts - size, (hface, hax, ts), "heading", "l", color=hcolor)
            y += ts * 1.3
        y += lh * 0.5
    if notes and rng.random() < 0.6:
        w.new_block()
        w.put(T.fmt_date(rng, *T.date_value(rng)), m, y, body, "paragraph", color=ink)
        y += lh * 1.5

    y_top = y
    row_id = 0
    for c in range(n_cols):
        x0 = m + c * (col_w + gutter)
        y = y_top
        last = "heading"
        while y < H - m - lh * 3:
            r = rng.random()
            kind = "heading" if (r < 0.18 and last != "heading") else "paragraph" if r < 0.65 else "list" if r < 0.83 else "table"
            w.new_block()
            if kind == "heading":
                hs = int(size * rng.uniform(1.1, 1.6))
                y += lh * 0.4
                w.put(T.sentence(rng, 1, 6).rstrip(".!?:"), x0, y + hs - size, (hface, hax, hs), "heading", color=hcolor, max_w=col_w)
                y += lh * hs / size * 1.1
            elif kind == "paragraph":
                indent = size * 2 if rng.random() < 0.3 else 0
                words = " ".join(T.sentence(rng) for _ in range(int(rng.integers(1, 6)))).split()
                for i, line in enumerate(_wrap(words, font, col_w - indent)):
                    if y > H - m:
                        break
                    w.put(line, x0 + (indent if i == 0 else 0), y, body, "paragraph", color=ink)
                    y += lh
                y += lh * rng.uniform(0.3, 1.0)
            elif kind == "list":
                bullet = T.pick(rng, ["•", "-", "*", "–", "num", "alpha"])
                for i in range(int(rng.integers(2, 7))):
                    w.new_block()   # every list item is its own block, as in DocLayNet
                    b = f"{i + 1}." if bullet == "num" else f"{'abcdefg'[i]})" if bullet == "alpha" else bullet
                    for j, line in enumerate(_wrap(T.sentence(rng, 2, 16).split(), font, col_w - size * 2.5)):
                        if y > H - m:
                            break
                        text = f"{b} {line}" if j == 0 else line
                        w.put(text, x0 + (size * 0.5 if j == 0 else size * 2.5), y, body, "list_item", color=ink)
                        y += lh
                y += lh * 0.5
            else:
                ncol = int(rng.integers(2, 6))
                nrow = int(rng.integers(2, 9))
                cw = col_w / ncol
                grid = rng.random() < 0.6
                ty0 = y - size * 1.2
                for ri in range(nrow):
                    if y > H - m:
                        break
                    spec = (hface, hax, size) if ri == 0 else body
                    for ci in range(ncol):
                        txt = T.word(rng).capitalize() if ri == 0 else (T.number(rng) if ci > 0 and rng.random() < 0.7 else " ".join(T.word(rng) for _ in range(int(rng.integers(1, 3)))))
                        num = ci > 0 and ri > 0
                        xx = x0 + cw * (ci + 1) - size * 0.5 if num else x0 + cw * ci + size * 0.4
                        w.put(txt, xx, y, spec, "table_cell", "r" if num else "l", row=row_id, color=ink, max_w=cw - size)
                    row_id += 1
                    if grid or ri == 0:
                        L.rules.append((x0, y + size * 0.6, x0 + col_w, y + size * 0.6, 1, ink))
                    y += lh * 1.15
                if grid:
                    for ci in range(ncol + 1):
                        L.rules.append((x0 + cw * ci, ty0, x0 + cw * ci, y - lh * 1.15 + size * 0.6, 1, ink))
                y += lh * 0.6
            last = kind
    if rng.random() < 0.5:
        w.new_block()
        w.put(str(rng.integers(1, 300)), W / 2, H - m * 0.5, (body[0], body[1], int(size * 0.85)), "other", "c", color=ink)
    return L


# ------------------------------------------------------------------------------- multi-column pages
def _flow(w, words, font_spec, label, x0, col_w, y, y_max, lh, ink, indent=0.0, justify_last=False):
    """Set words as wrapped lines from y; stop at y_max. Returns (y after, words that did not fit)."""
    font = w.bank.font(font_spec[0], font_spec[2], font_spec[1])
    lines = _wrap(words, font, col_w - indent)
    for i, line in enumerate(lines):
        if y > y_max:
            rest = " ".join(lines[i:]).split()
            return y, rest
        w.put(line, x0 + (indent if i == 0 else 0), y, font_spec, label, color=ink, max_w=col_w)
        y += lh
    return y, []


def _image_placeholder(L, rng, x0, y0, x1, y1):
    """A photo / chart stand-in: shaded rectangle with a few shapes (its pixels carry no text)."""
    base = tuple(int(v) for v in rng.integers(120, 225, 3))
    L.fills.append((x0, y0, x1, y1, base))
    for _ in range(int(rng.integers(2, 7))):
        cx0, cy0 = rng.uniform(x0, x1), rng.uniform(y0, y1)
        cx1, cy1 = min(x1, cx0 + rng.uniform(0.1, 0.5) * (x1 - x0)), min(y1, cy0 + rng.uniform(0.1, 0.5) * (y1 - y0))
        L.fills.append((cx0, cy0, cx1, cy1, tuple(int(v) for v in rng.integers(40, 250, 3))))


def build_columns(rng, bank: FontBank, style="newspaper") -> Layout:
    """Newspaper (3-6 narrow columns, small type, stories in bands) or magazine (2-3 columns, photos, pull quotes).
    Each story band spans a group of columns; text flows down one column and continues in the next, and a
    paragraph broken by a column change becomes a new block (as DocLayNet annotates it).
    """
    news = style == "newspaper"
    W = int(rng.integers(1400, 2000) if news else rng.integers(1100, 1600))
    H = int(W * rng.uniform(1.25, 1.5))
    size = int(W * (rng.uniform(0.0075, 0.0105) if news else rng.uniform(0.010, 0.0145)))
    body = _spec(bank, rng, mono=False, size=size)
    hface, hax = bank.pick(rng, bold=True)
    ink = _ink(rng)
    L = Layout(W, H, "general", blocks_known=True)
    m = W * rng.uniform(0.03, 0.06)
    w = Writer(L, bank, rng, m)
    lh = size * rng.uniform(1.12, 1.35)
    n_cols = int(rng.integers(3, 7) if news else rng.integers(2, 4))
    gutter = size * rng.uniform(0.6, 1.8)       # newspapers often have very narrow gutters
    col_w = (W - 2 * m - (n_cols - 1) * gutter) / n_cols
    col_x = [m + c * (col_w + gutter) for c in range(n_cols)]
    rules = rng.random() < (0.5 if news else 0.15)
    y = m + size

    # masthead / section header
    if rng.random() < (0.7 if news else 0.4):
        ts = int(W * rng.uniform(0.035, 0.07))
        w.new_block()
        name = " ".join(T.word(rng).capitalize() for _ in range(int(rng.integers(1, 4))))
        w.put(name, W / 2, y + ts * 0.8, (hface, hax, ts), "heading", "c", color=ink, max_w=W - 2 * m)
        y += ts * 1.25
        L.rules.append((m, y - size * 0.4, W - m, y - size * 0.4, 2, ink))
        w.new_block()
        w.put(T.fmt_date(rng, *T.date_value(rng)), m, y + size * 0.6, body, "other", color=ink)
        if rng.random() < 0.7:
            w.new_block()
            w.put(f"{T.pick(rng, ['Vol.', 'No.', 'Issue', 'Edition'])} {rng.integers(1, 400)}", W - m, y + size * 0.6, body, "other", "r", color=ink)
        y += lh * 1.6
        L.rules.append((m, y - size * 0.6, W - m, y - size * 0.6, 1, ink))
        y += lh * 0.5
    elif rng.random() < 0.5:
        w.new_block()
        w.put(T.sentence(rng, 1, 4).rstrip(".").upper(), m, y, (body[0], body[1], int(size * 0.9)), "other", color=ink)
        y += lh * 1.5

    y_bottom = H - m - lh
    while y < y_bottom - lh * 6:
        band_h = min(y_bottom - y, H * rng.uniform(0.22, 0.6))
        if y_bottom - (y + band_h) < lh * 8:
            band_h = y_bottom - y
        # split the columns into story groups
        groups, c = [], 0
        while c < n_cols:
            span = int(rng.integers(1, n_cols - c + 1)) if rng.random() < 0.6 else n_cols - c
            groups.append((c, c + span))
            c += span
        band_end = y
        for g0, g1 in groups:
            x0, x1 = col_x[g0], col_x[g1 - 1] + col_w
            gy = y
            hs = int(size * rng.uniform(1.4, 3.2 if (g1 - g0) > 1 else 2.0))
            w.new_block()
            for line in _wrap(T.sentence(rng, 3, 12).rstrip(".").split(), bank.font(hface, hs, hax), x1 - x0)[:3]:
                w.put(line, x0, gy + hs * 0.8, (hface, hax, hs), "heading", color=ink)
                gy += hs * 1.15
            gy += lh * 0.3
            if rng.random() < 0.4:   # deck / standfirst
                w.new_block()
                ds = (body[0], body[1], int(size * rng.uniform(1.1, 1.4)))
                gy, _ = _flow(w, T.sentence(rng, 8, 22).split(), ds, "paragraph", x0, x1 - x0, gy + size, y + band_h, ds[2] * 1.25, ink)
                gy += lh * 0.2
            if rng.random() < 0.5:   # byline
                w.new_block()
                w.put(f"By {T.person_name(rng)}", x0, gy + size, (body[0], body[1], int(size * 0.9)), "other", color=ink)
                gy += lh * 1.3
            if rng.random() < (0.35 if news else 0.6) and (y + band_h) - gy > lh * 14:
                ih = min((y + band_h - gy) * rng.uniform(0.3, 0.55), (x1 - x0) * 0.8)
                _image_placeholder(L, rng, x0, gy, x1, gy + ih)
                gy += ih + size * 0.8
                w.new_block()
                cs = (body[0], body[1], int(size * 0.85))
                gy, _ = _flow(w, T.sentence(rng, 6, 20).split(), cs, "paragraph", x0, x1 - x0, gy + size * 0.8, y + band_h, cs[2] * 1.2, ink)
                gy += lh * 0.5
            # body text flows down each column of the group
            words, block_open = [], False
            for c in range(g0, g1):
                cy = gy + size
                block_open = False
                while cy < y + band_h:
                    if not words:
                        words = " ".join(T.sentence(rng) for _ in range(int(rng.integers(2, 7)))).split()
                        if c > g0 or cy > gy + size:
                            cy += lh * rng.uniform(0.0, 0.5)
                        indent = size * rng.uniform(1, 2) if rng.random() < 0.7 else 0
                        if not news and rng.random() < 0.08 and (y + band_h) - cy > lh * 5:   # pull quote
                            w.new_block()
                            qs = (hface, hax, int(size * rng.uniform(1.4, 1.9)))
                            L.rules.append((col_x[c], cy - size, col_x[c] + col_w, cy - size, 2, ink))
                            cy, _ = _flow(w, T.sentence(rng, 6, 14).split(), qs, "paragraph", col_x[c], col_w, cy + qs[2] * 0.6, y + band_h, qs[2] * 1.2, ink)
                            cy += lh
                            words = []
                            continue
                        w.new_block()
                        block_open = True
                    elif not block_open:
                        w.new_block()   # paragraph continues in the next column: a new box
                        block_open = True
                        indent = 0
                    cy, words = _flow(w, words, body, "paragraph", col_x[c], col_w, cy, y + band_h, lh, ink, indent)
                    if words:
                        break
                band_end = max(band_end, cy)
                if rules and c < g1 - 1:
                    L.rules.append((col_x[c] + col_w + gutter / 2, gy, col_x[c] + col_w + gutter / 2, y + band_h, 1, ink))
            band_end = max(band_end, gy)
        y = band_end + lh
        if y < y_bottom - lh * 6:
            L.rules.append((m, y - lh * 0.6, W - m, y - lh * 0.6, int(rng.integers(1, 3)), ink))
    if rng.random() < 0.6:
        w.new_block()
        w.put(str(rng.integers(1, 80)), W - m if rng.random() < 0.5 else W / 2, H - m * 0.5, (body[0], body[1], int(size * 0.9)), "other",
              "r" if rng.random() < 0.5 else "c", color=ink)
    return L


MATH_VARS = list("xyzabcnkmtuvwfgpq") + ["\u03b1", "\u03b2", "\u03b8", "\u03bb", "\u03c3", "\u03bc"]
MATH_OPS = ["+", "\u2212", "=", "\u2264", "\u2248", "\u00b7", "\u00d7", "/"]


def put_formula(w, x_center, baseline, spec, max_w, number=None, x_right=None, color=(20, 20, 20)):
    """Draw a display formula from pieces (variables, operators, sub/superscripts, a fraction) and add ONE
    segment labelled "formula" covering it, like DocLayNet's formula boxes. Returns the segment's bottom."""
    rng, L = w.rng, w.L
    face, bax, size = spec
    font = w.bank.font(face, size, bax)
    small = w.bank.font(face, max(8, int(size * 0.7)), bax)
    pieces = []      # (dx, dy, text, font)

    def term():
        out = [(FontBank.supported(face, str(rng.choice(MATH_VARS))) or "x", 0, font)]
        r = rng.random()
        if r < 0.35:
            out.append((str(rng.choice(["2", "n", "k", "i", "2k", "-1", "T"])), -0.45 * size, small))
        elif r < 0.6:
            out.append((str(rng.choice(["i", "j", "0", "1", "k", "t", "ij"])), 0.25 * size, small))
        return out

    def seq(n):
        out = []
        for k in range(n):
            if k:
                out.append((" " + (FontBank.supported(face, str(rng.choice(MATH_OPS))) or "+") + " ", 0, font))
            if rng.random() < 0.15:
                out.append(("(", 0, font))
                out += term() + [(" + ", 0, font)] + term()
                out.append((")", 0, font))
            else:
                out += term()
        return out

    lhs = term() + [(" = ", 0, font)]
    frac = rng.random() < 0.4
    x = 0.0
    for t, dy, f in lhs:
        pieces.append((x, dy, t, f))
        x += f.getlength(t)
    if frac:
        num, den = seq(int(rng.integers(1, 4))), seq(int(rng.integers(1, 3)))
        wn = sum(f.getlength(t) for t, _, f in num)
        wd = sum(f.getlength(t) for t, _, f in den)
        fw = max(wn, wd) + size * 0.4
        for parts, width, dy0 in ((num, wn, -0.75 * size), (den, wd, 0.85 * size)):
            xx = x + (fw - width) / 2
            for t, dy, f in parts:
                pieces.append((xx, dy0 + dy, t, f))
                xx += f.getlength(t)
        bar_x = x
        x += fw
    for t, dy, f in seq(int(rng.integers(1, 4))) if (not frac or rng.random() < 0.5) else []:
        pieces.append((x, dy, t, f))
        x += f.getlength(t)
    total_w = min(x, max_w)
    x0 = x_center - total_w / 2
    boxes = []
    for dx, dy, t, f in pieces:
        if dx > max_w:
            break
        l, tp, r, b = f.getbbox(t, anchor="ls")
        ox, oy = x0 + dx, baseline + dy
        L.draws.append(((ox, oy), t, f, color))
        boxes.append((ox + l, oy + tp, ox + r, oy + b))
    if frac:
        bx = x0 + bar_x
        L.rules.append((bx, baseline - 0.3 * size, bx + fw, baseline - 0.3 * size, max(1, size // 14), color))
        boxes.append((bx, baseline - 0.3 * size - 1, bx + fw, baseline - 0.3 * size + 1))
    if number is not None and x_right is not None:
        l, tp, r, b = font.getbbox(number, anchor="ls")
        L.draws.append(((x_right - r, baseline), number, font, color))
        boxes.append((x_right - r + l, baseline + tp, x_right, baseline + b))
    boxes = [b for b in boxes if b[2] > b[0]]
    if not boxes:
        return baseline
    box = (min(b[0] for b in boxes), min(b[1] for b in boxes), max(b[2] for b in boxes), max(b[3] for b in boxes))
    text = normalize("".join(t for _, _, t, _ in pieces) + (f" {number}" if number else "")).strip() or "x"
    L.segs.append(Seg(text, "formula", box, box[:2], None, color, None, w.block))
    return box[3]


def build_academic(rng, bank: FontBank) -> Layout:
    """Research-paper page: title block, abstract, two (or one) columns with sections, display formulas,
    figures with captions, tables, footnotes and numbered references."""
    W = int(rng.integers(1150, 1550))
    H = int(W * rng.choice([1.294, 1.414]))
    two_col = rng.random() < 0.7
    size = int(W * (rng.uniform(0.0105, 0.013) if two_col else rng.uniform(0.012, 0.015)))
    serif = _spec(bank, rng, mono=False, size=size)
    hface, hax = bank.pick(rng, bold=True)
    ink = _ink(rng) if rng.random() < 0.3 else (15, 15, 15)
    L = Layout(W, H, "general", blocks_known=True)
    m = W * rng.uniform(0.07, 0.11)
    w = Writer(L, bank, rng, m)
    lh = size * rng.uniform(1.15, 1.35)
    font = bank.font(serif[0], serif[2], serif[1])
    y = m
    first_page = rng.random() < 0.45
    if rng.random() < 0.6:   # running header
        w.new_block()
        w.put(T.sentence(rng, 3, 8).rstrip("."), W / 2 if rng.random() < 0.5 else m, y, (serif[0], serif[1], int(size * 0.85)),
              "other", "c" if rng.random() < 0.5 else "l", color=ink, max_w=W - 2 * m)
        y += lh * 2
    if first_page:
        ts = int(size * rng.uniform(1.6, 2.3))
        w.new_block()
        for line in _wrap(T.sentence(rng, 5, 16).rstrip(".").split(), bank.font(hface, ts, hax), W - 2 * m)[:3]:
            w.put(line, W / 2, y + ts, (hface, hax, ts), "heading", "c", color=ink)
            y += ts * 1.25
        y += lh * 0.8
        w.new_block()
        authors = ", ".join(T.person_name(rng) for _ in range(int(rng.integers(1, 6))))
        for line in _wrap(authors.split(), font, W - 2 * m):
            w.put(line, W / 2, y, serif, "paragraph", "c", color=ink)
            y += lh
        if rng.random() < 0.8:
            w.new_block()
            aff = f"{T.pick(rng, ['Department of', 'School of', 'Institute for'])} {T.word(rng).capitalize()} {T.pick(rng, ['Science', 'Engineering', 'Studies'])}, {T.pick(rng, T.CITIES)}"
            w.put(aff, W / 2, y, (serif[0], serif[1], int(size * 0.9)), "paragraph", "c", color=ink, max_w=W - 2 * m)
            y += lh * 1.8
        w.new_block()
        w.put("Abstract", W / 2, y, (hface, hax, size), "heading", "c", color=ink)
        y += lh * 1.2
        w.new_block()
        ax0 = m + W * rng.uniform(0.04, 0.1)
        y, _ = _flow(w, " ".join(T.sentence(rng, 8, 20) for _ in range(int(rng.integers(3, 7)))).split(), serif, "paragraph",
                     ax0, W - 2 * ax0, y, H * 0.5, lh, ink)
        y += lh
    n_cols = 2 if two_col else 1
    gutter = W * rng.uniform(0.025, 0.04)
    col_w = (W - 2 * m - (n_cols - 1) * gutter) / n_cols
    foot = rng.random() < 0.3
    y_max = H - m - lh * (4 if foot else 1.5)
    sec = int(rng.integers(1, 7))
    eq = int(rng.integers(1, 20))
    refs = rng.random() < 0.2
    ref_n = int(rng.integers(1, 40))
    y_top = y
    for c in range(n_cols):
        x0 = m + c * (col_w + gutter)
        y = y_top + size
        while y < y_max - lh * 2:
            r = rng.random()
            w.new_block()
            if refs:
                entry = f"[{ref_n}] {T.person_name(rng)}, {T.person_name(rng)}. {T.sentence(rng, 5, 12)} {T.word(rng).capitalize()} {T.pick(rng, ['Conference', 'Journal', 'Symposium', 'Review'])}, {rng.integers(1990, 2026)}."
                ref_n += 1
                y, _ = _flow(w, entry.split(), (serif[0], serif[1], int(size * 0.9)), "list_item", x0, col_w, y, y_max, lh * 0.95, ink)
                y += lh * 0.15
            elif r < 0.12:
                y += lh * 0.6
                title = f"{sec} {T.sentence(rng, 1, 5).rstrip('.')}" if rng.random() < 0.8 else T.sentence(rng, 1, 4).rstrip(".")
                w.put(title, x0, y, (hface, hax, int(size * rng.uniform(1.0, 1.3))), "heading", color=ink, max_w=col_w)
                sec += 1
                y += lh * 1.4
            elif r < 0.62:
                indent = size * rng.uniform(1, 2) if rng.random() < 0.6 else 0
                words = " ".join(T.sentence(rng, 6, 20) for _ in range(int(rng.integers(2, 7)))).split()
                y, rest = _flow(w, words, serif, "paragraph", x0, col_w, y, y_max, lh, ink, indent)
                y += lh * rng.uniform(0.1, 0.6)
            elif r < 0.8:
                y += lh * 0.5
                bottom = put_formula(w, x0 + col_w / 2, y + size * 0.6, serif, col_w * 0.8, number=f"({eq})" if rng.random() < 0.8 else None,
                                     x_right=x0 + col_w, color=ink)
                eq += 1
                y = bottom + lh * 1.2
            elif r < 0.9 and y_max - y > lh * 12:
                ih = min((y_max - y) * rng.uniform(0.3, 0.6), col_w * 0.8)
                _image_placeholder(L, rng, x0 + col_w * 0.05, y, x0 + col_w * 0.95, y + ih)
                y += ih + lh
                w.new_block()
                cap = f"{T.pick(rng, ['Figure', 'Fig.'])} {rng.integers(1, 12)}: {T.sentence(rng, 5, 20)}"
                y, _ = _flow(w, cap.split(), (serif[0], serif[1], int(size * 0.9)), "paragraph", x0, col_w, y, y_max, lh * 0.95, ink)
                y += lh * 0.6
            elif y_max - y > lh * 8:
                w.put(f"Table {rng.integers(1, 9)}: {T.sentence(rng, 3, 10)}", x0, y, (serif[0], serif[1], int(size * 0.9)), "paragraph", color=ink, max_w=col_w)
                y += lh * 1.3
                w.new_block()
                ncol, nrow = int(rng.integers(2, 5)), int(rng.integers(3, 8))
                cw = col_w / ncol
                L.rules.append((x0, y - size * 1.2, x0 + col_w, y - size * 1.2, 2, ink))
                row_id = len(L.segs) * 100
                for ri in range(nrow):
                    if y > y_max:
                        break
                    for ci in range(ncol):
                        num = ci > 0 and ri > 0
                        txt = T.word(rng).capitalize() if ri == 0 else (T.number(rng) if num else T.word(rng))
                        xx = x0 + cw * (ci + 1) - size * 0.4 if num else x0 + cw * ci + size * 0.3
                        w.put(txt, xx, y, serif, "table_cell", "r" if num else "l", row=row_id + ri, color=ink, max_w=cw - size * 0.6)
                    if ri == 0:
                        L.rules.append((x0, y + size * 0.5, x0 + col_w, y + size * 0.5, 1, ink))
                    y += lh * 1.1
                L.rules.append((x0, y - lh * 0.5, x0 + col_w, y - lh * 0.5, 2, ink))
                y += lh * 0.6
            else:
                y += lh
    if foot:
        fy = H - m - lh * 2.6
        L.rules.append((m, fy - size, m + W * 0.15, fy - size, 1, ink))
        w.new_block()
        _flow(w, f"{rng.integers(1, 9)} {T.sentence(rng, 6, 18)}".split(), (serif[0], serif[1], int(size * 0.8)), "paragraph",
              m, W - 2 * m, fy, H - m, size * 0.95, ink)
    if rng.random() < 0.7:
        w.new_block()
        w.put(str(rng.integers(1, 30)), W / 2, H - m * 0.45, (serif[0], serif[1], int(size * 0.9)), "other", "c", color=ink)
    return L


BUILDERS = {"receipt": build_receipt, "invoice": build_invoice, "general": build_general,
            "newspaper": lambda rng, bank: build_columns(rng, bank, "newspaper"),
            "magazine": lambda rng, bank: build_columns(rng, bank, "magazine"),
            "academic": build_academic}
KIND_WEIGHTS = {"receipt": 0.3, "invoice": 0.2, "general": 0.2, "newspaper": 0.1, "magazine": 0.08, "academic": 0.12}


def build(rng, bank, kind=None) -> Layout:
    kind = kind or rng.choice(list(KIND_WEIGHTS), p=list(KIND_WEIGHTS.values()))
    return BUILDERS[kind](rng, bank)


# --------------------------------------------------------------------------------------------- render
def render(L: Layout, rng) -> np.ndarray:
    paper = rng.random()
    if paper < 0.6:
        bg = (255, 255, 255)
    elif paper < 0.85:
        bg = tuple(int(v) for v in rng.integers(232, 252, 3))
    else:
        base = int(rng.integers(215, 245))
        bg = (base, base - int(rng.integers(0, 12)), base - int(rng.integers(0, 25)))
    img = Image.new("RGB", (L.width, L.height), bg)
    d = ImageDraw.Draw(img)
    for x0, y0, x1, y1, col in L.fills:
        d.rectangle([x0, y0, x1, y1], fill=col)
    for x0, y0, x1, y1, lw, col in L.rules:
        d.line([x0, y0, x1, y1], fill=col, width=int(lw))
    for x0, y0, x1, y1 in L.barcodes:
        x = x0
        while x < x1:
            bw = float(rng.choice([1, 1, 2, 3])) * max(1.0, (x1 - x0) / 200)
            if rng.random() < 0.55:
                d.rectangle([x, y0, max(x, x + bw - 1), y1], fill=(10, 10, 10))
            x += bw
    for s in L.segs:
        if s.font is not None:
            d.text(s.origin, s.text, font=s.font, fill=s.color, anchor="ls")
    for origin, text, font, col in L.draws:
        d.text(origin, text, font=font, fill=col, anchor="ls")
    return np.array(img)