#!/usr/bin/env python3
"""训练 PHPWind 纯数字验证码识别小模型（CTC），导出 ONNX。
输入: labels.json {filename: {label, ...}} + 图片目录
模型: 小 CNN -> 高度聚合 -> (T=20, C=11) 时序 -> CTC 解码(10数字+blank)
输出: onnx 模型 + 训练日志

预处理(Go 侧需完全一致):
  灰度 -> resize 到 (160,64) BILINEAR -> float32 /255 -> (1,1,64,160)
"""
import json
import os
import random
import sys

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

W, H = 160, 64
NUM_CLASSES = 11  # 0-9 + blank(10)
DIGITS = "0123456789"


class Dataset:
    def __init__(self, items, aug=False, w=W, h=H):
        self.items = items
        self.aug, self.w, self.h = aug, w, h

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

    def __getitem__(self, i):
        path, lbl = self.items[i]
        im = Image.open(path).convert("L").resize((self.w, self.h), Image.BILINEAR)
        if self.aug:
            im = augment(im)
        a = np.asarray(im, dtype=np.float32) / 255.0
        x = torch.from_numpy(a).unsqueeze(0)  # (1,h,w)
        target = torch.tensor([DIGITS.index(c) for c in lbl], dtype=torch.long)
        return x, target, len(lbl)


def augment(im):
    """纯 PIL 数据增强：旋转/缩放/平移 + 高斯噪声。背景保持白色。"""
    if random.random() < 0.8:
        im = im.rotate(random.uniform(-6, 6), resample=Image.BILINEAR, fillcolor=255)
    scale = random.uniform(0.92, 1.08)
    tx, ty = random.uniform(-3, 3), random.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)
    arr = np.asarray(im, dtype=np.float32)
    if random.random() < 0.6:
        arr += np.random.normal(0, 10, arr.shape)
    return Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8), "L")


class CaptchaNet(nn.Module):
    def __init__(self, num_classes=NUM_CLASSES):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 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(),
        )
        self.head = nn.Conv2d(256, num_classes, 1)

    def forward(self, x):
        x = self.features(x)          # (B,256,8,20)
        x = self.head(x)              # (B,11,8,20)
        x = x.mean(dim=2)             # 高度聚合 (B,11,20)
        x = x.permute(0, 2, 1)        # (B,20,11)
        return x


def ctc_decode(logits):
    """logits: (B,T,C) -> 字符串列表"""
    out = []
    idxs = logits.argmax(dim=-1)  # (B,T)
    for row in idxs:
        prev = -1
        s = []
        for t in row.tolist():
            if t != prev and t != 10:  # 去重+跳过blank
                s.append(DIGITS[t])
            prev = t
        out.append("".join(s))
    return out


def main():
    labels_json = sys.argv[1]
    imgdir = sys.argv[2]
    out_model = sys.argv[3] if len(sys.argv) > 3 else "/tmp/sp_captest/captcha.onnx"
    labels = {k: v.get("label") for k, v in json.load(open(labels_json)).items()}

    all_items = [(os.path.join(imgdir, fn), lbl) for fn, lbl in labels.items() if lbl]
    random.Random(42).shuffle(all_items)
    val_frac = float(sys.argv[4]) if len(sys.argv) > 4 else 0.15
    n_val = int(len(all_items) * val_frac)
    val_items, train_items = all_items[:n_val], all_items[n_val:]
    print(f"train={len(train_items)} val={len(val_items)} total={len(all_items)}")

    train_ds, val_ds = Dataset(train_items, aug=True), Dataset(val_items, aug=False)
    tr = torch.utils.data.DataLoader(train_ds, batch_size=16, shuffle=True,
                                     collate_fn=lambda b: collate(b))
    va = torch.utils.data.DataLoader(val_ds, batch_size=16, shuffle=False,
                                     collate_fn=lambda b: collate(b))

    device = "cpu"
    model = CaptchaNet().to(device)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    sched = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, factor=0.5, patience=8)
    crit = nn.CTCLoss(blank=10, zero_infinity=True)

    best = 0.0
    epochs = int(sys.argv[5]) if len(sys.argv) > 5 else 200
    for epoch in range(epochs):
        model.train()
        tot = 0.0
        for x, target, tl in tr:
            x, target = x.to(device), target.to(device)
            tl = torch.tensor(tl)
            logits = model(x)                     # (B,T,C)
            lp = F.log_softmax(logits, dim=2)
            input_lengths = torch.full((x.size(0),), logits.size(1), dtype=torch.long)
            loss = crit(lp.permute(1, 0, 2), target, input_lengths, tl)
            opt.zero_grad(); loss.backward(); opt.step()
            tot += loss.item() * x.size(0)
        # eval
        model.eval()
        acc = 0.0; n = 0
        if len(va) > 0:
            with torch.no_grad():
                for x, target, tl in va:
                    preds = ctc_decode(model(x.to(device)))
                    for p, (_, lbl) in zip(preds, [(None, lbl) for _, lbl in val_ds.items[n:n + x.size(0)]]):
                        acc += (p == lbl)
                        n += 1
            acc /= max(1, n)
            sched.step(acc)
            if acc > best:
                best = acc
                torch.save(model.state_dict(), out_model + ".pt")
        if epoch % 10 == 0 or epoch == 199:
            print(f"epoch {epoch} loss={tot/len(train_ds):.3f} val_acc={acc:.2%} best={best:.2%}", flush=True)
    if len(va) == 0:  # 无验证集：保存最后权重
        torch.save(model.state_dict(), out_model + ".pt")
    print(f"DONE best_val_acc={best:.2%}")

    # 加载最优权重导出 ONNX
    model.load_state_dict(torch.load(out_model + ".pt"))
    model.eval()
    dummy = torch.randn(1, 1, H, W)
    torch.onnx.export(model, dummy, out_model,
                      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_model}")
    # 导出标签 JSON（校验用）
    with open(out_model + ".labels.json", "w") as f:
        json.dump({"charset": list(DIGITS), "blank": 10, "w": W, "h": H}, f)


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


if __name__ == "__main__":
    main()