"""
Multiscale feature extraction for satellite imagery.

Combines patch-level and global features for richer representations.
Uses DINOv2 for patch features and CLIP for global alignment.
"""

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


@dataclass
class MultiscaleFeatures:
    """Container for multiscale features."""
    global_feature: torch.Tensor  # (embed_dim,) - CLIP-style global
    patch_features: torch.Tensor  # (num_patches, patch_dim) - DINOv2-style
    patch_grid: Tuple[int, int]  # (H, W) grid of patches
    combined: torch.Tensor       # (combined_dim,) - fused feature


class PatchAggregator(nn.Module):
    """
    Aggregates patch features into a single representation.
    
    Supports multiple aggregation strategies:
    - mean: Average pooling
    - max: Max pooling
    - attention: Learnable attention pooling
    """
    
    def __init__(self, patch_dim: int, strategy: str = "attention"):
        super().__init__()
        
        self.strategy = strategy
        
        if strategy == "attention":
            self.attention = nn.Sequential(
                nn.Linear(patch_dim, patch_dim // 4),
                nn.Tanh(),
                nn.Linear(patch_dim // 4, 1),
            )
        elif strategy == "cls":
            self.cls_token = nn.Parameter(torch.randn(1, 1, patch_dim))
    
    def forward(self, patch_features: torch.Tensor) -> torch.Tensor:
        """
        Aggregate patch features.
        
        Args:
            patch_features: (B, num_patches, patch_dim)
            
        Returns:
            Aggregated feature (B, patch_dim)
        """
        if self.strategy == "mean":
            return patch_features.mean(dim=1)
        
        elif self.strategy == "max":
            return patch_features.max(dim=1)[0]
        
        elif self.strategy == "attention":
            # (B, num_patches, 1)
            attn_weights = self.attention(patch_features)
            attn_weights = F.softmax(attn_weights, dim=1)
            # (B, patch_dim)
            return (patch_features * attn_weights).sum(dim=1)
        
        elif self.strategy == "cls":
            B = patch_features.shape[0]
            cls_tokens = self.cls_token.expand(B, -1, -1)
            # Prepend CLS token
            x = torch.cat([cls_tokens, patch_features], dim=1)
            return x[:, 0]
        
        else:
            raise ValueError(f"Unknown strategy: {self.strategy}")


class MultiscaleExtractor(nn.Module):
    """
    Extracts features at multiple scales from satellite imagery.
    
    Combines:
    - Global features from CLIP (semantic alignment)
    - Patch features from DINOv2 (spatial details)
    - Cross-scale attention for feature fusion
    """
    
    def __init__(
        self,
        clip_model: nn.Module,
        dinov2_model: Optional[nn.Module] = None,
        embed_dim: int = 768,
        patch_dim: int = 768,
        fusion_dim: int = 512,
        use_cross_attention: bool = True
    ):
        super().__init__()
        
        self.clip_model = clip_model
        self.dinov2_model = dinov2_model
        
        self.embed_dim = embed_dim
        self.patch_dim = patch_dim
        self.fusion_dim = fusion_dim
        
        # Patch aggregation
        self.patch_aggregator = PatchAggregator(patch_dim, strategy="attention")
        
        # Cross-scale attention (fuses global + patch features)
        self.use_cross_attention = use_cross_attention
        if use_cross_attention:
            self.cross_attn = nn.MultiheadAttention(
                embed_dim=embed_dim,
                num_heads=8,
                dropout=0.1,
                batch_first=True
            )
            self.fusion_proj = nn.Linear(embed_dim + patch_dim, fusion_dim)
        else:
            # Simple concatenation + projection
            self.fusion_proj = nn.Linear(embed_dim + patch_dim, fusion_dim)
        
        # Final normalization
        self.layer_norm = nn.LayerNorm(fusion_dim)
    
    @torch.no_grad()
    def extract_clip_global(self, x: torch.Tensor) -> torch.Tensor:
        """Extract global features from CLIP."""
        # Assuming CLIP vision model
        if hasattr(self.clip_model, 'vision_model'):
            output = self.clip_model.vision_model(pixel_values=x)
            pooled = output.last_hidden_state[:, 0, :]  # CLS token
            global_feat = self.clip_model.visual_projection(pooled)
        else:
            # Fallback for other architectures
            global_feat = self.clip_model(x)
        
        return F.normalize(global_feat, dim=-1)
    
    @torch.no_grad()
    def extract_dinov2_patches(self, x: torch.Tensor) -> torch.Tensor:
        """Extract patch features from DINOv2."""
        if self.dinov2_model is None:
            # Return dummy features
            B = x.shape[0]
            num_patches = 196  # 14x14 for 224x224 input
            return torch.randn(B, num_patches, self.patch_dim, device=x.device)
        
        # DINOv2 forward pass
        output = self.dinov2_model(x)
        
        # Handle different output formats
        if hasattr(output, 'last_hidden_state'):
            patch_features = output.last_hidden_state[:, 1:]  # Remove CLS token
        elif isinstance(output, torch.Tensor):
            patch_features = output[:, 1:]  # Remove CLS token if present
        else:
            # Assume output is the patch features directly
            patch_features = output
        
        return patch_features
    
    def fuse_features(
        self,
        global_feat: torch.Tensor,
        patch_feat: torch.Tensor
    ) -> torch.Tensor:
        """
        Fuse global and patch features.
        
        Args:
            global_feat: (B, embed_dim)
            patch_feat: (B, patch_dim)
            
        Returns:
            Fused feature (B, fusion_dim)
        """
        if self.use_cross_attention:
            # Use global as query, patches as keys/values
            B = global_feat.shape[0]
            global_seq = global_feat.unsqueeze(1)  # (B, 1, embed_dim)
            patch_seq = patch_feat.unsqueeze(1)    # (B, 1, patch_dim) - simplified
            
            # Cross attention
            attn_out, _ = self.cross_attn(
                query=global_seq,
                key=patch_seq,
                value=patch_seq
            )
            attn_out = attn_out.squeeze(1)  # (B, embed_dim)
            
            # Concatenate and project
            combined = torch.cat([attn_out, patch_feat], dim=-1)
        else:
            combined = torch.cat([global_feat, patch_feat], dim=-1)
        
        # Project to fusion dim
        fused = self.fusion_proj(combined)
        fused = self.layer_norm(fused)
        
        return F.normalize(fused, dim=-1)
    
    def forward(
        self,
        x: torch.Tensor,
        return_separate: bool = False
    ) -> MultiscaleFeatures:
        """
        Extract multiscale features.
        
        Args:
            x: Input image tensor (B, C, H, W)
            return_separate: If True, return separate features instead of fused
            
        Returns:
            MultiscaleFeatures container
        """
        # Extract features
        global_feat = self.extract_clip_global(x)
        patch_feat = self.extract_dinov2_patches(x)
        
        # Aggregate patches
        patch_agg = self.patch_aggregator(patch_feat)
        
        # Compute patch grid
        B = x.shape[0]
        num_patches = patch_feat.shape[1]
        patch_grid = (int(num_patches ** 0.5), int(num_patches ** 0.5))
        
        # Fuse features
        combined = self.fuse_features(global_feat, patch_agg)
        
        return MultiscaleFeatures(
            global_feature=global_feat.squeeze(0) if B == 1 else global_feat,
            patch_features=patch_feat.squeeze(0) if B == 1 else patch_feat,
            patch_grid=patch_grid,
            combined=combined.squeeze(0) if B == 1 else combined
        )


class MultiscaleRetrievalHead(nn.Module):
    """
    Retrieval head that combines multiscale features.
    
    Projects fused features to the final embedding space
    used for similarity search.
    """
    
    def __init__(
        self,
        input_dim: int,
        output_dim: int = 768,
        hidden_dim: int = 256
    ):
        super().__init__()
        
        self.projection = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, output_dim),
        )
    
    def forward(self, features: MultiscaleFeatures) -> torch.Tensor:
        """
        Project multiscale features to retrieval space.
        
        Args:
            features: MultiscaleFeatures container
            
        Returns:
            Projected embedding (output_dim,)
        """
        return self.projection(features.combined)


# Convenience function
def create_multiscale_extractor(
    clip_model: nn.Module,
    dinov2_model: Optional[nn.Module] = None,
    embed_dim: int = 768,
    fusion_dim: int = 512
) -> MultiscaleExtractor:
    """
    Create a multiscale feature extractor.
    
    Args:
        clip_model: CLIP vision model for global features
        dinov2_model: Optional DINOv2 model for patch features
        embed_dim: CLIP embedding dimension
        fusion_dim: Output fusion dimension
        
    Returns:
        MultiscaleExtractor instance
    """
    return MultiscaleExtractor(
        clip_model=clip_model,
        dinov2_model=dinov2_model,
        embed_dim=embed_dim,
        patch_dim=768,  # DINOv2 default
        fusion_dim=fusion_dim,
        use_cross_attention=True
    )


# Self-check
if __name__ == "__main__":
    print("Testing MultiscaleExtractor...")
    
    # Test without actual models (dummy)
    class DummyModel(nn.Module):
        def __init__(self, output_dim=768):
            super().__init__()
            self.linear = nn.Linear(3, output_dim)
        
        def forward(self, x):
            B = x.shape[0]
            return torch.randn(B, 197, 768)  # 196 patches + CLS
    
    dummy_clip = DummyModel(768)
    dummy_dinov2 = DummyModel(768)
    
    extractor = MultiscaleExtractor(
        clip_model=dummy_clip,
        dinov2_model=dummy_dinov2,
        embed_dim=768,
        patch_dim=768,
        fusion_dim=512
    )
    
    # Test forward pass
    x = torch.randn(1, 3, 224, 224)
    features = extractor(x)
    
    print(f"Global feature shape: {features.global_feature.shape}")
    print(f"Patch features shape: {features.patch_features.shape}")
    print(f"Patch grid: {features.patch_grid}")
    print(f"Combined feature shape: {features.combined.shape}")
    
    # Test retrieval head
    head = MultiscaleRetrievalHead(input_dim=512, output_dim=768)
    embedding = head(features)
    print(f"Final embedding shape: {embedding.shape}")
    print(f"Embedding norm: {torch.norm(embedding).item():.4f}")
    
    print("\nMultiscaleExtractor test passed!")