"""H3 ⑤ 오디오 구동 립싱크 — 공급 오디오 대비 산출물을 관측값으로 뽑는다.

이 발주 전체의 전제는 **DoD 2 (오디오 동일성)** 다 — "립싱크가 잘 되더라" 는 인상은 모델이
공급 오디오를 무시하고 그럴듯한 음성을 새로 만들었을 때도 똑같이 나온다. 그래서 축 3개를 잰다:

  ⓐ **오디오 동일성** — 공급 트랙 ↔ 산출 트랙의 정규화 상호상관(lag 탐색).
     ④ 에서 pass-through(Wan)=0.995 / 재생성(H3)=0.023 이 나온 그 지표다. 판정선이 아니라
     **두 모드가 2자릿수 갈리는** 지표라서 쓴다.
  ⓑ **립싱크 정합** — 입 영역 모션 에너지 ↔ 오디오 RMS 포락선의 상호상관(lag=프레임).
     포즈/랜드마크 추정이 아니라 "입 근처가 언제 움직였나" 대리지표다. 얼굴 위치는
     --mouth-box 로 명시(자동 검출 없음 — 틀린 자동검출이 조용히 결론을 만든다).
  ⓒ **발화 구간** — 산출 오디오의 유성 구간과 중앙 f0(자기상관·옥타브 보정). 화자 귀속용.

★계측이 못 보는 축은 못 본다고 적는다 — 음소 뭉개짐·가사 정확도는 여기서 안 보인다(청취/STT 몫).
"""
import argparse
import json
import subprocess

import numpy as np

SR = 16000


def pcm(path, sr=SR):
    """mono float32 로 디코드. 오디오 스트림이 없으면 빈 배열."""
    p = subprocess.run(["ffmpeg", "-v", "error", "-i", path, "-ac", "1", "-ar", str(sr),
                        "-f", "f32le", "-"], capture_output=True)
    return np.frombuffer(p.stdout, dtype=np.float32).copy()


def envelope(x, hop=400):
    """25ms RMS 포락선."""
    n = len(x) // hop
    if n == 0:
        return np.zeros(0)
    return np.sqrt((x[:n * hop].reshape(n, hop) ** 2).mean(axis=1))


def zmean(v):
    return (v - v.mean()) / (v.std() + 1e-9)


def xcorr(a, b, max_lag, step=1):
    """a 기준 b 를 밀며 최대 상관과 그 lag. 반환 = (best_corr, best_lag, corr_at_lag0).

    ★**겹침 하한 = 짧은 쪽의 50%**. 이걸 안 걸면 lag 를 끝까지 밀었을 때 남는 몇백 ms 짜리
      꼬리끼리 우연히 맞아 `best_corr` 가 신호처럼 솟는다(실측: 무관 트랙 쌍이 겹침 1.9s 에서 0.033,
      공급 트랙 쌍이 겹침 1.81s 에서 0.28 — 둘 다 **판정에 못 쓰는 가장자리 인공물**이었다).
    """
    n0 = min(len(a), len(b))
    if n0 < 16:
        return None, None, None
    lag0 = float(np.corrcoef(zmean(a[:n0]), zmean(b[:n0]))[0, 1])
    best = (-2.0, 0)
    for lag in range(-max_lag, max_lag + 1, step):
        x, y = (a[lag:], b[:len(b) - lag]) if lag >= 0 else (a[:len(a) + lag], b[-lag:])
        n = min(len(x), len(y))
        if n < max(16, n0 // 2):
            continue
        c = float(np.corrcoef(zmean(x[:n]), zmean(y[:n]))[0, 1])
        if c > best[0]:
            best = (c, lag)
    return round(best[0], 3), best[1], round(lag0, 3)


def waveform_identity(src, out, slide=False):
    """공급 ↔ 산출 파형 동일성. ★포락선이 아니라 **샘플 파형**으로 잰다 —
    포락선은 서로 다른 음성이 같은 리듬만 가져도 높게 나와 pass-through 와 재생성을 못 가른다.

    slide=True: 공급이 산출보다 **훨씬 길 때**(ⓐ 30s ref vs 5.2s 산출) 쓴다. ±0.5s lag 로는
    긴 소스의 뒷부분을 아예 안 본다 — 그러면 "안 맞는다" 가 "안 찾아봤다" 와 구분되지 않는다.
    산출 전체 길이를 창으로 삼아 공급 전 구간을 10ms 격자로 훑고, 최대 상관과 그 오프셋을 낸다.
    """
    a, b = pcm(src), pcm(out)
    if len(a) == 0 or len(b) == 0:
        return {"src_samples": int(len(a)), "out_samples": int(len(b)), "verdict": "오디오 스트림 결손"}
    if slide:
        n = len(b)
        if len(a) <= n:
            raise SystemExit("❌ --slide 는 공급이 산출보다 길 때만 의미가 있다")
        best = (-2.0, 0)
        for off in range(0, len(a) - n, SR // 100):     # 10ms 격자
            c = float(np.corrcoef(zmean(b), zmean(a[off:off + n]))[0, 1])
            if c > best[0]:
                best = (c, off)
        lo = max(0, best[1] - SR // 100)
        for off in range(lo, min(len(a) - n, best[1] + SR // 100)):   # 1샘플 정밀 재탐색
            c = float(np.corrcoef(zmean(b), zmean(a[off:off + n]))[0, 1])
            if c > best[0]:
                best = (c, off)
        return {
            "src_duration_s": round(len(a) / SR, 3), "out_duration_s": round(len(b) / SR, 3),
            "mode": "slide (공급 전 구간 탐색)",
            "waveform_best_corr": round(best[0], 3),
            "waveform_best_src_offset_s": round(best[1] / SR, 3),
        }
    corr, lag, lag0 = xcorr(a, b, max_lag=SR // 2)   # ±0.5s
    return {
        "src_duration_s": round(len(a) / SR, 3), "out_duration_s": round(len(b) / SR, 3),
        "waveform_best_corr": corr, "waveform_best_lag_s": round((lag or 0) / SR, 4),
        "waveform_corr_at_lag0": lag0,
    }


def gray(path, w, h):
    raw = subprocess.run(["ffmpeg", "-v", "error", "-i", path, "-vf", f"scale={w}:{h},format=gray",
                          "-f", "rawvideo", "-"], capture_output=True, check=True).stdout
    return np.frombuffer(raw, dtype=np.uint8).reshape(-1, h, w).astype(np.float32)


def probe(path):
    out = subprocess.run(["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries",
                          "stream=width,height,nb_frames,r_frame_rate,duration", "-of", "json",
                          path], capture_output=True, text=True).stdout
    return json.loads(out)["streams"][0]


def mouth_series(path, box):
    """box=(x0,y0,x1,y1) 정규화 좌표. 두 시계열을 낸다.

    ⓐ openness — 그 영역의 **어두운 픽셀 비율**(입이 열리면 구강 내부가 어둡다). 프레임과 1:1 대응이라
       "지금 열려 있나" 를 직접 잰다.
    ⓑ motion — 프레임간 변화량. 열림/닫힘 **전이**에 반응한다.
    두 축이 필요한 이유: 계속 벌린 채 정지한 입은 ⓑ 로는 안 보이고, 천천히 도는 고개는 ⓐ 를 흔든다.
    """
    v = probe(path)
    W, H = int(v["width"]), int(v["height"])
    f = gray(path, W // 4, H // 4)
    h, w = f.shape[1], f.shape[2]
    x0, y0, x1, y1 = box
    crop = f[:, int(y0 * h):int(y1 * h), int(x0 * w):int(x1 * w)]
    # 임계는 그 영역 자체의 분포에서 잡는다(캐릭터 피부톤·조명이 클립마다 달라 절대값 고정이 위험).
    thr = np.percentile(crop, 12)
    openness = (crop < thr).mean(axis=(1, 2))
    motion = np.abs(np.diff(crop, axis=0)).mean(axis=(1, 2))
    return openness, motion, len(f)


def voiced_segments(x, hop=400, thresh_db=-40, min_len=4):
    """유성 구간 + 구간별 중앙 f0(자기상관). f0 는 화자 동일성 판정용으로는 약한 지표다."""
    env = envelope(x, hop)
    db = 20 * np.log10(env + 1e-9)
    on = db > thresh_db
    segs, i = [], 0
    while i < len(on):
        if on[i]:
            j = i
            while j < len(on) and on[j]:
                j += 1
            if j - i >= min_len:
                segs.append((i, j))
            i = j
        else:
            i += 1
    out = []
    for i, j in segs:
        a, b = i * hop, j * hop
        vals = []
        win = 1024
        for k in range(a, b - win, 512):
            fr = x[k:k + win]
            if np.sqrt((fr ** 2).mean()) < 0.02:
                continue
            fr = fr - fr.mean()
            ac = np.correlate(fr, fr, "full")[win - 1:]
            lo, hi = int(SR / 400), int(SR / 70)
            if hi >= len(ac):
                continue
            p = lo + int(np.argmax(ac[lo:hi]))
            if ac[p] > 0.3 * ac[0]:
                vals.append(SR / p)
        out.append({"start_s": round(a / SR, 2), "end_s": round(b / SR, 2),
                    "median_f0_hz": round(float(np.median(vals)), 1) if vals else None,
                    "n_f0_frames": len(vals)})
    return out


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--src-audio", required=True, help="공급한 오디오 레퍼런스")
    ap.add_argument("--out-video", required=True, help="산출 영상")
    ap.add_argument("--mouth-box", default=None,
                    help="입 영역 정규화 좌표 x0,y0,x1,y1 (미지정 시 립싱크 축 생략). "
                         "★자동 얼굴검출을 안 쓰는 이유 = 틀린 검출이 조용히 결론을 만든다")
    ap.add_argument("--extra-box", default=None, help="2인 구도용 2번째 입 영역(화자 귀속)")
    ap.add_argument("--slide", action="store_true",
                    help="공급이 산출보다 훨씬 길 때 — 공급 전 구간을 훑어 최대 상관 지점을 찾는다")
    a = ap.parse_args()

    res = {"src_audio": a.src_audio, "out_video": a.out_video}
    res["video"] = probe(a.out_video)
    res["audio_identity"] = waveform_identity(a.src_audio, a.out_video, slide=a.slide)

    src, out = pcm(a.src_audio), pcm(a.out_video)
    res["src_voiced_segments"] = voiced_segments(src)
    res["out_voiced_segments"] = voiced_segments(out)

    if a.mouth_box:
        fps = eval(res["video"]["r_frame_rate"])  # noqa: S307 — "24/1" 형식
        # ★두 기준 오디오를 **둘 다** 댄다. 산출 오디오와의 정합은 모델 내부의 AV 일관성일 뿐이고,
        #   이 발주가 묻는 것은 **공급 트랙**과의 정합이다. 하나만 재면 그 둘이 섞여 읽힌다.
        env = {"vs_output_audio": envelope(out, hop=int(SR / fps)),
               "vs_src_audio": envelope(src, hop=int(SR / fps))}
        for name, boxs in (("mouth", a.mouth_box), ("mouth_2", a.extra_box)):
            if not boxs:
                continue
            box = tuple(float(v) for v in boxs.split(","))
            openness, motion, nfr = mouth_series(a.out_video, box)
            ent = {"box": box, "frames": int(nfr)}
            # ★lag 탐색폭을 **±3 프레임(±125ms)** 으로 좁힌다. ±1s 로 열어 두면 무관한 트랙끼리도
            #   0.6 을 찍는다(네거티브 컨트롤 실측: 무관 트랙 best_corr 0.604@lag−22 vs lag0 0.063).
            #   짧은 시계열에 lag 49개를 훑으면 우연 최대값이 신호처럼 보인다 — 판정은 lag0 로 한다.
            for ref, e in env.items():
                for lbl, series, off in (("openness", openness, 0), ("motion", motion, 1)):
                    n = min(len(series), len(e) - off)
                    corr, lag, lag0 = xcorr(series[:n], e[off:n + off], max_lag=3)
                    ent[f"{lbl}_{ref}"] = {"best_corr": corr, "best_lag_frames": lag,
                                           "best_lag_s": round((lag or 0) / fps, 3),
                                           "corr_at_lag0": lag0}
            res[f"lipsync_{name}"] = ent
    print(json.dumps(res, ensure_ascii=False, indent=2))
