# -*- coding: utf-8 -*-
"""lapian.py - 拉片自动分析: 视频 -> 切点/关键帧/色板/音频节拍/转写/overview.png/shots.json/csv/blueprint_auto.md
用法: python lapian.py input.mp4 [outdir] [--whisper small] [--lang zh]
依赖: scenedetect[opencv] librosa faster-whisper matplotlib scikit-learn, ffmpeg 在 PATH
"""
import sys, os, json, csv, re, subprocess, argparse
import numpy as np, cv2

def tc(sec, fps):
    f = int(round(sec * fps)); s, fr = divmod(f, int(round(fps)))
    return "%02d:%02d:%02d:%02d" % (s // 3600, s // 60 % 60, s % 60, fr)

def detect_cuts(video, fps):
    from scenedetect import open_video, SceneManager
    from scenedetect.detectors import ContentDetector, AdaptiveDetector
    res = {}
    for name, det in [("content", ContentDetector(threshold=27)), ("adaptive", AdaptiveDetector())]:
        v = open_video(video); sm = SceneManager(); sm.add_detector(det); sm.detect_scenes(v)
        res[name] = [s[0].get_frames() for s in sm.get_scene_list()][1:]
    # ffmpeg scdet
    p = subprocess.run(["ffmpeg", "-hide_banner", "-i", video, "-vf", "scdet=threshold=10", "-an", "-f", "null", "-"],
                       capture_output=True, text=True, encoding="utf-8", errors="ignore")
    res["scdet"] = [int(round(float(t) * fps)) for t in re.findall(r"lavfi\.scd\.time:\s*([\d.]+)", p.stderr)]
    # reconcile: cluster within 3 frames, keep if >=2 detectors agree (or single strong content hit)
    allc = sorted((f, k) for k, l in res.items() for f in l)
    clusters = []
    for f, k in allc:
        if clusters and f - clusters[-1][-1][0] <= 3: clusters[-1].append((f, k))
        else: clusters.append([(f, k)])
    cuts = []
    for c in clusters:
        ks = {k for _, k in c}
        if len(ks) >= 2 or "content" in ks:
            fr = [f for f, k in c if k == "content"] or [f for f, _ in c]
            cuts.append(min(fr))
    return res, cuts

def palette(img, k=5):
    from sklearn.cluster import KMeans
    px = cv2.cvtColor(cv2.resize(img, (120, 120)), cv2.COLOR_BGR2RGB).reshape(-1, 3).astype(float)
    km = KMeans(k, n_init=4, random_state=0).fit(px)
    order = np.argsort(-np.bincount(km.labels_, minlength=k))
    return ["#%02x%02x%02x" % tuple(int(v) for v in km.cluster_centers_[i]) for i in order]

def frame_stats(img):
    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB).astype(float); hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(float)
    b, g, r = [img[..., i].mean() for i in range(3)]
    return dict(luma=round(lab[..., 0].mean() / 2.55, 1), contrast=round(lab[..., 0].std() / 2.55, 1),
                sat=round(hsv[..., 1].mean() / 2.55, 1), warmth_RminusB=round(r - b, 1),
                tint_a=round(lab[..., 1].mean() - 128, 1), tint_b=round(lab[..., 2].mean() - 128, 1))

def audio(video, out, cuts_sec, model_size, lang):
    import librosa
    wav = os.path.join(out, "audio.wav")
    subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", video, "-ac", "1", "-ar", "22050", wav], check=True)
    y, sr = librosa.load(wav, sr=22050)
    tempo, beats = librosa.beat.beat_track(y=y, sr=sr)
    bt = librosa.frames_to_time(beats, sr=sr).tolist()
    on = librosa.onset.onset_detect(y=y, sr=sr, units="time").tolist()
    rms = librosa.feature.rms(y=y)[0]; rt = librosa.times_like(rms, sr=sr)
    # phrases: 4-beat groups / downbeat guess
    phrases = [round(bt[i], 3) for i in range(0, len(bt), 8)]
    align = []
    for c in cuts_sec:
        nb = min(bt, key=lambda b: abs(b - c)) if bt else None
        no = min(on, key=lambda b: abs(b - c)) if on else None
        align.append(dict(cut=round(c, 3), nearest_beat=round(nb, 3) if nb is not None else None,
                          beat_off_ms=round((c - nb) * 1000) if nb is not None else None,
                          nearest_onset=round(no, 3) if no is not None else None,
                          onset_off_ms=round((c - no) * 1000) if no is not None else None))
    segs = []
    try:
        from faster_whisper import WhisperModel
        m = WhisperModel(model_size, device="cpu", compute_type="int8")
        y16 = librosa.resample(y, orig_sr=sr, target_sr=16000).astype(np.float32)  # 绕过 PyAV
        ss, _ = m.transcribe(y16, language=lang, word_timestamps=False, vad_filter=False)
        segs = [dict(start=round(s.start, 2), end=round(s.end, 2), text=s.text.strip()) for s in ss]
    except Exception as e:
        segs = [dict(error=str(e))]
    return dict(tempo=float(np.atleast_1d(tempo)[0]), beats=[round(b, 3) for b in bt], onsets=[round(o, 3) for o in on],
                phrases=phrases, cut_alignment=align, transcript=segs), (y, sr, rt, rms)

def main():
    ap = argparse.ArgumentParser(); ap.add_argument("video"); ap.add_argument("out", nargs="?")
    ap.add_argument("--whisper", default="small"); ap.add_argument("--lang", default="zh")
    a = ap.parse_args(); video = os.path.abspath(a.video)
    out = os.path.abspath(a.out or os.path.join(os.path.dirname(video), "lapian_out")); kf = os.path.join(out, "keyframes")
    os.makedirs(kf, exist_ok=True)
    cap = cv2.VideoCapture(video); fps = cap.get(cv2.CAP_PROP_FPS); n = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    W, H = int(cap.get(3)), int(cap.get(4))
    dets, cuts = detect_cuts(video, fps)
    bounds = [0] + cuts + [n]
    shots = []
    for i in range(len(bounds) - 1):
        s, e = bounds[i], bounds[i + 1] - 1
        sh = dict(shot=i + 1, start_frame=s, end_frame=e, frames=e - s + 1, start_tc=tc(s / fps, fps), end_tc=tc(e / fps, fps),
                  start_s=round(s / fps, 3), end_s=round((e + 1) / fps, 3), dur_s=round((e - s + 1) / fps, 3), keyframes={})
        for tag, f in [("first", s), ("mid", (s + e) // 2), ("last", e)]:
            cap.set(cv2.CAP_PROP_POS_FRAMES, f); ok, img = cap.read()
            if not ok: continue
            p = os.path.join(kf, "s%02d_%s_f%04d.jpg" % (i + 1, tag, f)); cv2.imencode(".jpg", img)[1].tofile(p)
            sh["keyframes"][tag] = os.path.basename(p)
            if tag == "mid": sh["palette"] = palette(img); sh.update(frame_stats(img)); mid = img
        shots.append(sh)
    au, (y, sr, rt, rms) = audio(video, out, [c / fps for c in cuts], a.whisper, a.lang)
    meta = dict(video=video, fps=fps, frames=n, size=[W, H], duration=round(n / fps, 3), detectors=dets, cuts=cuts)
    json.dump(dict(meta=meta, shots=shots, audio=au), open(os.path.join(out, "shots.json"), "w", encoding="utf-8"), ensure_ascii=False, indent=1)
    keys = ["shot", "start_frame", "end_frame", "frames", "start_tc", "end_tc", "dur_s", "luma", "contrast", "sat", "warmth_RminusB", "tint_a", "tint_b"]
    with open(os.path.join(out, "shots.csv"), "w", newline="", encoding="utf-8-sig") as f:
        w = csv.writer(f); w.writerow(keys + ["palette"])
        for s in shots: w.writerow([s.get(k) for k in keys] + [" ".join(s.get("palette", []))])
    overview(out, kf, shots, meta, au, y, sr, rt, rms)
    with open(os.path.join(out, "blueprint_auto.md"), "w", encoding="utf-8") as f:
        f.write("# 自动拉片骨架\n\n%s | %dx%d | %.3ffps | %d帧 | %d镜 | BPM≈%.1f\n\n" % (os.path.basename(video), W, H, fps, n, len(shots), au["tempo"]))
        f.write("|镜|起止帧|TC|时长|色板|亮度/对比/饱和|\n|-|-|-|-|-|-|\n")
        for s in shots: f.write("|%d|%d-%d|%s|%.2fs|%s|%s/%s/%s|\n" % (s["shot"], s["start_frame"], s["end_frame"], s["start_tc"], s["dur_s"], " ".join(s.get("palette", [])), s.get("luma"), s.get("contrast"), s.get("sat")))
        f.write("\n## 转写\n" + "\n".join("- [%s-%s] %s" % (t.get("start"), t.get("end"), t.get("text", t)) for t in au["transcript"]))
    print(json.dumps(dict(cuts=[(c, tc(c / fps, fps)) for c in cuts], detectors=dets, tempo=au["tempo"], shots=len(shots)), ensure_ascii=False))

def overview(out, kf, shots, meta, au, y, sr, rt, rms):
    import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
    plt.rcParams["font.sans-serif"] = ["Microsoft YaHei", "SimHei", "Noto Sans CJK SC", "DejaVu Sans"]; plt.rcParams["axes.unicode_minus"] = False
    n = len(shots); cols = min(n, 6); rows = -(-n // cols)
    fig = plt.figure(figsize=(cols * 3, rows * 3.6 + 3))
    gs = fig.add_gridspec(rows * 2 + 1, cols, height_ratios=[3, 0.5] * rows + [2.5])
    for i, s in enumerate(shots):
        r, c = divmod(i, cols); ax = fig.add_subplot(gs[r * 2, c])
        img = cv2.imdecode(np.fromfile(os.path.join(kf, s["keyframes"]["mid"]), np.uint8), 1)
        ax.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)); ax.axis("off")
        ax.set_title("S%02d %s %.2fs" % (s["shot"], s["start_tc"], s["dur_s"]), fontsize=9)
        pa = fig.add_subplot(gs[r * 2 + 1, c]); pa.axis("off")
        for j, h in enumerate(s.get("palette", [])):
            pa.add_patch(plt.Rectangle((j, 0), 1, 1, color=h)); pa.text(j + .5, .5, h, ha="center", va="center", fontsize=5, color="w" if sum(int(h[k:k + 2], 16) for k in (1, 3, 5)) < 380 else "k")
        pa.set_xlim(0, 5); pa.set_ylim(0, 1)
    ax = fig.add_subplot(gs[-1, :]); t = np.arange(len(y)) / sr
    ax.plot(t[::20], y[::20], lw=.3, color="#888"); ax.plot(rt, rms / rms.max() * np.abs(y).max(), color="#c33", lw=.8)
    for b in au["beats"]: ax.axvline(b, color="#39f", lw=.4, alpha=.5)
    for s in shots[1:]: ax.axvline(s["start_s"], color="k", lw=1.2, ls="--")
    for s in shots: ax.text((s["start_s"] + s["end_s"]) / 2, ax.get_ylim()[1] * .9, "S%d" % s["shot"], ha="center", fontsize=7)
    ax.set_xlim(0, meta["duration"]); ax.set_xlabel("秒 (黑虚线=切点, 蓝线=节拍, 红=RMS) BPM≈%.1f" % au["tempo"])
    fig.tight_layout(); fig.savefig(os.path.join(out, "overview.png"), dpi=110); plt.close(fig)

if __name__ == "__main__": main()