"""Encodeur audio from scratch (Phase 7).

Pipeline entièrement maison :

1. **Mel-spectrogramme** calculé à la main (`src/data/audio_features.py`) : FFT de
   PyTorch, banc de filtres mel triangulaires, log. Aucun modèle pré-entraîné.
2. **Patch embedding** : découpe le mel-spectrogramme en patchs (fréquence × temps)
   et projette chacun en `n_embd` — comme un ViT mais sur une image 1 canal.
3. **Positional embedding appris** + blocs transformer **bidirectionnels** (l'audio
   n'est pas causal : on voit tout le clip).
4. **Pooling** (CLS ou moyenne) → vecteur projeté dans l'espace texte.

Justification : docs/architecture.md §6.
"""

from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn as nn

from src.model.block import TransformerBlock


@dataclass
class AudioConfig:
    n_mels: int = 64
    n_frames: int = 96
    patch_mels: int = 16
    patch_frames: int = 16
    n_layer: int = 2
    n_head: int = 4
    n_embd: int = 128
    dropout: float = 0.0
    bias: bool = False
    ffn_mult: int = 4
    pool: str = "mean"  # mean | cls


class MelPatchEmbedding(nn.Module):
    """Découpe un mel-spectrogramme (B, 1, n_mels, n_frames) en patchs."""

    def __init__(self, n_mels: int, n_frames: int, patch_mels: int, patch_frames: int,
                 n_embd: int):
        super().__init__()
        if n_mels % patch_mels or n_frames % patch_frames:
            raise ValueError("n_mels/n_frames doivent être divisibles par la taille de patch")
        self.grid_mels = n_mels // patch_mels
        self.grid_frames = n_frames // patch_frames
        self.n_patches = self.grid_mels * self.grid_frames
        self.proj = nn.Conv2d(1, n_embd, kernel_size=(patch_mels, patch_frames),
                              stride=(patch_mels, patch_frames))

    def forward(self, mel: torch.Tensor) -> torch.Tensor:
        x = self.proj(mel)                       # (B, D, gm, gf)
        return x.flatten(2).transpose(1, 2)      # (B, n_patches, D)


class AudioEncoder(nn.Module):
    """Encodeur de mel-spectrogrammes : patch embedding + position + blocs."""

    def __init__(self, config: AudioConfig):
        super().__init__()
        self.config = config
        self.patch_embed = MelPatchEmbedding(
            config.n_mels, config.n_frames, config.patch_mels, config.patch_frames,
            config.n_embd,
        )
        n_tokens = self.patch_embed.n_patches + (1 if config.pool == "cls" else 0)
        self.cls_token = (
            nn.Parameter(torch.zeros(1, 1, config.n_embd)) if config.pool == "cls" else None
        )
        self.pos_embedding = nn.Parameter(torch.zeros(1, n_tokens, config.n_embd))
        self.dropout = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList([
            TransformerBlock(config.n_embd, config.n_head, dropout=config.dropout,
                             bias=config.bias, ffn_mult=config.ffn_mult,
                             use_rope=False, causal=False)
            for _ in range(config.n_layer)
        ])
        self.ln_f = nn.LayerNorm(config.n_embd, bias=config.bias)
        self._init_weights()

    def _init_weights(self) -> None:
        nn.init.normal_(self.pos_embedding, std=0.02)
        if self.cls_token is not None:
            nn.init.normal_(self.cls_token, std=0.02)
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, std=0.02)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Conv2d):
                nn.init.normal_(module.weight, std=0.02)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)

    @property
    def n_audio_tokens(self) -> int:
        return self.patch_embed.n_patches

    def forward(self, mel: torch.Tensor) -> torch.Tensor:
        x = self.patch_embed(mel)
        if self.cls_token is not None:
            x = torch.cat([self.cls_token.expand(x.size(0), -1, -1), x], dim=1)
        x = self.dropout(x + self.pos_embedding)
        for block in self.blocks:
            x = block(x)
        return self.ln_f(x)

    def pool(self, tokens: torch.Tensor) -> torch.Tensor:
        if self.config.pool == "cls":
            return tokens[:, 0]
        return tokens.mean(dim=1)