"""Modèle CLIP maison (Phase 5) : alignement image-texte from scratch.

Deux encodeurs (texte + image) produisent chacun un vecteur, projeté dans un espace
commun puis **normalisé L2**. La similarité est un produit scalaire. Aucun modèle
CLIP pré-entraîné n'est utilisé : encodeur image = `VisionEncoder` maison (Phase 4),
encodeur texte = transformer maison (Phase 2).

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

from __future__ import annotations

from dataclasses import dataclass, field

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

from src.model.projectors import MLPProjector
from src.model.text_model import CortexTextLM, TextConfig
from src.model.vit import VisionEncoder, ViTConfig


def _default_text_config() -> TextConfig:
    return TextConfig(vocab_size=32000, n_layer=2, n_head=4, n_embd=128,
                      ctx_len=32, tie_embeddings=True)


def _default_vision_config() -> ViTConfig:
    return ViTConfig(image_size=32, patch_size=8, n_layer=2, n_head=4, n_embd=128)


@dataclass
class ClipConfig:
    embed_dim: int = 64          # dimension de l'espace commun (image/texte)
    text: TextConfig = field(default_factory=_default_text_config)
    vision: ViTConfig = field(default_factory=_default_vision_config)
    init_temperature: float = 0.07
    learn_temperature: bool = True


class TextEncoder(nn.Module):
    """Transformer texte (Phase 2) -> pooling masqué -> vecteur ``(B, n_embd)``."""

    def __init__(self, config: TextConfig):
        super().__init__()
        self.backbone = CortexTextLM(config)

    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None):
        hidden = self.backbone.encode(input_ids)  # (B, T, n_embd), avant la tête
        if attention_mask is None:
            return hidden.mean(dim=1)
        mask = attention_mask.unsqueeze(-1).float()
        return (hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-6)


class CortexClip(nn.Module):
    """CLIP maison : encodeurs texte + image, projecteurs, `logit_scale` apprise."""

    def __init__(self, config: ClipConfig):
        super().__init__()
        self.config = config
        self.text_encoder = TextEncoder(config.text)
        self.vision_encoder = VisionEncoder(config.vision)
        self.text_proj = MLPProjector(config.text.n_embd, config.embed_dim)
        self.image_proj = MLPProjector(config.vision.n_embd, config.embed_dim)
        init_logit_scale = float(torch.log(torch.tensor(1.0 / config.init_temperature)))
        if config.learn_temperature:
            self.logit_scale = nn.Parameter(torch.tensor(init_logit_scale))
        else:
            self.register_buffer("logit_scale", torch.tensor(init_logit_scale))

    def encode_text(self, input_ids, attention_mask=None):
        pooled = self.text_proj(self.text_encoder(input_ids, attention_mask))
        return F.normalize(pooled, dim=-1)

    def encode_image(self, images):
        feat = self.vision_encoder.pool(self.vision_encoder(images))
        return F.normalize(self.image_proj(feat), dim=-1)

    def forward(self, images, input_ids, attention_mask=None):
        img = self.encode_image(images)
        txt = self.encode_text(input_ids, attention_mask)
        scale = self.logit_scale.exp().clamp(max=100.0)
        logits_per_image = scale * img @ txt.t()
        return logits_per_image, logits_per_image.t(), img, txt


def contrastive_loss(logits_per_image: torch.Tensor, logits_per_text: torch.Tensor) -> torch.Tensor:
    """Perte contrastive symétrique (écrite à la main).

    La bonne paire est sur la diagonale : image i ↔ texte i.
    """
    batch = logits_per_image.size(0)
    targets = torch.arange(batch, device=logits_per_image.device)
    return 0.5 * (F.cross_entropy(logits_per_image, targets)
                  + F.cross_entropy(logits_per_text, targets))


def clip_loss(model: CortexClip, images, input_ids, attention_mask=None) -> torch.Tensor:
    logits_per_image, logits_per_text, _, _ = model(images, input_ids, attention_mask)
    return contrastive_loss(logits_per_image, logits_per_text)


@torch.no_grad()
def retrieval_accuracy(model: CortexClip, images, input_ids, attention_mask=None) -> float:
    """Précision de retrieval image->texte (diagonale) sur un lot **sans doublon**."""
    img = model.encode_image(images)
    txt = model.encode_text(input_ids, attention_mask)
    preds = (img @ txt.t()).argmax(dim=-1)
    targets = torch.arange(images.size(0), device=images.device)
    return (preds == targets).float().mean().item()


@torch.no_grad()
def class_retrieval_accuracy(
    model: CortexClip,
    images: torch.Tensor,
    bank_ids: torch.Tensor,
    bank_mask: torch.Tensor,
    label_ids: torch.Tensor,
) -> float:
    """Retrieval image->texte **au niveau des classes**.

    ``bank_ids``/``bank_mask`` décrivent la banque de légendes **uniques** (ici les
    9 classes) ; ``label_ids[i]`` est l'indice de la vraie légende de l'image ``i``
    dans cette banque. C'est la bonne mesure quand les légendes se répètent (une
    légende partagée par 10 images ne peut pas être « la » bonne par un argmax
    sur 90 colonnes).
    """
    img = model.encode_image(images)
    txt = model.encode_text(bank_ids, bank_mask)
    preds = (img @ txt.t()).argmax(dim=-1)
    return (preds == label_ids).float().mean().item()