"""In-repo CNN vision encoder for multimodal training."""

from typing import Sequence
import torch
import torch.nn as nn


class CNNEncoder(nn.Module):
    """A compact convolutional encoder for single-image multimodal inputs."""

    def __init__(
        self,
        image_size: int,
        output_dim: int,
        channels: Sequence[int] = (32, 64, 128),
        kernel_size: int = 3,
    ):
        """Initialize the CNN encoder."""
        super().__init__()

        if image_size < 8:
            raise ValueError("image_size must be at least 8")
        if output_dim < 1:
            raise ValueError("output_dim must be positive")
        if not channels:
            raise ValueError("channels must contain at least one stage")

        layers: list[nn.Module] = []
        in_channels = 3
        stride = 2
        padding = kernel_size // 2

        for channel_dim in channels:
            layers.extend([
                nn.Conv2d(in_channels, channel_dim, kernel_size=kernel_size, stride=stride, padding=padding),
                nn.BatchNorm2d(channel_dim),
                nn.GELU(),
            ])
            in_channels = channel_dim

        self.backbone = nn.Sequential(*layers)
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.projection = nn.Linear(in_channels, output_dim)

    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
        """Encode pixel values into a single feature vector per image."""
        features = self.backbone(pixel_values)
        pooled = self.pool(features).flatten(1)
        return self.projection(pooled)