#!/usr/bin/env python
"""Easy4K SFDP Siamese NET AI pipeline (compact)."""
import argparse, csv, hashlib, http.cookiejar, json, math, os, random, re, subprocess, sys, time
import urllib.error, urllib.parse, urllib.request
from functools import lru_cache

import cv2, numpy as np, torch, torch.nn as nn
from torch.utils.data import DataLoader, Dataset, Sampler

class _Tee:
    """把 stdout/stderr 同时写到一份可由外部读取的实时日志。"""
    def __init__(self, *streams):
        self.streams = streams

    def write(self, s):
        for st in self.streams:
            try:
                st.write(s)
            except Exception:
                pass
        return len(s)

    def flush(self):
        for st in self.streams:
            try:
                st.flush()
            except Exception:
                pass


_live = os.environ.get("SFDP_LIVE_LOG")
if _live:
    try:
        _lf = open(_live, "a", encoding="utf-8", buffering=1)
        sys.stdout = _Tee(sys.stdout, _lf)
        sys.stderr = _Tee(sys.stderr, _lf)
    except Exception:
        pass


DATA = r"D:\sfdp-ai\siamese\data"
TILES = os.environ.get("SFDP_TILES") or os.path.join(DATA, "tiles")
LABELS = os.path.join(DATA, "labels")
TEACHER = r"D:\sfdp-ai\bin\SfdpTeacher.exe"
RAW = os.path.join(DATA, "raw")
BILI_RAW = os.path.join(DATA, "raw_bili")
TEST_TILES = os.path.join(DATA, "tiles_test")
TEST_LABELS = os.path.join(DATA, "labels_test")


class DSConv(nn.Module):
    def __init__(self, cin, cout, stride=1):
        super().__init__()
        self.dw = nn.Conv2d(cin, cin, 3, stride, 1, groups=cin, bias=False)
        self.bn1 = nn.BatchNorm2d(cin)
        self.pw = nn.Conv2d(cin, cout, 1, bias=False)
        self.bn2 = nn.BatchNorm2d(cout)
        self.act = nn.ReLU6(inplace=True)

    def forward(self, x):
        x = self.act(self.bn1(self.dw(x)))
        return self.act(self.bn2(self.pw(x)))


class Encoder(nn.Module):
    def __init__(self, emb=512):
        super().__init__()
        self.stem = nn.Sequential(nn.Conv2d(1, 32, 3, 2, 1, bias=False),
                                  nn.BatchNorm2d(32), nn.ReLU6(inplace=True))
        self.blocks = nn.Sequential(
            DSConv(32, 48, 2), DSConv(48, 48), DSConv(48, 96, 2), DSConv(96, 96),
            DSConv(96, 160, 2), DSConv(160, 160), DSConv(160, 224, 2), DSConv(224, 224),
            DSConv(224, 320, 2), DSConv(320, 320))
        self.proj = nn.Conv2d(320, emb, 1, bias=False)
        self.bn = nn.BatchNorm2d(emb)
        self.act = nn.ReLU6(inplace=True)

    def forward(self, x):
        x = self.act(self.bn(self.proj(self.blocks(self.stem(x)))))
        return x.mean(dim=(2, 3))


class SiameseSFDP(nn.Module):
    def __init__(self, emb=512, hidden=1024):
        super().__init__()
        self.encoder = Encoder(emb)
        self.head = nn.Sequential(nn.Linear(emb * 4, hidden), nn.ReLU(inplace=True),
                                  nn.Dropout(0.1), nn.Linear(hidden, 1), nn.Sigmoid())

    def forward(self, a, b):
        ea = self.encoder(a); eb = self.encoder(b)
        return self.head(torch.cat([ea, eb, (ea - eb).abs(), ea * eb], dim=1))


def count_params(m):
    return sum(p.numel() for p in m.parameters())


@lru_cache(maxsize=4096)
def load_gray(p):
    img = cv2.imread(p, cv2.IMREAD_GRAYSCALE)
    if img is None:
        raise FileNotFoundError(p)
    return img


_GRAY_CACHE = {}


def load_gray_cached(path):
    """同一切片会被同片段的多个样本反复读取：按 worker 进程做一个小 LRU，
    省掉重复 PNG 解码（320x320 约 100KB，512 条 ≈ 50MB/worker）。"""
    a = _GRAY_CACHE.get(path)
    if a is not None:
        return a
    try:
        a = load_gray(path)
    except FileNotFoundError:
        # 切片可能正被搬运/重写：短暂重试，实在没有就返回黑图，绝不让整轮训练崩掉
        a = None
        for _ in range(3):
            time.sleep(0.2)
            try:
                a = load_gray(path)
                break
            except FileNotFoundError:
                continue
        if a is None:
            a = np.zeros((320, 320), np.uint8)
    if len(_GRAY_CACHE) >= 512:
        _GRAY_CACHE.clear()
    _GRAY_CACHE[path] = a
    return a


FALLBACK_TILES = os.path.join(DATA, "tiles")


def _resolve_tile(root, rel):
    """Hot tiles (SSD) may still be copying: fall back to the original tiles dir."""
    p = os.path.join(root, rel)
    if os.path.exists(p):
        return p
    fb = os.path.join(FALLBACK_TILES, rel)
    return fb if os.path.exists(fb) else p


class PairDataset(Dataset):
    def __init__(self, manifest, tiles_root, crop=128, augment=False):
        self.items = [json.loads(l) for l in open(manifest, "r", encoding="utf-8") if l.strip()]
        self.tiles_root = tiles_root; self.crop = crop; self.augment = augment

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

    def __getitem__(self, i):
        it = self.items[i]
        a = load_gray_cached(_resolve_tile(self.tiles_root, it["a"]))
        b = load_gray_cached(_resolve_tile(self.tiles_root, it["b"]))
        c = self.crop; h, w = a.shape
        x = min(max(0, it["x"]), w - c); y = min(max(0, it["y"]), h - c)
        pa, pb = a[y:y + c, x:x + c], b[y:y + c, x:x + c]
        if self.augment:
            if np.random.rand() < 0.5:
                pa, pb = pa[:, ::-1], pb[:, ::-1]
            if np.random.rand() < 0.3:
                g = 0.8 + np.random.rand() * 0.45
                pa = np.clip((pa / 255.0) ** g * 255.0, 0, 255).astype(np.uint8)
                pb = np.clip((pb / 255.0) ** g * 255.0, 0, 255).astype(np.uint8)
        ta = torch.from_numpy(np.ascontiguousarray(pa)).float().div_(255.0).unsqueeze(0)
        tb = torch.from_numpy(np.ascontiguousarray(pb)).float().div_(255.0).unsqueeze(0)
        return (ta, tb, torch.tensor([it["soft"]], dtype=torch.float32),
                torch.tensor([it["hard"]], dtype=torch.float32))


# ---------------- bilibili ----------------
BUA = ("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
       "(KHTML, like Gecko) Chrome/124.0 Safari/537.36")
BILI_META = os.path.join(BILI_RAW, "sources_bili.json")
_jar = http.cookiejar.CookieJar()
_op = urllib.request.build_opener(urllib.request.HTTPCookieProcessor(_jar))
_op.addheaders = [("User-Agent", BUA), ("Referer", "https://www.bilibili.com/")]


def bjson(url, params=None):
    full = url + ("?" + urllib.parse.urlencode(params) if params else "")
    last = None
    for i in range(3):
        try:
            with _op.open(full, timeout=25) as r:
                return json.loads(r.read().decode("utf-8"))
        except urllib.error.HTTPError as e:
            if e.code == 412:
                raise
            last = e; time.sleep(2 + i * 2)
        except Exception as e:
            last = e; time.sleep(2 + i * 2)
    raise last


def bview(bvid):
    d = bjson("https://api.bilibili.com/x/web-interface/view", {"bvid": bvid})
    return (d.get("data") or None) if d.get("code") == 0 else None


def bplay(bvid, cid, qn):
    d = bjson("https://api.bilibili.com/x/player/playurl",
              {"bvid": bvid, "cid": cid, "qn": qn, "fnval": 0, "fourk": 0})
    if d.get("code") != 0:
        return None
    dd = d.get("data") or {}; durl = dd.get("durl") or []
    return {"quality": dd.get("quality"), "url": durl[0].get("url")} if durl else None


BILI_MAX_BPS = 6_000_000   # 下载限速 ~48 Mbps（千兆宽带留给家里其他人）


def bdownload(url, dst):
    tmp = dst + ".part"
    req = urllib.request.Request(url, headers={"User-Agent": BUA, "Referer": "https://www.bilibili.com/"})
    for _ in range(3):
        try:
            with urllib.request.urlopen(req, timeout=30) as r, open(tmp, "wb") as f:
                t0 = time.time()
                got = 0
                while True:
                    ch = r.read(262144)
                    if not ch:
                        break
                    f.write(ch)
                    got += len(ch)
                    # 令牌桶式限速：按已下载字节应花的时间补 sleep
                    expected = got / BILI_MAX_BPS
                    elapsed = time.time() - t0
                    if expected > elapsed:
                        time.sleep(expected - elapsed)
            os.replace(tmp, dst); return True
        except Exception:
            if os.path.exists(tmp):
                os.remove(tmp)
            time.sleep(3)
    return False


def cmd_rank(args):
    root = getattr(args, "outdir", None) or BILI_RAW
    meta_path = os.path.join(root, "sources_bili.json")
    os.makedirs(root, exist_ok=True)
    try:
        _op.open("https://www.bilibili.com/", timeout=20).read(2048)
    except Exception:
        pass
    meta = json.load(open(meta_path, "r", encoding="utf-8")) if os.path.exists(meta_path) else []
    done = {m["bvid"] for m in meta}
    seen, cands = set(), []

    def add(items):
        for it in items:
            b = it.get("bvid"); dur = int(it.get("duration") or 0); title = it.get("title", "")
            if not b or b in seen or b in done or not (args.min_sec <= dur <= args.max_sec):
                continue
            if any(x in title for x in ("直播", "回放", "合集", "教程", "reaction", "解说")):
                continue
            seen.add(b)
            cands.append({"bvid": b, "title": title, "author": (it.get("owner") or {}).get("name", ""),
                          "duration": dur})

    for rid in (3, 129, 1, 181, 71, 160, 4):
        try:
            d = bjson("https://api.bilibili.com/x/web-interface/ranking/v2", {"rid": rid})
            add((d.get("data") or {}).get("list") or [])
        except Exception as e:
            print("RANK_FAIL", rid, e, flush=True)
        time.sleep(1.2)
    for pn in range(1, args.pages + 1):
        try:
            d = bjson("https://api.bilibili.com/x/web-interface/popular", {"ps": 20, "pn": pn})
            add((d.get("data") or {}).get("list") or [])
        except Exception as e:
            print("POPULAR_FAIL", pn, e, flush=True)
        time.sleep(1.2)
    for rid in (3, 129, 1, 181):
        for pn in range(1, 6):
            try:
                d = bjson("https://api.bilibili.com/x/web-interface/dynamic/region",
                          {"rid": rid, "ps": 20, "pn": pn})
                add((d.get("data") or {}).get("archives") or [])
            except Exception as e:
                print("REGION_FAIL", rid, pn, e, flush=True); break
            time.sleep(1.0)
    existing = list(done); random.shuffle(existing)
    seeds = ([c["bvid"] for c in cands] + existing[:args.seed_existing])[:args.seeds + args.seed_existing]
    for b in seeds:
        try:
            d = bjson("https://api.bilibili.com/x/web-interface/archive/related", {"bvid": b})
            add(d.get("data") or [])
        except Exception as e:
            print("RELATED_FAIL", b, e, flush=True)
        time.sleep(0.5)
    print("CANDIDATES", len(cands), flush=True)
    saved = 0
    for it in cands:
        if saved >= args.target:
            break
        v = bview(it["bvid"])
        if not v:
            continue
        pu = None
        for qn in (80, 64, 32):
            pu = bplay(it["bvid"], v["cid"], qn)
            if pu:
                break
        if not pu or not pu.get("url"):
            continue
        name = re.sub(r"[^0-9A-Za-z._-]+", "_", it["title"])[:60] or it["bvid"]
        ext = ".flv" if ".flv" in pu["url"] else ".mp4"
        fname = f"{it['bvid']}_{name}{ext}"
        dst = os.path.join(root, fname)
        if os.path.exists(dst):
            continue
        if not bdownload(pu["url"], dst):
            print("DL_FAIL", it["bvid"], flush=True); continue
        meta.append({**it, "file": fname, "cid": v["cid"], "qn": pu["quality"],
                     "license": "research use only, not redistributed"})
        saved += 1
        print("OK", fname, flush=True)
        json.dump(meta, open(meta_path, "w", encoding="utf-8"), ensure_ascii=False, indent=2)
    print("TOTAL_FILES", len(meta), flush=True)


# ---------------- tiles ----------------
PATTERNS = {"1-1": [1], "2-2": [2], "3-3": [3], "mix": [2, 3, 2, 1, 3, 2]}
VEXT = (".webm", ".ogg", ".ogv", ".mp4", ".mov", ".mkv", ".flv")


def clip_id(stem, seg, pat):
    s = re.sub(r"[^0-9A-Za-z._-]+", "_", stem).strip("_")[:40]
    if any(ord(c) > 127 for c in stem) or len(s) < 4:
        s = (s + "_" if s else "mv_") + hashlib.sha1(stem.encode("utf-8")).hexdigest()[:8]
    return f"{s}_s{seg:02d}_{pat}"


def pulldown(frames, rounds, n):
    out, i = [], 0
    while len(out) < n:
        out.extend([frames[i % len(frames)]] * rounds[i % len(rounds)]); i += 1
    return out[:n]


def cmd_tiles(args):
    rng = random.Random(args.seed); os.makedirs(args.out, exist_ok=True)
    idx_path = os.path.join(args.out, "index.jsonl")
    done = set()
    if os.path.exists(idx_path):
        for line in open(idx_path, "r", encoding="utf-8"):
            try:
                done.add(json.loads(line)["clip"])
            except Exception:
                pass
    videos = [n for n in sorted(os.listdir(args.raw)) if os.path.splitext(n)[1].lower() in VEXT]
    print("VIDEOS", len(videos), "ALREADY", len(done), flush=True)
    total = 0
    with open(idx_path, "a", encoding="utf-8") as idx:
        for vi, name in enumerate(videos):
            cap = cv2.VideoCapture(os.path.join(args.raw, name))
            if not cap.isOpened():
                print("OPEN_FAIL", name, flush=True); continue
            nf = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
            fps = float(cap.get(cv2.CAP_PROP_FPS) or 0)
            w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0)
            h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0)
            if nf < args.frames * 3 or w < max(args.tile, args.min_width) or h < args.tile:
                print("SKIP_DIM", name, w, h, nf, flush=True); cap.release(); continue
            max_start = max(0, nf - args.frames * 4)
            for si in range(args.segments):
                start = rng.randint(0, max_start)
                x = rng.randint(0, w - args.tile); y = rng.randint(0, h - args.tile)
                cap.set(cv2.CAP_PROP_POS_FRAMES, start)
                frames = []
                for _ in range(max(4, args.frames // 2)):
                    ok, fr = cap.read()
                    if not ok:
                        break
                    frames.append(cv2.cvtColor(fr[y:y + args.tile, x:x + args.tile], cv2.COLOR_BGR2GRAY))
                if len(frames) < 4:
                    continue
                for pat, rounds in PATTERNS.items():
                    clip = clip_id(os.path.splitext(name)[0], si, pat)
                    if clip in done:
                        continue
                    cdir = os.path.join(args.out, clip); os.makedirs(cdir, exist_ok=True)
                    for k, img in enumerate(pulldown(frames, rounds, args.frames)):
                        cv2.imwrite(os.path.join(cdir, f"f{k:04d}.png"), img)
                    idx.write(json.dumps({"clip": clip, "src": name, "pattern": pat,
                                          "start_frame": start, "x": x, "y": y,
                                          "tile": args.tile, "frames": args.frames, "fps": fps},
                                         ensure_ascii=False) + "\n")
                    idx.flush(); total += 1
            cap.release()
            if (vi + 1) % 10 == 0:
                print("PROGRESS", vi + 1, "/", len(videos), "clips", total, flush=True)
    print("TOTAL_CLIPS", total, flush=True)


def cmd_teacher(args):
    os.makedirs(args.labels, exist_ok=True)
    print("SCAN_START", args.tiles, flush=True)
    dirs = [d for d in sorted(os.listdir(args.tiles)) if os.path.isdir(os.path.join(args.tiles, d))]
    print("SCAN_DIRS", len(dirs), flush=True)
    todo = []
    for i, d in enumerate(dirs):
        if not os.path.exists(os.path.join(args.labels, d + ".csv")):
            todo.append(d)
        if (i + 1) % 2000 == 0:
            print("SCAN", i + 1, "/", len(dirs), flush=True)
    print("DIRS", len(dirs), "TODO", len(todo), flush=True)
    for i, d in enumerate(todo):
        subprocess.run([TEACHER, os.path.join(args.tiles, d), os.path.join(args.labels, d + ".csv"),
                        "--mode", "uhd", "--crop", "128", "--per-pair", "4", "--limit", "0"],
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        if (i + 1) % 10 == 0:
            print("PROGRESS", i + 1, "/", len(todo), flush=True)
    print("DONE", flush=True)


# ---------------- label / train / eval / pack ----------------
def soft_from_flow(flow):
    return max(0.0, min(1.0, 1.0 - flow / 2.0))


def cmd_label(args):
    rng = random.Random(args.seed)
    clip_src = {}
    p = os.path.join(TILES, "index.jsonl")
    if os.path.exists(p):
        for line in open(p, "r", encoding="utf-8"):
            try:
                r = json.loads(line); clip_src[r["clip"]] = r["src"]
            except Exception:
                pass
    rows = []
    for name in sorted(os.listdir(LABELS)):
        if not name.endswith(".csv"):
            continue
        clip = name[:-4]; src = clip_src.get(clip, "unknown")
        with open(os.path.join(LABELS, name), "r", encoding="utf-8", newline="") as f:
            for r in csv.DictReader(f):
                try:
                    mad, flow, var = float(r["mad"]), float(r["flow"]), float(r["var"])
                    stage = r.get("stage", "") or ""
                    dup = int(float(r.get("dup", 0) or 0))
                    x, y = int(float(r["x"])), int(float(r["y"]))
                except (KeyError, ValueError, TypeError):
                    continue   # 被中断写坏的 CSV 行（缺列/None）直接跳过
                if dup == 1:
                    soft = 1.0 if (flow == 0.0 and var == 0.0) else soft_from_flow(flow)
                elif stage == "Flow":
                    if 0.1 <= flow <= 0.4:
                        continue
                    soft = soft_from_flow(flow)
                else:
                    soft = 0.0
                rows.append({"a": os.path.join(clip, r["frame_a"]), "b": os.path.join(clip, r["frame_b"]),
                             "x": x, "y": y, "crop": 128,
                             "soft": round(soft, 4), "hard": 1 if (stage == "Flow" and mad < 1.0) else 0,
                             "stage": stage, "src": src, "clip": clip})
    pos = [r for r in rows if r["soft"] >= 0.9]
    neg = [r for r in rows if r["soft"] <= 0.1]
    mid = [r for r in rows if 0.1 < r["soft"] < 0.9]
    rng.shuffle(pos); rng.shuffle(neg); rng.shuffle(mid)
    neg.sort(key=lambda r: -r["hard"])
    m = int(len(pos) * args.max_neg_ratio)
    neg_keep = neg[:m] if len(neg) > m else neg
    if neg_keep and len(pos) > len(neg_keep):
        pos = pos[:len(neg_keep)]
    mid_keep = mid[:int((len(pos) + len(neg_keep)) * args.mid_ratio)]
    data = pos + neg_keep + mid_keep; rng.shuffle(data)
    val_f = lambda r: sum(ord(c) for c in r["src"]) % 7 == 0
    train = [r for r in data if not val_f(r)]; val = [r for r in data if val_f(r)]
    for path, items in ((os.path.join(DATA, "train.jsonl"), train), (os.path.join(DATA, "val.jsonl"), val)):
        with open(path, "w", encoding="utf-8") as f:
            for it in items:
                f.write(json.dumps(it, ensure_ascii=False) + "\n")
    st = {"raw": len(rows), "kept_pos": len(pos), "kept_neg": len(neg_keep),
          "kept_mid": len(mid_keep), "hard_mined": sum(1 for r in data if r["hard"]),
          "train": len(train), "val": len(val)}
    json.dump(st, open(os.path.join(DATA, "label_stats.json"), "w", encoding="utf-8"),
              ensure_ascii=False, indent=2)
    print(json.dumps(st, ensure_ascii=False), flush=True)


def eval_model(model, loader, device):
    model.eval(); se = n = correct = pe = pn = ne = nn_ = 0
    with torch.no_grad():
        for a, b, y, _ in loader:
            a, b, y = a.to(device), b.to(device), y.to(device)
            p = model(a, b)
            se += torch.abs(p - y).sum().item(); n += y.numel()
            correct += ((p >= 0.5) == (y >= 0.5)).sum().item()
            mm = y >= 0.9
            if mm.any():
                pe += torch.abs(p[mm] - y[mm]).sum().item(); pn += int(mm.sum())
            mm = y <= 0.1
            if mm.any():
                ne += torch.abs(p[mm] - y[mm]).sum().item(); nn_ += int(mm.sum())
    return {"mae": se / max(1, n), "acc": correct / max(1, n),
            "pos_mae": pe / max(1, pn), "neg_mae": ne / max(1, nn_), "n": n}


class ClipGroupSampler(Sampler):
    """按片段分组取样：同一 clip 的切片连续读，机械盘上把随机读变成顺序读。
    片段之间每个 epoch 重新打乱、片段内部也打乱，样本分布不变。"""

    def __init__(self, dataset, shuffle=True, seed=1234):
        self.groups = {}
        for i, it in enumerate(dataset.items):
            self.groups.setdefault(it["clip"], []).append(i)
        self.clips = sorted(self.groups.keys())
        self.shuffle = shuffle
        self.seed = seed
        self.epoch = 0
        self.n = len(dataset)

    def set_epoch(self, e):
        self.epoch = e

    def __len__(self):
        return self.n

    def __iter__(self):
        rng = random.Random(self.seed + self.epoch)
        clips = self.clips[:]
        if self.shuffle:
            rng.shuffle(clips)
        for c in clips:
            idx = self.groups[c]
            if self.shuffle:
                idx = idx[:]
                rng.shuffle(idx)
            for i in idx:
                yield i


def _worker_init(_):
    # 每个 DataLoader worker 只用 1 个 torch 线程：默认多线程在 8 个 worker 里会互相抢核，
    # 实测反而拖慢取数（GPU 在等 batch）。解码/切片本身是单线程活。
    try:
        torch.set_num_threads(1)
    except Exception:
        pass


def cmd_train(args):
    random.seed(args.seed); np.random.seed(args.seed)
    torch.manual_seed(args.seed); torch.cuda.manual_seed_all(args.seed)
    torch.backends.cudnn.benchmark = True
    os.makedirs(args.out, exist_ok=True)
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    print("DEVICE", dev, torch.cuda.get_device_name(0) if dev == "cuda" else "", flush=True)
    tr = PairDataset(os.path.join(DATA, "train.jsonl"), TILES, augment=bool(args.augment))
    va = PairDataset(os.path.join(DATA, "val.jsonl"), TILES)
    print("SAMPLES", len(tr), len(va), flush=True)
    tl = DataLoader(tr, batch_size=args.batch, sampler=ClipGroupSampler(tr, shuffle=True),
                    num_workers=args.workers, worker_init_fn=_worker_init, prefetch_factor=2,
                    pin_memory=(dev == "cuda"), drop_last=True, persistent_workers=args.workers > 0)
    vl = DataLoader(va, batch_size=args.batch, sampler=ClipGroupSampler(va, shuffle=False),
                    num_workers=max(2, args.workers // 2), worker_init_fn=_worker_init,
                    pin_memory=(dev == "cuda"))
    model = SiameseSFDP().to(dev); print("PARAMS", count_params(model), flush=True)
    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
    steps = max(1, len(tl)); total = args.epochs * steps

    def lr_at(step):
        warm = args.warmup * steps
        if step < warm:
            return (step + 1) / max(1, warm)
        return 0.5 * (1 + math.cos(math.pi * (step - warm) / max(1, total - warm)))

    sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_at)
    best = 1e9
    if args.resume and os.path.exists(args.resume):
        ck = torch.load(args.resume, map_location=dev, weights_only=True)
        model.load_state_dict(ck["model"]); opt.load_state_dict(ck["opt"])
        best = ck.get("best", best)
        print("RESUME", args.resume, "prev", ck.get("epoch"), flush=True)
    log = os.path.join(args.out, "train_log.csv")
    if not os.path.exists(log):
        with open(log, "w", newline="", encoding="utf-8") as f:
            csv.writer(f).writerow(["epoch", "lr", "loss", "val_mae", "val_acc", "pos_mae", "neg_mae", "sec"])
    for ep in range(args.epochs):
        if hasattr(tl.sampler, "set_epoch"):
            tl.sampler.set_epoch(ep)
        model.train(); t0 = time.time(); run = cnt = 0
        se = 0
        step_i = 0
        for a, b, y, hard in tl:
            a, b, y, hard = a.to(dev), b.to(dev), y.to(dev), hard.to(dev)
            p = model(a, b); w = 1.0 + (args.hard_weight - 1.0) * hard
            loss = (w * (p - y) ** 2).mean()
            opt.zero_grad(set_to_none=True); loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            opt.step(); sched.step()
            run += loss.item() * a.size(0); cnt += a.size(0)
            se += torch.abs(p - y).sum().item()
            step_i += 1
            if step_i % 50 == 0:
                print(json.dumps({"step": step_i, "steps": len(tl), "ep": ep,
                                  "loss": round(run / max(1, cnt), 5),
                                  "pmae": round(se / max(1, cnt), 5),
                                  "elapsed": round(time.time() - t0, 1)}), flush=True)
        m = eval_model(model, vl, dev); sec = time.time() - t0; lr = opt.param_groups[0]["lr"]
        print(json.dumps({"epoch": ep, "loss": round(run / max(1, cnt), 5), "lr": round(lr, 6),
                          **{k: round(v, 5) for k, v in m.items()}, "sec": round(sec, 1)}), flush=True)
        with open(log, "a", newline="", encoding="utf-8") as f:
            csv.writer(f).writerow([ep, round(lr, 6), round(run / max(1, cnt), 5), round(m["mae"], 5),
                                    round(m["acc"], 5), round(m["pos_mae"], 5), round(m["neg_mae"], 5),
                                    round(sec, 1)])
        ck = {"model": model.state_dict(), "opt": opt.state_dict(), "epoch": ep, "best": best}
        torch.save(ck, os.path.join(args.out, "last.pt"))
        if m["mae"] < best:
            best = m["mae"]; ck["best"] = best
            torch.save(ck, os.path.join(args.out, "best.pt"))
            print("BEST", round(best, 5), "epoch", ep, flush=True)
    print("TRAIN_DONE best_mae", round(best, 5), flush=True)
    ck_best = os.path.join(args.out, "best.pt")
    if not os.path.exists(ck_best):
        ck_best = os.path.join(args.out, "last.pt")   # 本轮没刷新最优时用 last.pt 导出，避免导出报 FileNotFoundError
    print("EXPORT_SRC", ck_best, flush=True)
    export_onnx(ck_best, os.path.join(args.out, "easy4k_sfdp_siamese_fp32.onnx"))


def export_onnx(ckpt, out, opset=17):
    import onnx
    ck = torch.load(ckpt, map_location="cpu", weights_only=True)
    m = SiameseSFDP(); m.load_state_dict(ck["model"]); m.eval()
    a = torch.rand(2, 1, 128, 128); b = torch.rand(2, 1, 128, 128)
    torch.onnx.export(m, (a, b), out, input_names=["frame_a", "frame_b"], output_names=["score"],
                      dynamic_axes={"frame_a": {0: "batch"}, "frame_b": {0: "batch"}, "score": {0: "batch"}},
                      opset_version=opset, do_constant_folding=True)
    o = onnx.load(out); onnx.checker.check_model(o)
    meta = {"model_name": "Easy4K SFDP - Siamese NET AI", "model_version": "0.6.0",
            "task": "frame-duplicate similarity regression",
            "input": "two grayscale 128x128 patches, float32, [0,1]",
            "output": "score in [0,1]; 1 = duplicate/same picture",
            "training": "knowledge distillation from Easy4K SFDP four-stage judge",
            "runtime": "onnxruntime DirectML (NVIDIA / AMD / Intel) or CPU",
            "license": "cc-by-nc-4.0", "params": str(count_params(m))}
    have = {p.key for p in o.metadata_props}
    for k, v in meta.items():
        if k not in have:
            p = o.metadata_props.add(); p.key, p.value = k, v
    onnx.save(o, out)
    print("EXPORTED", out, "size_mb", round(os.path.getsize(out) / 1048576, 2), flush=True)

    # FP16 版：内部权重/算子转半精度，输入输出仍保持 float32，调用方不用改代码
    try:
        from onnxconverter_common import float16 as _f16
        o16 = _f16.convert_float_to_float16(onnx.load(out), keep_io_types=True)
        p16 = out.replace("_fp32.onnx", "_fp16.onnx") if out.endswith("_fp32.onnx") else (out[:-5] + "_fp16.onnx")
        onnx.save(o16, p16)
        print("EXPORTED_FP16", p16, "size_mb", round(os.path.getsize(p16) / 1048576, 2), flush=True)
    except Exception as e:
        print("EXPORTED_FP16_SKIP", repr(e), flush=True)


def cmd_eval(args):
    ds = PairDataset(os.path.join(DATA, "val.jsonl"), TILES)
    dl = DataLoader(ds, batch_size=512, shuffle=False, num_workers=6)
    ck = torch.load(args.ckpt, map_location="cpu", weights_only=True)
    net = SiameseSFDP(); net.load_state_dict(ck["model"]); net.eval()
    S, Y, ST = [], [], []
    with torch.no_grad():
        off = 0
        for a, b, y, _ in dl:
            p = net(a, b).squeeze(1).numpy(); S.append(p); Y.append(y.squeeze(1).numpy())
            ST.extend([ds.items[off + i]["stage"] for i in range(len(p))]); off += len(p)
    s = np.concatenate(S); y = np.concatenate(Y); t = (y >= 0.5).astype(int); p = (s >= 0.5).astype(int)
    tp = int(((p == 1) & (t == 1)).sum()); fp = int(((p == 1) & (t == 0)).sum())
    fn = int(((p == 0) & (t == 1)).sum()); tn = int(((p == 0) & (t == 0)).sum())
    print(json.dumps({"n": len(s), "acc": round((tp + tn) / len(s), 5),
                      "precision": round(tp / max(1, tp + fp), 5),
                      "recall": round(tp / max(1, tp + fn), 5)}), flush=True)
    by = {}
    for st in sorted(set(ST)):
        mm = np.array([x == st for x in ST])
        by[st] = {"n": int(mm.sum()), "mean_score": round(float(s[mm].mean()), 4),
                  "mean_target": round(float(y[mm].mean()), 4)}
    print("BY_STAGE", json.dumps(by, ensure_ascii=False), flush=True)
    for th in (0.4, 0.5, 0.6, 0.7, 0.8):
        pp = (s >= th).astype(int)
        tp_ = int(((pp == 1) & (t == 1)).sum()); fp_ = int(((pp == 1) & (t == 0)).sum())
        fn_ = int(((pp == 0) & (t == 1)).sum())
        pr = tp_ / max(1, tp_ + fp_); rc = tp_ / max(1, tp_ + fn_)
        print(json.dumps({"t": th, "acc": round(float((pp == t).mean()), 5), "precision": round(pr, 5),
                          "recall": round(rc, 5), "f1": round(2 * pr * rc / max(1e-9, pr + rc), 5)}), flush=True)


def cmd_compare(args):
    items = []
    for name in sorted(os.listdir(TEST_LABELS)):
        if not name.endswith(".csv"):
            continue
        clip = name[:-4]
        for r in csv.DictReader(open(os.path.join(TEST_LABELS, name), "r", encoding="utf-8", newline="")):
            flow, var = float(r["flow"]), float(r["var"]); dup = int(float(r.get("dup", 0) or 0))
            stage = r.get("stage", "")
            soft = (1.0 if (flow == 0.0 and var == 0.0) else soft_from_flow(flow)) if dup == 1 else \
                   (soft_from_flow(flow) if stage == "Flow" else 0.0)
            items.append((os.path.join(clip, r["frame_a"]), os.path.join(clip, r["frame_b"]),
                          int(float(r["x"])), int(float(r["y"])), soft))
    print("TEST_ITEMS", len(items), flush=True)
    A, B = [], []
    for pa, pb, x, y, _ in items:
        a = load_gray(os.path.join(TEST_TILES, pa)); b = load_gray(os.path.join(TEST_TILES, pb))
        h, w = a.shape; x = min(max(0, x), w - 128); y = min(max(0, y), h - 128)
        A.append(torch.from_numpy(np.ascontiguousarray(a[y:y + 128, x:x + 128])).float().div_(255.0).unsqueeze(0))
        B.append(torch.from_numpy(np.ascontiguousarray(b[y:y + 128, x:x + 128])).float().div_(255.0).unsqueeze(0))
    A = torch.stack(A); B = torch.stack(B)
    y = torch.tensor([it[4] for it in items], dtype=torch.float32).unsqueeze(1)
    t = (y >= 0.5).int().numpy()
    for path in args.ckpts.split(","):
        path = path.strip()
        if not os.path.exists(path):
            print("MISSING", path, flush=True); continue
        ck = torch.load(path, map_location="cpu", weights_only=True)
        net = SiameseSFDP(); net.load_state_dict(ck["model"]); net.eval()
        outs = []
        with torch.no_grad():
            for i in range(0, len(A), 512):
                outs.append(net(A[i:i + 512], B[i:i + 512]).numpy())
        s = np.concatenate(outs)
        rec = {"ckpt": os.path.basename(os.path.dirname(path)), "epoch": ck.get("epoch"),
               "mae": round(float(np.abs(s - y.numpy()).mean()), 5)}
        for th in (0.5, 0.6, 0.7):
            pp = (s >= th).astype(int)
            tp = int(((pp == 1) & (t == 1)).sum()); fp = int(((pp == 1) & (t == 0)).sum())
            fn = int(((pp == 0) & (t == 1)).sum()); tn = int(((pp == 0) & (t == 0)).sum())
            rec[f"acc@{th}"] = round((tp + tn) / len(s), 4)
        print(rec, flush=True)


def cmd_pack(args):
    import shutil
    card = open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "card.md"), "r", encoding="utf-8").read()
    os.makedirs(args.out, exist_ok=True)
    shutil.copy2(args.onnx, os.path.join(args.out, "model.onnx"))
    _p16 = args.onnx.replace("_fp32.onnx", "_fp16.onnx")
    if os.path.exists(_p16):
        shutil.copy2(_p16, os.path.join(args.out, "model_fp16.onnx"))
    size = round(os.path.getsize(args.onnx) / 1048576, 2)
    best = {}
    if os.path.exists(args.log):
        for row in csv.DictReader(open(args.log, "r", encoding="utf-8")):
            if not best or float(row["val_mae"]) < float(best["val_mae"]):
                best = row
    ev = {}
    if os.path.exists(args.eval):
        for enc in ("utf-8-sig", "utf-16", "utf-16-le"):
            try:
                for line in open(args.eval, "r", encoding=enc, errors="ignore"):
                    line = line.strip()
                    if line.startswith("{") and '"precision"' in line:
                        ev = json.loads(line); break
            except Exception:
                pass
            if ev:
                break
    st = json.load(open(os.path.join(DATA, "label_stats.json"), "r", encoding="utf-8"))
    clips = len([d for d in os.listdir(TILES) if os.path.isdir(os.path.join(TILES, d))])
    pct = lambda x: ("%.2f%%" % (float(x) * 100)) if str(x) not in ("", "n/a") else "n/a"
    vals = {"params": args.params, "opset": 17, "size_mb": size, "epochs": args.epochs, "clips": clips,
            "samples": st.get("train", 0) + st.get("val", 0), "train": st.get("train"), "val": st.get("val"),
            "ratio": "1 : %.2f" % (st.get("kept_neg", 0) / max(1, st.get("kept_pos", 1))),
            "hard": st.get("hard_mined"), "mae": best.get("val_mae", "n/a"),
            "acc": pct(best.get("val_acc", "n/a")), "precision": pct(ev.get("precision", "n/a")),
            "recall": pct(ev.get("recall", "n/a")), "pos_mae": best.get("pos_mae", "n/a"),
            "neg_mae": best.get("neg_mae", "n/a"), "amd_ms": args.amd_ms, "cpu_ms": args.cpu_ms}
    for k, v in vals.items():
        card = card.replace("{" + k + "}", str(v))
    cfg = {"model_name": "Easy4K SFDP - Siamese NET AI", "model_version": "0.6.0",
           "architecture": "siamese depthwise-separable CNN + difference/product fusion head",
           "input": {"names": ["frame_a", "frame_b"], "shape": ["batch", 1, 128, 128],
                     "dtype": "float32", "range": [0, 1], "color": "grayscale"},
           "output": {"name": "score", "shape": ["batch", 1], "meaning": "1 = duplicate"},
           "opset": 17, "providers": ["DmlExecutionProvider", "CPUExecutionProvider"],
           "recommended_threshold": 0.5,
           "metrics": {"val_mae": float(best.get("val_mae", 0)), "val_acc": float(best.get("val_acc", 0)),
                       "precision": float(ev.get("precision", 0)), "recall": float(ev.get("recall", 0))},
           "license": "cc-by-nc-4.0"}
    json.dump(cfg, open(os.path.join(args.out, "config.json"), "w", encoding="utf-8"),
              ensure_ascii=False, indent=2)
    open(os.path.join(args.out, "README.md"), "w", encoding="utf-8").write(card)
    open(os.path.join(args.out, "LICENSE"), "w", encoding="utf-8").write(
        "Easy4K SFDP - Siamese NET AI\nCopyright (c) 2026 Easy4K\nCC BY-NC 4.0\n"
        "https://creativecommons.org/licenses/by-nc/4.0/legalcode\n")
    open(os.path.join(args.out, "requirements.txt"), "w", encoding="utf-8").write(
        "onnxruntime-directml>=1.18\nnumpy>=1.26\n")
    src = os.path.join(RAW, "sources.json")
    if os.path.exists(src):
        shutil.copy2(src, os.path.join(args.out, "sources.json"))
    print("PACKED", args.out, flush=True)


def main():
    ap = argparse.ArgumentParser(); sub = ap.add_subparsers(dest="cmd", required=True)
    p = sub.add_parser("rank")
    p.add_argument("--target", type=int, default=40); p.add_argument("--pages", type=int, default=6)
    p.add_argument("--min-sec", type=int, default=60); p.add_argument("--max-sec", type=int, default=420)
    p.add_argument("--seeds", type=int, default=60); p.add_argument("--seed-existing", type=int, default=300)
    p.add_argument("--outdir", default=BILI_RAW); p.set_defaults(func=cmd_rank)
    p = sub.add_parser("tiles")
    p.add_argument("--raw", required=True); p.add_argument("--out", default=TILES)
    p.add_argument("--tile", type=int, default=320); p.add_argument("--frames", type=int, default=32)
    p.add_argument("--segments", type=int, default=4); p.add_argument("--min-width", type=int, default=640)
    p.add_argument("--seed", type=int, default=1234); p.set_defaults(func=cmd_tiles)
    p = sub.add_parser("teacher")
    p.add_argument("--tiles", default=TILES); p.add_argument("--labels", default=LABELS)
    p.set_defaults(func=cmd_teacher)
    p = sub.add_parser("label")
    p.add_argument("--max-neg-ratio", type=float, default=2.0); p.add_argument("--mid-ratio", type=float, default=0.25)
    p.add_argument("--seed", type=int, default=2026); p.set_defaults(func=cmd_label)
    p = sub.add_parser("train")
    p.add_argument("--out", required=True); p.add_argument("--epochs", type=int, default=8)
    p.add_argument("--batch", type=int, default=256); p.add_argument("--lr", type=float, default=2e-3)
    p.add_argument("--workers", type=int, default=8); p.add_argument("--hard-weight", type=float, default=2.0)
    p.add_argument("--warmup", type=int, default=2); p.add_argument("--augment", type=int, default=0)
    p.add_argument("--seed", type=int, default=2026); p.add_argument("--resume", default="")
    p.set_defaults(func=cmd_train)
    p = sub.add_parser("eval"); p.add_argument("--ckpt", required=True); p.set_defaults(func=cmd_eval)
    p = sub.add_parser("compare"); p.add_argument("--ckpts", required=True); p.set_defaults(func=cmd_compare)
    p = sub.add_parser("pack")
    p.add_argument("--onnx", required=True); p.add_argument("--out", required=True)
    p.add_argument("--log", required=True); p.add_argument("--eval", required=True)
    p.add_argument("--params", default="2602017"); p.add_argument("--epochs", default="8")
    p.add_argument("--amd-ms", default="0.56"); p.add_argument("--cpu-ms", default="0.317")
    p.set_defaults(func=cmd_pack)
    args = ap.parse_args(); args.func(args)


if __name__ == "__main__":
    main()