"""Runnable inference for diffusion-gemma-asr-small.

  python inference.py audio.wav
"""
import sys, numpy as np, torch, soundfile as sf
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer, WhisperFeatureExtractor
from model import AudioDiffusionGemma, AudioDiffusionConfig
from audio import audio_out_len, real_encoder_frames, WHISPER_ENC_FRAMES

MODEL_ID   = "google/diffusiongemma-26B-A4B-it"
WHISPER_ID = "openai/whisper-small"
D_MODEL, VOCAB = 2816, 262144
BOA, AUD, EOA  = 256000, 258881, 258883          # <|audio>, <|audio|>, <audio|>
EOS_IDS = {1, 106, 50}
SR, SUBSAMPLE = 16000, 8
INSTRUCTION = "Transcribe this audio to text."
# Per-segment limits (encoder window = 30 s). Segment longer audio at silence.
CANVAS_LEN, TARGET_SEG, MAX_SEG = 192, 13.0, 18.0


def load(ckpt="diffusion_asr_small.pt", device="cuda"):
    model_dir = snapshot_download(MODEL_ID)
    cfg = AudioDiffusionConfig(model_dir=model_dir, whisper_id=WHISPER_ID, whisper_dim=768,
                               d_model=D_MODEL, vocab_size=VOCAB, boa_token_id=BOA,
                               audio_token_id=AUD, eoa_token_id=EOA, subsample_factor=SUBSAMPLE)
    model = AudioDiffusionGemma.from_pretrained(cfg, dtype=torch.bfloat16, device=device)
    st = torch.load(ckpt, map_location=device)
    if "lora" in st:
        from peft import set_peft_model_state_dict
        model.apply_lora(); set_peft_model_state_dict(model.base, st["lora"])
    model.projector.load_state_dict(st["projector"]); model.projector.to(device, torch.float32)
    model.eval()
    tok = AutoTokenizer.from_pretrained(model_dir)
    fe = WhisperFeatureExtractor.from_pretrained(WHISPER_ID)
    return model, tok, fe


def _segment(wav):
    target, search, fw = int(TARGET_SEG * SR), int(3.0 * SR), int(0.08 * SR)
    segs, i = [], 0
    while i < len(wav):
        if len(wav) - i <= int(MAX_SEG * SR):
            segs.append(wav[i:]); break
        lo, hi = i + target - search, min(len(wav), i + target + search)
        region = wav[lo:hi]; nf = max(1, (len(region) - fw) // fw)
        energy = np.array([float(np.dot(region[k*fw:k*fw+fw], region[k*fw:k*fw+fw])) for k in range(nf)])
        cut = lo + int(np.argmin(energy)) * fw + fw // 2
        segs.append(wav[i:cut]); i = cut
    return segs


def _collapse(text):
    out = []
    for w in text.split():
        if not out or out[-1] != w:
            out.append(w)
    return " ".join(out)


@torch.no_grad()
def transcribe(wav, model, tok, fe, max_steps=16, device="cuda"):
    wav = np.asarray(wav, dtype=np.float32)
    instr = tok(INSTRUCTION, add_special_tokens=False)["input_ids"]
    prefix = [2] + instr + [BOA]
    Ta = audio_out_len(WHISPER_ENC_FRAMES, SUBSAMPLE)
    base = prefix + [AUD] * Ta + [EOA]
    texts = []
    segs = _segment(wav) if len(wav) > MAX_SEG * SR else [wav]
    for s in range(0, len(segs), 8):
        sub = segs[s:s + 8]
        mel = torch.stack([fe(c, sampling_rate=SR, return_tensors="pt").input_features[0] for c in sub]).to(device)
        B = len(sub)
        prompt_ids = torch.tensor([base] * B, dtype=torch.long, device=device)
        prompt_mask = torch.zeros(B, len(prefix) + Ta + 1, dtype=torch.long, device=device)
        for i, c in enumerate(sub):
            n_real = min(Ta, audio_out_len(real_encoder_frames(len(c)), SUBSAMPLE))
            prompt_mask[i, :len(prefix)] = 1
            prompt_mask[i, len(prefix):len(prefix) + max(1, n_real)] = 1
            prompt_mask[i, len(prefix) + Ta] = 1
        canvas, _ = model.generate(prompt_ids, prompt_mask, mel, canvas_len=CANVAS_LEN, max_steps=max_steps)
        for i in range(B):
            ids = []
            for t in canvas[i].tolist():
                if t in EOS_IDS:
                    break
                ids.append(t)
            texts.append(_collapse(tok.decode(ids, skip_special_tokens=True)))
    return " ".join(t for t in texts if t).strip()


if __name__ == "__main__":
    path = sys.argv[1] if len(sys.argv) > 1 else "audio.wav"
    wav, sr = sf.read(path)
    if wav.ndim > 1:
        wav = wav.mean(axis=1)
    if sr != SR:
        import librosa
        wav = librosa.resample(wav.astype(np.float32), orig_sr=sr, target_sr=SR)
    model, tok, fe = load()
    print(transcribe(wav, model, tok, fe))