"""Xavante - image_encoder.py - Encoder de imagem (Teorema 16.1/16.2)."""
from __future__ import annotations

import logging

import torch
import torch.nn as nn

logger = logging.getLogger(__name__)


class ImageEncoder(nn.Module):
    """CNN leve -> features em d_model."""

    def __init__(self, d_model: int = 512, in_channels: int = 3):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, 32, 4, 2, 1),
            nn.GELU(),
            nn.Conv2d(32, 64, 4, 2, 1),
            nn.GELU(),
            nn.Conv2d(64, 128, 4, 2, 1),
            nn.GELU(),
            nn.AdaptiveAvgPool2d((1, 1)),
        )
        self.proj = nn.Linear(128, d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [B, C, H, W]
        feat = self.conv(x).flatten(1)  # [B, 128]
        return self.proj(feat)  # [B, d_model]


__all__ = ["ImageEncoder"]