"""One-time feature extraction for ALL experiments (run once; every exp_*.py reads this cache).

    python experiments/extract_features.py

Extracts, per audio clip, one [71,1024] array (layout in common.py) from a frozen WavLM-large:
all 25 layers' mean+std, plus layer-9 fine-segment / max / velocity / speech-segment / 3-segment rows.
Covers the good+bad audio of all 4,892 training triples, all test candidates, and the utterance audio
needed by the utterance-conditioning experiment (~13.8k clips; ~4 GB cache; ~1 h on an 8 GB GPU).

Deterministic settings (cuDNN deterministic, no benchmark autotuning, deterministic algorithms where
available) so a re-run on the same machine reproduces the cache — and therefore every experiment number.
Resumable: existing clips are skipped, so an interrupted run just continues.
"""
import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
import numpy as np
from experiments import common as C
from empathyeval.data.release import build_train_items, build_index
from empathyeval.data.audio import cached_load

os.makedirs(C.CACHE, exist_ok=True)
_M = {}


def _model():
    if not _M:
        import torch
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        try:
            torch.use_deterministic_algorithms(True, warn_only=True)
        except TypeError:
            pass
        from transformers import AutoFeatureExtractor, WavLMModel
        dev = "cuda" if torch.cuda.is_available() else "cpu"
        _M["torch"] = torch
        _M["fe"] = AutoFeatureExtractor.from_pretrained(C.MID)
        _M["m"] = WavLMModel.from_pretrained(C.MID).to(dev).eval()
        _M["dev"] = dev
        print(f"loaded {C.MID} on {dev} (deterministic mode)", flush=True)
    return _M


def seg_means(h, K):
    T = h.shape[0]
    idx = np.linspace(0, T, K + 1).astype(int)
    return [h[idx[i]:max(idx[i] + 1, idx[i + 1])].mean(0) for i in range(K)]


def ensure(wav):
    if os.path.exists(C.cpath(wav)):
        return False
    y = cached_load(wav, C.cfg)[:16000 * C.CAP_S]
    m = _model()
    inp = m["fe"](y, sampling_rate=16000, return_tensors="pt").input_values.to(m["dev"])
    with m["torch"].no_grad():
        hs = m["m"](inp, output_hidden_states=True).hidden_states       # 25 x [1,T,1024]
    rows = []
    for h in hs:                                                        # rows 0..49: per-layer mean+std
        h = h[0]
        rows += [h.mean(0).cpu().numpy(), h.std(0).cpu().numpy()]
    h9 = hs[C.LAYER][0].cpu().numpy()                                   # layer-9 rows
    T = h9.shape[0]
    rows += seg_means(h9, C.N_FINE)                                     # 50..61 fine segments
    rows.append(h9.max(0))                                              # 62 max
    d = h9[1:] - h9[:-1] if T > 1 else h9 * 0.0
    rows += [np.abs(d).mean(0), d.std(0)]                               # 63,64 velocity stats
    e = np.array([np.sqrt((y[i * len(y) // T:(i + 1) * len(y) // T] ** 2).mean() + 1e-9) for i in range(T)])
    sp = h9[e > 0.12 * e.max()]                                         # 65..67 speech segments
    rows += seg_means(sp if sp.shape[0] >= 3 else h9, 3)
    s3 = [h9[:T // 3], h9[T // 3:2 * T // 3], h9[2 * T // 3:]] if T >= 3 else [h9, h9, h9]
    rows += [s.mean(0) for s in s3]                                     # 68..70 direct 3 segments
    np.save(C.cpath(wav), np.stack(rows).astype(np.float32))
    return True


def all_wavs():
    items = build_train_items(C.cfg)                                    # ALL 4,892 (train-size sweep needs them)
    qs, _ = build_index(C.cfg)
    wavs = []
    for it in items:
        wavs += [it.good_wav, it.bad_wav]
    for q in qs:
        wavs += [o.wav for o in q.options]
    for it in C.train_items(2500):                                      # utterances (conditioning experiment)
        wavs.append(it.utterance_wav)
    for q in qs:
        wavs.append(q.utterance_wav)
    return list(dict.fromkeys(wavs))


def main():
    wavs = all_wavs()
    done = 0
    for n, w in enumerate(wavs, 1):
        try:
            done += ensure(w)
        except Exception as e:
            print(f"  skip {w}: {e}", flush=True)
        if n % 300 == 0:
            print(f"  {n}/{len(wavs)} scanned, {done} new", flush=True)
    missing = sum(0 if os.path.exists(C.cpath(w)) else 1 for w in wavs)
    print(f"extracted {done}; {missing} missing of {len(wavs)}")
    if missing == 0:
        print("ALL-CACHED")


if __name__ == "__main__":
    main()