"""
GENERAL captcha trainer (FREE / local).
1) Train on dataset/synth (held-out val) — real generalization
2) Light fine-tune on few real samples (≤20× each) — no overfit
3) Export models/captcha.onnx + charsets.json with honest synth val_acc
"""
from __future__ import annotations

import json
import random
import re
import sys
from pathlib import Path

import numpy as np
from PIL import Image, ImageOps, ImageEnhance, ImageFilter

ROOT = Path(__file__).resolve().parent
SYNTH = ROOT / "dataset" / "synth"
LABELED = ROOT / "dataset" / "labeled"
SAMPLES = ROOT / "samples"
MODELS = ROOT / "models"

CHARS = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
CHAR2IDX = {c: i for i, c in enumerate(CHARS)}
NUM_CLASSES = len(CHARS)
CAPTCHA_LEN = 6
H, W = 64, 192


def parse_label(name: str) -> str:
    stem = Path(name).stem
    m = re.match(r"^(.+?)_(?:real)?\d+$", stem, re.I)
    lab = m.group(1) if m else stem
    return re.sub(r"[^A-Za-z0-9]", "", lab)


def list_folder(folder: Path):
    items = []
    if not folder.is_dir():
        return items
    for p in folder.glob("*.*"):
        if p.suffix.lower() not in {".png", ".jpg", ".jpeg", ".webp"}:
            continue
        lab = parse_label(p.name)
        if p.name.startswith("sarathi_h5Ow7y") or p.stem.startswith("h5Ow7y"):
            lab = "h5Ow7y"
        if p.stem.startswith("Kt89i0") or p.stem == "frr":
            lab = "Kt89i0"
        if len(lab) == CAPTCHA_LEN and all(c in CHAR2IDX for c in lab):
            items.append((p, lab))
    return items


def prep(img: Image.Image) -> np.ndarray:
    rgb = np.asarray(img.convert("RGB"), dtype=np.int16)
    r, g, b = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]
    blue = ((b > r + 25) & (b > g + 15) & (b > 90)).mean(axis=0)
    w = rgb.shape[1]
    cut = w
    for x in range(w - 1, w // 2, -1):
        if blue[x] >= 0.25:
            cut = x
        elif cut < w and blue[x] < 0.05:
            break
    if cut < w - 8:
        rgb = rgb[:, : max(w // 2, cut - 2), :]
    r, g, b = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]
    green = (g > r + 8) & (g > b + 8) & (g > 70)
    gray = (0.299 * r + 0.587 * g + 0.114 * b).astype(np.float32)
    ink = np.where(green, 255.0 - np.clip(g.astype(np.float32) * 1.1, 0, 255), gray)
    im = ImageOps.autocontrast(Image.fromarray(ink.astype(np.uint8), mode="L"))
    im = im.resize((W, H), Image.Resampling.BILINEAR)
    return np.asarray(im, dtype=np.float32) / 255.0


def augment(img: Image.Image) -> Image.Image:
    img = img.convert("RGB")
    if random.random() < 0.45:
        img = ImageEnhance.Contrast(img).enhance(random.uniform(0.85, 1.35))
    if random.random() < 0.35:
        img = ImageEnhance.Brightness(img).enhance(random.uniform(0.9, 1.12))
    if random.random() < 0.4:
        dx, dy = random.randint(-3, 3), random.randint(-2, 2)
        img = img.transform(img.size, Image.AFFINE, (1, 0, dx, 0, 1, dy), fillcolor=(220, 225, 220))
    if random.random() < 0.25:
        img = img.filter(ImageFilter.SMOOTH)
    return img


def main():
    try:
        sys.stdout.reconfigure(line_buffering=True)
    except Exception:
        pass

    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    from torch.utils.data import Dataset, DataLoader

    synth = list_folder(SYNTH)
    if len(synth) < 500:
        print("Need dataset/synth — run: python generate_dataset.py --clear --count 8000")
        raise SystemExit(2)

    random.shuffle(synth)
    n_val = max(200, int(len(synth) * 0.1))
    val_s, train_s = synth[:n_val], synth[n_val:]
    print(f"FREE local train | Synth train={len(train_s)} val={len(val_s)}", flush=True)

    reals = []
    for folder in (SAMPLES, LABELED):
        for p, lab in list_folder(folder):
            if folder == SAMPLES or "real" in p.stem.lower():
                reals.append((p, lab))
    # dedupe by path
    seen = set()
    real_u = []
    for p, lab in reals:
        if str(p) in seen:
            continue
        seen.add(str(p))
        real_u.append((p, lab))
    print(f"Real samples for light FT: {len(real_u)}", flush=True)

    class CapDS(Dataset):
        def __init__(self, items, aug=False):
            self.items = items
            self.aug = aug

        def __len__(self):
            return len(self.items)

        def __getitem__(self, i):
            path, lab = self.items[i]
            img = Image.open(path).convert("RGB")
            if self.aug:
                img = augment(img)
            x = torch.from_numpy(prep(img)[None, ...])
            y = torch.tensor([CHAR2IDX[c] for c in lab], dtype=torch.long)
            return x, y

    class CaptchaNet(nn.Module):
        def __init__(self):
            super().__init__()
            self.cnn = nn.Sequential(
                nn.Conv2d(1, 64, 3, padding=1),
                nn.BatchNorm2d(64),
                nn.ReLU(True),
                nn.MaxPool2d(2),
                nn.Conv2d(64, 128, 3, padding=1),
                nn.BatchNorm2d(128),
                nn.ReLU(True),
                nn.MaxPool2d(2),
                nn.Conv2d(128, 256, 3, padding=1),
                nn.BatchNorm2d(256),
                nn.ReLU(True),
                nn.MaxPool2d((2, 1)),
                nn.Conv2d(256, 256, 3, padding=1),
                nn.BatchNorm2d(256),
                nn.ReLU(True),
            )
            self.pool = nn.AdaptiveAvgPool2d((1, CAPTCHA_LEN))
            self.drop = nn.Dropout(0.25)
            self.classifier = nn.Sequential(
                nn.Linear(256, 256),
                nn.ReLU(True),
                nn.Linear(256, NUM_CLASSES),
            )

        def forward(self, x):
            f = self.pool(self.cnn(x)).squeeze(2).permute(0, 2, 1)
            f = self.drop(f)
            return self.classifier(f)

    model = CaptchaNet()
    opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
    epochs = 35
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)

    train_loader = DataLoader(CapDS(train_s, True), batch_size=64, shuffle=True)
    val_loader = DataLoader(CapDS(val_s, False), batch_size=64)

    def run_epoch(loader, train: bool):
        model.train(train)
        tot = n = ok = ch_ok = ch_n = 0
        for x, y in loader:
            logits = model(x)
            loss = sum(F.cross_entropy(logits[:, i], y[:, i]) for i in range(6)) / 6
            if train:
                opt.zero_grad()
                loss.backward()
                nn.utils.clip_grad_norm_(model.parameters(), 5)
                opt.step()
            pred = logits.argmax(-1)
            tot += float(loss.detach().item()) * x.size(0)
            n += x.size(0)
            ok += int((pred == y).all(1).sum())
            ch_ok += int((pred == y).sum())
            ch_n += y.numel()
        return tot / max(1, n), ok / max(1, n), ch_ok / max(1, ch_n)

    MODELS.mkdir(parents=True, exist_ok=True)
    best_path = MODELS / "captcha_best.pt"
    best_acc = -1.0

    print("=== Phase 1: general synth training ===", flush=True)
    for ep in range(1, epochs + 1):
        tr = run_epoch(train_loader, True)
        va = run_epoch(val_loader, False)
        sched.step()
        print(
            f"Epoch {ep:02d}/{epochs} train_loss={tr[0]:.3f} seq={tr[1]:.3f} char={tr[2]:.3f}  "
            f"val_loss={va[0]:.3f} seq={va[1]:.3f} char={va[2]:.3f}",
            flush=True,
        )
        if va[1] >= best_acc:
            best_acc = va[1]
            torch.save(
                {"model": model.state_dict(), "acc": best_acc, "height": H, "width": W, "len": 6},
                best_path,
            )
            print("  saved", best_path, f"synth_val_seq={best_acc:.3f}", flush=True)

    # Phase 2: light real fine-tune
    if real_u:
        print("=== Phase 2: light real fine-tune (≤20× each) ===", flush=True)
        ft_items = []
        for p, lab in real_u:
            for _ in range(20):
                ft_items.append((p, lab))
        # keep some synth so we don't forget
        ft_items += random.sample(train_s, min(800, len(train_s)))
        random.shuffle(ft_items)
        ft_loader = DataLoader(CapDS(ft_items, True), batch_size=32, shuffle=True)
        opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
        ckpt = torch.load(best_path, map_location="cpu", weights_only=False)
        model.load_state_dict(ckpt["model"])
        for ep in range(1, 7):
            tr = run_epoch(ft_loader, True)
            va = run_epoch(val_loader, False)
            print(
                f"FT {ep}/6 train_seq={tr[1]:.3f} char={tr[2]:.3f}  synth_val_seq={va[1]:.3f} char={va[2]:.3f}",
                flush=True,
            )
            # keep best by synth val (generalization)
            if va[1] >= best_acc * 0.95:  # allow tiny drop
                if va[1] > best_acc:
                    best_acc = va[1]
                torch.save(
                    {"model": model.state_dict(), "acc": best_acc, "height": H, "width": W, "len": 6},
                    best_path,
                )
                print("  saved after FT", f"synth_val_seq={va[1]:.3f}", flush=True)
            model.eval()
            with torch.no_grad():
                for p, lab in real_u[:8]:
                    x = torch.from_numpy(prep(Image.open(p).convert("RGB"))[None, None])
                    pred = "".join(CHARS[i] for i in model(x).argmax(-1)[0].tolist())
                    print(f"    REAL {p.name}: {pred} vs {lab}", flush=True)

    # Export
    ckpt = torch.load(best_path, map_location="cpu", weights_only=False)
    model.load_state_dict(ckpt["model"])
    model.eval()
    onnx_path = MODELS / "captcha.onnx"
    for p in (onnx_path, MODELS / "captcha.onnx.data"):
        if p.is_file():
            p.unlink()
    torch.onnx.export(
        model,
        torch.randn(1, 1, H, W),
        str(onnx_path),
        input_names=["image"],
        output_names=["logits"],
        dynamic_axes={"image": {0: "batch"}, "logits": {0: "batch"}},
        opset_version=17,
        dynamo=False,
    )
    meta = {
        "charset": list(CHARS),
        "height": H,
        "width": W,
        "len": 6,
        "engine": "portal_multihead",
        "arch": "portal_multihead_cols",
        "val_acc": float(best_acc),
        "note": "synth_val_sequence_accuracy_free_local",
        "preprocess": "sarathi_green",
    }
    (MODELS / "charsets.json").write_text(json.dumps(meta, indent=2), encoding="utf-8")
    print("Exported", onnx_path, "synth_val_seq=", round(best_acc * 100, 2), "%", flush=True)
    if best_acc < 0.45:
        print("WARNING: val_acc low — engine may fall back to builtin until better.", flush=True)
    else:
        print("OK to serve as portal_multihead (FREE local).", flush=True)


if __name__ == "__main__":
    main()