"""Convert downloaded real datasets into one page format.

Each split becomes data/real/<source>/<split>/pages.jsonl plus images/, and (where transcripts keep
their spaces) lines.jsonl plus lines/ with recognizer crops. A page record:

    {"image": "images/0001.jpg", "width": W, "height": H, "source": "cord-v2", "doc_type": "receipt",
     "segments": [{"quad": [[x, y] x4], "text": str, "label": str | None, "row": int | None, "ignore": bool,
                   "block": int | None}],      # layout block id (paragraph, list item, table...) when known
     "kv_pairs": [[key_idx, value_idx], ...],
     "ignore_regions": [[[x, y] x4], ...],
     "use": {"detector": bool, "recognizer": bool, "labels": bool, "relations": bool, "blocks": bool},
     "fields": {...}}   # ground-truth values for end-to-end evaluation, when the source has them

    python -m tinydoc.data.convert --raw data/raw --out data/real
"""

import argparse
import io
import json
import os
import re
import shutil
from multiprocessing import Pool
from pathlib import Path

import cv2
import numpy as np
import pyarrow.parquet as pq
from PIL import Image

from ..charset import latex_scripts_to_unicode, normalize
from ..geometry import crop_quad, expand_quad, merge_line_cells, order_quad, quad_bbox, quad_size

LINE_STORE_HEIGHT = 48     # stored crop height of the tight text box (margins are extra)
MIN_LINE_HEIGHT = 9        # skip crops whose source text is smaller than this (unreadable)

# ---------------------------------------------------------------------------------------------- CORD
CORD_LABELS = {
    "menu.nm": "item_desc", "menu.cnt": "item_qty", "menu.unitprice": "item_unit_price",
    "menu.price": "item_amount", "menu.discountprice": "discount",
    "sub_total.subtotal_price": "subtotal", "sub_total.tax_price": "tax",
    "sub_total.discount_price": "discount",
    "total.total_price": "total", "total.cashprice": "amount_paid", "total.creditcardprice": "amount_paid",
    "total.emoneyprice": "amount_paid", "total.changeprice": "change",
}
ROW_CATEGORIES = {"menu.nm", "menu.cnt", "menu.unitprice", "menu.price", "menu.discountprice"}


def _merge_words(words):
    """words: list of (quad, text) on one row, sorted left to right -> list of (quad, text) segments.

    Words join into one segment when the gap between them is under 1.2x the text height, matching
    how the synthetic generator separates columns.
    """
    segs, cur = [], None
    for q, t in words:
        if cur is None:
            cur = [q, q, [t]]
            continue
        _, h = quad_size(q)
        gap = q[0][0] - cur[1][1][0]
        if gap < 1.2 * h:
            cur[1] = q
            cur[2].append(t)
        else:
            segs.append(cur)
            cur = [q, q, [t]]
    if cur:
        segs.append(cur)
    out = []
    for first, last, texts in segs:
        quad = np.stack([first[0], last[1], last[2], first[3]]).astype(np.float32)
        out.append((quad, " ".join(texts)))
    return out


def _cord_fields(gp):
    def first(d, k):
        v = d.get(k) if isinstance(d, dict) else None
        return v if isinstance(v, str) else None
    menu = gp.get("menu", [])
    menu = menu if isinstance(menu, list) else [menu]
    st, tt = gp.get("sub_total", {}) or {}, gp.get("total", {}) or {}
    return {
        "subtotal": first(st, "subtotal_price"), "tax": first(st, "tax_price"),
        "total": first(tt, "total_price"), "amount_paid": first(tt, "cashprice") or first(tt, "creditcardprice"),
        "change": first(tt, "changeprice"),
        "items": [{"description": first(m, "nm"), "quantity": first(m, "cnt"),
                   "unit_price": first(m, "unitprice"), "amount": first(m, "price")} for m in menu if isinstance(m, dict)],
    }


def convert_cord(raw: Path, out: Path):
    splits = {"train": "train-*.parquet", "val": "validation-*.parquet", "test": "test-*.parquet"}
    for split, pat in splits.items():
        files = sorted((raw / "data").glob(pat))
        writer = SplitWriter(out / split, "cord-v2")
        for f in files:
            table = pq.read_table(f)
            for row in table.to_pylist():
                gt = json.loads(row["ground_truth"])
                img = _decode(row["image"]["bytes"])
                segments, kv = [], []
                for li, vl in enumerate(gt["valid_line"]):
                    cat = vl["category"]
                    label = CORD_LABELS.get(cat, "other")
                    row_id = (vl["group_id"] if cat in ROW_CATEGORIES else None)
                    key_idx, val_idx = [], []
                    for is_key in (1, 0):
                        ws = [w for w in vl["words"] if w["is_key"] == is_key and w["text"].strip()]
                        by_row = {}
                        for w in ws:
                            q = w["quad"]
                            quad = order_quad([[q["x1"], q["y1"]], [q["x2"], q["y2"]], [q["x3"], q["y3"]], [q["x4"], q["y4"]]])
                            by_row.setdefault(w["row_id"], []).append((quad, w["text"]))
                        for r in by_row.values():
                            r.sort(key=lambda x: x[0][0][0])
                            for quad, text in _merge_words(r):
                                (key_idx if is_key else val_idx).append(len(segments))
                                segments.append({"quad": quad, "text": normalize(text),
                                                 "label": "key" if is_key else label,
                                                 "row": row_id, "ignore": False})
                    kv += [[k, v] for k in key_idx for v in val_idx]
                ignore = []
                for dc in gt.get("dontcare", []):  # list of lists of bare quads
                    for q in (dc if isinstance(dc, list) else [dc]):
                        if isinstance(q, dict) and "x1" in q:
                            ignore.append(order_quad([[q["x1"], q["y1"]], [q["x2"], q["y2"]], [q["x3"], q["y3"]], [q["x4"], q["y4"]]]))
                writer.add(img, {
                    "doc_type": "receipt", "segments": segments, "kv_pairs": kv, "ignore_regions": ignore,
                    # Store names/addresses are blurred and some readable text is unannotated, so CORD
                    # must not teach the detector that unboxed text is background.
                    "use": {"detector": False, "recognizer": True, "labels": True, "relations": True, "blocks": False},
                    "fields": _cord_fields(gt["gt_parse"]),
                })
        writer.close()


# ----------------------------------------------------------------------------------------- WildReceipt
WILD_LABELS = {
    1: "merchant_name", 3: "merchant_address", 5: "merchant_phone", 7: "date", 9: "time",
    11: "item_desc", 13: "item_qty", 15: "item_amount", 17: "subtotal", 19: "tax", 23: "total",
}


def convert_wildreceipt(raw: Path, out: Path):
    root = raw / "wildreceipt"
    train = [json.loads(l) for l in open(root / "train.txt")]
    test = [json.loads(l) for l in open(root / "test.txt")]
    rng = np.random.default_rng(0)
    is_val = rng.random(len(train)) < 0.1
    parts = {"train": [d for d, v in zip(train, is_val) if not v], "val": [d for d, v in zip(train, is_val) if v], "test": test}
    for split, docs in parts.items():
        writer = SplitWriter(out / split, "wildreceipt")
        for d in docs:
            img = cv2.imread(str(root / d["file_name"]), cv2.IMREAD_COLOR)
            if img is None:
                continue
            segments = []
            for a in d["annotations"]:
                b = a["box"]
                lab = a["label"]
                label = None if lab == 0 else ("key" if lab % 2 == 0 and lab < 25 else WILD_LABELS.get(lab, "other"))
                segments.append({"quad": order_quad(list(zip(b[0::2], b[1::2]))), "text": normalize(a["text"]),
                                 "label": label, "row": None, "ignore": lab == 0})
            fields = {}
            for lab, name in ((23, "total"), (17, "subtotal"), (19, "tax"), (7, "date"), (9, "time")):
                vals = [a["text"] for a in d["annotations"] if a["label"] == lab]
                if vals:
                    fields[name] = vals[0]
            writer.add(img[:, :, ::-1], {
                "doc_type": "receipt", "segments": segments, "kv_pairs": [], "ignore_regions": [],
                # Transcripts have their spaces removed ("THANKYOUFORSHOPPING"), so no recognizer crops.
                "use": {"detector": True, "recognizer": False, "labels": True, "relations": False, "blocks": False},
                "fields": fields,
            })
        writer.close()


# ------------------------------------------------------------------------------------------ DocLayNet
DLN_LABELS = {1: "paragraph", 2: "paragraph", 3: "formula", 4: "list_item", 5: "other", 6: "other", 7: "other",
              8: "heading", 9: "table_cell", 10: "paragraph", 11: "heading"}
DLN_FORMULA, DLN_PICTURE = 3, 7
DLN_DPI = 150
DLN_JPEG_QUALITY = 85          # keeps ~20k pages inside Kaggle's 20 GB output limit
DLN_LINES_PER_PAGE = 20


def _doclaynet_file(job):
    """Convert one parquet shard (optionally capped) into its own part directory. Returns (part dir, pages)."""
    import pypdfium2 as pdfium

    path, part_dir, cap = job
    writer = SplitWriter(Path(part_dir), "doclaynet", max_lines_per_page=DLN_LINES_PER_PAGE,
                         jpeg_quality=DLN_JPEG_QUALITY, quiet=True)
    n = 0
    pf = pq.ParquetFile(path)
    for rg in range(pf.num_row_groups):
        if cap and n >= cap:
            break
        for row in pf.read_row_group(rg, columns=["pdf", "bboxes", "category_id", "pdf_cells", "metadata"]).to_pylist():
            if cap and n >= cap:
                break
            meta = row["metadata"]
            if not row["pdf"] or not any(row["pdf_cells"]):
                continue
            doc = pdfium.PdfDocument(row["pdf"])
            img = doc[0].render(scale=DLN_DPI / 72).to_numpy()[:, :, :3].copy()
            doc.close()
            sx = img.shape[1] / meta["coco_width"]
            sy = img.shape[0] / meta["coco_height"]
            segments, ignore = [], []
            for block, (box, cat, cells) in enumerate(zip(row["bboxes"], row["category_id"], row["pdf_cells"])):
                bx, by, bw, bh = box
                region = np.array([[bx * sx, by * sy], [(bx + bw) * sx, by * sy],
                                   [(bx + bw) * sx, (by + bh) * sy], [bx * sx, (by + bh) * sy]], np.float32)
                if cat == DLN_PICTURE:
                    ignore.append(region)  # text inside photos and charts is unannotated
                    continue
                merged = merge_line_cells([(c["bbox"][0] * sx, c["bbox"][1] * sy, (c["bbox"][0] + c["bbox"][2]) * sx,
                                            (c["bbox"][1] + c["bbox"][3]) * sy, c["text"]) for c in cells if c["text"].strip()])
                if cat == DLN_FORMULA:
                    # Maths is out of scope: the detector ignores it, the extractor learns to label it "formula"
                    # so assembly can skip it. Its text is never used as a recognizer target.
                    ignore.append(region)
                    if not merged:
                        continue
                    x0, y0 = min(m[0] for m in merged), min(m[1] for m in merged)
                    x1, y1 = max(m[2] for m in merged), max(m[3] for m in merged)
                    text = re.sub(r"\$[_^]?\{([^}]*)\}\$|\$", r"\1", " ".join(m[4] for m in merged))   # $_{q}$ -> q
                    text = normalize(text)[:60] or "x"
                    segments.append({"quad": np.array([[x0, y0], [x1, y0], [x1, y1], [x0, y1]], np.float32),
                                     "text": text, "label": "formula", "row": None, "ignore": False, "block": block})
                    continue
                for x0, y0, x1, y1, text in merged:
                    # PDF text marks sub/superscripts as $_{2}$ / $^{2}$: spell them in Unicode (H₂O, x²), and keep
                    # lines whose markup cannot be converted out of the recognizer targets.
                    text, clean = latex_scripts_to_unicode(text)
                    text = normalize(text)
                    if not text:
                        continue
                    segments.append({"quad": np.array([[x0, y0], [x1, y0], [x1, y1], [x0, y1]], np.float32),
                                     "text": text, "label": DLN_LABELS[cat], "row": None,
                                     "ignore": (y1 - y0) < 6, "block": block, **({} if clean else {"rec_skip": True})})
            if not segments:
                continue
            writer.add(img, {
                "doc_type": "general", "segments": segments, "kv_pairs": [], "ignore_regions": ignore,
                "use": {"detector": True, "recognizer": True, "labels": True, "relations": False, "blocks": True},
                "fields": {}, "meta": {"category": meta["doc_category"], "file": meta["original_filename"], "page": meta["page_no"]},
            })
            n += 1
    writer.close()
    return part_dir, n


def convert_doclaynet(raw: Path, out: Path, max_pages=None, workers=None, splits=None):
    """Shards are converted in parallel into part directories, then merged into one split directory."""
    workers = workers or os.cpu_count() or 4
    for split, pat in (("train", "train-*.parquet"), ("val", "validation-*.parquet"), ("test", "test-*.parquet")):
        files = sorted((raw / "data").glob(pat))
        if not files or (splits and split not in splits):
            continue
        cap = max_pages or (None if split == "train" else 600)
        dst = out / split
        shutil.rmtree(dst, ignore_errors=True)
        parts = out / f"_parts_{split}"
        shutil.rmtree(parts, ignore_errors=True)
        per_file = -(-cap // len(files)) if cap else None
        jobs = [(str(f), str(parts / f"{k:03d}"), per_file) for k, f in enumerate(files)]
        with Pool(min(workers, len(jobs))) as pool:
            done = sorted(pool.imap_unordered(_doclaynet_file, jobs))
        merge_parts([Path(d) for d, _ in done], dst, cap)
        shutil.rmtree(parts, ignore_errors=True)


def merge_parts(part_dirs, dst: Path, cap=None):
    """Concatenate part directories written by SplitWriter, renumbering images and line crops."""
    (dst / "images").mkdir(parents=True, exist_ok=True)
    n = n_lines = 0
    with open(dst / "pages.jsonl", "w") as pages, open(dst / "lines.jsonl", "w") as lines:
        for part in part_dirs:
            prefix = part.name
            for sub in ("images", "lines"):
                if (part / sub).is_dir():
                    (dst / sub).mkdir(exist_ok=True)
                    for f in (part / sub).iterdir():
                        os.replace(f, dst / sub / f"{prefix}_{f.name}")
            for line in open(part / "pages.jsonl"):
                if cap and n >= cap:
                    break
                rec = json.loads(line)
                rec["image"] = f"images/{prefix}_{Path(rec['image']).name}"
                pages.write(json.dumps(rec, ensure_ascii=False) + "\n")
                n += 1
            if (part / "lines.jsonl").exists():
                for line in open(part / "lines.jsonl"):
                    rec = json.loads(line)
                    rec["path"] = f"lines/{prefix}_{Path(rec['path']).name}"
                    lines.write(json.dumps(rec, ensure_ascii=False) + "\n")
                    n_lines += 1
    print(f"  doclaynet        {dst.name:5s} pages={n:6d} line_crops={n_lines}", flush=True)


# --------------------------------------------------------------------------------------------- Katana
def convert_katana(raw: Path, out: Path):
    for split, pat in (("train", "train-*"), ("val", "validation-*"), ("test", "test-*")):
        writer = SplitWriter(out / split, "katana-invoices")
        for f in sorted((raw / "data").glob(pat)):
            for row in pq.read_table(f).to_pylist():
                gp = json.loads(row["ground_truth"])["gt_parse"]
                hd, sm = gp.get("header", {}), gp.get("summary", {})
                writer.add(_decode(row["image"]["bytes"]), {
                    "doc_type": "invoice", "segments": [], "kv_pairs": [], "ignore_regions": [],
                    "use": {"detector": False, "recognizer": False, "labels": False, "relations": False, "blocks": False},
                    "fields": {
                        "invoice_number": hd.get("invoice_no"), "date": hd.get("invoice_date"),
                        "subtotal": sm.get("total_net_worth"), "tax": sm.get("total_vat"), "total": sm.get("total_gross_worth"),
                        "items": [{"description": it.get("item_desc"), "quantity": it.get("item_qty"),
                                   "unit_price": it.get("item_net_price"), "amount": it.get("item_gross_worth")}
                                  for it in gp.get("items", [])],
                    },
                })
        writer.close()


# ---------------------------------------------------------------------------------------------- output
def _decode(b: bytes) -> np.ndarray:
    return np.array(Image.open(io.BytesIO(b)).convert("RGB"))


class SplitWriter:
    def __init__(self, root: Path, source: str, max_lines_per_page=None, jpeg_quality=93, quiet=False):
        self.root, self.source, self.n, self.n_lines = root, source, 0, 0
        self.max_lines, self.quality, self.quiet = max_lines_per_page, jpeg_quality, quiet
        self.rng = np.random.default_rng(0)
        (root / "images").mkdir(parents=True, exist_ok=True)
        self.pages = open(root / "pages.jsonl", "w")
        self.lines = None

    def add(self, img: np.ndarray, rec: dict):
        name = f"images/{self.n:06d}.jpg"
        cv2.imwrite(str(self.root / name), np.ascontiguousarray(img[:, :, ::-1]), [cv2.IMWRITE_JPEG_QUALITY, self.quality])
        rec = {"image": name, "width": img.shape[1], "height": img.shape[0], "source": self.source, **rec}
        if rec["use"]["recognizer"]:
            self._export_lines(img, rec)
        rec["segments"] = [{**s, "quad": np.round(np.asarray(s["quad"]), 1).tolist()} for s in rec["segments"]]
        rec["ignore_regions"] = [np.round(np.asarray(q), 1).tolist() for q in rec["ignore_regions"]]
        self.pages.write(json.dumps(rec, ensure_ascii=False) + "\n")
        self.n += 1

    def _export_lines(self, img, rec):
        if self.lines is None:
            (self.root / "lines").mkdir(exist_ok=True)
            self.lines = open(self.root / "lines.jsonl", "w")
        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        order = np.array([i for i, s in enumerate(rec["segments"]) if s["label"] != "formula"], int)
        if self.max_lines and len(order) > self.max_lines:
            order = np.sort(self.rng.choice(order, self.max_lines, replace=False))
        for i in order:
            s = rec["segments"][i]
            if s["ignore"] or s.get("rec_skip") or not s["text"].strip():
                continue
            w, h = quad_size(s["quad"])
            if h < MIN_LINE_HEIGHT or w < h * 0.3:
                continue
            # Keep generous margins so training can jitter the crop like an imperfect detector would.
            mx, my = 0.6 * h, 0.35 * h
            crop = crop_quad(gray, expand_quad(s["quad"], mx, my, mx, my))
            scale = LINE_STORE_HEIGHT / h
            crop = cv2.resize(crop, (max(1, round(crop.shape[1] * scale)), max(1, round(crop.shape[0] * scale))), interpolation=cv2.INTER_AREA)
            name = f"lines/{self.n:06d}_{i:03d}.jpg"
            cv2.imwrite(str(self.root / name), crop, [cv2.IMWRITE_JPEG_QUALITY, 92])
            bx, by = mx * scale, my * scale
            self.lines.write(json.dumps({"path": name, "text": s["text"],
                                         "box": [round(bx, 1), round(by, 1), round(crop.shape[1] - bx, 1), round(crop.shape[0] - by, 1)]},
                                        ensure_ascii=False) + "\n")
            self.n_lines += 1

    def close(self):
        self.pages.close()
        if self.lines:
            self.lines.close()
        if not self.quiet:
            print(f"  {self.source:16s} {self.root.name:5s} pages={self.n:6d} line_crops={self.n_lines}")


CONVERTERS = {"cord-v2": convert_cord, "wildreceipt": convert_wildreceipt, "doclaynet": convert_doclaynet,
              "katana-invoices": convert_katana}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--raw", default="data/raw")
    ap.add_argument("--out", default="data/real")
    ap.add_argument("--only", nargs="*", default=list(CONVERTERS))
    ap.add_argument("--max-pages", type=int, default=None, help="cap per split (DocLayNet only), for quick tests")
    ap.add_argument("--workers", type=int, default=None, help="parallel DocLayNet shards (default: all cores)")
    ap.add_argument("--splits", nargs="*", default=None, help="DocLayNet splits to convert (default: all present)")
    args = ap.parse_args()
    for name in args.only:
        raw = Path(args.raw) / name
        if not raw.exists():
            print(f"skip {name}: {raw} not found")
            continue
        print(f"converting {name}")
        if name == "doclaynet":
            convert_doclaynet(raw, Path(args.out) / name, args.max_pages, args.workers, args.splits)
        else:
            CONVERTERS[name](raw, Path(args.out) / name)


if __name__ == "__main__":
    main()