# -*- coding: utf-8 -*-
"""关键帧字幕区二值掩码哈希聚类 + 拼字幕长图，供多模态批量认字。
用法：python cluster_subs.py --frames <frames目录> --outdir <长图输出目录> [--thresh 8] [--crop 0.15] [--rows 10]
生成：<outdir>/sheet_001.jpg ... 每张含 rows 句字幕区（原分辨率竖排+时间戳），及 reps.txt 映射。
"""
import argparse, glob, os, re
from PIL import Image, ImageDraw, ImageFont


def crop_subtitle(im, crop):
    w, h = im.size
    sh = max(1, int(h * crop))
    return im.crop((0, h - sh, w, h))


def bin_hash(img, size=24):
    """固定高阈值只抓亮色字幕文字（背景归零）→ 同句字幕跨帧掩码稳定可合并。"""
    g = img.convert("L").resize((size + 1, size), Image.LANCZOS)
    px = list(g.getdata())
    thr = 160  # 字幕文字通常近白(>200)，背景暗(<100)，固定阈值隔离文字
    bits = []
    for y in range(size):
        row = px[y * (size + 1):(y + 1) * (size + 1)]
        for x in range(size):
            bits.append(1 if row[x] > thr else 0)
    return bits


def hamming(a, b):
    return sum(1 for i in range(len(a)) if a[i] != b[i])


def ts_from_name(name):
    m = re.search(r"_([0-9]+\.[0-9]+)s\.jpg$", name)
    return float(m.group(1)) if m else 0.0


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--frames", required=True)
    ap.add_argument("--outdir", required=True)
    ap.add_argument("--thresh", type=int, default=25)
    ap.add_argument("--crop", type=float, default=0.15)
    ap.add_argument("--rows", type=int, default=10, help="每张长图放几句")
    args = ap.parse_args()

    os.makedirs(args.outdir, exist_ok=True)
    files = sorted(glob.glob(os.path.join(args.frames, "*.jpg")))
    print("总帧数: %d" % len(files))

    # 聚类
    clusters = []
    for f in files:
        sub = crop_subtitle(Image.open(f), args.crop)
        hsh = bin_hash(sub)
        best, bestd = None, 999
        for ci, c in enumerate(clusters):
            d = hamming(hsh, c["hash"])
            if d < bestd:
                bestd, best = d, ci
        ts = ts_from_name(os.path.basename(f))
        if best is not None and bestd <= args.thresh:
            clusters[best]["items"].append((f, ts))
        else:
            clusters.append({"hash": hsh, "items": [(f, ts)]})

    reps = []
    for c in clusters:
        it = sorted(c["items"], key=lambda x: x[1])
        mid = it[len(it) // 2]
        reps.append((mid[0], mid[1], len(it)))
    reps.sort(key=lambda x: x[1])
    print("聚类数(代表帧): %d" % len(reps))

    # 拼长图
    try:
        font = ImageFont.load_default()
    except Exception:
        font = None
    sheets = [reps[i:i + args.rows] for i in range(0, len(reps), args.rows)]
    map_lines = []
    for si, grp in enumerate(sheets, 1):
        crops = [crop_subtitle(Image.open(p), args.crop) for p, _, _ in grp]
        wmax = max(c.width for c in crops)
        ch = max(c.height for c in crops)
        label_h = 28
        canvas = Image.new("RGB", (wmax, (ch + label_h) * len(grp)), (255, 255, 255))
        d = ImageDraw.Draw(canvas)
        yy = 0
        for (p, ts, n), c in zip(grp, crops):
            d.rectangle([0, yy, wmax, yy + label_h], fill=(230, 230, 230))
            d.text((6, yy + 6), "%.3fs  (同句帧数=%d)" % (ts, n), fill=(0, 0, 0), font=font)
            canvas.paste(c, (0, yy + label_h))
            yy += ch + label_h
            map_lines.append("%03d\t%.3f" % (si, ts))
        canvas.save(os.path.join(args.outdir, "sheet_%03d.jpg" % si))
    with open(os.path.join(args.outdir, "reps.txt"), "w", encoding="utf-8") as fo:
        fo.write("\n".join(map_lines))
    print("长图张数: %d → %s" % (len(sheets), args.outdir))


if __name__ == "__main__":
    main()