"""Train one component.

    python -m tinydoc.train recognizer            # CNN + Transformer + CTC line reader
    python -m tinydoc.train detector              # text-segment detector
    python -m tinydoc.train orientation           # 0/90/180/270 page classifier
    python -m tinydoc.train extractor             # field / relation / doc-type model

Checkpoints go to runs/<component>/{last,best}.pt, metrics to runs/<component>/log.jsonl.
Runs on CUDA (with fp16 mixed precision), Apple MPS, or CPU. --max-minutes saves and exits cleanly before
a time limit (e.g. a Kaggle session); --resume continues from last.pt; --init <runs dir> fine-tunes from an
earlier best.pt (heads that grew, e.g. a new label, keep their old rows and initialise the new ones).
"""

import argparse
import json
import math
import os
import time
from collections import defaultdict
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F

from .charset import BLANK, CHARS
from .data import datasets as D
from .models import build, count_params
from .geometry import rect_quad, transform_quads
from .pipeline.detect import detector_input, kernels_to_quads, match_quads
from .synth.augment import geometric, photometric
from .schema import FIELD_LABELS, REL_TYPES

DEFAULTS = {
    "recognizer": dict(steps=60_000, batch=96, lr=8e-4, eval_every=2_000, warmup=1_500),
    "detector": dict(steps=40_000, batch=16, lr=2e-3, eval_every=2_000, warmup=1_000),
    "orientation": dict(steps=6_000, batch=16, lr=2e-3, eval_every=1_000, warmup=300),   # batch = pages x 8 patches
    "extractor": dict(steps=25_000, batch=32, lr=1e-3, eval_every=1_000, warmup=1_000),
}


# ------------------------------------------------------------------------------------------- losses
def detector_loss(logits, kernel, mask, neg_ratio=3.0):
    """logits/kernel/mask: (B, C, H, W). Channel 0 = line kernels, channel 1 = block kernels (weighted 0.5;
    its mask is zero on pages without block labels)."""
    loss = _kernel_loss(logits[:, :1], kernel[:, :1], mask[:, :1], neg_ratio)
    if logits.shape[1] > 1 and kernel.shape[1] > 1:
        loss = loss + 0.5 * _kernel_loss(logits[:, 1:2], kernel[:, 1:2], mask[:, 1:2], neg_ratio)
    return loss


def _kernel_loss(logits, kernel, mask, neg_ratio=3.0):
    bce = F.binary_cross_entropy_with_logits(logits, kernel, reduction="none") * mask
    pos = (kernel * mask).flatten()
    neg = ((1 - kernel) * mask).flatten()
    n_pos = int(pos.sum().item())
    n_neg = int(min(neg.sum().item(), max(neg_ratio * n_pos, 4096)))
    flat = bce.flatten()
    hard_neg = torch.topk(flat * neg, n_neg).values if n_neg > 0 else flat.new_zeros(1)
    bce_loss = ((flat * pos).sum() + hard_neg.sum()) / (n_pos + n_neg + 1e-6)
    prob = torch.sigmoid(logits) * mask
    inter = (prob * kernel).sum()
    dice = 1 - 2 * inter / (prob.sum() + (kernel * mask).sum() + 1e-6)
    return bce_loss + dice


def extractor_loss(out, b):
    field_logits, rel_logits, doc_logits = out
    lf = F.cross_entropy(field_logits.flatten(0, 1), b["labels"].flatten(), ignore_index=-100, label_smoothing=0.05)
    valid = b["valid"]
    pair = valid[:, None, :, None] * valid[:, None, None, :] * b["rel_known"][:, :, None, None]
    eye = torch.eye(valid.shape[1], device=valid.device)[None, None]
    pair = pair * (1 - eye)
    lr_ = F.binary_cross_entropy_with_logits(rel_logits, b["rel"], pos_weight=torch.tensor(4.0, device=valid.device), reduction="none")
    lr_ = (lr_ * pair).sum() / pair.sum().clamp(min=1)
    ld = F.cross_entropy(doc_logits, b["doc"])
    return lf + 2.0 * lr_ + 0.3 * ld


# -------------------------------------------------------------------------------------------- evals
def cer(pred, gt):
    if not gt:
        return float(len(pred) > 0)
    prev = list(range(len(gt) + 1))
    for i, pc in enumerate(pred, 1):
        cur = [i] + [0] * len(gt)
        for j, gc in enumerate(gt, 1):
            cur[j] = min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (pc != gc))
        prev = cur
    return prev[-1] / len(gt)


def decode_batch(logits):
    ids = logits.argmax(-1).cpu().numpy()
    out = []
    for row in ids:
        s, prev = [], BLANK
        for k in row:
            if k != BLANK and k != prev:
                s.append(CHARS[k - 1])
            prev = k
        out.append("".join(s))
    return out


@torch.no_grad()
def eval_recognizer(model, device, args):
    model.eval()
    ds = D.LineDataset("val", length=args.eval_size, seed=12345)
    dl = torch.utils.data.DataLoader(ds, batch_size=64, num_workers=args.workers, collate_fn=D.collate_lines)
    errs, chars, exact, n = defaultdict(float), defaultdict(int), defaultdict(int), defaultdict(int)
    for x, mask, steps, _, _, srcs, targets in dl:
        preds = decode_batch(model(x.to(device), mask.to(device)))
        for p, t, s in zip(preds, targets, srcs):
            gt = "".join(CHARS[i - 1] for i in t.tolist())
            errs[s] += cer(p, gt) * len(gt)
            chars[s] += len(gt)
            exact[s] += p == gt
            n[s] += 1
    m = {f"cer/{s}": errs[s] / max(chars[s], 1) for s in errs}
    m.update({f"line_acc/{s}": exact[s] / n[s] for s in n})
    m["cer"] = sum(errs.values()) / max(sum(chars.values()), 1)
    model.train()
    return m, -m["cer"]


@torch.no_grad()
def eval_detector(model, device, args):
    model.eval()
    m = {}
    f1s = []
    inter, denom = {}, {}
    for src in ("synth", "doclaynet", "wildreceipt"):
        pool = D.PagePool("val", "detector", {src: 1.0})
        if src not in pool.weights:
            continue
        tp = fp = fn = 0
        for i in range(args.eval_pages):
            rng = np.random.default_rng([999, i])
            img, quads, ign, regions, _, geo, strength, thermal, blocks = pool.sample(rng)
            img, hm = geometric(img, rng, geo)
            quads, regions = transform_quads(quads, hm), transform_quads(regions, hm)
            img = photometric(img, rng, strength, thermal)
            x, scale = detector_input(np.ascontiguousarray(img))
            out = torch.sigmoid(model(torch.from_numpy(x)[None].to(device)))[0].cpu().numpy()
            prob = out[0]
            if blocks is not None and len(blocks[0]) and out.shape[0] > 1:
                # block-kernel dice on this page (targets built like the training ones, in model coordinates)
                bq = [np.asarray(q) * scale for q in transform_quads([rect_quad(*r) for r in blocks[0]], hm)]
                pq = [np.asarray(q) * scale for q in transform_quads([rect_quad(r[0], r[1], r[0] + lh, r[1] + lh)
                                                                      for r, lh in zip(*blocks)], hm)]
                bk, _ = D.block_targets(out.shape[1], out.shape[2], bq, pq, True, [])
                pb = out[1] > 0.5
                inter[src] = inter.get(src, 0) + float((pb & (bk > 0)).sum())
                denom[src] = denom.get(src, 0) + float(pb.sum() + (bk > 0).sum())
            pred = [q / scale for q, _ in kernels_to_quads(prob)]
            gt = [q for q, g in zip(quads, ign) if not g]
            # predictions overlapping ignore regions/boxes are neither right nor wrong
            ign_q = [q for q, g in zip(quads, ign) if g] + list(regions)
            matches, up, ug = match_quads(pred, gt)
            if ign_q:
                _, up_after, _ = match_quads([pred[i] for i in up], ign_q, 0.1)
                up = up_after
            tp, fp, fn = tp + len(matches), fp + len(up), fn + len(ug)
        p, r = tp / max(tp + fp, 1), tp / max(tp + fn, 1)
        f1 = 2 * p * r / max(p + r, 1e-6)
        m.update({f"precision/{src}": p, f"recall/{src}": r, f"f1/{src}": f1})
        f1s.append(f1)
    m["f1"] = float(np.mean(f1s))
    for src in inter:
        m[f"block_dice/{src}"] = 2 * inter[src] / max(denom[src], 1)
    score = m["f1"]
    if inter:
        m["block_dice"] = float(np.mean([m[f"block_dice/{s_}"] for s_ in inter]))
        score = m["f1"] + 0.25 * m["block_dice"]   # lines first; blocks break ties between line-equal checkpoints
    model.train()
    return m, score


@torch.no_grad()
def eval_orientation(model, device, args):
    model.eval()
    ds = D.OrientationDataset("val", length=max(1, args.eval_size // 8), seed=4242)
    dl = torch.utils.data.DataLoader(ds, batch_size=16, num_workers=args.workers)
    correct = total = 0
    for x, y in dl:
        x, y = x.flatten(0, 1), y.flatten()
        pred = model(x.to(device)).argmax(-1).cpu()
        correct += (pred == y).sum().item()
        total += len(y)
    model.train()
    acc = correct / total
    return {"patch_acc": acc}, acc


@torch.no_grad()
def eval_extractor(model, device, args):
    model.eval()
    ds = D.LayoutDataset("val", length=args.eval_size, seed=777)
    dl = torch.utils.data.DataLoader(ds, batch_size=32, num_workers=args.workers, collate_fn=D.collate_layouts)
    R = len(REL_TYPES)
    by_src = defaultdict(lambda: {"tp": np.zeros(len(FIELD_LABELS)), "fp": np.zeros(len(FIELD_LABELS)), "fn": np.zeros(len(FIELD_LABELS)),
                                  "rel_tp": np.zeros(R), "rel_fp": np.zeros(R), "rel_fn": np.zeros(R), "doc": 0, "n": 0})
    for b in dl:
        t = {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in b.items()}
        fl, rl, dl_ = model(t["char_ids"], t["feats"], t["rel_boxes"], t["valid"])
        pred = fl.argmax(-1).cpu().numpy()
        rel = (torch.sigmoid(rl) > 0.5).cpu().numpy()
        doc = dl_.argmax(-1).cpu().numpy()
        for i, src in enumerate(b["src"]):
            s = by_src[src]
            n = int(b["valid"][i].sum())
            lab = b["labels"][i, :n].numpy()
            pr = pred[i, :n]
            known = lab >= 0
            for c in range(len(FIELD_LABELS)):
                s["tp"][c] += ((pr == c) & (lab == c) & known).sum()
                s["fp"][c] += ((pr == c) & (lab != c) & known).sum()
                s["fn"][c] += ((pr != c) & (lab == c) & known).sum()
            off_diag = ~np.eye(n, dtype=bool)
            for r in range(R):
                if b["rel_known"][i, r] > 0:
                    g = b["rel"][i, r, :n, :n].numpy() > 0.5
                    p = rel[i, r, :n, :n] & off_diag
                    s["rel_tp"][r] += int((p & g).sum())
                    s["rel_fp"][r] += int((p & ~g).sum())
                    s["rel_fn"][r] += int((~p & g).sum())
            s["doc"] += int(doc[i] == b["doc"][i])
            s["n"] += 1
    m, f1s = {}, []
    for src, s in by_src.items():
        present = (s["tp"] + s["fn"]) > 0
        present[0] = False  # "other" is not a field
        f1 = 2 * s["tp"] / np.maximum(2 * s["tp"] + s["fp"] + s["fn"], 1)
        mf1 = float(f1[present].mean()) if present.any() else 0.0
        m[f"field_macro_f1/{src}"] = mf1
        f1s.append(mf1)
        m[f"doc_acc/{src}"] = s["doc"] / s["n"]
        for r, name in enumerate(REL_TYPES):
            if s["rel_tp"][r] + s["rel_fn"][r] > 0:
                m[f"{name}_f1/{src}"] = float(2 * s["rel_tp"][r] / max(2 * s["rel_tp"][r] + s["rel_fp"][r] + s["rel_fn"][r], 1))
    m["field_macro_f1"] = float(np.mean(f1s))
    blocks = [v for k, v in m.items() if k.startswith("same_block_f1/")]
    score = m["field_macro_f1"]
    if blocks:
        m["same_block_f1"] = float(np.mean(blocks))
        score = 0.5 * (score + m["same_block_f1"])   # fields and page structure count equally
    model.train()
    return m, score


# --------------------------------------------------------------------------------------------- loop
BUCKET_GROUP = 4  # batches drawn together and regrouped by length


def identity_collate(samples):
    return samples


def bucketed(loader, size_of, collate, batch, rng):
    """Regroup each mega-batch so that samples of similar length share a batch (less padding)."""
    for mega in loader:
        mega = sorted(mega, key=size_of)
        chunks = [mega[i:i + batch] for i in range(0, len(mega) - batch + 1, batch)]
        for k in rng.permutation(len(chunks)):
            yield collate(chunks[k])


def make_loader(component, args):
    n = args.steps * args.batch
    if component in ("recognizer", "extractor"):   # whole mega-batches, so the stream never ends before the last step
        mega = args.batch * BUCKET_GROUP
        n = -(-n // mega) * mega
    if component == "recognizer":
        ds, collate = D.LineDataset("train", length=n, seed=args.seed), D.collate_lines
    elif component == "detector":
        ds, collate = D.DetectorDataset("train", length=n, seed=args.seed), None
    elif component == "orientation":
        ds, collate = D.OrientationDataset("train", length=n, seed=args.seed), None
    else:
        ds, collate = D.LayoutDataset("train", length=n, seed=args.seed), D.collate_layouts
    sampler = range(args.start_step * args.batch, n)  # sequential indices: resuming continues the stream
    size_of = {"recognizer": lambda s: s[0].shape[-1], "extractor": lambda s: s["feats"].shape[0]}.get(component)
    if size_of is None:
        return torch.utils.data.DataLoader(ds, batch_size=args.batch, sampler=sampler, num_workers=args.workers,
                                           collate_fn=collate, persistent_workers=args.workers > 0,
                                           prefetch_factor=4 if args.workers > 0 else None, drop_last=True)
    mega = torch.utils.data.DataLoader(ds, batch_size=args.batch * BUCKET_GROUP, sampler=sampler, num_workers=args.workers,
                                       collate_fn=identity_collate, persistent_workers=args.workers > 0,
                                       prefetch_factor=2 if args.workers > 0 else None, drop_last=True)
    return bucketed(mega, size_of, collate, args.batch, np.random.default_rng(args.seed + args.start_step))


def step_loss(component, model, batch, device):
    if component == "recognizer":
        x, mask, steps, targets, tlens, _, _ = batch
        logits = model(x.to(device), mask.to(device))
        logp = logits.float().log_softmax(-1).transpose(0, 1)
        return F.ctc_loss(logp, targets.to(device), steps.to(device), tlens.to(device), blank=BLANK, zero_infinity=True)
    if component == "detector":
        x, k, m = (t.to(device) for t in batch)
        return detector_loss(model(x).float(), k, m)
    if component == "orientation":
        x, y = batch
        x, y = x.flatten(0, 1), y.flatten()
        return F.cross_entropy(model(x.to(device)).float(), y.to(device), label_smoothing=0.05)
    b = {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in batch.items()}
    return extractor_loss(model(b["char_ids"], b["feats"], b["rel_boxes"], b["valid"]), b)


def init_from(model, path):
    """Load matching weights from a checkpoint. A tensor that grew (more labels or relations, appended at
    the end) gets its old values in the leading slice; the remainder keeps its fresh initialisation."""
    old = torch.load(path, map_location="cpu", weights_only=False)["model"]
    state = model.state_dict()
    full = partial = 0
    for k, v in state.items():
        o = old.get(k)
        if o is None or o.ndim != v.ndim:
            continue
        if o.shape == v.shape:
            state[k] = o
            full += 1
        elif all(a <= b for a, b in zip(o.shape, v.shape)):
            v = v.clone()
            v[tuple(slice(0, a) for a in o.shape)] = o
            state[k] = v
            partial += 1
    model.load_state_dict(state)
    print(f"initialised from {path}: {full} tensors copied, {partial} grown, {len(state) - full - partial} new")


EVALS = {"recognizer": eval_recognizer, "detector": eval_detector, "orientation": eval_orientation, "extractor": eval_extractor}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("component", choices=list(DEFAULTS))
    ap.add_argument("--steps", type=int)
    ap.add_argument("--batch", type=int)
    ap.add_argument("--lr", type=float)
    ap.add_argument("--eval-every", type=int)
    ap.add_argument("--warmup", type=int)
    ap.add_argument("--eval-size", type=int, default=2000, help="samples per eval (lines / patches / docs)")
    ap.add_argument("--eval-pages", type=int, default=40, help="detector eval pages per source")
    ap.add_argument("--workers", type=int, default=min(6, max(1, (os.cpu_count() or 8) - 3)),
                    help="each worker holds its own PyTorch import (~300 MB); keep modest on 24 GB machines")
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
    ap.add_argument("--amp", action=argparse.BooleanOptionalAction, default=None, help="fp16 mixed precision (default: on for CUDA)")
    ap.add_argument("--max-minutes", type=float, default=None, help="save a checkpoint and stop after this long")
    ap.add_argument("--real-root", help="override data/real (sets TINYDOC_REAL_ROOT)")
    ap.add_argument("--synth-root", help="override data/synth (sets TINYDOC_SYNTH_ROOT)")
    ap.add_argument("--font-root", help="override fonts/ (sets TINYDOC_FONT_ROOT)")
    ap.add_argument("--out", default="runs")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--resume", action="store_true")
    ap.add_argument("--init", help="runs dir whose <component>/best.pt initialises the model (fine-tuning)")
    ap.add_argument("--arch", default="", help="model options saved with the checkpoint, e.g. cnn=separable")
    args = ap.parse_args()
    arch = {k: int(v) if v.isdigit() else v for k, v in (kv.split("=", 1) for kv in args.arch.split(",") if kv)}
    for flag, env in (("real_root", "TINYDOC_REAL_ROOT"), ("synth_root", "TINYDOC_SYNTH_ROOT"), ("font_root", "TINYDOC_FONT_ROOT")):
        if getattr(args, flag):
            os.environ[env] = getattr(args, flag)  # inherited by DataLoader workers
    for k, v in DEFAULTS[args.component].items():
        if getattr(args, k) is None:
            setattr(args, k, v * 0.5 if (k == "lr" and args.init) else v)   # gentler when fine-tuning

    device = torch.device(args.device)
    use_amp = (device.type == "cuda") if args.amp is None else args.amp
    if device.type == "cuda":
        torch.backends.cudnn.benchmark = args.component in ("detector", "orientation")  # fixed input shapes
    scaler = torch.amp.GradScaler("cuda", enabled=use_amp and device.type == "cuda")
    deadline = time.time() + args.max_minutes * 60 if args.max_minutes else None
    out = Path(args.out) / args.component
    out.mkdir(parents=True, exist_ok=True)
    model = build(args.component, **arch).to(device)
    decay = [p for n, p in model.named_parameters() if p.ndim >= 2]
    no_decay = [p for n, p in model.named_parameters() if p.ndim < 2]
    opt = torch.optim.AdamW([{"params": decay, "weight_decay": 0.01}, {"params": no_decay, "weight_decay": 0.0}], lr=args.lr)
    args.start_step, best = 0, -math.inf
    if args.resume and (out / "last.pt").exists():
        ck = torch.load(out / "last.pt", map_location="cpu", weights_only=False)  # our own checkpoint
        model.load_state_dict(ck["model"])
        opt.load_state_dict(ck["optim"])
        args.start_step, best = ck["step"], ck.get("best", -math.inf)
        if "scaler" in ck and scaler.is_enabled():
            scaler.load_state_dict(ck["scaler"])
        print(f"resumed from step {args.start_step}")
    elif args.init:
        init_from(model, Path(args.init) / args.component / "best.pt")
    if args.start_step >= args.steps:
        print("already finished")
        return
    print(f"{args.component}: {count_params(model):,} params, device={device}, workers={args.workers}, "
          f"steps={args.steps}, batch={args.batch}, lr={args.lr}, amp={use_amp}")

    def lr_at(s):
        if s < args.warmup:
            return args.lr * (s + 1) / args.warmup
        t = (s - args.warmup) / max(1, args.steps - args.warmup)
        return args.lr * (0.05 + 0.95 * 0.5 * (1 + math.cos(math.pi * t)))

    log = open(out / "log.jsonl", "a")
    loader = make_loader(args.component, args)
    model.train()
    step, t0, running = args.start_step, time.time(), []
    for batch in loader:
        for g in opt.param_groups:
            g["lr"] = lr_at(step)
        with torch.autocast(device.type, dtype=torch.float16, enabled=use_amp):
            loss = step_loss(args.component, model, batch, device)
        opt.zero_grad(set_to_none=True)
        scaler.scale(loss).backward()
        scaler.unscale_(opt)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
        scaler.step(opt)
        scaler.update()
        running.append(loss.item())
        step += 1
        if step % 50 == 0:
            dt = time.time() - t0
            print(f"step {step:6d}  loss {np.mean(running):.4f}  lr {lr_at(step):.2e}  {50 * args.batch / dt:.0f} samples/s", flush=True)
            log.write(json.dumps({"step": step, "loss": float(np.mean(running)), "lr": lr_at(step)}) + "\n")
            running, t0 = [], time.time()
        out_of_time = deadline is not None and time.time() > deadline
        if step % args.eval_every == 0 or step == args.steps or out_of_time:
            metrics, score = EVALS[args.component](model, device, args)
            is_best = score > best
            best = max(best, score)
            ck = {"model": model.state_dict(), "optim": opt.state_dict(), "step": step, "best": best,
                  "component": args.component, "metrics": metrics, "scaler": scaler.state_dict(), "arch": arch}
            torch.save(ck, out / "last.pt")
            if is_best:
                torch.save({k: v for k, v in ck.items() if k != "optim"}, out / "best.pt")
            print(f"eval @ {step}: " + "  ".join(f"{k}={v:.4f}" for k, v in sorted(metrics.items())) + ("  *best*" if is_best else ""), flush=True)
            log.write(json.dumps({"step": step, "eval": metrics, "best": is_best}) + "\n")
            log.flush()
            t0 = time.time()
        if out_of_time:
            print(f"time budget reached at step {step}; resume with --resume", flush=True)
            break
        if step >= args.steps:
            break


if __name__ == "__main__":
    main()