#!/usr/bin/env python3
"""训练 PHPWind 定长 4 位数字验证码分类器，导出 ONNX。
模型: CNN 骨干 + 4 个位置头，每头 10 类（0-9）。
输入 (1,1,64,160) 灰度/255，输出 (B,4,10) logits。
用法: train_fixed.py <labels.json> <图片目录> <输出.onnx> [val_frac] [epochs]
"""
import json
import os
import random
import sys

import numpy as np
import torch
import torch.nn as nn
from PIL import Image

W, H = 160, 64
NDIGITS = 4
RGB = "rgb" in sys.argv


class 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, lbl = self.items[i]
        mode = "RGB" if RGB else "L"
        im = Image.open(path).convert(mode).resize((W, H), Image.BILINEAR)
        if self.aug:
            im = augment(im)
        a = np.asarray(im, dtype=np.float32) / 255.0
        if RGB:
            x = torch.from_numpy(a).permute(2, 0, 1)  # (C,H,W)
        else:
            x = torch.from_numpy(a).unsqueeze(0)
        t = torch.tensor([int(c) for c in lbl], dtype=torch.long)
        return x, t


def augment(im):
    import random as R
    from PIL import Image
    if R.random() < 0.8:
        im = im.rotate(R.uniform(-6, 6), resample=Image.BILINEAR, fillcolor=255)
    scale = R.uniform(0.92, 1.08)
    tx, ty = R.uniform(-3, 3), R.uniform(-2, 2)
    w, h = im.size
    a, b, c, d, e, f = (scale, 0, -w * scale / 2 + w / 2 + tx,
                        0, scale, -h * scale / 2 + h / 2 + ty)
    im = im.transform((w, h), Image.AFFINE, (a, b, c, d, e, f),
                      resample=Image.BILINEAR, fillcolor=255)
    # 可选轻微噪声（默认关，避免盖掉验证码原有干扰线）
    if os.environ.get("AUG_NOISE", "0") == "1":
        arr = np.asarray(im, dtype=np.float32)
        arr += np.random.normal(0, 5, arr.shape)
        return Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8), im.mode)
    return im


class FixedNet(nn.Module):
    def __init__(self):
        super().__init__()
        in_ch = 3 if RGB else 1
        self.features = nn.Sequential(
            nn.Conv2d(in_ch, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(128, 256, 3, padding=1), nn.BatchNorm2d(256), nn.ReLU(),
            nn.Conv2d(256, 256, 3, padding=1), nn.BatchNorm2d(256), nn.ReLU(),
        )
        # 保留空间位置：20列特征按4个数字位置各池化
        self.heads = nn.ModuleList([nn.Linear(256, 10) for _ in range(NDIGITS)])

    def forward(self, x):
        x = self.features(x)       # (B,256,8,20)
        x = x.mean(dim=2)          # 高度池化 (B,256,20)
        # 20列 -> 4个位置（每位置5列），保持左右顺序
        B, C, T = x.shape
        per = T // NDIGITS
        x = x.view(B, C, NDIGITS, per).mean(dim=3)  # (B,256,4)
        x = x.permute(0, 2, 1)     # (B,4,256)
        return torch.stack([self.heads[p](x[:, p]) for p in range(NDIGITS)], dim=1)  # (B,4,10)


def collate(batch):
    xs = torch.stack([b[0] for b in batch])
    ts = torch.stack([b[1] for b in batch])
    return xs, ts


def main():
    labels = json.load(open(sys.argv[1]))
    imgdir = sys.argv[2]
    out = sys.argv[3] if len(sys.argv) > 3 else "/tmp/sp_captest/captcha_fixed.onnx"
    val_frac = float(sys.argv[4]) if len(sys.argv) > 4 else 0.1
    epochs = int(sys.argv[5]) if len(sys.argv) > 5 else 400
    use_aug = sys.argv[6] != "0" if len(sys.argv) > 6 else True

    all_items = [(f"{imgdir}/{fn}", v["label"]) for fn, v in labels.items()
                 if v.get("label") and len(v["label"]) == NDIGITS]
    random.Random(42).shuffle(all_items)
    n_val = int(len(all_items) * val_frac)
    train_ds = Dataset(all_items[n_val:], aug=use_aug)
    val_ds = Dataset(all_items[:n_val], aug=False)
    tr = torch.utils.data.DataLoader(train_ds, batch_size=16, shuffle=True, collate_fn=collate)
    va = torch.utils.data.DataLoader(val_ds, batch_size=16, shuffle=False, collate_fn=collate)
    print(f"train={len(train_ds)} val={len(val_ds)} total={len(all_items)}")

    device = "cpu"
    model = FixedNet().to(device)
    # 断点续训: 从已有 checkpoint 加载权重
    resume_pt = out + ".pt"
    start_ep = 0
    if os.path.exists(resume_pt):
        model.load_state_dict(torch.load(resume_pt))
        start_ep = int(sys.argv[7]) if len(sys.argv) > 7 else 0
        print(f"从 checkpoint 续训, 起始 epoch={start_ep}")
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    sched = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, factor=0.5, patience=10)
    crit = nn.CrossEntropyLoss()

    best = 0.0
    for ep in range(start_ep, start_ep + epochs):
        model.train()
        tot = 0.0
        for x, t in tr:
            x, t = x.to(device), t.to(device)
            logits = model(x)  # (B,4,10)
            loss = sum(crit(logits[:, p], t[:, p]) for p in range(NDIGITS))
            opt.zero_grad(); loss.backward(); opt.step()
            tot += loss.item() * x.size(0)
        model.eval()
        acc = n = 0
        with torch.no_grad():
            for x, t in va:
                p = model(x).argmax(-1)  # (B,4)
                acc += (p == t).all(dim=1).sum().item()
                n += x.size(0)
        acc /= max(1, n)
        sched.step(acc)
        if acc > best:
            best = acc
            torch.save(model.state_dict(), out + ".pt")
        if ep % 10 == 0 or ep == epochs - 1:
            print(f"epoch {ep} loss={tot/len(train_ds):.3f} val_acc={acc:.2%} best={best:.2%}", flush=True)
    if not os.path.exists(out + ".pt"):  # 从未有更好验证精度时，保存最后权重
        torch.save(model.state_dict(), out + ".pt")
    print(f"DONE best_val_acc={best:.2%}")

    model.load_state_dict(torch.load(out + ".pt"))
    model.eval()
    dummy = torch.randn(1, 3 if RGB else 1, H, W)
    torch.onnx.export(model, dummy, out,
                      input_names=["input"], output_names=["logits"],
                      dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}},
                      opset_version=13, external_data=False)
    print(f"ONNX exported -> {out}")


if __name__ == "__main__":
    main()