from __future__ import annotations

"""Asymmetric dual encoder for CVGL caption test on GTA-UAV.

Architecture:
  Query:   DINOv3 ViT-L/16 (LVD, frozen) + LRSCLIP text (L1/L2/L3) -> GatedFusion -> query
  Gallery: DINOv3 ViT-L/16 (SAT, frozen) -> gallery
  Loss:    InfoNCE(query, gallery)

DINOv3 checkpoints use a custom key layout (not HuggingFace transformers).
LRSCLIP (DGTRS-CLIP ViT-L-14) uses open_clip layout with KPS positional embeddings.
"""

import logging
import math
import warnings
from pathlib import Path

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

LOGGER = logging.getLogger("caption_test.model")
coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
from safetensors.torch import load_file as load_safetensors

from src.models.adapters import inject_lora_into_dgtrs, inject_mona_into_dinov3
from src.models.dgtrs.model import DGTRSTextEncoder, load_dgtrs_text_encoder, tokenize_dgtrs
from src.models.dual_encoder import GatedFusion, ProjectionHead
from src.models.stripnet import inject_conv_mona_into_stripnet
from src.models.stripnet_encoder import StripNetEncoder


# ---------------------------------------------------------------------------
# DINOv3 ViT-L/16 — minimal implementation matching checkpoint key layout
# ---------------------------------------------------------------------------

class DINOv3Attention(nn.Module):
    """Multi-head self-attention with separate Q/K/V projections."""

    def __init__(self, dim: int = 1024, num_heads: int = 16) -> None:
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim ** -0.5
        self.q_proj = nn.Linear(dim, dim)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim)
        self.o_proj = nn.Linear(dim, dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, N, C = x.shape
        q = self.q_proj(x).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        k = self.k_proj(x).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        v = self.v_proj(x).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        attn = F.scaled_dot_product_attention(q, k, v)
        x = attn.permute(0, 2, 1, 3).reshape(B, N, C)
        return self.o_proj(x)


class DINOv3LayerScale(nn.Module):
    """Per-channel learnable scale (lambda)."""

    def __init__(self, dim: int) -> None:
        super().__init__()
        self.lambda1 = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * self.lambda1


class DINOv3MLP(nn.Module):
    """SwiGLU-like MLP: up_proj + GELU + down_proj."""

    def __init__(self, dim: int = 1024, mlp_dim: int = 4096) -> None:
        super().__init__()
        self.up_proj = nn.Linear(dim, mlp_dim)
        self.down_proj = nn.Linear(mlp_dim, dim)
        self.act = nn.GELU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(self.act(self.up_proj(x)))


class DINOv3Block(nn.Module):
    """Single DINOv3 transformer block."""

    def __init__(self, dim: int = 1024, num_heads: int = 16, mlp_dim: int = 4096) -> None:
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attention = DINOv3Attention(dim, num_heads)
        self.layer_scale1 = DINOv3LayerScale(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = DINOv3MLP(dim, mlp_dim)
        self.layer_scale2 = DINOv3LayerScale(dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.layer_scale1(self.attention(self.norm1(x)))
        x = x + self.layer_scale2(self.mlp(self.norm2(x)))
        return x


class DINOv3Embeddings(nn.Module):
    """Patch embedding + CLS token + register tokens."""

    def __init__(
        self,
        dim: int = 1024,
        patch_size: int = 16,
        num_registers: int = 4,
    ) -> None:
        super().__init__()
        self.patch_embeddings = nn.Conv2d(3, dim, patch_size, patch_size)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, dim))
        self.register_tokens = nn.Parameter(torch.zeros(1, num_registers, dim))
        self.mask_token = nn.Parameter(torch.zeros(1, 1, dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B = x.shape[0]
        patches = self.patch_embeddings(x).flatten(2).transpose(1, 2)  # [B, N, D]
        N = patches.shape[1]

        cls = self.cls_token.expand(B, -1, -1)
        reg = self.register_tokens.expand(B, -1, -1)

        # DINOv3: [CLS, registers, patches]
        x = torch.cat([cls, reg, patches], dim=1)

        # Positional embedding: interpolated sincos (RoPE applied in attention
        # in original, but pretrained checkpoints bake it into weights).
        # We use a simple learned-style pos embed computed on the fly.
        pos = self._get_pos_embed(N, x.device, x.dtype)
        # pos covers patches only, skip CLS + registers
        x[:, 1 + reg.shape[1]:] = x[:, 1 + reg.shape[1]:] + pos

        return x

    def _get_pos_embed(self, n_patches: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
        # DINOv3 uses RoPE internally — no additive pos embed needed.
        # Return zeros as placeholder (weights handle positioning via RoPE).
        return torch.zeros(1, n_patches, self.cls_token.shape[-1], device=device, dtype=dtype)


class DINOv3ViT(nn.Module):
    """DINOv3 ViT-L/16 matching the checkpoint key layout.

    Checkpoint keys:
        embeddings.cls_token, embeddings.patch_embeddings.{weight,bias},
        embeddings.register_tokens, embeddings.mask_token,
        layer.{i}.attention.{q,k,v,o}_proj.{weight,bias},
        layer.{i}.layer_scale{1,2}.lambda1,
        layer.{i}.mlp.{up,down}_proj.{weight,bias},
        layer.{i}.norm{1,2}.{weight,bias},
        norm.{weight,bias}
    """

    def __init__(
        self,
        dim: int = 1024,
        num_heads: int = 16,
        mlp_dim: int = 4096,
        num_layers: int = 24,
        patch_size: int = 16,
        num_registers: int = 4,
    ) -> None:
        super().__init__()
        self.embeddings = DINOv3Embeddings(dim, patch_size, num_registers)
        self.layer = nn.ModuleList([
            DINOv3Block(dim, num_heads, mlp_dim) for _ in range(num_layers)
        ])
        self.norm = nn.LayerNorm(dim)
        self.embed_dim = dim
        self.gradient_checkpointing = False

    def set_gradient_checkpointing(self, enable: bool = True) -> None:
        """Enable/disable gradient checkpointing to trade compute for VRAM."""
        self.gradient_checkpointing = enable

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass. Returns CLS token embedding [B, dim]."""
        x = self.embeddings(x)
        for block in self.layer:
            if self.gradient_checkpointing and self.training:
                x = torch.utils.checkpoint.checkpoint(
                    block, x, use_reentrant=False,
                )
            else:
                x = block(x)
        x = self.norm(x)
        return x[:, 0]  # CLS token

    @classmethod
    def from_pretrained(cls, path: str | Path) -> DINOv3ViT:
        """Load from .pth or .safetensors checkpoint."""
        model = cls()
        path = Path(path)
        LOGGER.info("🧊 Loading DINOv3 from %s", path.name)
        if path.suffix == ".safetensors":
            state = load_safetensors(str(path))
        else:
            state = torch.load(str(path), map_location="cpu", weights_only=False)
            if "model" in state:
                state = state["model"]
            elif "state_dict" in state:
                state = state["state_dict"]
        model.load_state_dict(state, strict=False)
        n_params = sum(p.numel() for p in model.parameters())
        LOGGER.info("🧊 DINOv3 loaded: %s params", f"{n_params:,}")
        return model


# LRSCLIPTextEncoder removed — replaced by official DGTRS architecture
# in src/models/dgtrs/model.py (DGTRSTextEncoder)


# ---------------------------------------------------------------------------
# Text fusion MLP: concat L1/L2/L3 -> project to D
# ---------------------------------------------------------------------------

class TextFusionMLP(nn.Module):
    """Fuse L1/L2/L3 text embeddings via concat + MLP.

    [B, 3*text_dim] -> [B, out_dim]
    """

    def __init__(
        self,
        text_dim: int = 768,
        out_dim: int = 1024,
    ) -> None:
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(3 * text_dim, out_dim),
            nn.GELU(),
            nn.Linear(out_dim, out_dim),
        )

    def forward(
        self,
        z_l1: torch.Tensor,
        z_l2: torch.Tensor,
        z_l3: torch.Tensor,
    ) -> torch.Tensor:
        """Fuse three text embeddings.

        Args:
            z_l1: L1 overview [B, text_dim].
            z_l2: L2 full description [B, text_dim].
            z_l3: L3 fingerprint [B, text_dim].

        Returns:
            Fused text embedding [B, out_dim].
        """
        cat = torch.cat([z_l1, z_l2, z_l3], dim=-1)
        return self.mlp(cat)


# ---------------------------------------------------------------------------
# Main model: AsymmetricEncoder
# ---------------------------------------------------------------------------

# ResidualGateFusin experiment
from .residual_fusions import ResidualGateType, GatedFusionResidual

class AsymmetricEncoder(nn.Module):
    """Dual encoder for CVGL with text fusion on both branches.

    Supports two modes:
      - **shared** (default): single DINOv3 WEB encoder for both drone and satellite,
        one set of MONA adapters. Saves ~4-5 GB VRAM and halves adapter params.
      - **asymmetric**: separate DINOv3 encoders (LVD for drone, SAT for satellite),
        each with their own MONA adapters (legacy mode).

    Query branch:   DINOv3 (drone) + text(L1/L2/L3) -> GatedFusion_q -> query [1024]
    Gallery branch: DINOv3 (sat)   + text(L1/L2/L3) -> GatedFusion_g -> gallery [1024]

    No projection layers — retrieval space is DINOv3 native 1024-dim.
    Text fusion MLP is shared between branches (same caption format).
    Two separate GatedFusion gates (drone/sat may weight text differently).

    For satellite images without captions, GatedFusion passes image features through
    (text_feat=None → gate acts as identity).

    Args:
        dino_web_path: Path to DINOv3 LVD checkpoint (used for both branches in shared mode).
        dino_sat_path: Path to DINOv3 SAT checkpoint (only used in asymmetric mode).
        lrsclip_path: Path to DGTRS-CLIP checkpoint (text encoder).
        init_gate: Initial fusion gate (image weight).
        baseline_mode: If True, gate = 1.0 (text ignored), DGTRS not loaded.
        shared_encoder: If True, use single DINOv3 WEB for both branches.
        device: Torch device string.
    """

    DINO_DIM = 1024
    TEXT_DIM = 768

    def __init__(
        self,
        gate_type: ResidualGateType = ResidualGateType.simple_residual_one_gate,
        dino_web_path: str = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth",
        dino_sat_path: str = "nn_models/DINO_SAT/model.safetensors",
        lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt",
        init_gate: float = 0.7,
        baseline_mode: bool = False,
        shared_encoder: bool = False,
        mona_bottleneck: int = 64,
        mona_last_n_blocks: int = 24,
        lora_rank: int = 4,
        device: str = "cuda",
        backbone: str = "dinov3",
        stripnet_path: str = "nn_models/STRIPNET/stripnet_s.pth",
        stripnet_mona_last_n_stages: int = 2,
        stripnet_freeze: bool = True,
    ) -> None:
        super().__init__()
        self.embed_dim = self.DINO_DIM  # native 1024 (StripNet projects 512 -> 1024)
        self.baseline_mode = baseline_mode
        self.shared_encoder = shared_encoder
        self.backbone = backbone
        self.device = device

        # Image encoder(s) (frozen + MONA adapters).
        if backbone == "stripnet":
            # StripNet always operates as shared encoder (one CNN for both branches).
            self.shared_encoder = True
            self.image_encoder = StripNetEncoder(checkpoint_path=stripnet_path, out_dim=self.DINO_DIM)
            if stripnet_freeze:
                self._freeze(self.image_encoder.backbone)
                LOGGER.info("StripNet backbone: frozen (Conv-MONA + projection trainable)")
            else:
                LOGGER.info("StripNet backbone: UNFROZEN — full fine-tune (use lower lr factor)")
            if stripnet_mona_last_n_stages > 0:
                inject_conv_mona_into_stripnet(
                    self.image_encoder.backbone,
                    bottleneck=mona_bottleneck,
                    last_n_stages=stripnet_mona_last_n_stages,
                )
            else:
                LOGGER.info("Conv-MONA disabled (stripnet_mona_last_n_stages=0)")
            LOGGER.info("StripNet backbone: shared encoder, projection 512 -> %d", self.DINO_DIM)
        elif shared_encoder:
            self.image_encoder = DINOv3ViT.from_pretrained(dino_web_path)
            self._freeze(self.image_encoder)
            inject_mona_into_dinov3(self.image_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks)
            LOGGER.info("Shared encoder mode: single DINOv3 WEB for drone + satellite")
        else:
            self.drone_encoder = DINOv3ViT.from_pretrained(dino_web_path)
            self.sat_encoder = DINOv3ViT.from_pretrained(dino_sat_path)
            self._freeze(self.drone_encoder)
            self._freeze(self.sat_encoder)
            inject_mona_into_dinov3(self.drone_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks)
            inject_mona_into_dinov3(self.sat_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks)
            LOGGER.info("Asymmetric encoder mode: DINOv3 WEB (drone) + DINOv3 SAT (satellite)")

        # Text encoder — official DGTRS architecture (frozen + LoRA).
        if not baseline_mode:
            self.text_encoder = load_dgtrs_text_encoder(lrsclip_path)
            self._freeze(self.text_encoder)
            inject_lora_into_dgtrs(self.text_encoder, rank=lora_rank)
        else:
            self.text_encoder = None

        # Shared text fusion MLP: 3×768 -> 1024 (native DINOv3 dim).
        if not baseline_mode:
            self.text_fusion = TextFusionMLP(
                text_dim=self.TEXT_DIM,
                out_dim=self.DINO_DIM,
            )

        # Separate gated fusion for query and gallery branches.
        #! Experimental Gated fusion on query branch.
        self.fusion_query = GatedFusionResidual(gate_type=gate_type,
            init_gate=init_gate, baseline_mode=baseline_mode)
        self.fusion_gallery = GatedFusionResidual(gate_type=gate_type,
            init_gate=init_gate, baseline_mode=baseline_mode)


    @staticmethod
    def _freeze(module: nn.Module) -> None:
        for p in module.parameters():
            p.requires_grad = False
        module.eval()

    def encode_drone(self, images: torch.Tensor) -> torch.Tensor:
        """Encode drone images with MONA adapters. Returns [B, 1024]."""
        if self.shared_encoder:
            return self.image_encoder(images)
        return self.drone_encoder(images)

    def encode_satellite(self, images: torch.Tensor) -> torch.Tensor:
        """Encode satellite images with MONA adapters. Returns [B, 1024]."""
        if self.shared_encoder:
            return self.image_encoder(images)
        return self.sat_encoder(images)

    def encode_text_levels(
        self,
        l1_texts: list[str],
        l2_texts: list[str],
        l3_texts: list[str],
    ) -> torch.Tensor | None:
        """Encode L1/L2/L3 captions and fuse. Returns [B, 1024] or None.

        Returns None if all captions are empty (no text available).
        For mixed batches (some have captions, some don't), encodes all
        texts (empty strings tokenize to pad+EOS — their outputs must be
        masked downstream, see `_fuse_with_mask`).
        """
        # Check if any caption is non-empty.
        if all(t == "" for t in l1_texts):
            return None

        z_l1 = self._encode_single_text(l1_texts)
        z_l2 = self._encode_single_text(l2_texts)
        z_l3 = self._encode_single_text(l3_texts)
        fused = self.text_fusion(z_l1, z_l2, z_l3)
        return F.normalize(fused, dim=-1)

    def _encode_single_text(self, texts: list[str]) -> torch.Tensor:
        """Tokenize and encode a list of strings using DGTRS tokenizer."""
        tokens = tokenize_dgtrs(list(texts)).to(self.device)
        return self.text_encoder(tokens)

    def _fuse_with_mask(
        self,
        img_feat: torch.Tensor,
        l1_texts: list[str] | None,
        l2_texts: list[str] | None,
        l3_texts: list[str] | None,
        fusion: GatedFusionResidual,
    ) -> torch.Tensor:
        """Fuse image features with optional text, respecting per-sample presence.

        For samples where caption is an empty string, output falls back to
        pure image features (avoiding noise contamination from empty-string
        text embeddings). For samples with captions, applies the standard
        gated fusion `σ(α)·img + (1-σ(α))·text`.

        Returns L2-normalized [B, D] embedding.
        """
        if (
            self.baseline_mode
            or l1_texts is None
            or l2_texts is None
            or l3_texts is None
        ):
            return F.normalize(fusion(img_feat, None), dim=-1)

        has_text = torch.tensor(
            [t != "" for t in l1_texts], dtype=torch.bool, device=img_feat.device,
        )
        if not has_text.any():
            return F.normalize(fusion(img_feat, None), dim=-1)

        z_text = self.encode_text_levels(l1_texts, l2_texts, l3_texts)
        if z_text is None:
            return F.normalize(fusion(img_feat, None), dim=-1)

        # Per-sample fusion: text-present samples use full gated fusion,
        # empty-caption samples pass through pure image features.
        fused_with_text = fusion(img_feat, z_text)
        out = torch.where(has_text.unsqueeze(-1), fused_with_text, img_feat)
        return F.normalize(out, dim=-1)

    def encode_query(
        self,
        drone_img: torch.Tensor,
        caption_l1: list[str] | None = None,
        caption_l2: list[str] | None = None,
        caption_l3: list[str] | None = None,
    ) -> torch.Tensor:
        """Encode drone → normalized query embedding with per-sample text mask."""
        drone_feat = self.encode_drone(drone_img)
        return self._fuse_with_mask(
            drone_feat, caption_l1, caption_l2, caption_l3, self.fusion_query,
        )

    def encode_gallery(
        self,
        sat_img: torch.Tensor,
        sat_caption_l1: list[str] | None = None,
        sat_caption_l2: list[str] | None = None,
        sat_caption_l3: list[str] | None = None,
    ) -> torch.Tensor:
        """Encode satellite → normalized gallery embedding with per-sample text mask."""
        sat_feat = self.encode_satellite(sat_img)
        return self._fuse_with_mask(
            sat_feat, sat_caption_l1, sat_caption_l2, sat_caption_l3, self.fusion_gallery,
        )

    def forward(
        self,
        drone_img: torch.Tensor,
        sat_img: torch.Tensor,
        caption_l1: list[str] | None = None,
        caption_l2: list[str] | None = None,
        caption_l3: list[str] | None = None,
        sat_caption_l1: list[str] | None = None,
        sat_caption_l2: list[str] | None = None,
        sat_caption_l3: list[str] | None = None,
    ) -> dict[str, torch.Tensor]:
        """Forward pass.

        Both branches use per-sample caption masking: samples with an empty
        caption string fall back to pure image features instead of being
        fused with noise from empty-string text embeddings.

        Args:
            drone_img: Drone images [B, 3, 256, 256].
            sat_img: Satellite images [B, 3, 256, 256].
            caption_l1/l2/l3: Drone L1/L2/L3 captions.
            sat_caption_l1/l2/l3: Satellite L1/L2/L3 captions.

        Returns:
            Dict with 'query' [B, embed_dim], 'gallery' [B, embed_dim],
            'gate_q', 'gate_g'.
        """
        query = self.encode_query(drone_img, caption_l1, caption_l2, caption_l3)
        gallery = self.encode_gallery(sat_img, sat_caption_l1, sat_caption_l2, sat_caption_l3)
        return {
            "query": query,
            "gallery": gallery,
            "gate_q": self.fusion_query.gate_value,
            "gate_g": self.fusion_gallery.gate_value,
        }

    def trainable_parameters(self) -> list[nn.Parameter]:
        """Return list of parameters that require grad."""
        return [p for p in self.parameters() if p.requires_grad]

    def save_checkpoint(self, path: str | Path, **extra) -> None:
        """Save model checkpoint with metadata."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        ckpt = {
            "model_state": self.state_dict(),
            "baseline_mode": self.baseline_mode,
            "shared_encoder": self.shared_encoder,
            **extra,
        }
        tmp = path.with_suffix(path.suffix + ".tmp")
        torch.save(ckpt, tmp)
        tmp.replace(path)
        LOGGER.info("💾 Model saved to %s", path)

    @classmethod
    def load_checkpoint(
        cls,
        path: str | Path,
        dino_web_path: str = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth",
        dino_sat_path: str = "nn_models/DINO_SAT/model.safetensors",
        lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt",
        device: str = "cuda",
    ) -> tuple[AsymmetricEncoder, dict]:
        """Load model from checkpoint.

        First builds the model (loading frozen backbones), then loads
        the saved trainable weights on top.

        Returns:
            (model, checkpoint_dict) — model ready for eval/resume,
            checkpoint_dict has optimizer_state, epoch, etc.
        """
        path = Path(path)
        LOGGER.info("📂 Loading checkpoint from %s", path)
        ckpt = torch.load(str(path), map_location="cpu", weights_only=False)

        model = cls(
            dino_web_path=dino_web_path,
            dino_sat_path=dino_sat_path,
            lrsclip_path=lrsclip_path,
            baseline_mode=ckpt.get("baseline_mode", False),
            shared_encoder=ckpt.get("shared_encoder", False),
            mona_bottleneck=ckpt.get("mona_bottleneck", 64),
            mona_last_n_blocks=ckpt.get("mona_last_n_blocks", 24),
            device=device,
        )
        model.load_state_dict(ckpt["model_state"], strict=False)
        model = model.to(device)
        LOGGER.info("✅ Checkpoint loaded (epoch=%s)", ckpt.get("epoch", "?"))
        return model, ckpt

    def train(self, mode: bool = True) -> AsymmetricEncoder:
        """Override to keep frozen encoders in eval mode."""
        super().train(mode)
        if self.shared_encoder:
            self.image_encoder.eval()
        else:
            self.drone_encoder.eval()
            self.sat_encoder.eval()
        if self.text_encoder is not None:
            self.text_encoder.train(mode)
        return self


# ---------------------------------------------------------------------------
# Image preprocessing (DINOv3: 256x256, ImageNet normalization)
# ---------------------------------------------------------------------------

_IMAGENET_MEAN = [0.485, 0.456, 0.406]
_IMAGENET_STD = [0.229, 0.224, 0.225]


def get_dino_transform(image_size: int = 256) -> torch.nn.Module:
    """Build eval/inference image transform for DINOv3 input."""
    from torchvision import transforms
    return transforms.Compose([
        transforms.Resize(image_size, interpolation=transforms.InterpolationMode.BICUBIC),
        transforms.CenterCrop(image_size),
        transforms.ToTensor(),
        transforms.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD),
    ])


def get_drone_train_transform(image_size: int = 256) -> torch.nn.Module:
    """Build training augmentation for drone images.

    Includes: RandomResizedCrop, HFlip, rotation, color jitter,
    grayscale, Gaussian blur.
    """
    from torchvision import transforms
    return transforms.Compose([
        transforms.RandomResizedCrop(
            image_size, scale=(0.7, 1.0),
            interpolation=transforms.InterpolationMode.BICUBIC,
        ),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.RandomRotation(degrees=15),
        transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.1),
        transforms.RandomGrayscale(p=0.05),
        transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 2.0)),
        transforms.ToTensor(),
        transforms.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD),
    ])


def get_satellite_train_transform(image_size: int = 256) -> torch.nn.Module:
    """Build training augmentation for satellite images.

    Lighter than drone: no rotation or blur (satellite is nadir/consistent).
    Includes: RandomResizedCrop, HFlip, color jitter, grayscale.
    """
    from torchvision import transforms
    return transforms.Compose([
        transforms.RandomResizedCrop(
            image_size, scale=(0.7, 1.0),
            interpolation=transforms.InterpolationMode.BICUBIC,
        ),
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.1),
        transforms.RandomGrayscale(p=0.05),
        transforms.ToTensor(),
        transforms.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD),
    ])