"""Training datasets: each mixes real pages (data/real, built by tinydoc.data.convert) with synthetic
documents, either pre-rendered (tinydoc.data.prerender) or generated on the fly. Every sample is a pure
function of (seed, index).

Data locations default to data/real, data/synth and fonts, and can be redirected with the environment
variables TINYDOC_REAL_ROOT, TINYDOC_SYNTH_ROOT and TINYDOC_FONT_ROOT (used on Kaggle).
"""

import json
import os
from functools import lru_cache
from pathlib import Path

import cv2
import numpy as np
import torch
from PIL import Image, ImageDraw

from ..charset import encode, normalize
from ..features import MAX_SEGMENTS, segment_features
from ..geometry import expand_quad, quad_bbox, quad_size, rect_quad, transform_quads
from ..models.detector import BLOCK_SHRINK, KERNEL_END_TRIM, KERNEL_HEIGHT_RATIO
from ..models.orientation import PATCH
from ..models.recognizer import LINE_HEIGHT, WIDTH_STRIDE, X_STRETCH
from ..schema import DOC_ID, LABEL_ID, REL_TYPES
from ..synth import text as T
from ..synth.augment import geometric, photometric
from ..synth.documents import KIND_WEIGHTS, build, render
from ..synth.fonts import FontBank, coverage

MAX_LINE_WIDTH = 4096       # inference cap
MAX_TRAIN_WIDTH = 1536      # training cap: wider samples are resampled (keeps attention cost bounded)
MATH_LINE_P = 0.15          # share of synthetic recognizer lines that carry inline maths
WIDTH_BUCKET = 64           # pad batch widths to multiples of this so compiled graphs get reused
SEGMENT_BUCKET = 32


def real_root():
    return os.environ.get("TINYDOC_REAL_ROOT", "data/real")


def synth_root():
    return os.environ.get("TINYDOC_SYNTH_ROOT", "data/synth")


def font_root():
    return os.environ.get("TINYDOC_FONT_ROOT", "fonts")


@lru_cache(maxsize=4)
def font_bank(root, split):
    return FontBank(root, "val" if split != "train" else "train")


@lru_cache(maxsize=32)
def load_records(path):
    p = Path(path)
    if not p.exists():
        return []
    return [json.loads(l) for l in open(p)]


def pick_weighted(rng, weights: dict):
    names = [k for k, v in weights.items() if v > 0]
    p = np.array([weights[k] for k in names], np.float64)
    return names[rng.choice(len(names), p=p / p.sum())]


# ---------------------------------------------------------------------------------------- page sources
# geometric mode probabilities (none, rotate, photo) and photometric strength per source
AUG = {
    "synth": ((0.35, 0.40, 0.25), 1.0),
    "doclaynet": ((0.60, 0.30, 0.10), 0.5),
    "wildreceipt": ((0.70, 0.30, 0.0), 0.5),
}
DEFAULT_AUG = ((0.6, 0.3, 0.1), 0.5)


class PagePool:
    """Samples un-augmented pages (image, quads, ignore flags, ignore regions) plus how to augment them.

    "synth" uses pre-rendered pages from TINYDOC_SYNTH_ROOT when they exist for the split, and renders
    on the fly otherwise.
    """

    def __init__(self, split, use_key, weights):
        self.split = split
        self.real = {}
        for name in weights:
            base = Path(synth_root() if name == "synth" else real_root()) / ("" if name == "synth" else name) / split
            recs = [r for r in load_records(str(base / "pages.jsonl")) if r["use"][use_key]]
            if recs:
                self.real[name] = (base, recs)
        self.weights = {k: v for k, v in weights.items() if k == "synth" or k in self.real}
        self._cache = {}

    def pick(self, rng):
        """Choose a page and its augmentation. Returns a dict; the image is loaded separately (load)."""
        src = pick_weighted(rng, self.weights)
        geo_p, strength = AUG.get(src, DEFAULT_AUG)
        geo = str(rng.choice(["none", "rotate", "photo"], p=geo_p))
        if src == "synth" and src not in self.real:
            kind = str(rng.choice(list(KIND_WEIGHTS), p=list(KIND_WEIGHTS.values())))
            L = build(rng, font_bank(font_root(), self.split), kind)
            img = np.ascontiguousarray(render(L, rng)[:, :, ::-1])  # BGR like cv2.imread
            segs = L.segments()
            quads = np.array([s["quad"] for s in segs], np.float32).reshape(-1, 4, 2)
            info = self._summarise(quads, np.array([detector_ignores(s) for s in segs], bool), [], L.width, L.height,
                                   block_rects(segs) if L.blocks_known else None)
            return {"src": src, "img": img, "info": info, "geo": geo, "strength": strength,
                    "thermal": kind == "receipt" and rng.random() < 0.4}
        base, recs = self.real[src]
        i = int(rng.integers(len(recs)))
        rec = recs[i]
        return {"src": src, "path": str(base / rec["image"]), "info": self.info(src, i), "geo": geo, "strength": strength,
                "thermal": rec.get("kind") == "receipt" and rng.random() < 0.4}

    @staticmethod
    def _summarise(quads, ignore, regions, W, H, blocks=None):
        if len(quads):
            v = quads[:, 3] - quads[:, 0]
            heights = np.linalg.norm(v, axis=1)
            ok = heights[~ignore] if (~ignore).any() else heights
            med = float(np.median(ok))
            bboxes = np.concatenate([quads.min(1), quads.max(1)], 1)
        else:
            med, bboxes = 20.0, np.zeros((0, 4), np.float32)
        # blocks: (rects (B, 4) x0 y0 x1 y1, line heights (B,)) when the page has block labels, else None
        return {"quads": quads, "ignore": ignore, "regions": regions, "med_h": med, "bboxes": bboxes, "W": W, "H": H,
                "blocks": blocks}

    def info(self, src, i):
        """Per-page box arrays and statistics, cached per worker (pages are revisited many times)."""
        key = (src, i)
        if key not in self._cache:
            rec = self.real[src][1][i]
            quads = np.array([s["quad"] for s in rec["segments"]], np.float32).reshape(-1, 4, 2)
            ignore = np.array([detector_ignores(s) for s in rec["segments"]], bool)
            regions = [np.asarray(q, np.float32) for q in rec["ignore_regions"]]
            blocks = block_rects(rec["segments"]) if rec["use"].get("blocks") else None
            self._cache[key] = self._summarise(quads, ignore, regions, rec["width"], rec["height"], blocks)
        return self._cache[key]

    @staticmethod
    def load(pick, reduce=1):
        """BGR page image, decoded at 1/reduce size (JPEG DCT scaling is much cheaper than resizing)."""
        if "img" in pick:
            return pick["img"]
        flag = {1: cv2.IMREAD_COLOR, 2: cv2.IMREAD_REDUCED_COLOR_2, 4: cv2.IMREAD_REDUCED_COLOR_4}[reduce]
        return cv2.imread(pick["path"], flag)

    def sample(self, rng):
        """Full page as RGB (orientation data and evaluation)."""
        p = self.pick(rng)
        info = p["info"]
        img = cv2.cvtColor(self.load(p), cv2.COLOR_BGR2RGB)
        return (img, list(info["quads"]), list(info["ignore"]), list(info["regions"]), p["src"], p["geo"],
                p["strength"], p["thermal"], info["blocks"])


def block_rects(segs):
    """Layout blocks of a page from its segments' block ids: (rects (B, 4), median line height per block)."""
    groups = {}
    for s in segs:
        if s.get("block") is not None and not s.get("ignore"):
            groups.setdefault(s["block"], []).append(np.asarray(s["quad"], np.float32))
    if not groups:
        return np.zeros((0, 4), np.float32), np.zeros(0, np.float32)
    rects, hs = [], []
    for qs in groups.values():
        q = np.stack(qs)
        rects.append([q[:, :, 0].min(), q[:, :, 1].min(), q[:, :, 0].max(), q[:, :, 1].max()])
        hs.append(float(np.median(q[:, :, 1].max(1) - q[:, :, 1].min(1))))
    return np.array(rects, np.float32), np.array(hs, np.float32)


def block_targets(h, w, block_quads, probe_quads, known, regions):
    """Block-kernel target and mask. probe_quads: one line-height square per block, transformed with the page,
    giving the local line height for the shrink."""
    kernel = np.zeros((h, w), np.uint8)
    mask = np.full((h, w), 1 if known else 0, np.uint8)
    if known:
        for bq, pq in zip(block_quads, probe_quads):
            lh = quad_size(pq)[1]
            d = BLOCK_SHRINK * lh
            bw, bh = quad_size(bq)
            if bw > 2 * d + 1 and bh > 2 * d + 1:
                cv2.fillPoly(kernel, [np.round(expand_quad(bq, -d, -d, -d, -d)).astype(np.int32)], 1)
        for r in regions:
            cv2.fillPoly(mask, [np.round(r).astype(np.int32)], 0)
    return kernel, mask


def detector_ignores(seg):
    """Segments the detector neither learns nor is scored on: unreadable text and display maths."""
    return bool(seg["ignore"]) or seg.get("label") == "formula"


def page_region(img, quads, regions, cx, cy, scale, side, fill):
    """Cut a side x side window centred on (cx, cy) of the page scaled by `scale`; move boxes with it.

    Cropping before resizing/warping keeps per-sample cost independent of page size.
    """
    H, W = img.shape[:2]
    src_side = side / scale
    x0, y0 = cx - src_side / 2, cy - src_side / 2
    sx0, sy0 = int(max(0, np.floor(x0))), int(max(0, np.floor(y0)))
    sx1, sy1 = int(min(W, np.ceil(x0 + src_side))), int(min(H, np.ceil(y0 + src_side)))
    out = np.empty((side, side, 3), np.uint8)
    out[:] = fill
    if sx1 <= sx0 or sy1 <= sy0:
        return out, [], []
    patch = img[sy0:sy1, sx0:sx1]
    ow, oh = max(1, round((sx1 - sx0) * scale)), max(1, round((sy1 - sy0) * scale))
    patch = cv2.resize(patch, (ow, oh), interpolation=cv2.INTER_AREA if scale < 1 else cv2.INTER_LINEAR)
    ox, oy = int(round((sx0 - x0) * scale)), int(round((sy0 - y0) * scale))
    px0, py0 = max(0, ox), max(0, oy)
    px1, py1 = min(side, ox + ow), min(side, oy + oh)
    out[py0:py1, px0:px1] = patch[py0 - oy:py1 - oy, px0 - ox:px1 - ox]
    sx, sy = ow / (sx1 - sx0), oh / (sy1 - sy0)
    m = np.array([[sx, 0, ox - sx0 * sx], [0, sy, oy - sy0 * sy], [0, 0, 1]])
    return out, transform_quads(quads, m), transform_quads(regions, m)


# -------------------------------------------------------------------------------------------- detector
def kernel_quad(q):
    w, h = quad_size(q)
    dy = h * (1 - KERNEL_HEIGHT_RATIO) / 2
    dx = min(KERNEL_END_TRIM * h, w * 0.25)
    return expand_quad(q, -dx, -dy, -dx, -dy)


def detector_targets(h, w, quads, ignore, regions, min_h=5):
    kernel = np.zeros((h, w), np.uint8)
    mask = np.ones((h, w), np.uint8)
    for q, ign in zip(quads, ignore):
        _, qh = quad_size(q)
        if ign or qh < min_h:
            cv2.fillPoly(mask, [np.round(q).astype(np.int32)], 0)
            continue
        cv2.fillPoly(kernel, [np.round(kernel_quad(q)).astype(np.int32)], 1)
    for r in regions:
        cv2.fillPoly(mask, [np.round(r).astype(np.int32)], 0)
    return kernel, mask


class DetectorDataset(torch.utils.data.Dataset):
    def __init__(self, split="train", length=100_000, crop=640, seed=0, weights=None):
        weights = weights or {"synth": 0.45, "doclaynet": 0.3, "wildreceipt": 0.25}
        self.pool = PagePool(split, "detector", weights)
        self.length, self.crop, self.seed = length, crop, seed

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        rng = np.random.default_rng([self.seed, idx])
        p = self.pool.pick(rng)
        info = p["info"]
        W, H = info["W"], info["H"]
        scale = np.exp(rng.uniform(np.log(7), np.log(56))) / info["med_h"]   # dense newspapers to large titles
        scale = float(np.clip(scale, 320 / max(H, W), 3000 / max(H, W)))
        reduce = 4 if scale <= 0.22 else 2 if scale <= 0.45 else 1
        img = self.pool.load(p, reduce)
        rs = img.shape[1] / W                        # actual decode scale
        c = self.crop
        quads, bboxes = info["quads"], info["bboxes"]
        if len(quads) and rng.random() < 0.75:
            cx, cy = quads[rng.integers(len(quads))].mean(0)
        else:
            cx, cy = rng.uniform(0, W), rng.uniform(0, H)
        cx += rng.uniform(-c / 4, c / 4) / scale
        cy += rng.uniform(-c / 4, c / 4) / scale
        side = int(c * 1.25)
        half = side / scale  # generous: rotation/perspective can pull in boxes from just outside the window
        near = (bboxes[:, 2] > cx - half) & (bboxes[:, 0] < cx + half) & (bboxes[:, 3] > cy - half) & (bboxes[:, 1] < cy + half)
        quads_n, ignore = quads[near] * rs, info["ignore"][near]
        regions = [r * rs for r in info["regions"]]
        known = info["blocks"] is not None
        bq = pq = np.zeros((0, 4, 2), np.float32)
        if known:
            rects, lhs = info["blocks"]
            sel = (rects[:, 2] > cx - half) & (rects[:, 0] < cx + half) & (rects[:, 3] > cy - half) & (rects[:, 1] < cy + half)
            bq = np.array([rect_quad(*r) for r in rects[sel]], np.float32).reshape(-1, 4, 2) * rs
            pq = np.array([rect_quad(r[0], r[1], r[0] + lh, r[1] + lh) for r, lh in zip(rects[sel], lhs[sel])],
                          np.float32).reshape(-1, 4, 2) * rs
        nq, nb = len(quads_n), len(bq)
        allq = np.concatenate([np.asarray(quads_n, np.float32).reshape(-1, 4, 2), bq, pq])
        fill = np.median(img[::16, ::16].reshape(-1, 3), 0).astype(np.uint8)
        region, allq, regions = page_region(img, allq, regions, cx * rs, cy * rs, scale / rs, side, fill)
        region = np.ascontiguousarray(region[:, :, ::-1])  # BGR -> RGB on the small region only
        region, m = geometric(region, rng, p["geo"])
        rh, rw = region.shape[:2]
        x0, y0 = (rw - c) // 2, (rh - c) // 2
        m = np.array([[1, 0, -x0], [0, 1, -y0], [0, 0, 1]], np.float64) @ m
        allq, regions = transform_quads(allq, m), transform_quads(regions, m)
        quads_n, bq, pq = allq[:nq], allq[nq:nq + nb], allq[nq + nb:]
        out = region[y0:y0 + c, x0:x0 + c]
        kernel, mask = detector_targets(c, c, quads_n, ignore, regions)
        bkernel, bmask = block_targets(c, c, bq, pq, known, regions)
        out = photometric(np.ascontiguousarray(out), rng, p["strength"], p["thermal"])
        return (torch.from_numpy(out).permute(2, 0, 1).float() / 255,
                torch.from_numpy(np.stack([kernel, bkernel])).float(), torch.from_numpy(np.stack([mask, bmask])).float())


# ------------------------------------------------------------------------------------------ recognizer
def synth_line(rng, bank):
    """Render one text line with margins. Returns (gray crop, tight box, text)."""
    for _ in range(10):
        face, bax = bank.pick(rng, bold=rng.random() < 0.2, italic_ok=True)
        text = FontBank.supported(face, normalize(T.line_text(rng))).strip()
        if text:
            break
    else:
        text, (face, bax) = "0", bank.pick(rng)
    size = int(rng.integers(18, 56))
    font = bank.font(face, size, bax)
    # Justified newspaper/magazine columns set words much tighter (or looser) than the font's own space.
    words = text.split(" ")
    space = font.getlength(" ") * (rng.uniform(0.45, 1.8) if len(words) > 1 and rng.random() < 0.35 else 1.0)
    placed, x = [], 0.0
    for wd in words:
        if wd:
            placed.append((x, wd, font, 0.0))
        x += font.getlength(wd) + space
    arr, box = _line_canvas(rng, placed, font)
    return arr, box, text


def synth_math_line(rng, bank):
    """A line with inline maths: sub/superscripts typeset small and raised or lowered (labelled x², H₂O), Greek and
    maths symbols drawn from a maths font when the text font lacks them. Returns (gray crop, tight box, text)."""
    for _ in range(10):
        runs = T.math_runs(T.math_line(rng))
        face, bax = bank.pick(rng, bold=rng.random() < 0.1, italic_ok=True)
        cov = coverage(face.path)
        if all(c == " " or c in cov or bank.fallback(c) for t, _ in runs for c in t):
            break
    else:
        return synth_line(rng, bank)
    size = int(rng.integers(20, 56))
    font = bank.font(face, size, bax)
    small_size = max(6, round(size * rng.uniform(0.58, 0.75)))
    rise, drop = size * rng.uniform(0.28, 0.45), size * rng.uniform(0.10, 0.22)
    placed, x = [], 0.0
    for text, kind in runs:
        spelled = T.scripted(text, kind) if kind else text
        if kind and spelled != text and rng.random() < 0.3 and all(c in cov for c in spelled):
            text, kind = spelled, ""     # typed with the font's own Unicode sub/superscript glyphs
        pt, dy = (size, 0.0) if not kind else (small_size, -rise if kind == "^" else drop)
        for ch in text:                  # one glyph at a time so each can fall back to the maths font
            if ch == " ":
                x += font.getlength(" ")
                continue
            f = bank.font(face, pt, bax) if ch in cov else bank.font(bank.fallback(ch), pt)
            placed.append((x, ch, f, dy))
            x += f.getlength(ch)
    if not placed:
        return synth_line(rng, bank)
    arr, box = _line_canvas(rng, placed, font)
    return arr, box, normalize(T.runs_label(runs))


def _line_canvas(rng, placed, font):
    """placed: [(x, text, font, dy)] along one baseline. Draws them on a page-coloured canvas with margins,
    sometimes with neighbouring lines, rules and a slight tilt. Returns (gray crop, tight box)."""
    boxes = [(x + bb[0], bb[1] + dy, x + bb[2], bb[3] + dy)
             for x, s, f, dy in placed for bb in [f.getbbox(s, anchor="ls")]]
    l, r = min(b[0] for b in boxes), max(b[2] for b in boxes)
    _, rt, _, rb = font.getbbox("Hg", anchor="ls")
    top, bot = min(rt, min(b[1] for b in boxes)), max(rb, max(b[3] for b in boxes))
    h, w = bot - top, max(r - l, 1)
    mx, my = 0.6 * h, 0.35 * h
    W, H = int(w + 2 * mx), int(h + 2 * my)
    bg = int(rng.integers(170, 256))
    ink = int(rng.integers(0, max(bg - 90, 1)))
    im = Image.new("L", (W, H), bg)
    d = ImageDraw.Draw(im)
    for x, s, f, dy in placed:
        d.text((mx - l + x, my - top + dy), s, font=f, fill=ink, anchor="ls")
    if rng.random() < 0.3:  # fragments of neighbouring lines peeking into the margins
        for sign in (-1, 1):
            if rng.random() < 0.6:
                d.text((mx - l + rng.uniform(-h, h), my - top + sign * h * rng.uniform(1.05, 1.5)),
                       T.line_text(rng), font=font, fill=ink, anchor="ls")
    if rng.random() < 0.06:
        yy = my + h * rng.uniform(0.85, 1.25)
        d.line([0, yy, W, yy], fill=ink, width=max(1, int(h / 20)))
    if rng.random() < 0.06:
        for xx in (mx * rng.uniform(0, 0.7), W - mx * rng.uniform(0, 0.7)):
            d.line([xx, 0, xx, H], fill=ink, width=max(1, int(h / 20)))
    arr = np.array(im)
    if rng.random() < 0.3:
        m = cv2.getRotationMatrix2D((W / 2, H / 2), rng.normal(0, 0.8), 1.0)
        m[0, 1] += rng.uniform(-0.15, 0.15)
        m[0, 2] -= m[0, 1] * H / 2
        arr = cv2.warpAffine(arr, m, (W, H), borderValue=bg)
    return arr, (mx, my, mx + w, my + h)


def jitter_crop(gray, box, rng):
    x0, y0, x1, y1 = box
    h = y1 - y0
    H, W = gray.shape[:2]
    cx0 = int(np.clip(x0 - rng.uniform(-0.05, 0.5) * h, 0, W - 1))
    cx1 = int(np.clip(x1 + rng.uniform(-0.05, 0.5) * h, cx0 + 1, W))
    cy0 = int(np.clip(y0 - rng.uniform(-0.08, 0.3) * h, 0, H - 1))
    cy1 = int(np.clip(y1 + rng.uniform(-0.08, 0.3) * h, cy0 + 1, H))
    return gray[cy0:cy1, cx0:cx1]


def line_to_uint8(gray) -> np.ndarray:
    """Grayscale crop -> uint8 (32, W), W a multiple of the CNN stride."""
    h, w = gray.shape[:2]
    nw = int(round(w * LINE_HEIGHT / max(h, 1) * X_STRETCH))
    nw = int(np.clip(nw, WIDTH_STRIDE * 2, MAX_LINE_WIDTH))
    nw = (nw + WIDTH_STRIDE - 1) // WIDTH_STRIDE * WIDTH_STRIDE
    return cv2.resize(gray, (nw, LINE_HEIGHT), interpolation=cv2.INTER_AREA if h > LINE_HEIGHT else cv2.INTER_LINEAR)


def line_to_input(gray) -> np.ndarray:
    """Grayscale crop -> float32 (32, W) in [0,1]."""
    return line_to_uint8(gray).astype(np.float32) / 255.0


class LineDataset(torch.utils.data.Dataset):
    def __init__(self, split="train", length=1_000_000, seed=0, weights=None):
        weights = weights or {"synth": 0.4, "cord-v2": 0.25, "doclaynet": 0.35}
        self.split, self.length, self.seed = split, length, seed
        self.real = {}
        for name in weights:
            if name != "synth":
                recs = load_records(str(Path(real_root()) / name / split / "lines.jsonl"))
                if recs:
                    self.real[name] = (Path(real_root()) / name / split, recs)
        self.weights = {k: v for k, v in weights.items() if k == "synth" or k in self.real}

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        rng = np.random.default_rng([self.seed, idx])
        for _ in range(20):
            out = self._sample(rng)
            if out[0].shape[-1] <= MAX_TRAIN_WIDTH:
                return out
        return out

    def _sample(self, rng):
        src = pick_weighted(rng, self.weights)
        tag = src
        if src == "synth":
            bank = font_bank(font_root(), self.split)
            maths = rng.random() < MATH_LINE_P
            gray, box, text = (synth_math_line if maths else synth_line)(rng, bank)
            strength, tag = 1.0, "synth-math" if maths else "synth"
        else:
            base, recs = self.real[src]
            r = recs[rng.integers(len(recs))]
            gray = cv2.imread(str(base / r["path"]), cv2.IMREAD_GRAYSCALE)
            box, text, strength = r["box"], r["text"], 0.5
        crop = jitter_crop(gray, box, rng)
        if rng.random() < 0.2:
            # small print: the text was only 7-13 px tall in the page image (dense newspapers, footnotes)
            th = max(box[3] - box[1], 1)
            f = min(1.0, rng.uniform(7, 13) / th)
            if f < 0.95:
                h0, w0 = crop.shape[:2]
                small = cv2.resize(crop, (max(1, round(w0 * f)), max(1, round(h0 * f))), interpolation=cv2.INTER_AREA)
                crop = cv2.resize(small, (w0, h0), interpolation=cv2.INTER_LINEAR if rng.random() < 0.5 else cv2.INTER_CUBIC)
        crop = photometric(crop, rng, strength, thermal=src == "synth" and rng.random() < 0.1)
        if rng.random() < 0.08:
            crop = 255 - crop
        x = line_to_uint8(crop)  # uint8 keeps worker->trainer transfer 4x smaller; scaled in collate
        return torch.from_numpy(x)[None], torch.tensor(encode(text), dtype=torch.long), tag   # tag: metrics source


def collate_lines(batch):
    imgs, targets, srcs = zip(*batch)
    wmax = max(i.shape[-1] for i in imgs)
    wmax = (wmax + WIDTH_BUCKET - 1) // WIDTH_BUCKET * WIDTH_BUCKET
    B = len(imgs)
    x = torch.ones(B, 1, LINE_HEIGHT, wmax)
    steps = torch.tensor([i.shape[-1] // WIDTH_STRIDE for i in imgs])
    for k, im in enumerate(imgs):
        x[k, :, :, : im.shape[-1]] = im.float() / 255.0
    step_mask = torch.arange(wmax // WIDTH_STRIDE)[None, :] >= steps[:, None]
    return x, step_mask, steps, torch.cat(targets), torch.tensor([len(t) for t in targets]), list(srcs), targets


# ----------------------------------------------------------------------------------------- orientation
def text_patches(gray, n, rng=None):
    """Pick up to n PATCH-sized windows dense in ink-like edges. Returns list of (x, y)."""
    H, W = gray.shape
    if H < PATCH or W < PATCH:
        return [(0, 0)]
    ink = cv2.adaptiveThreshold(gray, 1, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 25, 15).astype(np.float32)
    s = 8
    small = cv2.resize(ink, (W // s, H // s), interpolation=cv2.INTER_AREA)
    k = PATCH // s
    dens = cv2.boxFilter(small, -1, (k, k), normalize=True, borderType=cv2.BORDER_CONSTANT)
    dens[: k // 2, :] = dens[-(k // 2):, :] = 0
    dens[:, : k // 2] = dens[:, -(k // 2):] = 0
    # Text has moderate ink density; solid dark regions (photo backgrounds) are excluded.
    score = dens * (dens < 0.35)
    out = []
    for _ in range(n):
        if rng is not None:
            flat = score.ravel() ** 2
            if flat.sum() <= 0:
                break
            i = rng.choice(len(flat), p=flat / flat.sum())
        else:
            i = int(score.argmax())
            if score.ravel()[i] <= 0.01:
                break
        cy, cx = divmod(i, score.shape[1])
        out.append((int(np.clip(cx * s - PATCH // 2, 0, W - PATCH)), int(np.clip(cy * s - PATCH // 2, 0, H - PATCH))))
        score[max(0, cy - k // 2): cy + k // 2, max(0, cx - k // 2): cx + k // 2] = 0
    return out or [((W - PATCH) // 2, (H - PATCH) // 2)]


ORIENT_LONG_SIDE = 1280


def orientation_page(img_rgb, long_side=ORIENT_LONG_SIDE):
    gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY) if img_rgb.ndim == 3 else img_rgb
    s = long_side / max(gray.shape)
    return cv2.resize(gray, (max(1, round(gray.shape[1] * s)), max(1, round(gray.shape[0] * s))), interpolation=cv2.INTER_AREA)


class OrientationDataset(torch.utils.data.Dataset):
    """Each item is `per_page` patches from one page, each independently rotated by k*90 degrees."""

    def __init__(self, split="train", length=200_000, seed=0, weights=None, per_page=8):
        weights = weights or {"synth": 0.5, "doclaynet": 0.25, "wildreceipt": 0.25}
        self.pool = PagePool(split, "detector", weights)
        self.length, self.seed, self.per_page = length, seed, per_page

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        rng = np.random.default_rng([self.seed, idx])
        img, quads, ignore, _, _, _, strength, thermal, _ = self.pool.sample(rng)
        long_side = int(rng.integers(1000, 1600))
        gray = orientation_page(img, long_side)
        scale = long_side / max(img.shape[:2])
        centers = [q.mean(0) * scale for q, ig in zip(quads, ignore) if not ig]
        pos = text_patches(gray, self.per_page, rng)
        big = int(PATCH * 1.3)
        padded = cv2.copyMakeBorder(gray, big, big, big, big, cv2.BORDER_REPLICATE)
        patches, labels = [], []
        for i in range(self.per_page):
            # mostly centred on known text (the heuristic picker can land on wood grain or photos)
            if centers and rng.random() < 0.75:
                cx, cy = centers[rng.integers(len(centers))] + rng.uniform(-PATCH / 3, PATCH / 3, 2)
            else:
                x, y = pos[i % len(pos)]
                cx, cy = x + PATCH / 2, y + PATCH / 2
            cx, cy = int(np.clip(cx, 0, gray.shape[1])) + big, int(np.clip(cy, 0, gray.shape[0])) + big
            win = padded[cy - big // 2: cy - big // 2 + big, cx - big // 2: cx - big // 2 + big]
            if rng.random() < 0.6:  # residual skew
                m = cv2.getRotationMatrix2D((big / 2, big / 2), rng.normal(0, 3), 1.0)
                win = cv2.warpAffine(win, m, (big, big), borderMode=cv2.BORDER_REPLICATE)
            o = (big - PATCH) // 2
            patch = win[o:o + PATCH, o:o + PATCH]
            k = int(rng.integers(4))
            patch = np.ascontiguousarray(np.rot90(patch, -k))  # rotated k*90 degrees clockwise
            patches.append(photometric(patch, rng, strength, thermal))
            labels.append(k)
        x = torch.from_numpy(np.stack(patches).astype(np.float32) / 255)[:, None]
        return x, torch.tensor(labels)


# ------------------------------------------------------------------------------------------- extractor
CONFUSIONS = {"0": "O", "O": "0", "o": "0", "1": "l", "l": "1", "I": "1", "5": "S", "S": "5", "8": "B", "B": "8",
              ",": ".", ".": ",", "2": "Z", "6": "G", "g": "9", "9": "g", "e": "c", "c": "e", "u": "v", ":": ";"}


def ocr_noise(text, rng, rate):
    out = []
    for ch in text:
        r = rng.random()
        if r < rate * 0.6:
            out.append(CONFUSIONS.get(ch, ch if rng.random() < 0.5 else T.pick(rng, "abcdefghijklmnopqrstuvwxyz0123456789")))
        elif r < rate * 0.8:
            continue
        elif r < rate:
            out.extend([ch, T.pick(rng, ".,'-il")])
        else:
            out.append(ch)
    return "".join(out) or text


EXTRACTOR_KINDS = {"receipt": 0.3, "invoice": 0.35, "general": 0.1, "newspaper": 0.07, "magazine": 0.06, "academic": 0.12}


class LayoutDataset(torch.utils.data.Dataset):
    def __init__(self, split="train", length=200_000, seed=0, weights=None):
        weights = weights or {"synth": 0.45, "cord-v2": 0.17, "wildreceipt": 0.13, "doclaynet": 0.25}
        self.split, self.length, self.seed = split, length, seed
        self.real = {}
        for name in weights:
            if name != "synth":
                recs = [r for r in load_records(str(Path(real_root()) / name / split / "pages.jsonl")) if r["use"]["labels"]]
                if recs:
                    self.real[name] = recs
        self.prerendered = load_records(str(Path(synth_root()) / split / "pages.jsonl"))
        self.weights = {k: v for k, v in weights.items() if k == "synth" or k in self.real}

    def __len__(self):
        return self.length

    def sample_doc(self, rng):
        src = pick_weighted(rng, self.weights)
        # Layouts without rendering are cheap, so half the synthetic documents are generated here with a mix that
        # favours invoices (v2's field model lost invoice accuracy when invoices were 20% of the mix).
        if src == "synth" and (not self.prerendered or rng.random() < 0.5):
            kind = str(rng.choice(list(EXTRACTOR_KINDS), p=list(EXTRACTOR_KINDS.values())))
            L = build(rng, font_bank(font_root(), self.split), kind)
            segs = [{"bbox": list(s.box), "text": s.text, "label": s.label, "row": s.row, "block": s.block} for s in L.segs]
            return segs, L.kv, L.doc_type, L.width, L.height, (True, True, L.blocks_known), f"synth-{L.doc_type}"
        recs = self.prerendered if src == "synth" else self.real[src]
        rec = recs[rng.integers(len(recs))]
        segs = [{"bbox": quad_bbox(s["quad"]), "text": s["text"], "label": s["label"], "row": s["row"], "block": s.get("block")}
                for s in rec["segments"] if not s["ignore"]]
        keep = [i for i, s in enumerate(rec["segments"]) if not s["ignore"]]
        remap = {old: new for new, old in enumerate(keep)}
        kv = [(remap[k], remap[v]) for k, v in rec["kv_pairs"] if k in remap and v in remap]
        use = rec["use"]
        label = f"synth-{rec['doc_type']}" if src == "synth" else src   # per-kind scores in evaluation
        return segs, kv, rec["doc_type"], rec["width"], rec["height"], (use["relations"], use["relations"], use.get("blocks", False)), label

    def __getitem__(self, idx):
        rng = np.random.default_rng([self.seed, idx])
        segs, kv, doc_type, W, H, rel_known, src = self.sample_doc(rng)
        if self.split == "train":
            rate = rng.uniform(0, 0.08)
            drop_spaces = rng.random() < 0.08
            for s in segs:
                if rng.random() < 0.5:
                    s["text"] = ocr_noise(s["text"], rng, rate)
                if drop_spaces:
                    s["text"] = s["text"].replace(" ", "")
                h = s["bbox"][3] - s["bbox"][1]
                s["bbox"] = [v + rng.normal(0, 0.06 * h) for v in s["bbox"]]
                s["conf"] = float(np.clip(1 - rate * rng.uniform(2, 8), 0.3, 1))
            if rng.random() < 0.3:   # page embedded in a larger photo frame
                f = rng.uniform(1.0, 1.8)
                ox, oy = rng.uniform(0, W * (f - 1)), rng.uniform(0, H * (f - 1))
                for s in segs:
                    s["bbox"] = [s["bbox"][0] + ox, s["bbox"][1] + oy, s["bbox"][2] + ox, s["bbox"][3] + oy]
                W, H = W * f, H * f
        order = sorted(range(len(segs)), key=lambda i: (round(segs[i]["bbox"][1] / 8), segs[i]["bbox"][0]))
        if len(order) > MAX_SEGMENTS:
            start = int(rng.integers(0, len(order) - MAX_SEGMENTS + 1))
            order = order[start:start + MAX_SEGMENTS]
        pos = {old: new for new, old in enumerate(order)}
        segs = [segs[i] for i in order]
        n = len(segs)
        f = segment_features(segs, W, H)
        labels = np.array([LABEL_ID[s["label"]] if s["label"] is not None else -100 for s in segs], np.int64)
        rel = np.zeros((len(REL_TYPES), n, n), np.float32)
        rows = np.array([s["row"] if s["row"] is not None else -1 for s in segs])
        same = (rows[:, None] == rows[None, :]) & (rows[:, None] >= 0)
        np.fill_diagonal(same, False)
        rel[0] = same
        for k, v in kv:
            if k in pos and v in pos:
                rel[1, pos[k], pos[v]] = 1
        blocks = np.array([s["block"] if s.get("block") is not None else -1 for s in segs])
        same = (blocks[:, None] == blocks[None, :]) & (blocks[:, None] >= 0)
        np.fill_diagonal(same, False)
        rel[2] = same
        rel_mask = np.array(rel_known, np.float32)
        return {**{k: torch.from_numpy(v) for k, v in f.items()}, "labels": torch.from_numpy(labels),
                "rel": torch.from_numpy(rel), "rel_known": torch.from_numpy(rel_mask),
                "doc": DOC_ID[doc_type], "src": src}


def collate_layouts(batch):
    B = len(batch)
    n = max(max(b["feats"].shape[0] for b in batch), 1)
    n = (n + SEGMENT_BUCKET - 1) // SEGMENT_BUCKET * SEGMENT_BUCKET
    L = batch[0]["char_ids"].shape[1]
    F_ = batch[0]["feats"].shape[1]
    R = batch[0]["rel"].shape[0]
    out = {"char_ids": torch.zeros(B, n, L, dtype=torch.long), "feats": torch.zeros(B, n, F_),
           "rel_boxes": torch.zeros(B, n, 4), "valid": torch.zeros(B, n), "labels": torch.full((B, n), -100, dtype=torch.long),
           "rel": torch.zeros(B, R, n, n), "rel_known": torch.stack([b["rel_known"] for b in batch]),
           "doc": torch.tensor([b["doc"] for b in batch]), "src": [b["src"] for b in batch]}
    for i, b in enumerate(batch):
        k = b["feats"].shape[0]
        out["char_ids"][i, :k] = b["char_ids"]
        out["feats"][i, :k] = b["feats"]
        out["rel_boxes"][i, :k] = b["rel_boxes"]
        out["valid"][i, :k] = 1
        out["labels"][i, :k] = b["labels"]
        out["rel"][i, :, :k, :k] = b["rel"]
    return out