"""Pre-render synthetic pages to disk (clean renders; augmentation still happens at training time).

Rendering pages is the most CPU-expensive part of the detector/orientation data pipeline. Doing it once
lets training run on machines with few CPU cores (e.g. Kaggle's 4-core GPU sessions).

    python -m tinydoc.data.prerender --pages 10000 --out data/synth
"""

import argparse
import json
from multiprocessing import Pool
from pathlib import Path

import cv2
import numpy as np

from ..synth.documents import KIND_WEIGHTS, build, render
from ..synth.fonts import FontBank

_bank = None


def _init(font_root):
    global _bank
    _bank = FontBank(font_root, "train")


def _one(args):
    i, seed, out = args
    rng = np.random.default_rng([seed, i])
    kind = str(rng.choice(list(KIND_WEIGHTS), p=list(KIND_WEIGHTS.values())))
    L = build(rng, _bank, kind)
    img = render(L, rng)
    name = f"images/{i:06d}.jpg"
    cv2.imwrite(str(Path(out) / name), img[:, :, ::-1], [cv2.IMWRITE_JPEG_QUALITY, 90])
    segs = [{**s, "quad": np.round(s["quad"], 1).tolist()} for s in L.segments()]
    return {"image": name, "width": L.width, "height": L.height, "source": "synth", "kind": kind,
            "doc_type": L.doc_type, "segments": segs, "kv_pairs": [list(p) for p in L.kv], "ignore_regions": [],
            "use": {"detector": True, "recognizer": False, "labels": True, "relations": True, "blocks": L.blocks_known},
            "fields": L.fields}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--pages", type=int, default=10_000)
    ap.add_argument("--out", default="data/synth")
    ap.add_argument("--fonts", default="fonts")
    ap.add_argument("--seed", type=int, default=2024)
    ap.add_argument("--workers", type=int, default=8)
    args = ap.parse_args()
    out = Path(args.out) / "train"
    (out / "images").mkdir(parents=True, exist_ok=True)
    with Pool(args.workers, initializer=_init, initargs=(args.fonts,)) as pool, open(out / "pages.jsonl", "w") as fh:
        for k, rec in enumerate(pool.imap(_one, [(i, args.seed, str(out)) for i in range(args.pages)], chunksize=16)):
            fh.write(json.dumps(rec, ensure_ascii=False) + "\n")
            if (k + 1) % 1000 == 0:
                print(f"  {k + 1}/{args.pages}", flush=True)
    print(f"wrote {args.pages} pages to {out}")


if __name__ == "__main__":
    main()