"""Precompute v4 caches — same GPU-heavy pipeline as precompute_fast.py, but the
cache stores the RAW (unnormalized) 50 Hz RMS envelope instead of the old
min-max "conf-gated" intensity target. The v4 intensity target is derived in
train_v4 from (probs, rms_raw) with a global reference level.

Cache item: {states (T,768) f16, probs (T,9) f16, rms (T,) f16, file}
Usage: python3 precompute_v4.py <device> <shard_idx> <n_shards>
"""
import os, sys, threading, queue, warnings
warnings.filterwarnings("ignore")
import logging; logging.disable(logging.CRITICAL)

import torch, torchaudio
import torch.nn.functional as F

sys.path.insert(0, "/tmp/work/s2st_emotion")
from precompute_shard_v3 import gather_files, MAX_SEC
from surgical_emotion import SurgicalConfig, SurgicalEmotionEncoder

DEV = sys.argv[1] if len(sys.argv) > 1 else "cuda:0"
SHARD = int(sys.argv[2]) if len(sys.argv) > 2 else 0
N_SHARDS = int(sys.argv[3]) if len(sys.argv) > 3 else 4
N_LOAD = 4
OUT_DIR = "/tmp/work/teacher_cache_v4"
os.makedirs(OUT_DIR, exist_ok=True)

resamplers = {}

def to_16k(wav, sr):
    if sr == 16000:
        return wav
    if sr not in resamplers:
        resamplers[sr] = torchaudio.transforms.Resample(sr, 16000).to(DEV)
    return resamplers[sr](wav.to(DEV)).cpu()

def load_one(f):
    try:
        wav, sr = torchaudio.load(f)
        wav = wav.mean(0)
        if sr != 16000:
            wav = to_16k(wav, sr)
    except Exception:
        return None
    if wav.numel() < 8000:
        return None
    if wav.numel() > MAX_SEC * 16000:
        wav = wav[: int(MAX_SEC * 16000)]
    wav = wav / wav.abs().max().clamp(min=1e-4)
    return wav

def rms_raw_gpu(wav_dev, T_f):
    """RAW 50 Hz RMS envelope (9-tap box smoothed), NO min-max normalization."""
    L = wav_dev.numel()
    n_full = (L + 319) // 320
    w = torch.zeros(T_f * 320, device=wav_dev.device, dtype=wav_dev.dtype)
    w[: min(L, T_f * 320)] = wav_dev[: min(L, T_f * 320)]
    ms = w.view(T_f, 320).pow(2).mean(1).sqrt()
    if T_f > n_full:
        ms[n_full:] = ms[n_full - 1]
    if T_f >= 9:
        ms = F.conv1d(ms[None, None], torch.ones(9, device=DEV)[None, None] / 9.0,
                      padding=4)[0, 0]
    return ms

def main():
    files, counts = gather_files()
    mine = files[SHARD::N_SHARDS]
    print(f"[v4 shard {SHARD}/{N_SHARDS} on {DEV}] total={len(files)} mine={len(mine)}", flush=True)

    partial_path = f"{OUT_DIR}/shard_{SHARD}_partial.pt"
    cache = []
    if os.path.exists(partial_path):
        cache = torch.load(partial_path, map_location="cpu", weights_only=False)
        print(f"[v4 shard {SHARD}] resuming with {len(cache)} done", flush=True)
    mine = mine[len(cache):]

    from funasr import AutoModel
    fun = AutoModel(model="/tmp/work/emotion2vec_plus_base", trust_remote_code=False,
                    device=DEV, disable_update=True)
    enc = SurgicalEmotionEncoder(fun.model, SurgicalConfig()).to(DEV)
    enc.eval()
    torch.set_grad_enabled(False)

    q = queue.Queue(maxsize=32)
    def loader(paths, q):
        for p in paths:
            q.put((p, load_one(p)))
        q.put(None)
    th = threading.Thread(target=loader, args=(mine, q), daemon=True)
    th.start()

    skipped = 0
    done = len(cache)
    while True:
        item = q.get()
        if item is None:
            break
        path, wav = item
        if wav is None:
            skipped += 1
            continue
        wg = wav.to(DEV)
        src = F.layer_norm(wg[None], wg[None].shape)
        st = enc.teacher.extract_features(src, padding_mask=None, mask=False)["x"][0].float()
        Ti = st.size(0)
        probs = torch.softmax(enc.teacher.proj(st), dim=-1)
        rms = rms_raw_gpu(wg, Ti)
        cache.append({
            "states": st.half().cpu(),
            "probs": probs.half().cpu(),
            "rms": rms.half().cpu(),
            "file": os.path.basename(path),
        })
        done = len(cache)
        if done % 500 == 0:
            print(f"[v4 shard {SHARD}] {done}/{len(files)}", flush=True)
        if done % 2000 == 0:
            torch.save(cache, partial_path)

    torch.save(cache, f"{OUT_DIR}/shard_{SHARD}.pt")
    if os.path.exists(partial_path):
        os.remove(partial_path)
    print(f"[v4 shard {SHARD}] saved {len(cache)} utts (skipped {skipped})", flush=True)

if __name__ == "__main__":
    main()