#model.py
import torch
import torch.nn as nn
import open_clip

class CONFIG:
    # Model parameters
    clip_model_name = "ViT-B-32"
    pretrained_dataset = "laion400m_e31"
    embedding_dim = 512
    fusion_type = 'sum'  # 'sum', 'concat', or 'attention'
    device = "cuda" if torch.cuda.is_available() else "cpu"
    image_size = (224, 224)
    batch_size = 16
    device = "cuda" if torch.cuda.is_available() else "cpu"


class FusionModel(nn.Module):
    def __init__(self, clip_model, fusion_type='sum'):
        super(FusionModel, self).__init__()
        self.clip_model = clip_model
        self.fusion_type = fusion_type

        # Depending on the fusion type, define the combined embedding size
        if fusion_type == 'sum':
            self.combined_dim = CONFIG.embedding_dim
        elif fusion_type == 'concat':
            self.combined_dim = CONFIG.embedding_dim * 2
        elif fusion_type == 'attention':
            self.combined_dim = CONFIG.embedding_dim
            # Define the attention layer
            self.attention = nn.MultiheadAttention(embed_dim=CONFIG.embedding_dim, num_heads=8)
        else:
            raise ValueError("Invalid fusion type")

        # MLP to project the combined embedding to the final embedding space
        self.mlp = nn.Sequential(
            nn.Linear(self.combined_dim, CONFIG.embedding_dim),
            nn.ReLU(),
            nn.Dropout(p=0.1),
            nn.Linear(CONFIG.embedding_dim, CONFIG.embedding_dim)
        )

    def forward(self, image, text):
        # Encode image and text using CLIP encoders
        image_features = self.clip_model.encode_image(image)  # Shape: [batch_size, embedding_dim]
        text_tokens = open_clip.tokenize(text).to(CONFIG.device)  # Tokenize text
        text_features = self.clip_model.encode_text(text_tokens)  # Shape: [batch_size, embedding_dim]

        # Normalize the features
        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
        text_features = text_features / text_features.norm(dim=-1, keepdim=True)

        # Fusion strategies
        if self.fusion_type == 'sum':
            combined_features = image_features + text_features
        elif self.fusion_type == 'concat':
            combined_features = torch.cat([image_features, text_features], dim=-1)
        elif self.fusion_type == 'attention':
            combined_features = self.attention_fusion(image_features, text_features)
        else:
            raise ValueError("Invalid fusion type")

        # Pass through MLP
        output = self.mlp(combined_features)
        return output

    def attention_fusion(self, image_features, text_features):
        # Reshape features to (sequence_length, batch_size, embedding_dim)
        # Since we have two features per sample, sequence_length = 2
        features = torch.stack([image_features, text_features], dim=0)  # Shape: [2, batch_size, embedding_dim]

        # Multi-head attention expects inputs in shape (sequence_length, batch_size, embedding_dim)
        # Query, Key, Value are all the features in this case
        attn_output, _ = self.attention(features, features, features)    # Shape: [2, batch_size, embedding_dim]

        # Aggregate the attention outputs (e.g., take the mean over the sequence length)
        combined_features = attn_output.mean(dim=0)                      # Shape: [batch_size, embedding_dim]

        return combined_features