"""Table 1 context-window row / Finding 4 — context-window length: mean-pooled layer-9 performance with 15/30/45/60 s of audio.
Uses a 9-layer-truncated WavLM (last_hidden_state == layer-9 output) so 60 s fits an 8 GB GPU, and one
60 s forward per clip whose prefix means give every shorter cap. Own cache (~35 min extraction on GPU).

    python experiments/exp_window.py
"""
import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
import hashlib
import numpy as np
from experiments import common as C
from empathyeval.data.audio import cached_load

CAPS = [15, 30, 45, 60]
CD = "cache/exp_window"
os.makedirs(CD, exist_ok=True)


def cp(w):
    return os.path.join(CD, hashlib.md5(w.encode()).hexdigest()[:16] + ".npy")


def extract():
    import torch
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    from transformers import AutoFeatureExtractor, WavLMModel
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    fe = AutoFeatureExtractor.from_pretrained(C.MID)
    m = WavLMModel.from_pretrained(C.MID)
    m.encoder.layers = m.encoder.layers[:C.LAYER]          # truncate -> output == layer-9
    m = m.to(dev).eval()
    print(f"extracting 60 s layer-9 features on {dev}", flush=True)
    wavs = [w for it in C.train_items(2500) for w in (it.good_wav, it.bad_wav)]
    wavs += [o.wav for q in C.QS for o in q.options]
    for n, w in enumerate(dict.fromkeys(wavs), 1):
        if os.path.exists(cp(w)):
            continue
        y = cached_load(w, C.cfg)[:16000 * CAPS[-1]]
        inp = fe(y, sampling_rate=16000, return_tensors="pt").input_values.to(dev)
        with torch.no_grad():
            h = m(inp).last_hidden_state[0]                # [T,1024]
        np.save(cp(w), np.stack([h[:50 * s].mean(0).cpu().numpy() for s in CAPS]).astype(np.float32))
        if n % 500 == 0:
            print(f"  {n} clips", flush=True)


if not os.path.exists(cp(C.train_items(1)[0].good_wav)):
    extract()
mem = {}
for i, s in enumerate(CAPS):
    def f(w, i=i):
        a = mem.get(w)
        if a is None:
            a = mem[w] = np.load(cp(w))
        return a[i]
    C.report(f"cap {s:2d}s (mean pooling)", C.answers(f))