"""Reusable image encoder; AS2 itself does not require an image modality."""

import torch
from torch import nn


class CellCNN(nn.Module):
    """Shared-weight CNN that encodes a single 28×28 grayscale cell image.

    Args:
        d_model: Output embedding dimension.
        backbone: Architecture variant.
            ``"default"`` — lightweight 3-conv CNN (original SoftASP design).
            ``"resnet18"`` — ResNet-18 adapted for 1-channel 28×28 input;
            weights are loaded externally via ``pretrained_from`` or trained
            from scratch.  Requires ``torchvision``.
    """

    def __init__(self, d_model: int = 128, backbone: str = "default") -> None:
        super().__init__()
        self.backbone_name = backbone

        if backbone == "default":
            self.features = nn.Sequential(
                nn.Conv2d(1, 32, kernel_size=3, padding=1),
                nn.ReLU(inplace=True),
                nn.MaxPool2d(2),  # → 32 × 14 × 14
                nn.Conv2d(32, 64, kernel_size=3, padding=1),
                nn.ReLU(inplace=True),
                nn.MaxPool2d(2),  # → 64 × 7 × 7
                nn.Conv2d(64, 128, kernel_size=3, padding=1),
                nn.ReLU(inplace=True),
                nn.AdaptiveAvgPool2d(1),  # → 128 × 1 × 1
                nn.Flatten(),  # → 128
            )
            self.proj = nn.Linear(128, d_model)

        elif backbone == "resnet18":
            import torchvision.models as _tv  # lazy import — optional dep

            _r = _tv.resnet18(weights=None)
            # Adapt for 1-channel 28×28 input:
            #   • smaller first conv (3×3, stride 1) preserves 28→28 spatial size
            #   • remove the 3×3 max-pool so small images keep enough resolution
            _r.conv1 = nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1, bias=False)
            # Use all residual stages + avgpool → 512-d feature vector
            self.features = nn.Sequential(
                _r.conv1,
                _r.bn1,
                _r.relu,
                nn.Identity(),
                _r.layer1,
                _r.layer2,
                _r.layer3,
                _r.layer4,
                _r.avgpool,
                nn.Flatten(),  # → 512
            )
            self.proj = nn.Linear(512, d_model)

        else:
            raise ValueError(
                f"Unknown CellCNN backbone: {backbone!r}. Choose 'default' or 'resnet18'."
            )

        self.norm = nn.LayerNorm(d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (N, 1, 28, 28) float32

        Returns:
            (N, d_model)
        """
        h = self.features(x)
        return self.norm(self.proj(h))