#!/usr/bin/env python3
"""build-adn027-page.py —— 依据 measure.tsv 生成 ADN-027 各档画质对比页

用法：
  python3 build-adn027-page.py <样本目录> <measure.tsv> <输出目录> [变体id ...]
变体 id 就是文件名里 av1-<id>-<seg>.mp4 的 <id>；不给则自动选一组代表档位。
"""
import json, os, shutil, sys, re
from collections import defaultdict

SAMPLES, TSV, OUT = sys.argv[1], sys.argv[2], sys.argv[3]
WANT = sys.argv[4:]

# 变体说明：id -> (显示名, preset/crf 描述, 编码器)
META = {
    "svtnew-p2-c21": ("SVT-AV1 4.2.0 preset 2 / crf 21", "2", "21", "svt"),
    "svtnew-p3-c21": ("SVT-AV1 4.2.0 preset 3 / crf 21", "3", "21", "svt"),
    "svtnew-p4-c21": ("SVT-AV1 4.2.0 preset 4 / crf 21", "4", "21", "svt"),
    "svtnew-p6-c21": ("SVT-AV1 4.2.0 preset 6 / crf 21", "6", "21", "svt"),
    "svtnew-p4-c19": ("SVT-AV1 4.2.0 preset 4 / crf 19", "4", "19", "svt"),
    "p4c19": ("SVT-AV1 1.7.0 preset 4 / crf 19", "4", "19", "svt-old"),
    "p4c21": ("SVT-AV1 1.7.0 preset 4 / crf 21", "4", "21", "svt-old"),
    "nv-p3-cq26": ("NVENC p3 hq cq26", "p3", "26", "nvenc"),
    "nv-p3-cq30": ("NVENC p3 hq cq30", "p3", "30", "nvenc"),
    "nv-p3-cq33": ("NVENC p3 hq cq33", "p3", "33", "nvenc"),
    "nv-p3-cq36": ("NVENC p3 hq cq36", "p3", "36", "nvenc"),
    "nv-p5uhq-cq26": ("NVENC p5 UHQ cq26", "p5", "26", "nvenc"),
    "nv-p5uhq-cq32": ("NVENC p5 UHQ cq32", "p5", "32", "nvenc"),
    "nv-p7-cq26": ("NVENC p7 hq cq26", "p7", "26", "nvenc"),
    "nv-p7-cq32": ("NVENC p7 hq cq32", "p7", "32", "nvenc"),
}

rows = defaultdict(lambda: {"kbps": [], "ssim": []})
segs = set()
with open(TSV, encoding="utf-8") as f:
    next(f)
    for ln in f:
        p = ln.rstrip("\n").split("\t")
        if len(p) < 5:
            continue
        v, seg, _b, kbps, ssim = p[0], p[1], p[2], p[3], p[4]
        if ssim in ("", "?"):
            continue
        rows[v]["kbps"].append(float(kbps)); rows[v]["ssim"].append(float(ssim))
        segs.add(seg)

def avg(xs): return sum(xs) / len(xs) if xs else 0.0

have = {v: d for v, d in rows.items() if d["kbps"] and all(
    os.path.exists(os.path.join(SAMPLES, f"av1-{v}-{s}.mp4")) for s in sorted(segs))}

if not WANT:
    WANT = [v for v in ["svtnew-p4-c21", "svtnew-p6-c21", "p4c21",
                        "nv-p3-cq30", "nv-p3-cq33", "nv-p5uhq-cq32"] if v in have]
WANT = [v for v in WANT if v in have]
if not WANT:
    sys.exit("没有可用的变体（检查 measure.tsv 与样本文件）")

os.makedirs(OUT, exist_ok=True)
for s in sorted(segs):
    src = os.path.join(SAMPLES, f"src-{s}.mp4")
    if os.path.exists(src):
        shutil.copy2(src, os.path.join(OUT, f"src-{s}.mp4"))

vars_js = []
for v in WANT:
    for s in sorted(segs):
        shutil.copy2(os.path.join(SAMPLES, f"av1-{v}-{s}.mp4"), os.path.join(OUT, f"av1-{v}-{s}.mp4"))
    name, preset, crf, _enc = META.get(v, (v, "?", "?", "?"))
    vars_js.append({"id": v, "name": name, "preset": preset, "crf": crf,
                    "kbps": round(avg(rows[v]["kbps"])), "ssim": f"{avg(rows[v]['ssim']):.5f}"})

tpl = open(os.path.expanduser("~/Downloads/adn027-compare-template.html"), encoding="utf-8").read()
html = (tpl.replace("__VARS__", json.dumps(vars_js, ensure_ascii=False))
           .replace("__SRC_KBPS__", "1996")
           .replace("__SRC_BYTES__", "1812297098"))
# 下拉里显示可读名
html = html.replace("o.textContent = `${v.id}（preset ${v.preset} / crf ${v.crf}）`;",
                    "o.textContent = v.name;")
open(os.path.join(OUT, "index.html"), "w", encoding="utf-8").write(html)

print(f"✅ 对比页已生成：{OUT}/index.html")
print(f"   片段：{', '.join(sorted(segs))}")
print(f"   档位：{len(vars_js)} 个")
print(f"{'档位':<28}{'码率kbps':>10}{'相对源':>9}{'平均SSIM':>11}{'整片预估':>11}")
sb, sk = 1812297098, 1996
for v in vars_js:
    est = sb * v["kbps"] / sk / 1073741824
    print(f"{v['name']:<28}{v['kbps']:>10}{v['kbps']/sk*100:>8.0f}%{v['ssim']:>11}{est:>10.2f}G")
print(f"{'（原片 h264）':<28}{sk:>10}{100:>8.0f}%{'—':>11}{sb/1073741824:>10.2f}G")