"""Encodeur image ViT from scratch (Phase 4).

Patch embedding + blocs transformer + **positional embedding appris**.

Choix : un ViT « plat » (patchs × transformer, sans CNN), comme dans « An Image is
Worth 16x16 Words », mais **entièrement réimplémenté**.

* **Patch embedding** : une convolution de noyau `patch_size` et de pas
  `patch_size` découpe l'image en patchs non chevauchants et projette chacun en
  `n_embd`. C'est une projection linéaire par patch, exprimée en conv pour la
  simplicité.
* **Positional embedding appris** (et **pas** RoPE) : sur une grille 2D de taille
  fixe et connue, apprendre un vecteur par position est simple, direct et
  suffisant. RoPE 2D apporte peu ici. RoPE reste utilisé pour le texte.
* **CLS token** : un token appris additionnel résume l'image (utilisé par le
  projecteur / la tête de classification). Sinon `pool='mean'`.

Justification complète : docs/architecture.md §3.
"""

from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn as nn

from src.model.block import TransformerBlock


@dataclass
class ViTConfig:
    image_size: int = 32
    patch_size: int = 8
    in_channels: int = 3
    n_layer: int = 4
    n_head: int = 4
    n_embd: int = 128
    dropout: float = 0.0
    bias: bool = False
    ffn_mult: int = 4
    pool: str = "cls"  # cls | mean
    use_rope: bool = False  # positional embedding appris par défaut


class PatchEmbedding(nn.Module):
    """Découpe l'image en patchs et projette chaque patch en ``n_embd``."""

    def __init__(self, image_size: int, patch_size: int, in_channels: int, n_embd: int):
        super().__init__()
        if image_size % patch_size != 0:
            raise ValueError(
                f"image_size ({image_size}) doit être divisible par patch_size ({patch_size})"
            )
        self.image_size = image_size
        self.patch_size = patch_size
        self.grid = image_size // patch_size
        self.n_patches = self.grid * self.grid
        self.proj = nn.Conv2d(in_channels, n_embd, kernel_size=patch_size, stride=patch_size)

    def forward(self, images: torch.Tensor) -> torch.Tensor:
        # images: (B, C, H, W) -> (B, n_patches, n_embd)
        x = self.proj(images)              # (B, n_embd, grid, grid)
        x = x.flatten(2).transpose(1, 2)   # (B, n_patches, n_embd)
        return x


class VisionEncoder(nn.Module):
    """Encodeur d'images : patch embedding + position apprise + N blocs."""

    def __init__(self, config: ViTConfig):
        super().__init__()
        self.config = config
        self.patch_embed = PatchEmbedding(
            config.image_size, config.patch_size, config.in_channels, 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=config.use_rope,
                    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)

    def forward(self, images: torch.Tensor) -> torch.Tensor:
        """Retourne les features ``(B, n_tokens, n_embd)`` (CLS inclus si ``pool='cls'``)."""
        x = self.patch_embed(images)
        if self.cls_token is not None:
            cls = self.cls_token.expand(x.size(0), -1, -1)
            x = torch.cat([cls, 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:
        """Réduit les tokens en un vecteur ``(B, n_embd)``."""
        if self.config.pool == "cls":
            return tokens[:, 0]
        return tokens.mean(dim=1)


class ImageClassifier(nn.Module):
    """Tête de classification (utilisée pour l'overfit de test, Phase 4)."""

    def __init__(self, config: ViTConfig, n_classes: int):
        super().__init__()
        self.encoder = VisionEncoder(config)
        self.head = nn.Linear(config.n_embd, n_classes)

    def forward(self, images: torch.Tensor) -> torch.Tensor:
        tokens = self.encoder(images)
        return self.head(self.encoder.pool(tokens))


class ImageDecoder(nn.Module):
    """Tête de reconstruction (auto-encodeur jouet), pour prouver que l'encodeur
    capture de l'information sans utiliser d'étiquette."""

    def __init__(self, config: ViTConfig, hidden: int = 256):
        super().__init__()
        self.encoder = VisionEncoder(config)
        out_dim = config.in_channels * config.image_size * config.image_size
        self.net = nn.Sequential(
            nn.Linear(config.n_embd, hidden), nn.GELU(), nn.Linear(hidden, out_dim)
        )
        self.image_size = config.image_size
        self.in_channels = config.in_channels

    def forward(self, images: torch.Tensor) -> torch.Tensor:
        feat = self.encoder.pool(self.encoder(images))
        out = self.net(feat)
        return out.view(-1, self.in_channels, self.image_size, self.image_size)