#!/usr/bin/env python3
"""Free-running synthesis gate + ASR scoring for the base model or a LoRA checkpoint.

For every sentence x reference (x seed): synthesize with voice cloning (reference in the system
prompt, exactly like inference), save the wav, and measure
  * duration ratio   = generated seconds / (words / 2.5)          (babble or early stop)
  * active fraction  = share of 100 ms frames with RMS > 0.02      (silence collapse)
  * tail ratio       = RMS(last 25%) / RMS(first 75%)              (fade-out collapse, July failure)
then transcribe with the Egyptian-dialect Whisper judge and compute normalised WER / CER.

Usage (fish-speech root, uv env, HF_HOME=/root/work/hf):
  python scripts/gate_eval.py --tag base --sentences data/eval/quick.jsonl --out results/eval/base
  python scripts/gate_eval.py --tag s250 --lora results/<run>/checkpoints/step_000000250.ckpt \
        --lora-config r_32_alpha_16_slow --sentences data/eval/quick.jsonl --out results/eval/s250
"""
import argparse
import json
import os
import re
import sys
import time
from pathlib import Path

os.environ.setdefault("HF_HOME", "/root/work/hf")
os.environ.setdefault("OMP_NUM_THREADS", "8")

import numpy as np
import soundfile as sf
import torch
import torchaudio
from hydra import compose, initialize
from hydra.utils import instantiate

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from fish_speech.models.text2semantic.inference import (  # noqa: E402
    decode_one_token_ar, decode_to_audio, encode_audio, generate_long, load_codec_model,
)
from fish_speech.models.text2semantic.llama import BaseTransformer, DualARTransformer  # noqa: E402

ap = argparse.ArgumentParser()
ap.add_argument("--base", default="checkpoints/s2-pro")
ap.add_argument("--lora", default=None, help="Lightning .ckpt with LoRA weights (unmerged eval)")
ap.add_argument("--lora-config", default="r_32_alpha_16_slow")
ap.add_argument("--merged", default=None, help="merged model dir (alternative to --lora)")
ap.add_argument("--tag", required=True)
ap.add_argument("--sentences", default="data/eval/quick.jsonl")
ap.add_argument("--refs", default="data/eval/refs.jsonl")
ap.add_argument("--out", required=True)
ap.add_argument("--seeds", type=int, default=1)
ap.add_argument("--temperature", type=float, default=0.8)
ap.add_argument("--top-p", type=float, default=0.8)
ap.add_argument("--max-refs", type=int, default=99)
ap.add_argument("--merge-mode", choices=["unmerged", "bf16"], default="unmerged",
                help="unmerged: keep the LoRA as a separate fp-accurate path (default); bf16: merge into bf16 weights (old behaviour)")
ap.add_argument("--no-asr", action="store_true")
ap.add_argument("--prefix", default="", help="text prepended to every synthesized sentence and to the reference text (e.g. \"[egyptian accent] \")")
ap.add_argument("--judge", choices=["soniox", "whisper"], default=None,
                help="default: soniox when SONIOX_API_KEY is set, else whisper")
ap.add_argument("--asr-model", default="oddadmix/whisper-large-v3-turbo-arabic-dialectal-v2")
args = ap.parse_args()
judge = args.judge or ("soniox" if os.environ.get("SONIOX_API_KEY") else "whisper")
SONIOX_API = os.environ.get("SONIOX_API_BASE", "https://api.soniox.com").rstrip("/")   # EU keys: https://api.eu.soniox.com
if judge == "soniox":
    import requests
    try:
        code = requests.get(f"{SONIOX_API}/v1/models", timeout=20,
                            headers={"Authorization": f"Bearer {os.environ.get('SONIOX_API_KEY', '')}"}).status_code
    except Exception as e:
        code = repr(e)
    if code != 200:
        print(f"WARNING: Soniox key check failed ({code}); falling back to the Whisper judge")
        judge = "whisper"

device = "cuda"
out = Path(args.out); (out / "wav").mkdir(parents=True, exist_ok=True)
sentences = [json.loads(l) for l in Path(args.sentences).read_text(encoding="utf-8").splitlines() if l.strip()]
refs = [json.loads(l) for l in Path(args.refs).read_text(encoding="utf-8").splitlines() if l.strip()][: args.max_refs]

# ---------------------------------------------------------------- model
t0 = time.time()
if args.lora:
    with initialize(version_base="1.3", config_path="../fish_speech/configs/lora"):
        lc = instantiate(compose(config_name=args.lora_config))
    model = BaseTransformer.from_pretrained(args.base, load_weights=True, max_length=8192, lora_config=lc)
    model = model.to(device=device, dtype=torch.bfloat16).train()      # July trap 11: train() -> load -> eval()
    sd = torch.load(args.lora, map_location="cpu", weights_only=False)
    sd = sd.get("state_dict", sd)
    sd = {k[len("model."):] if k.startswith("model.") else k: v for k, v in sd.items() if "lora" in k}
    missing, unexpected = model.load_state_dict(sd, strict=False)
    n_loaded = len(sd)
    assert not unexpected, unexpected[:5]
    if args.merge_mode == "unmerged":
        # keep W and the LoRA term separate: merging a ~0.1% update into bf16 weights rounds most of it away
        for m in model.modules():
            if hasattr(m, "merge_weights"):
                m.merge_weights = False
    model.eval()                                                        # bf16 mode: loralib merges here
    print(f"loaded {n_loaded} LoRA tensors from {args.lora} (merge mode: {args.merge_mode})")
else:
    src = args.merged or args.base
    model = DualARTransformer.from_pretrained(src, load_weights=True, max_length=8192)
    model = model.to(device=device, dtype=torch.bfloat16).eval()
model.fixed_temperature = torch.tensor(0.7, device=device)
model.fixed_top_p = torch.tensor(0.7, device=device)
model.fixed_repetition_penalty = torch.tensor(1.5, device=device)
model._cache_setup_done = False
with torch.device(device):
    model.setup_caches(max_batch_size=1, max_seq_len=model.config.max_seq_len, dtype=torch.bfloat16)
codec = load_codec_model(Path(args.base) / "codec.pth", device)
sr = codec.sample_rate
print(f"model ready in {time.time()-t0:.0f}s | codec sr {sr}")

ref_codes = {r["name"]: encode_audio(r["wav"], codec, device).cpu() for r in refs}
for r in refs:
    print(f"ref {r['name']}: {tuple(ref_codes[r['name']].shape)} codes")


def synth(text, ref):
    chunks = []
    gen = generate_long(
        model=model, device=device, decode_one_token=decode_one_token_ar, text=text, num_samples=1,
        max_new_tokens=0, top_p=args.top_p, top_k=30, temperature=args.temperature,
        repetition_penalty=1.1, compile=False, iterative_prompt=True, chunk_length=300,
        prompt_text=[args.prefix + ref["text"]], prompt_tokens=[ref_codes[ref["name"]]],
    )
    for resp in gen:
        if resp.action == "sample" and resp.codes is not None:
            chunks.append(resp.codes)
    codes = torch.cat(chunks, dim=1) if len(chunks) > 1 else chunks[0]
    wav = decode_to_audio(codes.to(device), codec).float().cpu().numpy()
    return codes.shape[1], wav


def audio_stats(wav, words):
    dur = len(wav) / sr
    frame = int(sr * 0.1)
    frames = [wav[i:i + frame] for i in range(0, max(frame, len(wav) - frame + 1), frame)]
    rms = np.array([np.sqrt(np.mean(f ** 2)) for f in frames if len(f)])
    active = float((rms > 0.02).mean()) if len(rms) else 0.0
    k = max(1, int(0.75 * len(rms)))
    head = float(np.sqrt(np.mean(rms[:k] ** 2))) if k else 0.0
    tail = float(np.sqrt(np.mean(rms[k:] ** 2))) if len(rms) > k else 0.0
    return {"dur_s": round(dur, 2), "dur_ratio": round(dur / max(0.5, words / 2.5), 2),
            "active_frac": round(active, 3), "tail_ratio": round(tail / head, 3) if head > 1e-6 else 0.0}


rows = []
t1 = time.time()
for s in sentences:
    for ref in refs:
        for seed in range(args.seeds):
            torch.manual_seed(seed); torch.cuda.manual_seed(seed)
            name = f"{s['id']}__{ref['name']}__s{seed}"
            try:
                n_codes, wav = synth(args.prefix + s["text"], ref)
            except Exception as e:  # keep going; record the failure
                print("GEN ERROR", name, repr(e)[:200]); rows.append({"name": name, "error": repr(e)[:200]}); continue
            path = out / "wav" / f"{name}.wav"
            sf.write(str(path), np.clip(wav, -1.0, 1.0), sr, subtype="PCM_16")
            st = audio_stats(wav, s.get("words") or len(s["text"].split()))
            rows.append({"name": name, "id": s["id"], "src": s.get("src"), "ref": ref["name"], "seed": seed,
                         "text": s["text"], "asr_text": s.get("asr_text"), "n_codes": int(n_codes), "wav": str(path), **st})
            print(f"{len(rows):4d}/{len(sentences)*len(refs)*args.seeds} {name} {st['dur_s']}s ratio={st['dur_ratio']} act={st['active_frac']} tail={st['tail_ratio']}", flush=True)
print(f"synthesis done in {(time.time()-t1)/60:.1f} min")

# ---------------------------------------------------------------- ASR judge
DIAC = re.compile(r"[ؐ-ًؚ-ٰٟۖ-ۭـ]")
PUNCT = re.compile(r"[^\w\s]", re.UNICODE)
AR_DIGITS = str.maketrans("٠١٢٣٤٥٦٧٨٩", "0123456789")


def norm(t):
    t = DIAC.sub("", t).translate(AR_DIGITS)
    t = t.replace("أ", "ا").replace("إ", "ا").replace("آ", "ا").replace("ى", "ي").replace("ة", "ه").replace("ؤ", "و").replace("ئ", "ي")
    t = PUNCT.sub(" ", t)
    return " ".join(t.split())


def edit_distance(a, b):
    prev = list(range(len(b) + 1))
    for i, x in enumerate(a, 1):
        cur = [i]
        for j, y in enumerate(b, 1):
            cur.append(min(prev[j] + 1, cur[j - 1] + 1, prev[j - 1] + (x != y)))
        prev = cur
    return prev[-1]


def soniox_transcribe(path, retries=3):
    """Soniox stt-async-v5 (same engine as the July judge and the podcast transcripts)."""
    import requests
    api = SONIOX_API
    hdr = {"Authorization": f"Bearer {os.environ['SONIOX_API_KEY']}"}
    last = None
    for attempt in range(retries):
        try:
            with open(path, "rb") as f:
                r = requests.post(f"{api}/v1/files", headers=hdr, files={"file": (Path(path).name, f, "audio/wav")}, timeout=120)
            r.raise_for_status()
            file_id = r.json()["id"]
            try:
                r = requests.post(f"{api}/v1/transcriptions", headers=hdr,
                                  json={"file_id": file_id, "model": "stt-async-v5", "language_hints": ["ar"]}, timeout=30)
                r.raise_for_status()
                tid = r.json()["id"]
                try:
                    for _ in range(180):
                        s = requests.get(f"{api}/v1/transcriptions/{tid}", headers=hdr, timeout=30).json()
                        if s.get("status") == "completed":
                            break
                        if s.get("status") == "error":
                            raise RuntimeError(f"soniox error: {s.get('error_message')}")
                        time.sleep(1)
                    else:
                        raise TimeoutError("soniox poll timeout")
                    return requests.get(f"{api}/v1/transcriptions/{tid}/transcript", headers=hdr, timeout=30).json().get("text", "")
                finally:
                    requests.delete(f"{api}/v1/transcriptions/{tid}", headers=hdr, timeout=30)
            finally:
                requests.delete(f"{api}/v1/files/{file_id}", headers=hdr, timeout=30)
        except Exception as e:  # transient API errors: retry
            last = e
            time.sleep(3 * (attempt + 1))
    raise RuntimeError(f"soniox failed after {retries} attempts: {last!r}")


def score(r, hyp):
    if r.get("asr_text"):   # phoneme-tag sets: plain reference, digits verbalised on both sides (Soniox writes numerals)
        sys.path.insert(0, str(Path(__file__).resolve().parent))
        from egy_numbers import verbalize
        ref_n, hyp_n = norm(verbalize(r["asr_text"])), norm(verbalize(hyp))
    else:
        ref_n, hyp_n = norm(r["text"]), norm(hyp)
    rw, hw = ref_n.split(), hyp_n.split()
    r["hyp"] = hyp
    r["wer"] = round(edit_distance(rw, hw) / max(1, len(rw)), 4)
    r["cer"] = round(edit_distance(ref_n.replace(" ", ""), hyp_n.replace(" ", "")) / max(1, len(ref_n.replace(" ", ""))), 4)


if not args.no_asr and rows:
    todo = [r for r in rows if "wav" in r]
    print(f"ASR judge: {judge} on {len(todo)} files")
    if judge == "soniox":
        from concurrent.futures import ThreadPoolExecutor

        def one(r):
            try:
                return soniox_transcribe(r["wav"])
            except Exception as e:
                return e

        with ThreadPoolExecutor(max_workers=8) as ex:
            for r, hyp in zip(todo, ex.map(one, todo)):
                if isinstance(hyp, Exception):
                    r["asr_error"] = repr(hyp)[:200]; print("ASR ERROR", r["name"], r["asr_error"])
                else:
                    score(r, hyp)
    else:
        del model; torch.cuda.empty_cache()
        from transformers import pipeline
        asr = pipeline("automatic-speech-recognition", model=args.asr_model, torch_dtype=torch.float16, device=0)
        for r in todo:
            wav, s_r = sf.read(r["wav"], dtype="float32")
            wav16 = torchaudio.functional.resample(torch.from_numpy(wav), s_r, 16000).numpy()
            hyp = asr({"raw": wav16, "sampling_rate": 16000}, chunk_length_s=30,
                      generate_kwargs={"language": "ar", "task": "transcribe"})["text"]
            score(r, hyp)

# ---------------------------------------------------------------- report
ok = [r for r in rows if "wav" in r]
scored = [r for r in ok if "wer" in r]
def mean(k, rs=None): rs = ok if rs is None else rs; return round(float(np.mean([r[k] for r in rs])), 4) if rs else None
summary = {
    "tag": args.tag, "judge": judge if not args.no_asr else None, "n": len(ok), "errors": len(rows) - len(ok),
    "n_scored": len(scored), "wer": mean("wer", scored), "cer": mean("cer", scored),
    "dur_ratio_median": round(float(np.median([r["dur_ratio"] for r in ok])), 3) if ok else None,
    "active_frac_mean": mean("active_frac"), "tail_ratio_mean": mean("tail_ratio"),
    "collapsed": sum(1 for r in ok if r["tail_ratio"] < 0.3 or r["active_frac"] < 0.5),
    "too_long": sum(1 for r in ok if r["dur_ratio"] > 1.6), "too_short": sum(1 for r in ok if r["dur_ratio"] < 0.5),
}
for src in sorted({r["src"] for r in scored}):
    sub = [r for r in scored if r["src"] == src]
    summary[f"wer_{src}"] = round(float(np.mean([r["wer"] for r in sub])), 4)
for ref in sorted({r["ref"] for r in scored}):
    sub = [r for r in scored if r["ref"] == ref]
    summary[f"wer_ref_{ref}"] = round(float(np.mean([r["wer"] for r in sub])), 4)
(out / "rows.jsonl").write_text("\n".join(json.dumps(r, ensure_ascii=False) for r in rows) + "\n", encoding="utf-8")
(out / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
print("SUMMARY", json.dumps(summary, ensure_ascii=False))