"""
benchmark.py - word-level OCR benchmark on handwritten word crops (used by notebook CELL 6).

Recreated from how the notebook calls it; the original lived only in Colab's /content.

Data layout (bench_sample.zip unpacks to this):
    bench_sample/<lang>/<image>          lang = hi-IN | kn-IN | ml-IN | ta-IN | te-IN
    bench_sample/<lang>/labels.csv       file,text   (also .tsv / .txt with "file<TAB>text")

    python benchmark.py sample <dataset_root> --per-lang 500     # -> bench_sample.zip (same layout, random subset)
    python benchmark.py summary bench_direct.csv bench_direct_langhint_remap.csv

Document benchmark (full pages / PDFs against a ground-truth .txt with the same name):
    bench_docs/<doc>.<jpg|png|pdf|...>   +   bench_docs/<doc>.txt     (UTF-8, the whole expected text)
    bench_docs/<lang>/<doc>.* + .txt     optional: a language folder sets the language for script checks

In the notebook:
    run(predict_fn, "/content/bench_sample", tag="direct", limit=None)   # -> bench_direct.csv + table
    summary(["bench_direct.csv", "bench_direct_remap.csv"])
    upload_docs("/content/bench_docs")                                    # Colab file picker: docs + .txt
    run_docs(predict_doc, "/content/bench_docs", tag="docs", lang="kn-IN") # -> bench_docs.csv + table
"""
from __future__ import annotations

import csv, io, random, re, sys, time, unicodedata, zipfile
from pathlib import Path

from validate import cer, nfc, script_purity, wer

IMAGE_EXT = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp", ".webp"}
DOC_EXT = IMAGE_EXT | {".pdf"}
LANGS = {"hi-IN", "kn-IN", "ml-IN", "ta-IN", "te-IN"}
FIELDS = ["tag", "lang", "file", "gt", "pred", "cer", "cer_norm", "word_acc", "wa_norm", "wer",
          "empty", "halluc", "wrong_script", "seconds"]
ZERO_WIDTH = dict.fromkeys(map(ord, "​‌‍"))


# ---------------------------------------------------------------- scoring
def normalize(s: str) -> str:
    """Lenient form for CERnorm / WAnorm: NFC, no zero-width joiners, no punctuation or spaces,
    Indic digits -> ASCII. Separates reading errors from formatting differences."""
    s = nfc(s).translate(ZERO_WIDTH)
    s = "".join(str(unicodedata.digit(c)) if unicodedata.category(c) == "Nd" else c for c in s)
    return "".join(c for c in s if unicodedata.category(c)[0] not in ("P", "Z", "S", "C"))


def is_hallucination(gt: str, pred: str) -> bool:
    """Output that is not a reading of the word: a sentence about the image, LaTeX, or far too long."""
    letters = [c for c in pred if unicodedata.category(c)[0] == "L"]
    latin = sum("A" <= c <= "z" for c in letters)
    gt_latin = any("A" <= c <= "z" for c in gt)
    return bool(pred.strip()) and (
        (not gt_latin and letters and latin / len(letters) > 0.5)
        or "$$" in pred or "\\text" in pred
        or len(pred) > 3 * len(gt) + 10)


def score(gt: str, pred: str, lang: str) -> dict:
    gt, pred = nfc(gt).strip(), nfc(pred or "").strip()
    one = lambda s: re.sub(r"\s+", " ", s)
    cap = lambda x: round(min(x, 1.0), 4)    # a word is at most fully wrong; stops one long hallucination skewing the mean
    return {
        "cer": cap(cer(pred, gt)),
        "cer_norm": cap(cer(normalize(pred), normalize(gt))),
        "word_acc": int(one(pred) == one(gt)),
        "wa_norm": int(normalize(pred) == normalize(gt)),
        "wer": cap(wer(pred, gt)),
        "empty": int(not pred),
        "halluc": int(is_hallucination(gt, pred)),
        "wrong_script": int(bool(pred) and script_purity(pred, lang) < 0.5),
    }


# ---------------------------------------------------------------- data
def load_labels(lang_dir: Path) -> dict[str, str]:
    labels = {}
    for f in sorted(lang_dir.glob("*.csv")) + sorted(lang_dir.glob("*.tsv")) + sorted(lang_dir.glob("*.txt")):
        delim = "," if f.suffix == ".csv" else "\t"
        for row in csv.reader(open(f, encoding="utf-8"), delimiter=delim):
            if len(row) >= 2 and row[0] not in ("file", "filename", "image"):
                labels[Path(row[0]).name] = row[1]
    return labels


def items(root, limit=None):
    """(lang, image path, ground truth) for every labelled image, `limit` per language."""
    out = []
    for lang_dir in sorted(p for p in Path(root).iterdir() if p.is_dir()):
        labels = load_labels(lang_dir)
        imgs = sorted(p for p in lang_dir.iterdir() if p.suffix.lower() in IMAGE_EXT and p.name in labels)
        out += [(lang_dir.name, p, labels[p.name]) for p in imgs[:limit]]
    return out


def doc_items(root, lang=None):
    """(lang, document path, ground truth) for every document with a <stem>.txt next to it.
    Language: the parent folder if it is a language code, else `lang`, else "-" (script check skipped)."""
    out, missing = [], []
    for p in sorted(Path(root).rglob("*")):
        if p.suffix.lower() not in DOC_EXT or p.name.startswith("."):
            continue
        gt = p.with_suffix(".txt")
        if not gt.is_file():
            missing.append(p.name)
            continue
        out.append((p.parent.name if p.parent.name in LANGS else lang or "-", p,
                    gt.read_text(encoding="utf-8-sig")))
    if missing:
        print(f"no ground truth .txt for {len(missing)} file(s), skipped: {', '.join(missing)}")
    return out


def upload_docs(dest="/content/bench_docs"):
    """Colab file picker: select documents and their .txt ground truth together (same file names)."""
    from google.colab import files as colab_files
    dest = Path(dest)
    dest.mkdir(parents=True, exist_ok=True)
    for name, data in colab_files.upload().items():
        if name.lower().endswith(".zip"):                # a zip of the whole folder also works
            zipfile.ZipFile(io.BytesIO(data)).extractall(dest)
        else:
            (dest / Path(name).name).write_bytes(data)
    for lang, p, gt in doc_items(dest):
        print(f"{p.relative_to(dest)}  ({lang}, ground truth {len(gt)} chars)")
    return dest


# ---------------------------------------------------------------- run + report
def run(predict, root, tag="direct", limit=None, out=None):
    """predict(path) -> text. Writes bench_<tag>.csv and prints the summary table."""
    return _bench(predict, items(root, limit), tag, out)


def run_docs(predict, root, tag="docs", lang=None, out=None):
    """Same as run(), on whole documents: predict(path) -> the full text of the image or PDF."""
    return _bench(predict, doc_items(root, lang), tag, out)


def _bench(predict, todo, tag, out):
    out = Path(out or f"bench_{tag}.csv")
    with open(out, "w", encoding="utf-8", newline="") as fh:
        wr = csv.DictWriter(fh, fieldnames=FIELDS)
        wr.writeheader()
        for i, (lang, path, gt) in enumerate(todo, 1):
            t0 = time.time()
            try:
                pred = predict(str(path)) or ""
            except Exception as e:                      # one bad crop must not stop a 2,500-image run
                print(f"{path.name}: {e}")
                pred = ""
            row = {"tag": tag, "lang": lang, "file": path.name, "gt": gt, "pred": pred,
                   "seconds": round(time.time() - t0, 3)}
            wr.writerow({**row, **score(gt, pred, lang)})
            if i % 100 == 0 or i == len(todo):
                print(f"{i}/{len(todo)}")
                fh.flush()
    summary([out])
    return out


def summary(paths):
    rows = [r for p in paths for r in csv.DictReader(open(p, encoding="utf-8"))]
    cols = [("CER", "cer"), ("CERnorm", "cer_norm"), ("WordAcc", "word_acc"), ("WAnorm", "wa_norm"),
            ("WER", "wer"), ("empty", "empty"), ("halluc", "halluc"), ("wrongScr", "wrong_script")]
    print(f"{'tag':<22}{'lang':<8}{'n':>6} " + " ".join(f"{h:>8}" for h, _ in cols) + f" {'s/img':>8} {'p95 s':>8}")
    for tag in sorted({r["tag"] for r in rows}):
        tr = [r for r in rows if r["tag"] == tag]
        for lang in ["ALL"] + sorted({r["lang"] for r in tr}):
            g = tr if lang == "ALL" else [r for r in tr if r["lang"] == lang]
            mean = lambda k: sum(float(r[k]) for r in g) / len(g)
            secs = sorted(float(r["seconds"]) for r in g)
            p95 = secs[min(len(secs) - 1, int(0.95 * len(secs)))]
            print(f"{tag:<22}{lang:<8}{len(g):>6} " + " ".join(f"{mean(k):>8.3f}" for _, k in cols)
                  + f" {mean('seconds'):>8.3f} {p95:>8.3f}")


def sample(dataset_root, per_lang=500, out="bench_sample.zip", seed=0):
    """Random `per_lang` crops per language from a dataset in the same layout -> zip for Colab."""
    rnd = random.Random(seed)
    with zipfile.ZipFile(out, "w", zipfile.ZIP_DEFLATED) as z:
        for lang_dir in sorted(p for p in Path(dataset_root).iterdir() if p.is_dir()):
            labels = load_labels(lang_dir)
            imgs = sorted(p for p in lang_dir.iterdir() if p.suffix.lower() in IMAGE_EXT and p.name in labels)
            pick = rnd.sample(imgs, min(per_lang, len(imgs)))
            for p in pick:
                z.write(p, f"bench_sample/{lang_dir.name}/{p.name}")
            buf = io.StringIO()
            csv.writer(buf).writerows([("file", "text")] + [(p.name, labels[p.name]) for p in pick])
            z.writestr(f"bench_sample/{lang_dir.name}/labels.csv", buf.getvalue())
            print(f"{lang_dir.name}: {len(pick)}")
    print("wrote", out)


if __name__ == "__main__":
    cmd, *rest = sys.argv[1:] or ["help"]
    if cmd == "summary":
        summary(rest)
    elif cmd == "sample":
        n = int(rest[rest.index("--per-lang") + 1]) if "--per-lang" in rest else 500
        sample(rest[0], n)
    else:
        print(__doc__)