from pathlib import Path
from os.path import exists
import torch
import torch.nn as nn
from copy import deepcopy
from typing_extensions import Literal
from typing import Optional, List
from torchtyping import TensorType
from transformers.file_utils import ModelOutput
from config import MultimodalConfig

from adapters import (
    Adapter,
    ParallelAdapter,
    AdapterWrapper,
    ParallelAdapterWrapper,
)
from image_prefix import ImagePrefix
from image_input import ImageInput
from sampling import generate
from utils import build_labels 
from transforms import get_transforms 
import os


from transformers import GPTJForCausalLM
from transformers import GPT2TokenizerFast

def get_tokenizer(name, sequence_length=2048):
    """
    Gets tokenizer for LM
    """
    if name == "gpt2":
        tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
        tokenizer.pad_token_id = tokenizer.eos_token#_id #transformers 4.25.1 update. New version is: tokenizer.eos_token_id
        tokenizer.padding_side = "right"
        tokenizer.model_max_length = sequence_length
        # setup lm settings
        tokenizer.add_special_tokens(
            {"cls_token": "<|image|>"}
        )  # add special image token to tokenizer
    else:
        raise ValueError(f"Tokenizer {name} not recognized")
    return tokenizer


class LimberGPTJ(GPTJForCausalLM):
    

    def setup_multimodal(self, config, device='cuda:0'):
        
        if isinstance(config, (str, Path)):
            config = MultimodalConfig.from_yml(
                config
            )  # load config from yml file if config is a string
        else:
            assert isinstance(config, MultimodalConfig)
        self.multimodal_config = config
        self.seq_len = self.config.max_position_embeddings
        self.tokenizer = get_tokenizer('gpt2', sequence_length=self.seq_len)

        self.image_token = self.tokenizer.cls_token_id
        self.eos_token = self.tokenizer.eos_token_id
        self.resize_token_embeddings(len(self.tokenizer))
        self.config.pad_token_id = self.tokenizer.eos_token_id
        self.word_embedding = self.transformer.wte

        if config.freeze_lm:
            for name, param in self.named_parameters():  # freeze lm weights
                if config.adapter_config and "adapter" in name:
                    param.requires_grad = True
                else:
                    param.requires_grad = False


        # adapter settings
        self.mlp_adapter_added, self.attn_adapter_added = False, False
        self.image_prefix = ImagePrefix(
            config=config,
            out_dim=self.config.hidden_size,
        )

        self.image_prefix_seq_len = self.image_prefix.out_seq_len
        if 'identity' in config.encoder_name or 'nfresnet' in config.encoder_name:
            inp_rez = None
        else:
            inp_rez = self.image_prefix.enc.input_resolution
        self.transforms = get_transforms(
            config.image_size,
            config.encoder_name,
            input_resolution= inp_rez,
        )

        # add adapters
        if config.adapter_config:
            mlp_config = deepcopy(config.adapter_config.get("mlp", None))
            if mlp_config:
                assert mlp_config.get("adapter_type") is not None
                self.add_adapters(
                    location="mlp",
                    adapter_type=mlp_config.pop("adapter_type"),
                    downsample_factor=mlp_config.pop("downsample_factor", 4),
                    **mlp_config,
                )
            attn_config = deepcopy(config.adapter_config.get("attention", None))
            if attn_config:
                assert attn_config.get("adapter_type") is not None
                self.add_adapters(
                    location="attention",
                    adapter_type=attn_config.pop("adapter_type"),
                    **attn_config,
                )

        if config.freeze_img_encoder:
            for param in self.image_prefix.enc.parameters():
                param.requires_grad = False

    def preprocess_inputs(self, input_list: list, embed = True) -> List[torch.Tensor]:
        """
        Expects a list of strings and instances of ImageInput
        Converts them into a list of tensors and then optionally runs self.embed over it
        """
        for i in range(len(input_list)):
            inp = input_list[i]
            if isinstance(inp, str):
                input_list[i] = self.tokenizer.encode(inp, return_tensors="pt")
            elif isinstance(inp, ImageInput):
                input_list[i] = inp.get_transformed_image(transform_fn = self.transforms)
            else:
                raise Exception(f'Invalid input type:{type(inp)}')

        if embed == True:
            return self.embed(input_list)
        else: 
            return input_list

    def embed(self, inputs: List[torch.Tensor]) -> TensorType["b", "s", "d"]:
        """
        Embeds a list of tensors In the correct format to input into the LM (b, s, d).
        For each tensor, if it's 2d assume it's text and use word embedding,
        if it's 4d, assume it's an image, and use image_prefix to embed.
        """
        emb_list = []
        for x in inputs:
            if x.ndim == 2:
                x = x.to(self.device)
                emb_list.append(self.word_embedding(x))
            elif x.ndim == 4:
                x = x.to(self.device).half()
                image_embeddings = self.image_prefix(x)
                emb_list.append(image_embeddings)
            elif x.ndim == 3:
                x = x.unsqueeze(0).to(self.device).half()
                image_embeddings=self.image_prefix(x)
                emb_list.append(image_embeddings)
            else:
                raise ValueError(f"Expected 2d or 4d tensor, got {x.ndim}d")
        return torch.cat(emb_list, dim=1)

    @torch.no_grad()
    def generate(
        self,
        embeddings: TensorType["b", "s", "d"],
        max_steps: int = 100,
        temperature: float = 0.7,
        top_k: int = 0,
        top_p: float = 0.9,
        decode: bool = True,

    ):
        """
        Generates captions for a batch of embeddings.
        """
        return generate(
            self,
            embeddings=embeddings,
            max_steps=max_steps,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            decode=decode,
        )


    def forward(
        self,
        images: TensorType["b", "c", "h", "w"] = None,
        captions: Optional[TensorType["b", "seq"]] = None,
        output_hidden_states: bool = True,
        input_embeddings: TensorType["b", "s", "d"] = None,
        attention_mask=None,
        return_dict=None
    ) -> ModelOutput:
        assert captions is not None, "Must provide captions in training"
        assert any([i is not None for i in [images, input_embeddings]]) and not all(
            [i is not None for i in [images, input_embeddings]]
        ), "Pass in either images, or input embeddings, not both."


        if input_embeddings is None:
            input_embeddings = self.image_prefix(images)

        labels = build_labels(
            input_embeddings, captions, self.eos_token, self.device
        )  # build labels from input_embeddings

        word_embeddings = self.word_embedding(captions)
        input_embeddings = torch.cat(
            (
                input_embeddings,
                word_embeddings,
            ),
            dim=1,
        )
        
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        lm_outputs = super().forward(
            inputs_embeds=input_embeddings,
            labels=labels,
            output_hidden_states=output_hidden_states,
            attention_mask=attention_mask,
            use_cache=False
        )
        return lm_outputs

    def add_adapters(
        self,
        downsample_factor: int = 4,
        adapter_type: Literal["normal", "parallel", "scaled_parallel"] = "normal",
        location: Literal["mlp", "attention"] = "mlp",
        ff_attr: str = "mlp",
        attn_attr: str = "attn",
        **adapter_kwargs,
    ):
        #This function is based on the add_adapters method used to create the MAGMA model. Source here:
        # https://github.com/Aleph-Alpha/magma/blob/a69a575bcdbfa02c5c984f9c0058ac84b5d9c841/magma/magma.py#L102
        """
        Adds an adapter layer to `self` at the specified location
        """
        assert adapter_type in [
            "normal",
            "parallel",
            "scaled_parallel",
        ], "adapter_type must be one of 'normal', 'parallel', or 'scaled_parallel'"
        assert location in [
            "mlp",
            "attention",
        ], "location must be one of 'mlp' or 'attention'"

        for l in range(len(self.transformer.h)):
            if location == "mlp":
                if self.mlp_adapter_added:
                    raise ValueError("Adapter layer already added")
                mlp = getattr(self.transformer.h[l], ff_attr)
                if adapter_type in ["parallel", "scaled_parallel"]:
                    adapter_layer = ParallelAdapter(
                        module=mlp,
                        dim=self.config.hidden_size,
                        downsample_factor=downsample_factor,
                        scaled=adapter_type == "scaled_parallel",
                        **adapter_kwargs,
                    )
                else:
                    adpt = Adapter(
                        dim=self.config.hidden_size,
                        downsample_factor=downsample_factor,
                        **adapter_kwargs,
                    )
                    adapter_layer = nn.Sequential(
                        *[
                            mlp,
                            adpt,
                        ]
                    )
                setattr(self.transformer.h[l], ff_attr, adapter_layer)
            else:
                if self.attn_adapter_added:
                    raise ValueError("Adapter layer already added")
                attn = getattr(self.transformer.h[l], attn_attr)
                if adapter_type in ["parallel", "scaled_parallel"]:
                    adapter_layer = ParallelAdapterWrapper(
                        module=attn,
                        dim=self.config.hidden_size,
                        downsample_factor=downsample_factor,
                        scaled="scaled" in adapter_type,
                        **adapter_kwargs,
                    )
                else:
                    adapter_layer = AdapterWrapper(
                        attn_block=attn,
                        dim=self.config.hidden_size,
                        downsample_factor=downsample_factor,
                        **adapter_kwargs,
                    )
                setattr(self.transformer.h[l], attn_attr, adapter_layer)

        if location == "mlp":
            self.mlp_adapter_added = True
        else:
            self.attn_adapter_added = True
        
         

       


    
