"""Surgical emotion encoder built on emotion2vec_plus_base.

Surgery performed on the donor model (iic/emotion2vec_plus_base):
  1. The frozen backbone (conv frontend + prenet + 8 AltBlocks, 768-d) is kept
     intact and exposes its *frame-level* continuous states at 50 Hz
     (20 ms/frame, stride 320 @ 16 kHz). The stock pipeline only surfaces the
     temporally mean-pooled utterance vector, discarding within-word emotion
     dynamics - the surgery re-exposes them.
  2. A lightweight trainable implant ("EmotionAdapter") fuses groups of 4
     consecutive teacher frames into continuous emotion vectors at 12.5 Hz
     (80 ms/frame), preserving intensity dynamics across words.
  3. Two distilled heads provide interpretable signals per 12.5 Hz frame:
       - emotion probabilities (9 classes, same labels as the teacher)
       - a continuous intensity scalar (distilled from teacher confidence)

Output alignment (designed for later dots.tts injection):
  dots.tts LLM token rate = 48000 / 7680 = 6.25 Hz (160 ms/token)
  emotion frame rate      = 12.5 Hz (80 ms/frame)  -> exactly 2 emotion
  frames per LLM token, aligned on frame boundaries.

Interface:
  enc = SurgicalEmotionEncoder.from_pretrained(funasr_model_dir, ckpt_path)
  out = enc(wav_16k)            # or enc.from_states(...) for cached states
  out["frame_embs"]     (B, T, 768)  @ 12.5 Hz, continuous
  out["frame_probs"]    (B, T, 9)    per-frame emotion probabilities
  out["frame_intensity"](B, T, 1)    per-frame continuous intensity in [0,1]
  out["utt_emb"]        (B, 768)     utterance-level embedding (mean pooled)
  out["utt_probs"]      (B, 9)
  out["frame_lengths"]  (B,)         valid frames per sample
"""

from __future__ import annotations

import json
import math
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F

TEACHER_SAMPLE_RATE = 16000
TEACHER_FRAME_RATE = 50.0  # samples/frame stride 320 @ 16 kHz (verified empirically)
TEACHER_DIM = 768
N_EMO_LABELS = 9
EMO_LABELS = [
    "angry", "disgusted", "fearful", "happy",
    "neutral", "other", "sad", "surprised", "<unk>",
]


@dataclass
class SurgicalConfig:
    group_size: int = 4              # 50 Hz / 4 = 12.5 Hz
    d_model: int = 512               # adapter working width
    out_dim: int = 768               # continuous emotion vector dim
    n_emotions: int = N_EMO_LABELS
    n_conv_layers: int = 3
    conv_kernel: int = 9
    conv_dilations: tuple = (1, 2, 4)
    dropout: float = 0.1

    @property
    def frame_rate_hz(self) -> float:
        return TEACHER_FRAME_RATE / self.group_size


class ConvFusionBlock(nn.Module):
    """Depthwise-separable 1D conv block at teacher frame rate.

    Blends neighbouring 50 Hz frames so each pooled 12.5 Hz vector integrates
    short-range dynamics before temporal pooling.
    """

    def __init__(self, d_model: int, kernel: int, dilation: int, dropout: float):
        super().__init__()
        padding = dilation * (kernel - 1) // 2
        self.dwconv = nn.Conv1d(d_model, d_model, kernel, padding=padding,
                                dilation=dilation, groups=d_model)
        self.pwconv = nn.Conv1d(d_model, 2 * d_model, 1)
        self.out_proj = nn.Conv1d(d_model, d_model, 1)
        self.norm = nn.LayerNorm(d_model)
        self.drop = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        # x: (B, T, D); mask: (B, T) True where valid
        h = x.transpose(1, 2)                      # (B, D, T)
        h = h * mask.unsqueeze(1).to(h.dtype)
        h = self.dwconv(h) * mask.unsqueeze(1).to(h.dtype)
        h = self.pwconv(h).transpose(1, 2)         # (B, T, 2D)
        h = F.glu(h, dim=-1)
        h = self.norm(h)
        h = h * F.silu(self.out_proj(h.transpose(1, 2)).transpose(1, 2))
        h = self.drop(h)
        return x + h * mask.unsqueeze(-1).to(h.dtype)


class GroupedSoftPool(nn.Module):
    """Attention pooling over non-overlapping groups of G frames.

    The scoring linear is zero-initialised, so at initialisation the pooling
    degenerates to a uniform mean over each group (a safe starting point that
    distillation then sharpens into selective pooling).
    """

    def __init__(self, d_model: int, group_size: int):
        super().__init__()
        self.group_size = group_size
        self.score = nn.Linear(d_model, 1)
        nn.init.zeros_(self.score.weight)
        nn.init.zeros_(self.score.bias)

    def forward(self, x: torch.Tensor, lengths: torch.Tensor):
        # x: (B, T, D), lengths: (B,) valid frames at teacher rate
        B, T, D = x.shape
        G = self.group_size
        T_out = (T + G - 1) // G
        pad = T_out * G - T
        if pad:
            x = F.pad(x, (0, 0, 0, pad))
            mask = torch.arange(T_out * G, device=x.device)[None, :] < lengths[:, None]
        else:
            mask = torch.arange(T, device=x.device)[None, :] < lengths[:, None]
        x = x.view(B, T_out, G, D)
        mask = mask.view(B, T_out, G)
        scores = self.score(x).squeeze(-1)                       # (B, T_out, G)
        scores = scores.masked_fill(~mask, float("-inf"))
        attn = torch.softmax(scores, dim=-1)
        attn = torch.nan_to_num(attn)                            # all-masked groups
        out = torch.einsum("btg,btgd->btd", attn, x)             # (B, T_out, D)
        out_lengths = torch.clamp((lengths + G - 1) // G, min=1)
        return out, out_lengths


class EmotionAdapter(nn.Module):
    """The trainable implant: 50 Hz teacher states -> 12.5 Hz emotion vectors."""

    def __init__(self, cfg: SurgicalConfig):
        super().__init__()
        self.cfg = cfg
        self.in_norm = nn.LayerNorm(TEACHER_DIM)
        self.in_proj = nn.Linear(TEACHER_DIM, cfg.d_model)
        self.fusion = nn.ModuleList([
            ConvFusionBlock(cfg.d_model, cfg.conv_kernel, d, cfg.dropout)
            for d in cfg.conv_dilations[: cfg.n_conv_layers]
        ])
        self.pool = GroupedSoftPool(cfg.d_model, cfg.group_size)
        self.out_norm = nn.LayerNorm(cfg.d_model)
        self.out_proj = nn.Linear(cfg.d_model, cfg.out_dim)
        self.emotion_head = nn.Linear(cfg.out_dim, cfg.n_emotions)
        self.intensity_head = nn.Linear(cfg.out_dim, 1)

    def forward(self, states: torch.Tensor, lengths: torch.Tensor) -> dict:
        """states: (B, T, 768) teacher frame states; lengths: (B,) valid T."""
        mask = torch.arange(states.size(1), device=states.device)[None, :] < lengths[:, None]
        x = self.in_norm(states)
        x = self.in_proj(x)
        x = F.gelu(x) * mask.unsqueeze(-1).to(x.dtype)
        for block in self.fusion:
            x = block(x, mask)
        x, out_lengths = self.pool(x, lengths)
        x = self.out_norm(x)
        embs = self.out_proj(x)                                   # (B, T*, 768) @12.5Hz
        omask = torch.arange(embs.size(1), device=embs.device)[None, :] < out_lengths[:, None]
        embs = embs * omask.unsqueeze(-1).to(embs.dtype)
        probs = torch.softmax(self.emotion_head(embs), dim=-1)
        intensity = torch.sigmoid(self.intensity_head(embs))      # (B, T*, 1)
        return {
            "frame_embs": embs,
            "frame_probs": probs,
            "frame_intensity": intensity,
            "frame_lengths": out_lengths,
        }


class SurgicalEmotionEncoder(nn.Module):
    """Frozen emotion2vec_plus_base teacher + surgical adapter."""

    def __init__(self, teacher_model, cfg: SurgicalConfig):
        super().__init__()
        self.cfg = cfg
        self.teacher = teacher_model
        self.teacher.eval()
        for p in self.teacher.parameters():
            p.requires_grad = False
        self.adapter = EmotionAdapter(cfg)

    # ---------- teacher extraction ----------

    @torch.no_grad()
    def extract_teacher_states(self, wav: torch.Tensor, wav_lengths: torch.Tensor):
        """wav: (B, S) float @16 kHz already loudness-normalised per sample;
        wav_lengths: (B,). Returns states (B, T, 768) fp32 and lengths (B,)."""
        outs = []
        for i in range(wav.size(0)):
            src = F.layer_norm(wav[i : i + 1, : int(wav_lengths[i])], (int(wav_lengths[i]),))
            out = self.teacher.extract_features(src.to(wav.device), padding_mask=None, mask=False)
            outs.append(out["x"].squeeze(0))
        lengths = torch.tensor([o.size(0) for o in outs], device=wav.device)
        T = int(lengths.max().item())
        states = torch.zeros(wav.size(0), T, TEACHER_DIM, device=wav.device)
        for i, o in enumerate(outs):
            states[i, : o.size(0)] = o
        return states.float(), lengths

    @torch.no_grad()
    def teacher_frame_probs(self, states: torch.Tensor, lengths: torch.Tensor):
        """Per-frame teacher probabilities + intensity proxy (max prob)."""
        logits = self.teacher.proj(states)                        # (B, T, 9)
        probs = torch.softmax(logits, dim=-1)
        inten = probs[..., : N_EMO_LABELS - 1].max(dim=-1, keepdim=True).values
        return probs, inten

    # ---------- full pipeline ----------

    def forward(self, wav: torch.Tensor, wav_lengths: torch.Tensor) -> dict:
        states, lengths = self.extract_teacher_states(wav, wav_lengths)
        return self.from_states(states, lengths)

    def from_states(self, states: torch.Tensor, lengths: torch.Tensor) -> dict:
        out = self.adapter(states, lengths)
        omask = (
            torch.arange(out["frame_embs"].size(1), device=states.device)[None, :]
            < out["frame_lengths"][:, None]
        ).unsqueeze(-1).to(out["frame_embs"].dtype)
        embs = out["frame_embs"] * omask
        utt_emb = (embs.sum(dim=1) / out["frame_lengths"].unsqueeze(-1)).to(embs.dtype)
        utt_logits = self.adapter.emotion_head(utt_emb)
        utt_probs = torch.softmax(utt_logits, dim=-1)
        out["utt_emb"] = utt_emb
        out["utt_probs"] = utt_probs
        return out

    # ---------- serialization ----------

    def save_surgical(self, path: str):
        payload = {
            "format_version": 1,
            "module_type": "surgical_emotion_encoder",
            "teacher_source": "iic/emotion2vec_plus_base",
            "sample_rate": TEACHER_SAMPLE_RATE,
            "teacher_frame_rate_hz": TEACHER_FRAME_RATE,
            "frame_rate_hz": self.cfg.frame_rate_hz,
            "labels": EMO_LABELS,
            "config": {
                "group_size": self.cfg.group_size,
                "d_model": self.cfg.d_model,
                "out_dim": self.cfg.out_dim,
                "n_emotions": self.cfg.n_emotions,
                "n_conv_layers": self.cfg.n_conv_layers,
                "conv_kernel": self.cfg.conv_kernel,
                "conv_dilations": list(self.cfg.conv_dilations),
                "dropout": self.cfg.dropout,
            },
            "adapter_state": {k: v.cpu() for k, v in self.adapter.state_dict().items()},
        }
        torch.save(payload, path)

    @classmethod
    def from_pretrained(cls, funasr_model_dir: str, ckpt_path: str | None = None,
                        device: str = "cuda:0"):
        import logging
        logging.disable(logging.CRITICAL)
        from funasr import AutoModel

        teacher_wrapper = AutoModel(model=funasr_model_dir, trust_remote_code=False,
                                    device=device, disable_update=True)
        teacher = teacher_wrapper.model
        cfg = SurgicalConfig()
        enc = cls(teacher, cfg).to(device)
        if ckpt_path is not None:
            payload = torch.load(ckpt_path, map_location=device, weights_only=False)
            enc.adapter.load_state_dict(payload["adapter_state"])
        enc.eval()
        return enc


def resample_to_16k(wav: torch.Tensor, orig_sr: int) -> torch.Tensor:
    if orig_sr == TEACHER_SAMPLE_RATE:
        return wav
    import torchaudio

    return torchaudio.functional.resample(wav, orig_sr, TEACHER_SAMPLE_RATE)