"""StripNet image encoder wrapper for the caption-test pipeline.

Exposes the same interface as DINOv3ViT: `forward(images) -> [B, embed_dim]`.
StripNet's deepest stage produces [B, 512, H/32, W/32]; we apply global average
pooling (GAP) and project to the target retrieval dimension via Linear(512→1024)
to match DINOv3 native dim and keep GatedFusion unchanged.
"""
from __future__ import annotations

import logging
import torch
import torch.nn as nn

from src.models.stripnet import StripNet, get_stripnet_small, load_stripnet_small_pretrained

LOGGER = logging.getLogger("caption_test.stripnet_encoder")


class StripNetEncoder(nn.Module):
    """StripNet-small + GAP + projection to `out_dim`.

    Frozen backbone (BatchNorm in eval mode); only the projection head and
    any injected Conv-MONA adapters are trainable.
    """

    LAST_STAGE_DIM = 512  # StripNet-small last stage embed dim

    def __init__(
        self,
        checkpoint_path: str,
        out_dim: int = 1024,
        *,
        load_pretrained: bool = True,
    ) -> None:
        super().__init__()
        self.out_dim = out_dim
        if load_pretrained:
            self.backbone: StripNet = load_stripnet_small_pretrained(checkpoint_path)
        else:
            LOGGER.info("StripNet-small: random init (stripnet_load_pretrained=False)")
            self.backbone = get_stripnet_small()
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.projection = nn.Linear(self.LAST_STAGE_DIM, out_dim)
        nn.init.trunc_normal_(self.projection.weight, std=0.02)
        nn.init.zeros_(self.projection.bias)

    def set_gradient_checkpointing(self, enable: bool = True) -> None:
        self.backbone.set_gradient_checkpointing(enable)

    def train(self, mode: bool = True):
        """Override: keep frozen backbone in eval mode (BN running stats stable)."""
        super().train(mode)
        # Frozen backbone always in eval; trainable adapters/projection follow `mode`.
        if not any(p.requires_grad for p in self.backbone.parameters()):
            self.backbone.eval()
        return self

    def forward(self, images: torch.Tensor) -> torch.Tensor:
        feat = self.backbone.forward_last_features(images)  # [B, 512, H/32, W/32]
        pooled = self.pool(feat).flatten(1)                  # [B, 512]
        return self.projection(pooled)                        # [B, out_dim]
