#!/usr/bin/env python3
"""dubbing_2b_ml_s2full12_v4_cv3 断网自包含验证: 完整加载 + en/es/ja/ko/zh 各 dub 1 条 + whisper 回测 en(WER)/ko(CER)/zh(CER)。

用法: HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 MODELSCOPE_OFFLINE=1 \
      python verify_export.py [--model-dir DIR] [--langs en,es,ja,ko,zh]

验证点:
  - 断网完整加载 (HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE / MODELSCOPE_OFFLINE);
  - 每语种出一条非静音 wav, 时长合理 (2-15s, 超出仅 WARN);
  - 打印每语种判定源语/译文 (ja 额外打印 kata 化后文本) 与参考译文供人工核对;
  - input_zh.wav 为英语源 (en->zh 方向, 覆盖 en 源路径), 其余为中文源;
  - en 用 whisper-large-v3 回测算 WER; ko/zh 回测算字级 CER。
"""
import argparse
import json
import os
import sys

os.environ.setdefault('HF_HUB_OFFLINE', '1')
os.environ.setdefault('TRANSFORMERS_OFFLINE', '1')
os.environ.setdefault('MODELSCOPE_OFFLINE', '1')

_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _HERE)

WHISPER = os.environ.get('DUBBING_WHISPER',
                         '/mnt/pfs-write/users/lusheng/Index-MT/models/whisper-large-v3')


def cer(ref, hyp):
    """字级 CER (去空格与标点)。"""
    import string
    drop = set(string.punctuation) | set('。，、！？；：“”‘’（）《》…—')
    norm = lambda s: [c for c in s.replace(' ', '') if c not in drop]
    r, h = norm(ref), norm(hyp)
    if not r:
        return float('nan')
    prev = list(range(len(h) + 1))
    for i, rc in enumerate(r, 1):
        cur = [i]
        for j, hc in enumerate(h, 1):
            cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (rc != hc)))
        prev = cur
    return prev[-1] / len(r)


def wer(ref, hyp):
    import string
    norm = lambda s: [w.strip(string.punctuation) for w in s.lower().split()
                      if w.strip(string.punctuation)]
    r, h = norm(ref), norm(hyp)
    if not r:
        return float('nan')
    prev = list(range(len(h) + 1))
    for i, rc in enumerate(r, 1):
        cur = [i]
        for j, hc in enumerate(h, 1):
            cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (rc != hc)))
        prev = cur
    return prev[-1] / len(r)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--model-dir', default=_HERE)
    ap.add_argument('--langs', default='en,es,ja,ko,zh')
    args = ap.parse_args()
    langs = args.langs.split(',')

    assert not os.environ.get('DUBBING_HOME') or \
        os.path.abspath(os.environ['DUBBING_HOME']) == os.path.abspath(args.model_dir)
    from modeling_dubbing import DubbingBridgeModel
    model = DubbingBridgeModel.from_pretrained(args.model_dir)
    print('[verify] model loaded, single process', flush=True)

    refs = json.load(open(os.path.join(args.model_dir, 'samples', 'refs.json')))
    report = []
    for lang in langs:
        src = os.path.join(args.model_dir, 'samples', f'input_{lang}.wav')
        out = os.path.join(args.model_dir, 'samples', f'dub_{lang}.wav')
        wav, sr, info = model.dub(src, lang=lang, out_wav=out, return_info=True)
        dur = wav.shape[1] / sr
        rms = float(wav.float().pow(2).mean().sqrt())
        assert rms > 1e-4, f'{lang}: silent output (rms={rms})'
        status = 'OK' if 2.0 <= dur <= 15.0 else 'WARN dur out of 2-15s'
        rec = {'lang': lang, 'dur': round(dur, 2), 'rms': round(rms, 4),
               'src_lang': info.get('src_lang'), 'src_line': info['src_line'],
               'tgt_raw': info['tgt_raw'], 'tgt_cv': info['tgt_cv'],
               'ref': refs[lang]['ref'], 'status': status}
        report.append(rec)
        print(f'[verify:{lang}] dur={dur:.2f}s rms={rms:.4f} src_lang={info.get("src_lang")} {status}',
              flush=True)
        print(f'[verify:{lang}] src    : {info["src_line"]}', flush=True)
        print(f'[verify:{lang}] hyp_raw: {info["tgt_raw"]}', flush=True)
        print(f'[verify:{lang}] cv_text: {info["tgt_cv"]}' +
              ('  (ja kata)' if lang == 'ja' else ''), flush=True)
        print(f'[verify:{lang}] ref    : {refs[lang]["ref"]}', flush=True)

    wl = [l for l in ('en', 'ko', 'zh') if l in langs]
    if wl:
        from transformers import pipeline as hf_pipeline
        asr = hf_pipeline('automatic-speech-recognition', model=WHISPER, device=0)
        print('[verify] whisper loaded', flush=True)
    for lang in wl:
        t = asr(os.path.join(args.model_dir, 'samples', f'dub_{lang}.wav'),
                generate_kwargs={'language': lang, 'task': 'transcribe'})['text']
        m = wer(refs[lang]['ref'], t) if lang == 'en' else cer(refs[lang]['ref'], t)
        name = 'WER' if lang == 'en' else 'CER'
        print(f'[verify:{lang}] whisper backtest: {t.strip()!r}', flush=True)
        print(f'[verify:{lang}] {name} vs ref: {m:.3f}', flush=True)
        report.append({'lang': lang, 'whisper': t, name.lower() + '_vs_ref': round(m, 3)})

    json.dump(report, open(os.path.join(args.model_dir, 'samples', 'verify_report.json'), 'w'),
              indent=1, ensure_ascii=False)
    print('VERIFY EXPORT DONE', flush=True)


if __name__ == '__main__':
    main()