"""End-of-turn classifier: Whisper encoder (base for the released model) on the last 8 s of audio -> P(speaker is done)."""
import math, torch, torch.nn as nn
from transformers import WhisperConfig, WhisperModel
from transformers.models.whisper.modeling_whisper import WhisperEncoder
SR, CTX_S, N_MELS, HOP, NFFT = 16000, 8, 80, 160, 400
N_SAMPLES = SR * CTX_S; N_FRAMES = N_SAMPLES // HOP          # 128000 samples -> 800 mel frames -> 400 positions
class LogMel(nn.Module):
    """Whisper's log-mel front end in torch (matches WhisperFeatureExtractor up to float error)."""
    def __init__(self):
        super().__init__()
        from transformers import WhisperFeatureExtractor
        fe = WhisperFeatureExtractor(feature_size=N_MELS)
        self.register_buffer("window", torch.hann_window(NFFT), persistent=False)
        self.register_buffer("fb", torch.tensor(fe.mel_filters, dtype=torch.float32).T, persistent=False)   # [80, 201]
    def forward(self, x):                                          # x [B, N_SAMPLES] float
        st = torch.stft(x.float(), NFFT, HOP, window=self.window, return_complex=True)
        mag = st[..., :-1].abs() ** 2                               # [B, 201, 800]
        mel = torch.clamp(self.fb @ mag, min=1e-10).log10()
        mel = torch.maximum(mel, mel.amax(dim=(1, 2), keepdim=True) - 8.0)
        return (mel + 4.0) / 4.0
class EOT(nn.Module):
    def __init__(self, base="openai/whisper-tiny", pretrained=True):
        super().__init__()
        cfg = WhisperConfig.from_pretrained(base); cfg.max_source_positions = N_FRAMES // 2
        self.mel = LogMel(); self.enc = WhisperEncoder(cfg)
        if pretrained:
            sd = WhisperModel.from_pretrained(base).encoder.state_dict()
            sd["embed_positions.weight"] = sd["embed_positions.weight"][: cfg.max_source_positions]
            self.enc.load_state_dict(sd)
        d = cfg.d_model
        self.pool = nn.Linear(d, 1)
        self.head = nn.Sequential(nn.Linear(d, 256), nn.GELU(), nn.Dropout(0.1), nn.Linear(256, 1))
    def forward_mel(self, mel):
        h = self.enc(mel).last_hidden_state                        # [B, 400, d]
        a = self.pool(h).softmax(1)
        return self.head((a * h).sum(1)).squeeze(-1)              # logit
    def forward(self, wav):
        return self.forward_mel(self.mel(wav))


class ConvLogMel(LogMel):
    """Same as LogMel but the STFT is a fixed DFT conv1d (center=True, reflect pad): exports to any ONNX runtime."""
    def __init__(self):
        super().__init__()
        n = torch.arange(NFFT, dtype=torch.float64); k = torch.arange(NFFT // 2 + 1, dtype=torch.float64)[:, None]
        ang = 2 * math.pi * k * n / NFFT; win = torch.hann_window(NFFT, dtype=torch.float64)
        basis = torch.cat([torch.cos(ang) * win, -torch.sin(ang) * win])[:, None, :].float()      # [402, 1, 400]
        self.register_buffer("basis", basis, persistent=False)
    def forward(self, x):
        x = torch.nn.functional.pad(x.float()[:, None, :], (NFFT // 2, NFFT // 2), mode="reflect")
        st = torch.nn.functional.conv1d(x, self.basis, stride=HOP)                                   # [B, 402, 801]
        re, im = st[:, : NFFT // 2 + 1], st[:, NFFT // 2 + 1:]
        mag = (re ** 2 + im ** 2)[..., :-1]
        mel = torch.clamp(self.fb @ mag, min=1e-10).log10()
        mel = torch.maximum(mel, mel.amax(dim=(1, 2), keepdim=True) - 8.0)
        return (mel + 4.0) / 4.0