import logging
from omegaconf import OmegaConf
import copy

import spacy
import torch
from torch import nn
from torchvision import transforms as T
from torchvision.transforms.functional import InterpolationMode
from transformers import LlamaTokenizer, BertTokenizer, BitsAndBytesConfig

import sys
sys.path.append("./")
from model.knwl_model import KnwlModel, KnwlEncoder
from model.utils import drop_sequence_mask, cat_pad, disabled_train, download_cached_file
from model.eva_vit import create_eva_vit_g
from model.qformer import BertConfig, BertLMHeadModel
from model.llama import LlamaForCausalLM, LlamaConfig


class GPTK(nn.Module):
    def __init__(
        self,
        img_size=224,
        drop_path_rate=0,
        use_grad_checkpoint=True,
        vit_precision="fp16",
        num_query_token=32,
        llm_model="",
        prompt="",
        max_txt_len=128,
        max_output_txt_len=256,
        d_knwl=768,
        topk={},
        pc=0.1,
        pt=0.1,
        pv=0.1
    ):
        super().__init__()
        
        self.topk = {k: v for k, v in topk.items() if v > 0}
        self.pc = pc
        self.pt = pt
        self.pv = pv

        # LLM
        self.llm_tokenizer = LlamaTokenizer.from_pretrained(llm_model, use_fast=False, truncation_side="left")
        llm_config = LlamaConfig.from_pretrained(llm_model)
        llm_config.gradient_checkpointing = True
        llm_config.use_cache = True
        quantization_config = BitsAndBytesConfig(
            load_in_8bit=True,
            llm_int8_threshold=6.0,
            llm_int8_has_fp16_weight=False,
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type='nf4'
        )
        self.llm_model = LlamaForCausalLM.from_pretrained(
            llm_model, config=llm_config, torch_dtype=torch.float16, quantization_config=quantization_config
        )
        self.llm_tokenizer.add_special_tokens({'pad_token': '[PAD]'})
        self.llm_tokenizer.add_special_tokens({'bos_token': '</s>'})
        self.llm_tokenizer.add_special_tokens({'eos_token': '</s>'})
        self.llm_tokenizer.add_special_tokens({'unk_token': '</s>'})
        self.llm_model.resize_token_embeddings(len(self.llm_tokenizer))
        for name, param in self.llm_model.named_parameters():
            param.requires_grad = False

        # ViT image encoder
        self.visual_encoder, self.ln_vision = self.init_vision_encoder(
            img_size, drop_path_rate, use_grad_checkpoint, vit_precision
        )
        for name, param in self.visual_encoder.named_parameters():
            param.requires_grad = False
        self.visual_encoder = self.visual_encoder.eval()
        self.visual_encoder.train = disabled_train
        logging.info("freeze vision encoder")

        # Q-former
        self.tokenizer = self.init_tokenizer(truncation_side="left")
        self.Qformer, self.query_tokens = self.init_Qformer(
            num_query_token, self.visual_encoder.num_features
        )
        self.Qformer.resize_token_embeddings(len(self.tokenizer))
        self.Qformer.cls = None
        self.llm_proj = nn.Linear(
            self.Qformer.config.hidden_size, self.llm_model.config.hidden_size
        )

        # Knowledge modules
        if len(self.topk) > 0:  # all added modules must contain "knwl" in their names
            self.knwl_encoder = KnwlEncoder(self.visual_encoder.num_features)
            self.knwl_query = copy.deepcopy(self.query_tokens)
            for k in self.topk.keys():
                m = KnwlModel(d_knwl=d_knwl, d_out=self.knwl_encoder.d, pt=pt)
                setattr(self, f"knwl_{k}", m)

        self.max_txt_len = max_txt_len
        self.max_output_txt_len = max_output_txt_len
        self.prompt = prompt
        prompt_tokens = self.llm_tokenizer(self.prompt, return_tensors="pt")
        self.prompt_length = prompt_tokens.attention_mask.sum(1)

    @classmethod
    def init_tokenizer(cls, truncation_side="right"):
        tokenizer = BertTokenizer.from_pretrained("bert-base-uncased", truncation_side=truncation_side)
        tokenizer.add_special_tokens({"bos_token": "[DEC]"})
        return tokenizer

    @classmethod
    def init_Qformer(cls, num_query_token, vision_width, cross_attention_freq=2):
        encoder_config = BertConfig.from_pretrained("bert-base-uncased")
        encoder_config.encoder_width = vision_width
        # insert cross-attention layer every other block
        encoder_config.add_cross_attention = True
        encoder_config.cross_attention_freq = cross_attention_freq
        encoder_config.query_length = num_query_token
        
        logging.disable(logging.CRITICAL)
        Qformer = BertLMHeadModel.from_pretrained(
            "bert-base-uncased", config=encoder_config
        )
        logging.disable(logging.NOTSET)
        
        query_tokens = nn.Parameter(
            torch.zeros(1, num_query_token, encoder_config.hidden_size)
        )
        query_tokens.data.normal_(mean=0.0, std=encoder_config.initializer_range)
        return Qformer, query_tokens

    @classmethod
    def init_vision_encoder(cls, img_size, drop_path_rate, use_grad_checkpoint, precision):
        visual_encoder = create_eva_vit_g(
            img_size, drop_path_rate, use_grad_checkpoint, precision
        )
        ln_vision = nn.LayerNorm(visual_encoder.num_features)

        return visual_encoder, ln_vision

    def concat_text_input_output(self, input_ids, input_atts, output_ids, output_atts):
        input_part_targets_len = []
        llm_tokens = {"input_ids": [], "attention_mask": []}
        for i in range(input_ids.size(0)):
            this_input_ones = input_atts[i].sum()
            input_part_targets_len.append(this_input_ones)
            llm_tokens['input_ids'].append(
                torch.cat([
                    input_ids[i][:this_input_ones],
                    output_ids[i][1:],
                    input_ids[i][this_input_ones:]
                ])
            )
            llm_tokens['attention_mask'].append(
                torch.cat([
                    input_atts[i][:this_input_ones],
                    output_atts[i][1:],
                    input_atts[i][this_input_ones:]
                ])
            )
        llm_tokens['input_ids'] = torch.stack(llm_tokens['input_ids'])
        llm_tokens['attention_mask'] = torch.stack(llm_tokens['attention_mask'])
        
        return llm_tokens, input_part_targets_len
    
    def forward_qformer(self, image, knowledge, prompt):
        views = []
        
        # knowledge embeds
        if len(self.topk) > 0 and knowledge is not None:
            embeds, masks = [], []
            for k in knowledge.keys():
                embeds_k, masks_k = getattr(self, f"knwl_{k}")(knowledge[k])
                embeds.append(embeds_k)
                masks.append(masks_k)
            embeds = cat_pad(embeds, cat_dim=0, pad_dim=1)
            masks = cat_pad(masks, cat_dim=0, pad_dim=1)

            embeds = self.knwl_encoder(
                inputs_embeds=embeds, attention_mask=masks
            )
            embeds = nn.functional.dropout(
                embeds, p=self.pc, training=self.training
            )

            N, (S, d) = len(image), embeds.shape[1:]
            embeds = embeds.reshape(-1, N, S, d)
            embeds = embeds.transpose(0, 1).flatten(1, 2)
            masks = masks.reshape(-1, N, S)
            masks = masks.transpose(0, 1).flatten(1, 2)
            views.append((embeds, masks, self.knwl_query))

        # image embeds
        embeds = self.ln_vision(self.visual_encoder(image))
        embeds = nn.functional.dropout(
            embeds, p=self.pc, training=self.training
        )
        masks = drop_sequence_mask(
            *embeds.shape[:2], image.device, self.pt, self.training
        )
        views.append((embeds, masks, self.query_tokens))

        # Qformer forward
        text_Qformer = self.tokenizer(
            prompt,
            padding='longest',
            truncation=True,
            max_length=self.max_txt_len,
            return_tensors="pt",
        ).to(image.device)
        
        qfm_embeds, qfm_masks = [], []
        for embeds, masks, query in views:
            query = query.expand(image.shape[0], -1, -1)
            query_atts = torch.ones(query.size()[:-1], dtype=torch.long).to(embeds.device)
            Qformer_atts = torch.cat([query_atts, text_Qformer.attention_mask], dim=1)

            query_output = self.Qformer.bert(
                text_Qformer.input_ids,
                attention_mask=Qformer_atts,
                query_embeds=query,
                encoder_hidden_states=embeds,
                encoder_attention_mask=masks,
                return_dict=True,
            )
            embeds = self.llm_proj(query_output.last_hidden_state[:, :query.size(1),:])
            masks = torch.ones(embeds.size()[:-1], dtype=torch.long).to(image.device)
            qfm_embeds.append(embeds)
            qfm_masks.append(masks)

        # drop views
        if self.training:
            view_masks = drop_sequence_mask(len(image), len(qfm_embeds), image.device, self.pv)
            qfm_masks = [m * view_masks[:, i:(i+1)] for i, m in enumerate(qfm_masks)]
        llm_embeds = torch.cat(qfm_embeds, dim=1)
        llm_masks = torch.cat(qfm_masks, dim=1)

        return llm_embeds, llm_masks

    def forward(self, samples):
        inputs_llm, atts_llm = self.forward_qformer(
            samples["image"], samples["knowledge"], samples["prompt"]
        )
        device = inputs_llm.device

        self.llm_tokenizer.padding_side = "right"
        self.llm_tokenizer.truncation_side = 'left'
        text_input_tokens = self.llm_tokenizer(
            samples['prompt'],
            return_tensors="pt",
            padding="longest",
            truncation=True,
            max_length=self.max_txt_len,
        ).to(device)

        self.llm_tokenizer.truncation_side = 'right'
        text_output_tokens = self.llm_tokenizer(
            [t + self.llm_tokenizer.eos_token for t in samples['output']],
            return_tensors="pt",
            padding="longest",
            truncation=True,
            max_length=self.max_output_txt_len,
        ).to(device)

        llm_tokens, input_part_targets_len = self.concat_text_input_output(
            text_input_tokens.input_ids,
            text_input_tokens.attention_mask,
            text_output_tokens.input_ids,
            text_output_tokens.attention_mask,
        )

        # do not apply loss to the padding
        targets = llm_tokens['input_ids'].masked_fill(
            llm_tokens['input_ids'] == self.llm_tokenizer.pad_token_id, -100
        )

        # do not apply loss to the text input (i.e., instruction)
        for i, l in enumerate(input_part_targets_len):
            targets[i][:l] = -100

        # do not apply loss to the query tokens
        empty_targets = (
            torch.ones(atts_llm.size(), dtype=torch.long).to(device).fill_(-100)
        )
        targets = torch.cat([empty_targets, targets], dim=1)

        inputs_embeds = self.llm_model.get_input_embeddings()(llm_tokens['input_ids'])
        inputs_embeds = torch.cat([inputs_llm, inputs_embeds], dim=1)
        attention_mask = torch.cat([atts_llm, llm_tokens['attention_mask']], dim=1)

        outputs = self.llm_model(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            return_dict=True,
            labels=targets,
        )
        loss = outputs.loss

        return loss

    @torch.no_grad()
    def generate(
        self,
        samples,
        use_nucleus_sampling=False,
        num_beams=5,
        max_length=256,
        min_length=1,
        top_p=0.9,
        repetition_penalty=1.5,
        length_penalty=1,
        num_captions=1,
        temperature=1,
        streamer=None,
        auto_cast=False
    ):
        prompt = samples["prompt"] if "prompt" in samples.keys() else self.prompt
        if isinstance(prompt, str):
            prompt = [prompt] * samples["image"].size(0)
        else:
            assert len(prompt) == samples["image"].size(0), \
                "The number of prompts must be equal to the batch size."

        with torch.cuda.amp.autocast(auto_cast):
            inputs_llm, atts_llm = self.forward_qformer(
                samples["image"], samples["knowledge"], prompt
            )
            device = inputs_llm.device
        
            self.llm_tokenizer.padding_side = "left"
            llm_tokens = self.llm_tokenizer(
                prompt,
                padding="longest",
                return_tensors="pt"
            ).to(device)

            inputs_embeds = self.llm_model.get_input_embeddings()(llm_tokens.input_ids)
            inputs_embeds = torch.cat([inputs_llm, inputs_embeds], dim=1)
            attention_mask = torch.cat([atts_llm, llm_tokens.attention_mask], dim=1)

            outputs = self.llm_model.generate(
                inputs_embeds=inputs_embeds,
                attention_mask=attention_mask,
                do_sample=use_nucleus_sampling,
                top_p=top_p,
                temperature=temperature,
                num_beams=num_beams,
                max_length=max_length,
                min_length=min_length,
                repetition_penalty=repetition_penalty,
                length_penalty=length_penalty,
                num_return_sequences=num_captions,
                streamer=streamer
            )

        if streamer is None:
            outputs[outputs == 0] = 2 # convert output id 0 to 2 (eos_token_id)
            output_text = self.llm_tokenizer.batch_decode(outputs, skip_special_tokens=True)
            output_text = [text.strip() for text in output_text]

            return output_text
        else:
            return outputs


    @torch.no_grad()
    def predict_answers(
        self,
        samples,
        num_beams=5,
        max_len=10,
        min_len=1,
        length_penalty=0
    ):
        output_text = self.generate(
            samples,
            num_beams=num_beams,
            max_length=max_len,
            min_length=min_len,
            length_penalty=length_penalty
        )
        output_text = self._lemmatize(output_text)

        return output_text
    
    def _lemmatize(self, answers):
        lemmatizer = spacy.load("en_core_web_sm")
        def apply(answer):
            doc = lemmatizer(answer)

            words = []
            for token in doc:
                if token.pos_ in ["NOUN", "VERB"]:
                    words.append(token.lemma_)
                else:
                    words.append(token.text)
            answer = " ".join(words)

            return answer

        return [apply(answer) for answer in answers]

    @classmethod
    def from_config(cls, cfg):
        llm_model = cfg.get("llm_model")
        num_query_token = cfg.get("num_query_token", 32)

        img_size = cfg.get("image_size", 224)
        drop_path_rate = cfg.get("drop_path_rate", 0)
        use_grad_checkpoint = cfg.get("use_grad_checkpoint", True)
        vit_precision = cfg.get("vit_precision", "fp16")

        prompt = cfg.get("prompt", "")
        max_txt_len = cfg.get("max_txt_len", 128)
        max_output_txt_len = cfg.get("max_output_txt_len", 256)

        d_knwl = cfg.get("d_knwl", 768)
        topk = cfg.get("topk", {})
        pc = cfg.get("pc", 0.1)
        pt = cfg.get("pt", 0.1)
        pv = cfg.get("pv", 0.4)

        model = cls(
            img_size=img_size,
            drop_path_rate=drop_path_rate,
            use_grad_checkpoint=use_grad_checkpoint,
            vit_precision=vit_precision,
            num_query_token=num_query_token,
            llm_model=llm_model,
            prompt=prompt,
            max_txt_len=max_txt_len,
            max_output_txt_len=max_output_txt_len,
            d_knwl=d_knwl,
            topk=topk,
            pc=pc,
            pt=pt,
            pv=pv
        )

        pretrain_path = cfg.get("pretrained", None)
        assert pretrain_path is not None, "Pretrain_path is None."
        cached_file = download_cached_file(
            pretrain_path, check_hash=False, progress=True
        )
        checkpoint = torch.load(cached_file, map_location="cpu")
        state_dict = checkpoint["model"]
        model.load_state_dict(state_dict, strict=False)
        logging.info("load checkpoint from %s" % pretrain_path)
        
        return model


def get_gptk_image_transform(model_type: str = "gptk-7b"):
    assert model_type in ("gptk-7b", "gptk-13b")
    model_config = OmegaConf.load(f"model/{model_type}.yaml")
    size = model_config.get("image_size", 224)

    normalizer = T.Normalize(
        mean=(0.48145466, 0.4578275, 0.40821073),
        std=(0.26862954, 0.26130258, 0.27577711)
    )
    trans_train = T.Compose([
        T.RandomResizedCrop(
            size=(size, size), scale=(0.5, 1.0),
            interpolation=InterpolationMode.BICUBIC,
        ),
        T.RandomHorizontalFlip(),
        T.ToTensor(),
        normalizer
    ])

    trans_val = T.Compose([
        T.Resize(
            size=(size, size), interpolation=InterpolationMode.BICUBIC,
        ),
        T.RandomHorizontalFlip(),
        T.ToTensor(),
        normalizer
    ])

    return trans_train, trans_val


def get_gptk_model(
        model_type: str = "gptk-7b",
        d_knwl: int = 768,
        topk: dict = {"whole": 0, "five": 0, "nine": 0},
        pc: float = 0.1, pt: float = 0.1, pv: float = 0.4
    ):
    assert model_type in ("gptk-7b", "gptk-13b")

    model_config = OmegaConf.load(f"model/{model_type}.yaml")
    model_config.pv = pv
    model_config.pt = pt
    model_config.pc = pc
    model_config.topk = {k: v for k, v in topk.items() if v > 0}
    model_config.d_knwl = d_knwl

    model = GPTK.from_config(model_config)
    if sum(topk.values()) > 0:
        model.knwl_query.data.copy_(model.query_tokens.data.clone().detach())

    return model