"""Shared harness for the experiments reported in the technical report.

Every experiment reads ONE unified feature cache (built once by `python experiments/extract_features.py`
with deterministic torch settings), the ground-truth labels in inference/annotations.jsonl, and the frozen
metric Final = (accuracy + group bonus) / (#questions + #groups). Groups/subtasks derive from question ids.

Unified cache: one [71, 1024] float32 array per audio clip
  rows  0..49  layer L mean = row 2L, layer L std = row 2L+1              (L = 0..24)
  rows 50..61  layer-9: 12 fine equal-time segment means (K in {2,3,4,6,12} derivable)
  row  62      layer-9: elementwise max over time
  row  63      layer-9: |delta| mean (frame-to-frame velocity),  row 64: delta std
  rows 65..67  layer-9: 3 energy-gated speech-segment means (silence-trimmed thirds)
  rows 68..70  layer-9: 3 direct equal-time segment means (the shipped pooling's segments)
"""
import hashlib
import json
import os
import random
import sys

os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")   # deterministic cuBLAS (before torch loads)
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
os.chdir(ROOT)
sys.path.insert(0, ROOT)

import numpy as np
import yaml

cfg = yaml.safe_load(open("configs/phase1.yaml"))
CACHE = "cache/exp_unified"
MID, LAYER, CAP_S = "microsoft/wavlm-large", 9, 15

# ---- cache row indices -----------------------------------------------------
def R_MEAN(L=LAYER): return 2 * L
def R_STD(L=LAYER): return 2 * L + 1
R_FINE0, N_FINE = 50, 12
R_MAX, R_DVEL, R_DSTD = 62, 63, 64
R_SPEECH0, R_SEG3 = 65, 68
POOL_FINAL = [R_MEAN(), R_STD(), R_SEG3, R_SEG3 + 1, R_SEG3 + 2]     # the shipped pooling (5120-d)

def cpath(wav):
    return os.path.join(CACHE, hashlib.md5(wav.encode()).hexdigest()[:16] + ".npy")

_mem = {}
def arr(wav):
    a = _mem.get(wav)
    if a is None:
        a = _mem[wav] = np.load(cpath(wav))
    return a

def feat(rows):
    """featfn selecting fixed cache rows."""
    return lambda w: np.concatenate([arr(w)[r] for r in rows])

def feat_fineK(K, base_rows=(R_MEAN(), R_STD())):
    """featfn: base rows + K coarse segment means derived from the 12 fine segments."""
    b = np.linspace(0, N_FINE, K + 1).astype(int)
    def f(w):
        a = arr(w)
        segs = [a[R_FINE0 + b[i]:R_FINE0 + b[i + 1]].mean(0) for i in range(K)]
        return np.concatenate([a[r] for r in base_rows] + segs)
    return f

# ---- data ------------------------------------------------------------------
from empathyeval.data.release import build_index, build_train_items          # noqa: E402

QS, _GROUPS_IDX = build_index(cfg)

def train_items(n=2500, seed=0):
    items = build_train_items(cfg)
    random.Random(seed).shuffle(items)
    return items[:n]

def load_labels():
    lab = {}
    for line in open("inference/annotations.jsonl", encoding="utf-8"):
        r = json.loads(line)
        lab[r["question_id"]] = r["answer"]
    return lab

LABELS = load_labels()
def gid(q): return q.rsplit("_", 1)[0]
def subtask(q): return "tone" if q.split("_", 1)[0] == "emovdb" else "context"
GROUPS = {}
for q in LABELS:
    GROUPS.setdefault(gid(q), []).append(q)

# ---- metric ----------------------------------------------------------------
def final(preds, qids=None):
    """Frozen metric over the given qids (default: all)."""
    lab = LABELS if qids is None else {q: LABELS[q] for q in qids}
    gs = {g: qq for g, qq in GROUPS.items() if qq[0] in lab}
    acc = sum(preds.get(q) == lab[q] for q in lab)
    bon = sum(all(preds.get(q) == lab[q] for q in qq) for qq in gs.values())
    return (acc + bon) / (len(lab) + len(gs))

def report(name, preds, width=34):
    t = final(preds, [q for q in LABELS if subtask(q) == "tone"])
    c = final(preds, [q for q in LABELS if subtask(q) == "context"])
    print(f"  {name:{width}s} tone={t:.3f} context={c:.3f} Final={final(preds):.4f}", flush=True)
    return final(preds)

# ---- ranker ----------------------------------------------------------------
def train_ranker(featfn, n_items=2500, C=0.5):
    from sklearn.linear_model import LogisticRegression
    from sklearn.preprocessing import StandardScaler
    X, y = [], []
    for it in train_items(n_items):
        try:
            g, b = featfn(it.good_wav), featfn(it.bad_wav)
        except Exception:
            continue
        X.append(g - b); y.append(1)
        X.append(b - g); y.append(0)
    X = np.array(X)
    sc = StandardScaler().fit(X)
    clf = LogisticRegression(max_iter=3000, C=C).fit(sc.transform(X), y)
    return sc, clf

def answers(featfn, n_items=2500, C=0.5, model=None):
    """Train a ranker on the pairwise task and answer every test question (argmax)."""
    sc, clf = model if model else train_ranker(featfn, n_items, C)
    out = {}
    for q in QS:
        fv = sc.transform([featfn(o.wav) for o in q.options])
        out[q.qid] = [o.letter for o in q.options][int(np.argmax(clf.decision_function(fv)))]
    return out

# ---- paired bootstrap over groups -------------------------------------------
def paired_bootstrap(preds_a, preds_b, n=2000, seed=0):
    """Resample the groups; score BOTH systems on each resample. Returns delta = b - a stats."""
    glist = list(GROUPS.values())
    rng = random.Random(seed)
    deltas = []
    for _ in range(n):
        samp = [glist[rng.randrange(len(glist))] for _ in range(len(glist))]
        def fin(p):
            acc = sum(p.get(q) == LABELS[q] for qq in samp for q in qq)
            bon = sum(all(p.get(q) == LABELS[q] for q in qq) for qq in samp)
            nd = sum(len(qq) for qq in samp)
            return (acc + bon) / (nd + len(samp))
        deltas.append(fin(preds_b) - fin(preds_a))
    d = np.array(deltas)
    return {"median": float(np.median(d)), "lo": float(np.percentile(d, 2.5)),
            "hi": float(np.percentile(d, 97.5)), "p_better": float((d > 0).mean())}