import torch
import torch.nn as nn

from transformers import CLIPImageProcessor

try:
    from imagebind.models import imagebind_model
    from imagebind.models.imagebind_model import ModalityType
    from imagebind.data import load_and_transform_audio_data
except ImportError:
    pass


class ImageBindWrapper(nn.Module):
    def __init__(self, vision_tower, select_layer, select_feature="patch", delay_load=False):
        super().__init__()

        self.is_loaded = False

        self.vision_tower_name = vision_tower
        self.select_layer = select_layer
        self.select_feature = select_feature

        if not delay_load:
            self.load_model()

    def load_model(self):
        self.image_processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-large-patch14")
        self.vision_tower = imagebind_model.imagebind_huge(pretrained=True)
        for p in self.vision_tower.parameters():
            p.requires_grad = False
        self.vision_tower.eval()
        self.is_loaded = True

    def train(self, mode=True):
        self.training = mode

        if self.is_loaded:
            self.vision_tower.eval()

    @torch.no_grad()
    def forward(self, x):
        if type(x) == dict:
            if x["audios"] is not None:
                inputs = {ModalityType.AUDIO: load_and_transform_audio_data(x["audios"], device=self.device).half()}
                embeddings = self.vision_tower(inputs)
                audio_embedding = embeddings[ModalityType.AUDIO]
                return audio_embedding.unsqueeze(1)
        else:
            inputs = {ModalityType.VISION: x.to(dtype=self.dtype)}
            embeddings = self.vision_tower(inputs)
            vision_embedding = embeddings[ModalityType.VISION]
            if vision_embedding.ndim == 2:
                return vision_embedding.unsqueeze(1)
            if vision_embedding.shape[1] == 257:
                return vision_embedding[:, 1:]
            raise ValueError(f"Unexpected shape: {vision_embedding.shape}")

    @property
    def dummy_feature(self):
        return torch.zeros(1, 1024, device=self.device, dtype=self.dtype)

    @property
    def dtype(self):
        return self.vision_tower.modality_preprocessors.vision.cls_token.dtype

    @property
    def device(self):
        return self.vision_tower.modality_preprocessors.vision.cls_token.device

    @property
    def hidden_size(self):
        return 1024