"""
Cross-modal alignment for satellite imagery retrieval.

Implements multiple approaches:
1. Modality-specific projection heads
2. Contrastive cross-modal loss
3. Wavelength-aware encoding
4. Domain adaptation
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass


@dataclass
class CrossModalConfig:
    """Configuration for cross-modal alignment."""
    embed_dim: int = 768
    projection_dim: int = 256
    modalities: List[str] = None
    temperature: float = 0.07
    use_wavelength_encoding: bool = True
    use_domain_adaptation: bool = True
    
    def __post_init__(self):
        if self.modalities is None:
            self.modalities = ["optical", "sar", "multispectral"]


class ModalityProjectionHead(nn.Module):
    """Projection head for a single modality."""
    
    def __init__(self, input_dim: int, output_dim: int):
        super().__init__()
        self.projection = nn.Sequential(
            nn.Linear(input_dim, input_dim),
            nn.GELU(),
            nn.Linear(input_dim, output_dim),
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.normalize(self.projection(x), dim=-1)


class WavelengthEncoder(nn.Module):
    """Encode wavelength information for each modality."""
    
    def __init__(self, num_channels: int, output_dim: int):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(num_channels, output_dim),
            nn.GELU(),
            nn.Linear(output_dim, output_dim),
        )
    
    def forward(self, wavelengths: torch.Tensor) -> torch.Tensor:
        return self.encoder(wavelengths)


class CrossModalAligner(nn.Module):
    """
    Cross-modal alignment using modality-specific projections.
    
    Based on CLOSP and DOFA-CLIP approaches:
    - Each modality has its own projection head
    - Wavelength encoding for channel-aware processing
    - Contrastive loss for alignment
    """
    
    def __init__(self, config: CrossModalConfig):
        super().__init__()
        self.config = config
        
        # Modality-specific projection heads
        self.projection_heads = nn.ModuleDict({
            mod: ModalityProjectionHead(config.embed_dim, config.projection_dim)
            for mod in config.modalities
        })
        
        # Wavelength encoders (if enabled)
        if config.use_wavelength_encoding:
            self.wavelength_encoders = nn.ModuleDict({
                mod: WavelengthEncoder(3, config.projection_dim)  # 3 channels for wavelength
                for mod in config.modalities
            })
        
        # Domain adaptation layer (if enabled)
        if config.use_domain_adaptation:
            self.domain_adaptor = nn.Sequential(
                nn.Linear(config.projection_dim, config.projection_dim),
                nn.GELU(),
                nn.Linear(config.projection_dim, config.projection_dim),
            )
        
        # Learnable temperature
        self.logit_scale = nn.Parameter(torch.ones([]) * torch.log(torch.tensor(1.0 / config.temperature)))
    
    def project(self, features: torch.Tensor, modality: str) -> torch.Tensor:
        """Project features using modality-specific head."""
        return self.projection_heads[modality](features)
    
    def align_with_wavelength(
        self,
        features: torch.Tensor,
        modality: str,
        wavelengths: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """Align features using wavelength encoding."""
        if not self.config.use_wavelength_encoding or wavelengths is None:
            return self.project(features, modality)
        
        # Get wavelength embedding
        wave_emb = self.wavelength_encoders[modality](wavelengths)
        
        # Combine features with wavelength info
        combined = features + wave_emb
        return self.projection_heads[modality](combined)
    
    def contrastive_loss(
        self,
        features_a: torch.Tensor,
        features_b: torch.Tensor,
        temperature: Optional[float] = None
    ) -> torch.Tensor:
        """Compute contrastive loss between two sets of features."""
        if temperature is None:
            temperature = self.config.temperature
        
        # Normalize
        features_a = F.normalize(features_a, dim=-1)
        features_b = F.normalize(features_b, dim=-1)
        
        # Compute similarity
        logit_scale = self.logit_scale.exp()
        logits = logit_scale * features_a @ features_b.t()
        
        # Labels (diagonal is positive)
        labels = torch.arange(len(features_a), device=features_a.device)
        
        # Symmetric loss
        loss_a = F.cross_entropy(logits, labels)
        loss_b = F.cross_entropy(logits.t(), labels)
        
        return (loss_a + loss_b) / 2
    
    def cross_modal_retrieve(
        self,
        query_features: torch.Tensor,
        query_modality: str,
        gallery_features: Dict[str, torch.Tensor],
        k: int = 5
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Cross-modal retrieval.
        
        Args:
            query_features: Query features from source modality
            query_modality: Modality of query
            gallery_features: Dict of gallery features per modality
            k: Number of results
            
        Returns:
            (indices, scores) for top-k results
        """
        # Project query
        query_proj = self.project(query_features, query_modality)
        
        all_scores = []
        all_indices = []
        
        # Search across all target modalities
        offset = 0
        for mod, features in gallery_features.items():
            # Project gallery
            gallery_proj = self.project(features, mod)
            
            # Compute similarity
            scores = query_proj @ gallery_proj.t()
            all_scores.append(scores)
            all_indices.append(torch.arange(len(features), device=features.device) + offset)
            offset += len(features)
        
        # Concatenate
        all_scores = torch.cat(all_scores, dim=-1)
        all_indices = torch.cat(all_indices, dim=-1)
        
        # Top-k
        topk_scores, topk_idx = all_scores.topk(k, dim=-1)
        topk_indices = all_indices[topk_idx]
        
        return topk_indices, topk_scores


class ContrastiveCrossModalLoss(nn.Module):
    """
    Contrastive loss for cross-modal alignment.
    
    Based on CLOSP approach: align SAR and optical via shared text anchor.
    """
    
    def __init__(self, temperature: float = 0.07):
        super().__init__()
        self.temperature = temperature
        self.logit_scale = nn.Parameter(torch.ones([]) * torch.log(torch.tensor(1.0 / temperature)))
    
    def forward(
        self,
        features_a: torch.Tensor,
        features_b: torch.Tensor,
        features_text: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """
        Compute cross-modal contrastive loss.
        
        Args:
            features_a: Features from modality A (e.g., optical)
            features_b: Features from modality B (e.g., SAR)
            features_text: Optional text features for triple alignment
            
        Returns:
            Loss value
        """
        features_a = F.normalize(features_a, dim=-1)
        features_b = F.normalize(features_b, dim=-1)
        
        logit_scale = self.logit_scale.exp()
        
        # Image-image contrastive loss
        logits_ab = logit_scale * features_a @ features_b.t()
        logits_ba = logits_ab.t()
        
        labels = torch.arange(len(features_a), device=features_a.device)
        
        loss_a2b = F.cross_entropy(logits_ab, labels)
        loss_b2a = F.cross_entropy(logits_ba, labels)
        
        loss = (loss_a2b + loss_b2a) / 2
        
        # Text-image contrastive loss (if available)
        if features_text is not None:
            features_text = F.normalize(features_text, dim=-1)
            
            logits_t2a = logit_scale * features_text @ features_a.t()
            logits_t2b = logit_scale * features_text @ features_b.t()
            
            loss_t2a = F.cross_entropy(logits_t2a, labels)
            loss_t2b = F.cross_entropy(logits_t2b, labels)
            
            loss = loss + (loss_t2a + loss_t2b) / 2
        
        return loss


class DomainAdaptationLayer(nn.Module):
    """
    Domain adaptation for bridging modality gaps.
    
    Based on SARCLIP approach: transfer knowledge from optical to SAR.
    """
    
    def __init__(self, embed_dim: int, num_modalities: int = 3):
        super().__init__()
        
        # Modality-specific adapters
        self.adapters = nn.ModuleList([
            nn.Sequential(
                nn.Linear(embed_dim, embed_dim),
                nn.GELU(),
                nn.Linear(embed_dim, embed_dim),
            )
            for _ in range(num_modalities)
        ])
        
        # Shared adapter
        self.shared_adapter = nn.Sequential(
            nn.Linear(embed_dim, embed_dim),
            nn.GELU(),
            nn.Linear(embed_dim, embed_dim),
        )
    
    def forward(
        self,
        features: torch.Tensor,
        modality_idx: int
    ) -> torch.Tensor:
        """
        Apply domain adaptation.
        
        Args:
            features: Input features
            modality_idx: Index of the modality
            
        Returns:
            Adapted features
        """
        # Modality-specific adaptation
        adapted = self.adapters[modality_idx](features)
        
        # Shared adaptation
        shared = self.shared_adapter(features)
        
        # Combine
        return adapted + shared


# Self-check
if __name__ == "__main__":
    print("Testing Cross-Modal Alignment...")
    
    config = CrossModalConfig()
    aligner = CrossModalAligner(config)
    
    # Test projection
    features = torch.randn(8, 768)
    optical_proj = aligner.project(features, "optical")
    sar_proj = aligner.project(features, "sar")
    
    print(f"Optical projection shape: {optical_proj.shape}")
    print(f"SAR projection shape: {sar_proj.shape}")
    
    # Test contrastive loss
    loss = aligner.contrastive_loss(optical_proj, sar_proj)
    print(f"Contrastive loss: {loss.item():.4f}")
    
    # Test cross-modal retrieval
    gallery_features = {
        "optical": torch.randn(100, 768),
        "sar": torch.randn(100, 768),
        "multispectral": torch.randn(100, 768),
    }
    
    query = torch.randn(1, 768)
    indices, scores = aligner.cross_modal_retrieve(query, "optical", gallery_features, k=5)
    
    print(f"Retrieved indices: {indices}")
    print(f"Retrieved scores: {scores}")
    
    print("\nCross-Modal Alignment test passed!")