import torch
import torch.nn as nn
from transformers import SiglipVisionModel
from .configuration import TinyDocVLMConfig

class SigLIPVisionEncoder(nn.Module):
    """
    Wrapper around HuggingFace's SiglipVisionModel.
    Handles encoding of multiple image tiles and thumbnails.
    """
    def __init__(self, config: TinyDocVLMConfig):
        super().__init__()
        self.config = config
        
        # Load from config or create model
        vision_config = config.vision_config
        # We can initialize from config. If we are running pretraining, we load weights.
        # During runtime we might load a pretrained siglip model.
        self.encoder = SiglipVisionModel(vision_config)
        self.hidden_size = vision_config.hidden_size
        
        # Add special region classification or auxiliary layers if needed in future
        # For now, just a wrapper around the SigLIP vision encoder

    def forward(
        self, 
        pixel_values: torch.Tensor, 
        interpolate_pos_encoding: bool = False
    ) -> torch.Tensor:
        """
        Args:
            pixel_values: shape (batch_size, num_tiles, channels, height, width) 
                          or (batch_size * num_tiles, channels, height, width)
            interpolate_pos_encoding: whether to interpolate positional embeddings if resolution changes
            
        Returns:
            visual_features: shape (batch_size, num_tiles, num_patches, hidden_size)
        """
        # If input has shape (batch_size, num_tiles, channels, height, width)
        if len(pixel_values.shape) == 5:
            batch_size, num_tiles, channels, height, width = pixel_values.shape
            # Flatten batch and tiles for vision encoder
            pixel_values = pixel_values.view(batch_size * num_tiles, channels, height, width)
        else:
            # Assumed to be already flattened (batch_size * num_tiles, channels, height, width)
            batch_size = 1
            num_tiles = pixel_values.shape[0]
            channels, height, width = pixel_values.shape[1:]

        # Run through SigLIP Vision Model
        outputs = self.encoder(
            pixel_values=pixel_values,
            interpolate_pos_encoding=interpolate_pos_encoding
        )
        
        # Last hidden state: (batch_size * num_tiles, num_patches, hidden_size)
        # For SigLIP-B/16 with 384x384 input: num_patches = (384/16)^2 = 24^2 = 576
        last_hidden_state = outputs.last_hidden_state
        
        # Reshape back to batch format
        num_patches = last_hidden_state.shape[1]
        last_hidden_state = last_hidden_state.view(batch_size, num_tiles, num_patches, self.hidden_size)
        
        return last_hidden_state