"""Evaluate the full system on held-out *real* documents, plus robustness and unseen-font synthetic checks.

    python -m tinydoc.evaluate --weights models                    # all held-out real test sets
    python -m tinydoc.evaluate --pages my_docs/pages.jsonl         # your own annotated documents
    python -m tinydoc.evaluate --robustness 100 --synthetic 150    # + rotation/skew and unseen-font docs

Metrics: detection P/R/F1 (IoU 0.5), recognition CER on ground-truth boxes, per-field accuracy, exact numeric
accuracy, how often a field is wrong vs. left empty, how often a wrong value was NOT flagged for review,
orientation accuracy under rotation, page latency, and parameter count.
Annotation format: see tinydoc/data/convert.py (one JSON page record per line).
"""

import argparse
import json
import re
import tempfile
import time
from collections import defaultdict
from pathlib import Path

import cv2
import numpy as np

from .geometry import quad_size
from .pipeline.assemble import (assemble_page_structure, block_groups, block_order, groups_from_block_map, number_style,
                                parse_amount, parse_date, structure_from_groups, xy_cut)
from .pipeline.detect import match_quads
from .train import cer

NUMERIC = ("subtotal", "tax", "total", "amount_paid", "change")
TEXTUAL = ("invoice_number", "date", "merchant_name")


def _norm(s):
    return re.sub(r"[\s:#.,'-]", "", (s or "")).lower()


def field_scores(result, gt_fields, scores):
    """Compare one document's output with its ground truth; accumulate into scores[field][...]."""
    pred_fields = result.get("fields") or {}
    doc_flagged = bool((result.get("validation") or {}).get("needs_review"))
    gt_money = [gt_fields.get(k) for k in NUMERIC if gt_fields.get(k)]
    dec = number_style(gt_money + [i.get("amount") for i in gt_fields.get("items", []) if i.get("amount")])
    for k in NUMERIC + TEXTUAL:
        g = gt_fields.get(k)
        if not g:
            continue
        p = pred_fields.get(k)
        pv = p.get("value") if isinstance(p, dict) else None
        if k in NUMERIC:
            gv = parse_amount(g, dec)
            if gv is None:
                continue
            ok = pv is not None and abs(abs(pv) - abs(gv)) < 0.005
        elif k == "date":
            giso, gcands = parse_date(g)
            if giso is None and len(gcands) > 1:
                # genuinely ambiguous date (e.g. 07/05/22): the right answer is null plus both candidates
                ok = pv is None and isinstance(p, dict) and set(p.get("candidates") or []) == set(gcands)
            else:
                ok = pv is not None and (pv == giso if giso else _norm(p.get("raw")) == _norm(g))
            if pv is None and isinstance(p, dict) and p.get("candidates") and giso in p["candidates"]:
                scores[k]["flagged_ambiguous"] += 1
        else:
            ok = pv is not None and _norm(pv) == _norm(g)
        wrong = pv is not None and not ok
        s = scores[k]
        s["n"] += 1
        s["correct"] += ok
        s["empty"] += pv is None
        s["wrong"] += wrong
        # dangerous case: a wrong value that neither the field nor the document flags for review
        s["wrong_unflagged"] += wrong and not (p.get("needs_review") or doc_flagged)
    gi = [i for i in gt_fields.get("items", []) if i.get("amount")]
    if gi:
        pi = pred_fields.get("items") or []
        scores["item_count"]["n"] += 1
        scores["item_count"]["correct"] += len(pi) == len(gi)
        g_amts = sorted(parse_amount(i["amount"], dec) or 0 for i in gi)
        p_amts = sorted((i.get("amount") or {}).get("value") or 0 for i in pi)
        scores["item_amounts"]["n"] += 1
        scores["item_amounts"]["correct"] += len(g_amts) == len(p_amts) and all(abs(a - b) < 0.005 for a, b in zip(g_amts, p_amts))


def summarise_fields(fscores):
    out = {}
    for k, v in fscores.items():
        row = {"accuracy": round(v["correct"] / v["n"], 4), "n": v["n"]}
        if k not in ("item_count", "item_amounts"):
            row.update(wrong_rate=round(v["wrong"] / v["n"], 4), empty_rate=round(v["empty"] / v["n"], 4),
                       wrong_unflagged_rate=round(v["wrong_unflagged"] / v["n"], 4))
        if v.get("flagged_ambiguous"):
            row["flagged_ambiguous"] = v["flagged_ambiguous"]
        out[k] = row
    nums = [fscores[k] for k in NUMERIC if fscores[k]["n"]]
    return out, (round(sum(v["correct"] for v in nums) / sum(v["n"] for v in nums), 4) if nums else None)


def run_doc(eng, img_rgb):
    with tempfile.NamedTemporaryFile(suffix=".png") as tmp:
        cv2.imwrite(tmp.name, np.ascontiguousarray(img_rgb[:, :, ::-1]))
        t0 = time.time()
        out = eng.process(tmp.name)
        return out, time.time() - t0


def eval_real(eng, args):
    report = {}
    sources = args.pages or [str(p) for p in sorted(Path(args.real_root).glob(f"*/{args.split}/pages.jsonl"))]
    for pages_path in sources:
        base = Path(pages_path).parent
        name = base.parent.name if base.name in ("train", "val", "test") else base.name
        recs = [json.loads(l) for l in open(pages_path)]
        if args.limit:
            recs = recs[: args.limit]
        if not recs:
            continue
        det = {"tp": 0, "fp": 0, "fn": 0}
        rec_err = rec_chars = 0
        fscores = defaultdict(lambda: defaultdict(int))
        latency, stage_ms = [], defaultdict(list)
        for r in recs:
            img = cv2.imread(str(base / r["image"]), cv2.IMREAD_COLOR)[:, :, ::-1].copy()
            segs = r["segments"]
            scored = [not s["ignore"] and s.get("label") != "formula" for s in segs]   # maths is skipped by design
            gt_q = [np.asarray(s["quad"], np.float32) for s, ok in zip(segs, scored) if ok]
            if r["use"]["detector"] and gt_q:
                pred = [q for q, _ in eng.detect(img)]
                ign = [np.asarray(s["quad"], np.float32) for s, ok in zip(segs, scored) if not ok] + [np.asarray(q, np.float32) for q in r["ignore_regions"]]
                m, up, ug = match_quads(pred, gt_q)
                if ign and up:
                    _, up, _ = match_quads([pred[i] for i in up], ign, 0.1)
                det["tp"] += len(m)
                det["fp"] += len(up)
                det["fn"] += len(ug)
            if r["use"]["recognizer"]:
                ok = [s for s, sc in zip(segs, scored) if sc and s["text"] and not s.get("rec_skip")
                      and quad_size(s["quad"])[1] >= 9][: args.lines_per_page]
                if ok:
                    texts = eng.recognize(img, [np.asarray(s["quad"], np.float32) for s in ok])
                    for s, (t, _) in zip(ok, texts):
                        rec_err += cer(t, s["text"]) * len(s["text"])
                        rec_chars += len(s["text"])
            if r.get("fields"):
                out, dt = run_doc(eng, img)
                latency.append(dt)
                for k, v in out["timing_ms"].items():
                    stage_ms[k].append(v)
                field_scores(out, r["fields"], fscores)
        res = {"documents": len(recs)}
        if det["tp"] + det["fn"]:
            p = det["tp"] / max(det["tp"] + det["fp"], 1)
            rc = det["tp"] / max(det["tp"] + det["fn"], 1)
            res["detection"] = {"precision": round(p, 4), "recall": round(rc, 4), "f1": round(2 * p * rc / max(p + rc, 1e-9), 4)}
        if rec_chars:
            res["recognition_cer_on_gt_boxes"] = round(rec_err / rec_chars, 4)
        if fscores:
            res["fields"], res["exact_numeric_accuracy"] = summarise_fields(fscores)
        if latency:
            res["latency_ms_per_page"] = {"mean": round(1000 * np.mean(latency), 1), "p50": round(1000 * np.median(latency), 1),
                                          "p90": round(1000 * np.percentile(latency, 90), 1),
                                          "stages_mean": {k: round(float(np.mean(v)), 1) for k, v in stage_ms.items() if k != "total"}}
        report[name] = res
        print(f"{name}: {json.dumps(res)}", flush=True)
    return report


def _pair_f1(pred_blocks, gt_blocks):
    """Pairwise F1 of 'these two segments are in the same block' (pred and gt: lists of index lists)."""
    def pairs(blocks):
        return {(a, b) for g in blocks for i, a in enumerate(sorted(g)) for b in sorted(g)[i + 1:]}
    p, g = pairs(pred_blocks), pairs(gt_blocks)
    tp = len(p & g)
    return tp, len(p) - tp, len(g) - tp


def eval_blocks(eng, args):
    """Paragraph/block grouping on DocLayNet pages from ground-truth lines, three ways: geometric rules, the
    extractor's same_block relation, and the detector's block regions. Picks the one the engine should use
    (returned as "blocks"; the caller writes it to config.json). A learned method must beat the rules by 0.01."""
    from .schema import FIELD_LABELS, REL_TYPES
    base = Path(args.real_root) / "doclaynet" / args.split
    if not (base / "pages.jsonl").exists():
        return None
    recs = [json.loads(l) for l in open(base / "pages.jsonl")][: args.blocks]
    tot = {"geometric": [0, 0, 0], "model": [0, 0, 0], "detector": [0, 0, 0]}
    seen = set()

    def add(mode, blocks_or_groups, gt):
        seen.add(mode)
        tot[mode] = [a + b for a, b in zip(tot[mode], _pair_f1(blocks_or_groups, list(gt.values())))]

    for r in recs:
        keep = [s for s in r["segments"] if not s["ignore"] and s.get("block") is not None]
        if len(keep) < 2:
            continue
        segs = [{"id": i, "text": s["text"], "bbox": [min(x for x, _ in s["quad"]), min(y for _, y in s["quad"]),
                                                       max(x for x, _ in s["quad"]), max(y for _, y in s["quad"])], "conf": 1.0}
                for i, s in enumerate(keep)]
        gt = {}
        for i, s in enumerate(keep):
            gt.setdefault(s["block"], []).append(i)
        labels, rel, _ = eng.extract(segs, r["width"], r["height"])
        names = [FIELD_LABELS[int(x.argmax())] for x in labels]
        boxes = [s["bbox"] for s in segs]
        blocks, _ = assemble_page_structure(segs, xy_cut(boxes), names, rel[0])
        add("geometric", [b["source"] for b in blocks], gt)
        if rel.shape[0] > REL_TYPES.index("same_block"):
            groups = block_order(boxes, block_groups(boxes, names, rel[REL_TYPES.index("same_block")]))
            add("model", [b["source"] for b in structure_from_groups(segs, groups, names, rel[0])[0]], gt)
        img = cv2.imread(str(base / r["image"]), cv2.IMREAD_COLOR)
        if img is not None:
            _, block_map, scale = eng.detect_full(img[:, :, ::-1].copy())
            if block_map is not None:
                groups = block_order(boxes, groups_from_block_map(boxes, block_map, scale))
                add("detector", [b["source"] for b in structure_from_groups(segs, groups, names, rel[0])[0]], gt)
    out = {"pages": len(recs)}
    for k, (tp, fp, fn) in tot.items():
        if k in seen and tp + fp + fn:
            out[f"{k}_pair_f1"] = round(2 * tp / max(2 * tp + fp + fn, 1), 4)
    best = "geometric"
    for k in ("model", "detector"):
        if out.get(f"{k}_pair_f1", 0) > out.get(f"{best}_pair_f1", 0) + (0.01 if best == "geometric" else 0):
            best = k
    out["blocks"] = best
    out["use_model_blocks"] = best == "model"   # older readers
    print(f"block grouping: {json.dumps(out)}", flush=True)
    return out


def eval_robustness(eng, args):
    """CORD test receipts rotated by 90/180/270 degrees and skewed: orientation accuracy and field survival."""
    base = Path(args.real_root) / "cord-v2" / args.split
    recs = [json.loads(l) for l in open(base / "pages.jsonl")][: args.robustness]
    out = {}
    for cond in ("upright", "rot90", "rot180", "rot270", "skew+6", "skew-6"):
        fscores = defaultdict(lambda: defaultdict(int))
        orient_ok = 0
        for r in recs:
            img = cv2.imread(str(base / r["image"]), cv2.IMREAD_COLOR)[:, :, ::-1].copy()
            expect = 0
            if cond.startswith("rot"):
                k = int(cond[3:]) // 90
                img = np.ascontiguousarray(np.rot90(img, -k))  # page rotated k*90 degrees clockwise
                expect = 90 * k
            elif cond.startswith("skew"):
                ang = float(cond[4:])
                h, w = img.shape[:2]
                m = cv2.getRotationMatrix2D((w / 2, h / 2), ang, 1.0)
                img = cv2.warpAffine(img, m, (w, h), borderMode=cv2.BORDER_REPLICATE)
            res, _ = run_doc(eng, img)
            orient_ok += res["pages"][0].get("rotation_corrected_deg", 0) == expect
            field_scores(res, r["fields"], fscores)
        fs, num = summarise_fields(fscores)
        out[cond] = {"orientation_accuracy": round(orient_ok / len(recs), 4), "exact_numeric_accuracy": num,
                     "total_accuracy": fs.get("total", {}).get("accuracy")}
        print(f"robustness {cond}: {json.dumps(out[cond])}", flush=True)
    return out


def eval_synthetic(eng, args):
    """Receipts and invoices rendered with font families held out from training (all fields known)."""
    from .synth.augment import augment_page
    from .synth.documents import build, render
    from .synth.fonts import FontBank

    bank = FontBank(args.fonts, "val")
    out = {}
    for kind in ("receipt", "invoice"):
        fscores = defaultdict(lambda: defaultdict(int))
        doc_ok = 0
        for i in range(args.synthetic):
            rng = np.random.default_rng([31337, i, kind == "invoice"])
            L = build(rng, bank, kind)
            img = render(L, rng)
            img, _, _ = augment_page(img, [], [], rng, strength=0.7, geo_mode=str(rng.choice(["none", "rotate"])))
            res, _ = run_doc(eng, img)
            doc_ok += (res["document_type"] or {}).get("value") == kind
            field_scores(res, L.fields, fscores)
        fs, num = summarise_fields(fscores)
        out[kind] = {"documents": args.synthetic, "doc_type_accuracy": round(doc_ok / args.synthetic, 4),
                     "fields": fs, "exact_numeric_accuracy": num}
        print(f"synthetic {kind} (unseen fonts): {json.dumps(out[kind])}", flush=True)
    return out


def _math_counts(pred, gt):
    """Maths characters (sub/superscripts, Greek, operators) in gt that pred has, in gt, and extra ones in pred."""
    from collections import Counter
    from .charset import MATH_CHARS
    g = Counter(c for c in gt if c in MATH_CHARS)
    p = Counter(c for c in pred if c in MATH_CHARS)
    return sum(min(n, p[c]) for c, n in g.items()), sum(g.values()), sum(max(0, n - g[c]) for c, n in p.items())


def _broken_encoding(text):
    """Ground truth from a PDF with a broken font encoding ("> Ì > Ê μ Õ i Ì"): mostly single characters, many of
    them accented capitals. One DocLayNet test document is like this; its lines say nothing about OCR quality."""
    chars = [c for c in text if not c.isspace()]
    words = text.split()
    return (len(chars) >= 4 and sum(c in "ÀÁÂÄÃÅÇÈÉÊËÌÍÎÏÑÒÓÔÖÕØÙÚÛÜÝ" for c in chars) / len(chars) > 0.15
            and sum(len(w) == 1 for w in words) / len(words) > 0.5)


def eval_math(eng, args):
    """Inline maths read as text: synthetic lines in held-out font families, and real DocLayNet test lines that
    contain sub/superscripts, Greek or maths symbols (read at their ground-truth boxes)."""
    from .charset import MATH_CHARS
    from .data.datasets import synth_math_line
    from .synth.fonts import FontBank

    def summary(pairs):
        hit = tot = extra = 0
        for p, g in pairs:
            h, t, e = _math_counts(p, g)
            hit, tot, extra = hit + h, tot + t, extra + e
        chars = sum(len(g) for _, g in pairs)
        return {"lines": len(pairs), "cer": round(sum(cer(p, g) * len(g) for p, g in pairs) / max(chars, 1), 4),
                "exact": round(sum(p == g for p, g in pairs) / max(len(pairs), 1), 4),
                "math_char_recall": round(hit / max(tot, 1), 4), "math_chars": tot, "invented_math_chars": extra}

    bank, rng, pairs = FontBank(args.fonts, "val"), np.random.default_rng(4242), []
    for _ in range(args.math):
        gray, (x0, y0, x1, y1), text = synth_math_line(rng, bank)
        img = np.repeat(gray[:, :, None], 3, axis=2)
        quad = np.array([[x0, y0], [x1, y0], [x1, y1], [x0, y1]], np.float32)
        pairs.append((eng.recognize(img, [quad])[0][0], text))
    out = {"synthetic_unseen_fonts": summary(pairs)}
    print(f"maths lines, synthetic (unseen fonts): {json.dumps(out['synthetic_unseen_fonts'])}", flush=True)

    pairs, base = [], Path(args.real_root) / "doclaynet" / args.split
    for line in open(base / "pages.jsonl") if (base / "pages.jsonl").exists() else []:
        r = json.loads(line)
        segs = [s for s in r["segments"] if any(c in MATH_CHARS for c in s["text"]) and not s["ignore"]
                and not s.get("rec_skip") and s.get("label") != "formula" and quad_size(s["quad"])[1] >= 9
                and not _broken_encoding(s["text"])]
        if not segs:
            continue
        img = cv2.imread(str(base / r["image"]), cv2.IMREAD_COLOR)[:, :, ::-1].copy()
        texts = eng.recognize(img, [np.asarray(s["quad"], np.float32) for s in segs])
        pairs += [(t, s["text"]) for s, (t, _) in zip(segs, texts)]
        if len(pairs) >= args.math:
            break
    if pairs:
        out["doclaynet"] = summary(pairs)
        print(f"maths lines, DocLayNet {args.split}: {json.dumps(out['doclaynet'])}", flush=True)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--backend", choices=["onnx", "torch"], default="onnx")
    ap.add_argument("--weights", default="models")
    ap.add_argument("--device", default="cpu")
    ap.add_argument("--real-root", default="data/real")
    ap.add_argument("--fonts", default="fonts")
    ap.add_argument("--split", default="test")
    ap.add_argument("--pages", nargs="*", help="explicit pages.jsonl files (e.g. your own annotated set)")
    ap.add_argument("--limit", type=int, default=0, help="documents per source (0 = all)")
    ap.add_argument("--lines-per-page", type=int, default=40)
    ap.add_argument("--robustness", type=int, default=0, help="CORD test receipts for rotation/skew checks")
    ap.add_argument("--synthetic", type=int, default=0, help="unseen-font synthetic docs per kind")
    ap.add_argument("--blocks", type=int, default=0, help="DocLayNet pages for the block-grouping comparison")
    ap.add_argument("--math", type=int, default=0, help="lines for the inline-maths check (synthetic and DocLayNet)")
    ap.add_argument("--out", default="eval_report.json")
    args = ap.parse_args()

    from .pipeline.engine import Engine
    eng = Engine(args.weights, backend=args.backend, device=args.device)
    report = {"parameters": eng.backend.params, "total_parameters": int(sum(eng.backend.params.values())),
              "backend": args.backend, "sources": eval_real(eng, args)}
    if args.robustness:
        report["robustness"] = eval_robustness(eng, args)
    if args.synthetic:
        report["synthetic_unseen_fonts"] = eval_synthetic(eng, args)
    if args.blocks:
        report["block_grouping"] = eval_blocks(eng, args)
    if args.math:
        report["inline_maths"] = eval_math(eng, args)
    Path(args.out).write_text(json.dumps(report, indent=2))
    print(f"wrote {args.out}")


if __name__ == "__main__":
    main()