import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from transformers import (
    AutoModel,
    BertTokenizer,
    SiglipModel,
    SiglipProcessor,
    WhisperModel,
    WhisperFeatureExtractor,
)


class CompactEmbedder(nn.Module):
    def __init__(self, local_dir="./multi-modal-embed-small"):
        local_dir = os.path.abspath(local_dir)
        super().__init__()
        text_path = os.path.normpath(f"{local_dir}/../base-models/all-MiniLM-L6-v2")
        image_path = os.path.normpath(
            f"{local_dir}/../base-models/siglip-base-patch16-512"
        )
        audio_path = os.path.normpath(f"{local_dir}/../base-models/whisper-tiny")

        self.text_tokenizer = BertTokenizer.from_pretrained(
            text_path, local_files_only=True
        )
        self.text_encoder = AutoModel.from_pretrained(text_path, local_files_only=True)

        self.image_processor = SiglipProcessor.from_pretrained(
            image_path, local_files_only=True
        )
        self.image_encoder = SiglipModel.from_pretrained(
            image_path, local_files_only=True
        ).vision_model
        self.audio_processor = WhisperFeatureExtractor.from_pretrained(
            audio_path, local_files_only=True
        )
        self.image_proj = nn.Linear(768, 384)
        self.audio_processor = WhisperFeatureExtractor.from_pretrained(
            audio_path, local_files_only=True
        )
        self.audio_encoder = WhisperModel.from_pretrained(
            audio_path, local_files_only=True
        ).encoder

        checkpoint_path = f"{local_dir}/model.pt"
        state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True)

        self.text_encoder.load_state_dict(
            {
                k.replace("text_encoder.encoder.", ""): v
                for k, v in state_dict.items()
                if k.startswith("text_encoder.encoder.")
            }
        )

        self.image_encoder.load_state_dict(
            {
                k.replace("image_encoder.vision_encoder.", ""): v
                for k, v in state_dict.items()
                if k.startswith("image_encoder.vision_encoder.")
            }
        )

        self.image_proj.load_state_dict(
            {
                k.replace("image_encoder.projection.", ""): v
                for k, v in state_dict.items()
                if k.startswith("image_encoder.projection.")
            }
        )

        self.audio_encoder.load_state_dict(
            {
                k.replace("audio_encoder.encoder.", ""): v
                for k, v in state_dict.items()
                if k.startswith("audio_encoder.encoder.")
            }
        )

        self.eval()

    @torch.no_grad()
    def encode_text(self, texts):
        if isinstance(texts, str):
            texts = [texts]
        inputs = self.text_tokenizer(
            texts, padding=True, truncation=True, return_tensors="pt"
        )
        inputs = {k: v.to(next(self.parameters()).device) for k, v in inputs.items()}
        outputs = self.text_encoder(**inputs)
        embeddings = outputs.last_hidden_state.mean(dim=1)
        return F.normalize(embeddings, p=2, dim=-1)

    @torch.no_grad()
    def encode_image(self, images):
        if isinstance(images, Image.Image):
            images = [images]
        inputs = self.image_processor(images=images, return_tensors="pt")
        inputs = {k: v.to(next(self.parameters()).device) for k, v in inputs.items()}
        outputs = self.image_encoder(**inputs)
        embeddings = self.image_proj(outputs.pooler_output)
        return F.normalize(embeddings, p=2, dim=-1)

    @torch.no_grad()
    def encode_audio(self, waveform):
        if isinstance(waveform, torch.Tensor):
            waveform = waveform.squeeze().numpy()
        inputs = self.audio_processor(
            waveform, sampling_rate=16000, return_tensors="pt"
        )
        inputs = {k: v.to(next(self.parameters()).device) for k, v in inputs.items()}
        outputs = self.audio_encoder(**inputs)
        embeddings = outputs.last_hidden_state.mean(dim=1)
        return F.normalize(embeddings, p=2, dim=-1)


if __name__ == "__main__":
    model = CompactEmbedder(local_dir="./models/multi-modal-embed-small")

    text_vec = model.encode_text("example query")
    print(f"Text shape: {text_vec.shape}")