"""Full teacher-vs-student benchmark on identical inputs.

Paths measured per identical input:
  teacher_pure     : stock compute path  (layer_norm -> extract_features -> mean -> proj -> softmax)
  teacher_pipeline : stock funasr user path (AutoModel.generate, granularity=frame)
  student_total    : SurgicalEmotionEncoder end-to-end (frozen backbone + adapter, 12.5 Hz outputs)
  adapter_only     : adapter on precomputed teacher states (marginal surgical cost)

Quality (same inputs, both models):
  top-1 agreement, JS divergence of utterance probs, pooled-embedding cosine,
  frame-level cosine (student upsampled to 50 Hz vs teacher states),
  intensity stats + correlation with conf-gated energy target.

Scaling: durations 1/2/3/4/10 s, adapter batch 1/4/8, fp16 autocast variant,
small CPU comparison. Saves benchmark_report.json.
"""

import glob
import json
import os
import sys
import time
import warnings

warnings.filterwarnings("ignore")
import logging

logging.disable(logging.CRITICAL)

sys.path.insert(0, "/tmp/work/s2st_emotion")
import numpy as np
import torch
import torch.nn.functional as F
import torchaudio

from funasr import AutoModel
from surgical_emotion import SurgicalConfig, SurgicalEmotionEncoder

FUNASR_DIR = "/tmp/work/emotion2vec_plus_base"
CKPT = "/tmp/work/surgical_out/surgical_emotion_module.pt"
OUT = "/tmp/work/surgical_out/benchmark_report.json"
DEV = "cuda:0"
LABELS = ["angry", "disgusted", "fearful", "happy", "neutral", "other", "sad", "surprised", "unk"]
EMO_FIELD = {"01": "neutral", "02": "calm", "03": "happy", "04": "sad",
             "05": "angry", "06": "fearful", "07": "disgust", "08": "surprised"}

report = {"gpu": torch.cuda.get_device_name(0), "torch": torch.__version__,
          "device": DEV, "labels": LABELS, "paths": {}, "files": {}, "scaling": {}, "quality_summary": {}}


def load_wav(path, max_s=12.0):
    wav, sr = torchaudio.load(path)
    wav = wav.mean(0)
    if sr != 16000:
        wav = torchaudio.functional.resample(wav, sr, 16000)
    wav = wav[: int(max_s * 16000)]
    return wav / wav.abs().max().clamp(min=1e-4)


def js_div(p, q, eps=1e-8):
    p = p.clamp(min=eps)
    q = q.clamp(min=eps)
    m = 0.5 * (p + q)
    kl = lambda a, b: float((a * (a.log() - b.log())).sum(-1))
    return 0.5 * kl(p, m) + 0.5 * kl(q, m)


def rms_env(wav, T):
    frames = torch.zeros(T)
    for i in range(T):
        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 >= 9:
        frames = F.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 timeit(fn, warmup=3, reps=10):
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    ts = []
    for _ in range(reps):
        t0 = time.perf_counter()
        fn()
        torch.cuda.synchronize()
        ts.append(time.perf_counter() - t0)
    return {"mean_ms": 1000 * float(np.mean(ts)), "std_ms": 1000 * float(np.std(ts)),
            "min_ms": 1000 * float(np.min(ts))}


print("=== loading models (shared teacher backbone) ===", flush=True)
fun = AutoModel(model=FUNASR_DIR, trust_remote_code=False, device=DEV, disable_update=True)
teacher = fun.model.eval()
enc = SurgicalEmotionEncoder(teacher, SurgicalConfig()).to(DEV)
payload = torch.load(CKPT, map_location=DEV, weights_only=False)
enc.adapter.load_state_dict(payload["adapter_state"])
enc.eval()
n_teacher = sum(p.numel() for p in teacher.parameters())
n_adapter = sum(p.numel() for p in enc.adapter.parameters())
report["params"] = {"teacher_backbone": n_teacher, "student_adapter": n_adapter,
                    "adapter_overhead_pct": 100.0 * n_adapter / n_teacher}
print(f"teacher params {n_teacher/1e6:.2f}M, adapter {n_adapter/1e6:.2f}M", flush=True)
torch.set_grad_enabled(False)  # inference-only benchmark

torch.cuda.reset_peak_memory_stats()

# ---------------- collect identical inputs ----------------
files = sorted(glob.glob("/tmp/work/ravdess_kag/Actor_*/*.wav"))
files = [f for f in files if os.path.basename(f).split("-")[1] == "01"
         and 19 <= int(os.path.basename(f).split("-")[-1].split(".")[0]) <= 24]
by_emo = {}
for f in files:
    by_emo.setdefault(os.path.basename(f).split("-")[2], []).append(f)
bench_files = []
for e in sorted(by_emo):
    bench_files.extend(by_emo[e][:3])
print(f"bench inputs: {len(bench_files)} files", flush=True)


# ---------------- the four measured paths ----------------
def teacher_pure(wav):
    src = F.layer_norm(wav[None], wav[None].shape)
    out = teacher.extract_features(src.to(DEV), padding_mask=None, mask=False)
    x = out["x"]                                    # (1, T, 768) @50Hz
    probs = torch.softmax(teacher.proj(x.mean(1)), -1)
    return x, probs


def teacher_pipeline(path):
    return fun.generate(path, granularity="frame", extract_embedding=True)


def student_total(wav):
    return enc(wav[None].to(DEV), torch.tensor([wav.numel()], device=DEV))


def adapter_only(states, lengths):
    return enc.from_states(states, lengths)


# ---------------- per-file quality + timing ----------------
print("=== per-file run (teacher & student on identical wavs) ===", flush=True)
agrees, jss, pool_coss, frame_coss, corrs = [], [], [], [], []
per_path_times = {"teacher_pure": [], "student_total": [], "adapter_only": []}

for f in bench_files:
    wav = load_wav(f).cpu()
    dur = wav.numel() / 16000

    # warm + time each path on this identical input
    t0 = time.perf_counter(); tx, tprob = teacher_pure(wav); torch.cuda.synchronize()
    per_path_times["teacher_pure"].append(time.perf_counter() - t0)
    states = tx.detach()
    lengths = torch.tensor([states.size(1)], device=DEV)

    t0 = time.perf_counter(); sout = student_total(wav); torch.cuda.synchronize()
    per_path_times["student_total"].append(time.perf_counter() - t0)

    t0 = time.perf_counter(); aout = adapter_only(states, lengths); torch.cuda.synchronize()
    per_path_times["adapter_only"].append(time.perf_counter() - t0)

    # quality: student vs teacher on the same input
    sp = sout["utt_probs"][0]
    tp = tprob[0]
    agrees.append(int(sp.argmax().item() == tp.argmax().item()))
    jss.append(js_div(sp[None], tp[None]))
    pool_coss.append(float(F.cosine_similarity(sout["utt_emb"][0], states.mean(1)[0], dim=0)))
    up = F.interpolate(sout["frame_embs"].transpose(1, 2), size=states.size(1), mode="nearest").transpose(1, 2)
    frame_coss.append(float(F.cosine_similarity(up[0], states[0], dim=-1).mean()))

    inten = sout["frame_intensity"][0, :, 0].cpu()
    T = int(lengths[0])
    tgt = rms_env(wav, T)
    tmax = torch.softmax(teacher.proj(states[0]), -1)[:, :8].max(-1).values.cpu().clamp(0, 1).sqrt()
    tgt = tmax * tgt
    tgt_ds = F.interpolate(tgt[None, None], size=inten.size(0), mode="linear")[0, 0]
    c = float(np.corrcoef(inten.numpy(), tgt_ds.numpy())[0, 1]) if inten.var() > 1e-6 and tgt_ds.var() > 1e-6 else None
    if c is not None:
        corrs.append(c)

    report["files"][os.path.basename(f)] = {
        "emotion": EMO_FIELD.get(os.path.basename(f).split("-")[2], "?"),
        "dur_s": round(dur, 2),
        "teacher_top": LABELS[int(tp.argmax())], "teacher_p": round(float(tp.max()), 3),
        "student_top": LABELS[int(sp.argmax())], "student_p": round(float(sp.max()), 3),
        "js": round(jss[-1], 5),
        "pooled_cos": round(pool_coss[-1], 5),
        "frame_cos": round(frame_coss[-1], 5),
        "intensity": {"mean": round(float(inten.mean()), 3), "std": round(float(inten.std()), 3),
                      "min": round(float(inten.min()), 3), "max": round(float(inten.max()), 3),
                      "corr_target": round(c, 3) if c is not None else None},
    }

n = len(bench_files)
report["quality_summary"] = {
    "n_files": n,
    "top1_agreement": float(np.mean(agrees)),
    "mean_js": float(np.mean(jss)),
    "mean_pooled_cos": float(np.mean(pool_coss)),
    "mean_frame_cos": float(np.mean(frame_coss)),
    "mean_intensity_corr": float(np.mean(corrs)),
}
durs = np.array([load_wav(f).numel() / 16000 for f in bench_files])
for k, v in per_path_times.items():
    v = np.array(v)
    report["paths"][k] = {
        "mean_ms_per_utt": float(1000 * v.mean()), "rtf_mean": float((v / durs).mean()),
        "utt_per_s": float(1.0 / v.mean()), "audio_x_realtime": float((durs / v).mean()),
    }
print(json.dumps(report["quality_summary"], indent=1), flush=True)

# ---------------- duration scaling (identical content, trimmed) ----------------
print("=== duration scaling ===", flush=True)
base = load_wav(bench_files[0])
long_wav = base.repeat(4)[: 10 * 16000]
report["scaling"]["duration"] = {}
for sec in [1, 2, 3, 4, 10]:
    w = long_wav[: sec * 16000]
    row = {"dur_s": sec}
    row["teacher_pure"] = timeit(lambda: teacher_pure(w))
    row["student_total"] = timeit(lambda: student_total(w))
    st, ln = (lambda r: (r[0].detach(), torch.tensor([r[0].size(1)], device=DEV)))(teacher_pure(w))
    row["adapter_only"] = timeit(lambda: adapter_only(st, ln))
    for k in ["teacher_pure", "student_total", "adapter_only"]:
        row[k]["rtf"] = row[k]["mean_ms"] / 1000.0 / sec
    report["scaling"]["duration"][f"{sec}s"] = row
    print(f"{sec}s: teacher {row['teacher_pure']['mean_ms']:.1f}ms (RTF {row['teacher_pure']['rtf']:.3f}) | "
          f"student {row['student_total']['mean_ms']:.1f}ms (RTF {row['student_total']['rtf']:.3f}) | "
          f"adapter {row['adapter_only']['mean_ms']:.2f}ms", flush=True)

# ---------------- funasr pipeline (stock user path) on 3 files ----------------
t0 = time.perf_counter()
for f in bench_files[:3]:
    teacher_pipeline(f)
pipe_t = (time.perf_counter() - t0) / 3
d3 = float(np.mean([load_wav(f).numel() / 16000 for f in bench_files[:3]]))
report["paths"]["teacher_pipeline_funasr"] = {"mean_ms_per_utt": pipe_t * 1000, "rtf_mean": pipe_t / d3,
                                              "utt_per_s": 1.0 / pipe_t, "audio_x_realtime": d3 / pipe_t}

# ---------------- adapter batch scaling ----------------
print("=== adapter batch scaling ===", flush=True)
w3 = long_wav[: 3 * 16000]
x3, _ = teacher_pure(w3)
T3 = x3.size(1)
report["scaling"]["adapter_batch"] = {}
for B in [1, 4, 8]:
    st = x3.expand(B, T3, 768).contiguous()
    ln = torch.full((B,), T3, device=DEV, dtype=torch.long)
    r = timeit(lambda: adapter_only(st, ln))
    r["utt_per_s"] = 1000.0 / r["mean_ms"] * B
    report["scaling"]["adapter_batch"][f"B{B}"] = r
    print(f"B={B}: {r['mean_ms']:.2f}ms -> {r['utt_per_s']:.1f} utt/s", flush=True)

# ---------------- fp16 autocast student ----------------
with torch.autocast("cuda", dtype=torch.float16):
    r16 = timeit(lambda: student_total(w3))
report["scaling"]["student_fp16_autocast_3s"] = r16
print(f"student fp16 3s: {r16['mean_ms']:.1f}ms", flush=True)

# ---------------- memory ----------------
report["memory"] = {"peak_gpu_allocated_mb": torch.cuda.max_memory_allocated() / 1e6,
                    "peak_gpu_reserved_mb": torch.cuda.max_memory_reserved() / 1e6}

# ---------------- CPU comparison (3s input) ----------------
print("=== CPU comparison (3 s input, 3 reps) ===", flush=True)
enc_cpu = SurgicalEmotionEncoder(teacher.to("cpu"), SurgicalConfig())
enc_cpu.adapter.load_state_dict(payload["adapter_state"])
enc_cpu.eval()
w3c = w3.cpu()


def student_cpu():
    with torch.no_grad():
        enc_cpu(w3c[None], torch.tensor([w3c.numel()]))


with torch.no_grad():
    tc = timeit(student_cpu, warmup=1, reps=3)
report["scaling"]["cpu_3s"] = {"student_total": tc}
print(f"cpu student 3s: {tc['mean_ms']:.0f}ms", flush=True)
teacher.to(DEV)
enc.to(DEV)

with open(OUT, "w") as fp:
    json.dump(report, fp, indent=1)
print("BENCH_DONE", flush=True)