src/features/sar_adapter.py

"""
SAR-specific adapter layers for CLIP.

Adds lightweight adapter modules to improve SAR modality handling
without modifying the base CLIP weights.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional


class SARAdapter(nn.Module):
    """
    Lightweight adapter for SAR imagery.
    
    Applies learnable transformations to bridge the domain gap between
    optical and SAR imagery. Uses:
    1. Channel projection (2ch SAR → 3ch RGB-like)
    2. Learnable scaling to match CLIP input distribution
    3. Optional speckle noise reduction
    """
    
    def __init__(
        self,
        in_channels: int = 2,
        out_channels: int = 3,
        hidden_dim: int = 64,
        dropout: float = 0.1
    ):
        super().__init__()
        
        # Channel projection: 2ch (VV, VH) → 3ch (RGB-like)
        self.channel_proj = nn.Sequential(
            nn.Conv2d(in_channels, hidden_dim, kernel_size=1, bias=False),
            nn.BatchNorm2d(hidden_dim),
            nn.GELU(),
            nn.Conv2d(hidden_dim, out_channels, kernel_size=1, bias=False),
            nn.BatchNorm2d(out_channels),
        )
        
        # Learnable scaling per channel
        self.channel_scale = nn.Parameter(torch.ones(out_channels))
        self.channel_bias = nn.Parameter(torch.zeros(out_channels))
        
        # Optional speckle reduction
        self.speckle_reduction = nn.Sequential(
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, groups=out_channels),
            nn.Sigmoid(),
        )
        
        self.dropout = nn.Dropout2d(dropout)
        
        self._init_weights()
    
    def _init_weights(self):
        """Initialize weights with small values."""
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
    
    def forward(self, x: torch.Tensor, apply_speckle: bool = True) -> torch.Tensor:
        """
        Forward pass.
        
        Args:
            x: SAR image tensor, shape (B, 2, H, W) with VV, VH channels
            apply_speckle: Whether to apply speckle reduction
            
        Returns:
            Projected tensor, shape (B, 3, H, W)
        """
        # Channel projection
        x = self.channel_proj(x)
        
        # Apply speckle reduction
        if apply_speckle:
            mask = self.speckle_reduction(x)
            x = x * mask
        
        # Apply learnable scaling
        x = x * self.channel_scale.view(1, -1, 1, 1) + self.channel_bias.view(1, -1, 1, 1)
        
        x = self.dropout(x)
        
        return x


class SARCLIPWrapper(nn.Module):
    """
    Wraps a CLIP model with SAR adapter.
    
    Handles the preprocessing pipeline for SAR imagery:
    1. Log-scale transformation
    2. Speckle reduction
    3. Channel projection via SARAdapter
    """
    
    def __init__(
        self,
        clip_model: nn.Module,
        adapter: Optional[SARAdapter] = None,
        device: Optional[str] = None
    ):
        super().__init__()
        
        self.clip_model = clip_model
        self.adapter = adapter or SARAdapter()
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        
        self.adapter.to(self.device)
    
    def log_scale(self, x: torch.Tensor) -> torch.Tensor:
        """Apply log-scale transformation to SAR amplitude data."""
        return torch.log1p(x)
    
    def preprocess_sar(self, x: torch.Tensor) -> torch.Tensor:
        """
        Preprocess SAR imagery for CLIP.
        
        Args:
            x: Raw SAR tensor, shape (B, 2, H, W)
            
        Returns:
            Preprocessed tensor, shape (B, 3, H, W)
        """
        # Log-scale
        x = self.log_scale(x)
        
        # Normalize to [0, 1] range
        x = x - x.min(dim=-1, keepdim=True)[0].min(dim=-2, keepdim=True)[0]
        x = x / (x.max(dim=-1, keepdim=True)[0].max(dim=-2, keepdim=True)[0] + 1e-8)
        
        # Apply adapter
        x = self.adapter(x, apply_speckle=True)
        
        return x
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass through SAR adapter then CLIP.
        
        Args:
            x: SAR image tensor, shape (B, 2, H, W)
            
        Returns:
            CLIP embedding, shape (B, embed_dim)
        """
        x = self.preprocess_sar(x)
        return self.clip_model(x)


def create_sar_adapter_for_clip(
    clip_model: nn.Module,
    in_channels: int = 2,
    hidden_dim: int = 64
) -> SARCLIPWrapper:
    """
    Convenience function to create SAR adapter for existing CLIP model.
    
    Args:
        clip_model: Existing CLIP model
        in_channels: Number of SAR channels (default: 2 for VV/VH)
        hidden_dim: Hidden dimension in adapter
        
    Returns:
        SARCLIPWrapper with adapter attached
    """
    adapter = SARAdapter(
        in_channels=in_channels,
        out_channels=3,
        hidden_dim=hidden_dim
    )
    
    return SARCLIPWrapper(clip_model, adapter)


# Self-check
if __name__ == "__main__":
    print("Testing SARAdapter...")
    
    # Test adapter
    adapter = SARAdapter(in_channels=2, out_channels=3)
    
    # Dummy SAR input (2 channels: VV, VH)
    x = torch.randn(2, 2, 224, 224)
    
    # Forward pass
    out = adapter(x)
    print(f"Input shape: {x.shape}")
    print(f"Output shape: {out.shape}")
    
    # Verify output is 3 channels
    assert out.shape[1] == 3, f"Expected 3 channels, got {out.shape[1]}"
    
    # Test log scale
    wrapper = SARCLIPWrapper.__new__(SARCLIPWrapper)
    wrapper.adapter = adapter
    
    x_log = wrapper.log_scale(x.abs())  # abs() because SAR can have negative values
    print(f"Log-scaled shape: {x_log.shape}")
    print(f"Log-scaled range: [{x_log.min():.4f}, {x_log.max():.4f}]")
    
    print("\nSARAdapter test passed!")