from __future__ import annotations

import torch
import torch.nn as nn

PARAKEET_NAME = "nvidia/parakeet-tdt-0.6b-v3"
PARAKEET_DIM = 1024
LM_DIM = 5376
PROJ_HIDDEN = 4096
HOP_LENGTH = 160
SUBSAMPLE_LAYERS = 3


def mel_frames_for(n_samples: int) -> int:
    return int(n_samples) // HOP_LENGTH


def valid_frames_for(n_samples: int, sr: int = 16000) -> int:
    n = mel_frames_for(n_samples)
    for _ in range(SUBSAMPLE_LAYERS):
        n = (n - 1) // 2 + 1
    return max(1, n)


class SquaredReLU(nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.pow(torch.nn.functional.relu(x), 2)


class RMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float = 1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
        return (self.weight.to(torch.float32) * hidden_states).to(input_dtype)


class ParakeetSoundProjection(nn.Module):
    def __init__(
        self,
        in_dim: int = PARAKEET_DIM,
        out_dim: int = LM_DIM,
        hidden: int = PROJ_HIDDEN,
        bias: bool = False,
        eps: float = 1e-5,
        out_dtype: torch.dtype = torch.bfloat16,
        target_rms: float | None = None,
    ):
        super().__init__()
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.out_dtype = out_dtype

        self.norm = RMSNorm(in_dim, eps=eps)
        self.linear1 = nn.Linear(in_dim, hidden, bias=bias)
        self.activation = SquaredReLU()
        self.linear2 = nn.Linear(hidden, out_dim, bias=bias)

        self.register_buffer("target_rms", torch.tensor(
            float(target_rms) if target_rms else 0.0))
        self.out_gain = nn.Parameter(torch.ones(()))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        wdt = self.linear1.weight.dtype
        out = self.linear2(self.activation(self.linear1(self.norm(x.to(wdt)))))
        if float(self.target_rms) > 0:
            rms = out.pow(2).mean(dim=-1, keepdim=True).clamp_min(1e-12).sqrt()
            out = out / rms * self.target_rms * self.out_gain
        return out.to(self.out_dtype)


class ParakeetAudioFrontEnd(nn.Module):
    def __init__(
        self,
        parakeet_name: str = PARAKEET_NAME,
        hidden: int = PROJ_HIDDEN,
        out_dim: int = LM_DIM,
        param_dtype: torch.dtype = torch.float32,
        encoder_dtype: torch.dtype = torch.bfloat16,
        target_rms: float | None = None,
    ):
        super().__init__()
        self.encoder = _load_parakeet_encoder(parakeet_name, encoder_dtype)
        for p in self.encoder.parameters():
            p.requires_grad_(False)
        self.encoder.eval()
        self.encoder_dtype = encoder_dtype

        self.projector = ParakeetSoundProjection(
            in_dim=self.encoder.config.hidden_size,
            out_dim=out_dim,
            hidden=hidden,
            out_dtype=encoder_dtype,
            target_rms=target_rms,
        ).to(param_dtype)

    def train(self, mode: bool = True):
        super().train(mode)
        self.encoder.eval()
        return self

    def _encode(self, input_features, attention_mask=None) -> torch.Tensor:
        x = input_features.to(self.encoder_dtype)
        lengths = None
        if attention_mask is not None:
            lengths = attention_mask.sum(dim=-1).tolist()

        if lengths is None or (len(set(lengths)) == 1
                               and lengths[0] == x.shape[1]):
            with torch.no_grad():
                out = self.encoder(input_features=x,
                                   attention_mask=attention_mask)
            return out.last_hidden_state.detach()

        rows = []
        with torch.no_grad():
            for b, length in enumerate(lengths):
                length = int(length)
                ones = x.new_ones((1, length), dtype=attention_mask.dtype)
                out = self.encoder(
                    input_features=x[b:b + 1, :length].contiguous(),
                    attention_mask=ones)
                rows.append(out.last_hidden_state[0].detach())
        width = max(r.shape[0] for r in rows)
        return torch.stack([
            torch.nn.functional.pad(r, (0, 0, 0, width - r.shape[0]))
            for r in rows
        ])

    def forward(self, input_features, attention_mask=None) -> torch.Tensor:
        feats = self._encode(input_features, attention_mask)
        return self.projector(feats)


def _load_parakeet_encoder(name: str, dtype: torch.dtype):
    from transformers import ParakeetEncoder
    try:
        return ParakeetEncoder.from_pretrained(name, dtype=dtype)
    except Exception:
        from transformers import AutoModelForTDT
        full = AutoModelForTDT.from_pretrained(name, dtype=dtype)
        enc = getattr(full, "encoder", None)
        if enc is None:
            enc = full.model.encoder
        return enc


def merge_audio_into_embeds(
    model,
    frontend: ParakeetAudioFrontEnd,
    input_ids: torch.Tensor,
    input_features: torch.Tensor,
    valid_frames: torch.Tensor | list[int],
    audio_token_id: int,
    encoder_attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    embed = model.get_input_embeddings()
    inputs_embeds = embed(input_ids)
    audio_embeds = frontend(input_features, encoder_attention_mask)

    if isinstance(valid_frames, torch.Tensor):
        valid_frames = valid_frames.tolist()

    parts = []
    for b in range(input_ids.shape[0]):
        n = int(valid_frames[b])
        n_slots = int((input_ids[b] == audio_token_id).sum())
        if n_slots != n:
            raise ValueError(
                f"row {b}: {n_slots} audio tokens but {n} valid Parakeet frames"
            )
        if audio_embeds.shape[1] < n:
            raise ValueError(
                f"row {b}: encoder emitted {audio_embeds.shape[1]} frames "
                f"but {n} were predicted"
            )
        parts.append(audio_embeds[b, :n])
    audio_flat = torch.cat(parts, dim=0).to(inputs_embeds.dtype)
    mask = (input_ids == audio_token_id).unsqueeze(-1).expand_as(inputs_embeds)
    return inputs_embeds.masked_scatter(mask, audio_flat)