"""Precompute frozen-teacher caches for a shard of training files.

Usage: python3 precompute_shard.py <device> <shard_idx> <n_shards>
Files are enumerated deterministically (train actors 1-18, speech + song);
this process handles every n_shards-th file starting at shard_idx.
Saves /tmp/work/teacher_cache/shard_<idx>.pt
"""

import glob
import os
import sys
import warnings

warnings.filterwarnings("ignore")
import logging

logging.disable(logging.CRITICAL)

import torch
import torchaudio

sys.path.insert(0, "/tmp/work/s2st_emotion")
from funasr import AutoModel
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 2
MAX_SEC = 12.0

OUT_DIR = "/tmp/work/teacher_cache"
os.makedirs(OUT_DIR, exist_ok=True)


def gather_files():
    files = sorted(glob.glob("/tmp/work/ravdess_kag/**/*.wav", recursive=True))
    out = []
    for f in files:
        b = os.path.basename(f)
        parts = b.split("-")
        try:
            actor = int(parts[-1].split(".")[0])
            channel = parts[1]  # 01=speech, 02=song
        except Exception:
            continue
        if actor <= 18 and channel in ("01", "02"):
            out.append(f)
    return out


def rms_envelope_50hz(wav, T_frames):
    frames = torch.zeros(T_frames)
    for i in range(T_frames):
        seg = wav[i * 320: (i + 1) * 320]
        frames[i] = float(seg.pow(2).mean().sqrt()) if seg.numel() else (frames[i - 1] if i else 0.0)
    if T_frames >= 9:
        frames = torch.nn.functional.conv1d(
            frames[None, None], torch.ones(9)[None, None] / 9.0, padding=4)[0, 0]
    lo, hi = frames.min(), frames.max()
    return torch.full_like(frames, 0.5) if hi - lo < 1e-6 else (frames - lo) / (hi - lo)


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

    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)

    cache, skipped = [], 0
    for i, f in enumerate(mine):
        try:
            wav, sr = torchaudio.load(f)
            wav = wav.mean(0)
            if sr != 16000:
                wav = torchaudio.functional.resample(wav, sr, 16000)
        except Exception:
            skipped += 1
            continue
        if wav.numel() < 8000:
            skipped += 1
            continue
        if wav.numel() > MAX_SEC * 16000:
            wav = wav[: int(MAX_SEC * 16000)]
        wav = wav / wav.abs().max().clamp(min=1e-4)
        states, lengths = enc.extract_teacher_states(wav[None].to(DEV),
                                                     torch.tensor([wav.numel()], device=DEV))
        probs, inten = enc.teacher_frame_probs(states, lengths)
        T = int(lengths[0].item())
        rms = rms_envelope_50hz(wav, T)
        inten = (inten.squeeze(0)[:, 0].cpu().clamp(min=0, max=1).sqrt() * rms).unsqueeze(-1)
        cache.append({
            "states": states.squeeze(0).half().cpu(),
            "probs": probs.squeeze(0).half().cpu(),
            "inten": inten.half().cpu(),
            "file": os.path.basename(f),
        })
        if (i + 1) % 200 == 0:
            print(f"[shard {SHARD}] {i+1}/{len(mine)}", flush=True)

    torch.save(cache, f"{OUT_DIR}/shard_{SHARD}.pt")
    print(f"[shard {SHARD}] saved {len(cache)} utts (skipped {skipped}) -> shard_{SHARD}.pt", flush=True)


if __name__ == "__main__":
    main()