"""Self-distillation training for the surgical emotion adapter.

Teacher = frozen emotion2vec_plus_base (frame states @ 50 Hz + pooled classifier).
Student = EmotionAdapter producing 12.5 Hz continuous vectors + heads.

Losses (no emotion labels required):
  L_frame   : cosine + LN-MSE between up-sampled student vectors and teacher states
  L_utt_cos : cosine of pooled embeddings
  L_utt_kl  : KL(teacher pooled probs || student pooled probs)
  L_frm_kl  : per-12.5Hz-frame KL between teacher group-mean probs and student probs
  L_inten   : MSE between student intensity and teacher max-prob proxy

Data: RAVDESS Actor_01/02 (already downloaded) + any 16 kHz wav tree.
Teacher states are precomputed once; epochs iterate over the cache.
"""

from __future__ import annotations

import glob
import json
import math
import os
import random
import sys
import time

import numpy as np
import torch
import torch.nn.functional as F
import torchaudio

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from surgical_emotion import (
    EMO_LABELS, SurgicalConfig, SurgicalEmotionEncoder,
)

FUNASR_DIR = "/tmp/work/emotion2vec_plus_base"
RAVDESS_DIR = "/tmp/work/ravdess_kag"
OUT_DIR = "/tmp/work/surgical_out"
MAX_SEC = 12.0
BATCH = 8
LR = 3e-4
STEPS = 900
WARMUP = 50
SEED = 42
VAL_FRAC = 0.1
FP16 = True
MAX_TRAIN_FILES = 900
TRAIN_ACTORS = set(range(1, 19))      # actors 1..18 for train/val
TEST_ACTORS = set(range(19, 25))      # actors 19..24 held out for eval


def _actor_of(path: str) -> int:
    try:
        return int(os.path.basename(path).split("-")[-1].split(".")[0])
    except Exception:
        return 0


def _is_speech(path: str) -> bool:
    # filename: modality-channel-emotion-intensity-statement-repetition-actor
    try:
        return os.path.basename(path).split("-")[1] == "01"  # vocal channel = speech
    except Exception:
        return False


def gather_wavs():
    files = sorted(glob.glob(os.path.join(RAVDESS_DIR, "**", "*.wav"), recursive=True))
    files = [f for f in files if _is_speech(f) and _actor_of(f) in TRAIN_ACTORS]
    random.shuffle(files)
    files = files[:MAX_TRAIN_FILES]
    if not files:  # fallback: any wav on the machine for a smoke run
        files = [FUNASR_DIR + "/example/test.wav"]
    return files


def ravdess_label(path: str) -> str:
    # modality-vocality-emotion-intensity-statement-repetition-actor.wav
    emo = {"01": "neutral", "02": "calm", "03": "happy", "04": "sad",
           "05": "angry", "06": "fearful", "07": "disgust", "08": "surprised"}
    try:
        return emo.get(os.path.basename(path).split("-")[2], "?")
    except Exception:
        return "?"


def load_wav(path: str):
    wav, sr = torchaudio.load(path)
    wav = wav.mean(dim=0)
    if sr != 16000:
        wav = torchaudio.functional.resample(wav, sr, 16000)
    return wav


def rms_envelope_50hz(wav: torch.Tensor, T_frames: int) -> torch.Tensor:
    """Short-time RMS at ~50 Hz, smoothed and per-utterance min-maxed to [0,1]."""
    n = wav.numel()
    frames = torch.zeros(T_frames)
    for i in range(T_frames):
        seg = wav[i * 320 : (i + 1) * 320]
        if seg.numel() == 0:
            frames[i] = frames[i - 1] if i else 0.0
        else:
            frames[i] = float(seg.pow(2).mean().sqrt())
    if T_frames >= 9:
        k = torch.ones(9) / 9.0
        frames = torch.nn.functional.conv1d(
            frames[None, None], k[None, None], padding=4)[0, 0]
    lo, hi = frames.min(), frames.max()
    if hi - lo < 1e-6:
        return torch.full_like(frames, 0.5)
    return (frames - lo) / (hi - lo)


def main():
    random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
    os.makedirs(OUT_DIR, exist_ok=True)
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    log = open(os.path.join(OUT_DIR, "train.log"), "a", buffering=1)

    def P(*a):
        print(*a); print(*a, file=log)

    P(f"=== distill start {time.strftime('%H:%M:%S')} device={device} ===")

    enc = SurgicalEmotionEncoder.from_pretrained(FUNASR_DIR, device=device)
    adapter = enc.adapter
    opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01)
    sched = torch.optim.lr_scheduler.LambdaLR(
        opt, lambda s: min((s + 1) / WARMUP, 0.5 * (1 + math.cos(math.pi * s / STEPS)))
    )
    scaler = torch.amp.GradScaler("cuda", enabled=FP16)

    # ---------------- data ----------------
    files = gather_wavs()
    P(f"wav files: {len(files)}")
    meta = [(f, ravdess_label(f)) for f in files]

    # ---------------- precompute teacher states ----------------
    cache = []
    t0 = time.time()
    for i, (f, lab) in enumerate(meta):
        try:
            wav = load_wav(f)
        except Exception:
            continue
        if wav.numel() < 8000:            # <0.5s skip
            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(device), torch.tensor([wav.numel()], device=device))
        probs, inten = enc.teacher_frame_probs(states, lengths)
        T = int(lengths[0].item())
        rms = rms_envelope_50hz(wav, T)
        # composite intensity target: confidence-gated expression energy
        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(),
            "label": lab, "file": os.path.basename(f),
        })
        if (i + 1) % 48 == 0:
            P(f"  teacher {i+1}/{len(meta)}  ({time.time()-t0:.0f}s)")
    P(f"teacher cache: {len(cache)} utts, {time.time()-t0:.0f}s")

    random.shuffle(cache)
    n_val = max(8, int(len(cache) * VAL_FRAC))
    val, train = cache[:n_val], cache[n_val:]
    P(f"train={len(train)}  val={len(val)}")

    def make_batch(pool):
        items = random.sample(pool, min(BATCH, len(pool)))
        T = max(it["states"].size(0) for it in items)
        B = len(items)
        states = torch.zeros(B, T, 768)
        probs = torch.zeros(B, T, 9)
        inten = torch.zeros(B, T, 1)
        lengths = torch.zeros(B, dtype=torch.long)
        for j, it in enumerate(items):
            t = it["states"].size(0)
            states[j, :t] = it["states"].float()
            probs[j, :t] = it["probs"].float()
            inten[j, :t] = it["inten"].float()
            lengths[j] = t
        return states.to(device), probs.to(device), inten.to(device), lengths.to(device)

    def compute_loss(batch, train_mode):
        states, tprobs, tinten, lengths = batch
        out = adapter(states, lengths) if train_mode else adapter(states, lengths)
        G = SurgicalConfig().group_size
        embs, flen = out["frame_embs"], out["frame_lengths"]
        T = states.size(1)
        T_out = embs.size(1)

        # 1) frame reconstruction: upsample student to teacher rate
        up = F.interpolate(embs.transpose(1, 2), size=T, mode="nearest").transpose(1, 2)
        m = (torch.arange(T, device=states.device)[None, :] < lengths[:, None]).unsqueeze(-1)
        cos = F.cosine_similarity(up, states, dim=-1).unsqueeze(-1)
        l_frame = ((1 - cos) * m).sum() / m.sum().clamp(min=1)
        l_mse = F.mse_loss(F.layer_norm(up, (768,)), F.layer_norm(states, (768,)), reduction="none")
        l_mse = (l_mse * m).sum() / m.sum().clamp(min=1)

        # 2) utterance losses
        om = (torch.arange(T_out, device=states.device)[None, :] < flen[:, None]).unsqueeze(-1)
        utt_s = (embs * om).sum(1) / flen.unsqueeze(-1).float()
        utt_probs_s = torch.softmax(enc.adapter.emotion_head(utt_s.float()), dim=-1)
        utt_t = (states * m).sum(1) / lengths.unsqueeze(-1).float()
        l_utt_cos = (1 - F.cosine_similarity(utt_s, utt_t, dim=-1)).mean()
        utt_probs_t = torch.softmax(enc.teacher.proj(utt_t), dim=-1)
        l_utt_kl = F.kl_div(
            torch.log(utt_probs_s.clamp(min=1e-8)), utt_probs_t,
            reduction="batchmean")

        # 3) frame-level classification distillation (group-mean teacher probs)
        tprobs_g = F.interpolate(tprobs.transpose(1, 2), size=T_out, mode="area").transpose(1, 2)
        tprobs_g = tprobs_g / tprobs_g.sum(-1, keepdim=True).clamp(min=1e-8)
        l_frm = F.kl_div(
            torch.log(out["frame_probs"].clamp(min=1e-8)), tprobs_g,
            reduction="none").sum(-1)
        l_frm = (l_frm * om.squeeze(-1)).sum() / om.sum().clamp(min=1)

        # 4) intensity distillation
        tinten_g = F.interpolate(tinten.transpose(1, 2), size=T_out, mode="area").transpose(1, 2)
        l_int = (((out["frame_intensity"] - tinten_g) ** 2) * om).sum() / om.sum().clamp(min=1)

        loss = (1.0 * l_frame + 0.5 * l_mse + 0.5 * l_utt_cos
                + 1.0 * l_utt_kl + 1.0 * l_frm + 0.5 * l_int)
        return loss, {
            "frame": l_frame.item(), "mse": l_mse.item(), "utt_cos": l_utt_cos.item(),
            "utt_kl": l_utt_kl.item(), "frm_kl": l_frm.item(), "inten": l_int.item(),
        }

    # ---------------- train ----------------
    best_val = float("inf")
    step = 0
    while step < STEPS:
        adapter.train()
        batch = make_batch(train)
        with torch.autocast("cuda", dtype=torch.float16, enabled=FP16):
            loss, parts = compute_loss(batch, True)
        opt.zero_grad(set_to_none=True)
        scaler.scale(loss).backward()
        scaler.unscale_(opt)
        torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0)
        scaler.step(opt); scaler.update(); sched.step()
        if step % 20 == 0:
            P(f"step {step:4d}  loss {loss.item():.4f}  " +
              " ".join(f"{k}={v:.3f}" for k, v in parts.items()))
        if step % 100 == 0 or step == STEPS - 1:
            adapter.eval()
            with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16, enabled=FP16):
                vb = make_batch(val)
                vloss, vparts = compute_loss(vb, False)
            P(f"  >> val step {step}: {vloss.item():.4f} " +
              " ".join(f"{k}={v:.3f}" for k, v in vparts.items()))
            if vloss.item() < best_val:
                best_val = vloss.item()
                enc.save_surgical(os.path.join(OUT_DIR, "surgical_emotion_module.pt"))
                P("  >> saved best")
        step += 1

    enc.save_surgical(os.path.join(OUT_DIR, "surgical_emotion_module_last.pt"))
    P(f"=== done, best_val={best_val:.4f} ===")


if __name__ == "__main__":
    main()