import re

import torch
from torch import nn
from torchvision import transforms


from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
from accelerate import Accelerator
from models.opt import OPTModel, OPTConfig, OPTForCausalLM
import models.vit 

from PIL import Image
import json 
import numpy as np



import torch.nn.functional as F
from transformers.tokenization_utils_base import BatchEncoding

def rank_answer(model, image, question_input, answer_ids, answer_atts, k, tokenizer):

    num_ques = question_input.input_ids.size(0)
    start_ids = answer_ids[0,0].repeat(num_ques,1) # bos token

    start_ids = torch.cat((question_input.input_ids,  start_ids), dim=1)
    attention_mask = torch.cat((question_input.attention_mask,  torch.ones((num_ques, 1)).to(question_input.attention_mask.device)), dim=1)
    
    start_input = {'input_ids': start_ids, 'attention_mask': attention_mask}
    start_input = BatchEncoding(start_input)
    
    
    
    start_output = model(image, start_input, return_dict = True, mode='evaluate')     
    
    logits = start_output.logits[:,-1,:] # first token's logit

    # topk_probs: top-k probability 
    # topk_ids: [num_question, k]        
    answer_first_token = answer_ids[:,1]
    prob_first_token = F.softmax(logits,dim=1).index_select(dim=1, index=answer_first_token) 
    topk_probs, topk_ids = prob_first_token.topk(k,dim=1) 

    # answer input: [num_question*k, answer_len]                 
    input_ids = []
    input_atts = []
    for b, topk_id in enumerate(topk_ids):
        input_ids.append(answer_ids.index_select(dim=0, index=topk_id))
        input_atts.append(answer_atts.index_select(dim=0, index=topk_id))
    input_ids = torch.cat(input_ids,dim=0)  
    input_atts = torch.cat(input_atts,dim=0)  

    start_ids = tile(start_ids, 0, k)
    attention_mask = tile(attention_mask, 0, k)
    image = tile(image, 0, k)
    
    
        
    
    input_ids = torch.cat((start_ids, input_ids), dim=1) # include the  <s> ?
    input_atts = torch.cat((attention_mask, input_atts), dim=1)
        
    targets_ids = input_ids.masked_fill(input_ids == tokenizer.pad_token_id, -100)

    
    
    # repeat encoder's output for top-k answers


    inputs = {'input_ids': input_ids, 'attention_mask': input_atts}
    inputs = BatchEncoding(inputs)
    
    output = model(image, inputs, labels = targets_ids, return_dict = True, mode='train', reduction='none')                 

    answer_loss = output.loss 
    answer_loss = answer_loss.view(input_ids.size(0),-1)

    # topk_prob: first token probability

    topk_probs = topk_probs.view(-1,1)
    log_probs = torch.cat([topk_probs.log(), -answer_loss],dim=1)

    # re-calculate log probabilities for the answer sequences using chain rule
    log_probs_sum = log_probs.sum(1)
    log_probs_sum = log_probs_sum.view(num_ques,k)

    topk_probs = F.softmax(log_probs_sum, dim=-1)
    # get top-k after re-ranking
    topk_probs, rerank_id = topk_probs.topk(k,dim=1) 
    topk_ids = torch.gather(topk_ids, 1, rerank_id)    

    return topk_ids, topk_probs
    
def tile(x, dim, n_tile):
    init_dim = x.size(dim)
    repeat_idx = [1] * x.dim()
    repeat_idx[dim] = n_tile
    x = x.repeat(*(repeat_idx))
    order_index = torch.LongTensor(np.concatenate([init_dim * np.arange(n_tile) + i for i in range(init_dim)]))
    return torch.index_select(x, dim, order_index.to(x.device))    





class VisOPT(nn.Module):
    def __init__(self,                 
                 opt_model_name = 'facebook/opt-350m',
                 vision_model_name = 'vit_base_patch16_224',
                 use_vis_prefix = True,
                 start_layer_idx = 11,
                 end_layer_idx = 23,
                 return_hidden_state_vision = True,
                 injected_hidden_states = 1,
                 
                 ):
        super().__init__()
        print("Loading VisOPT ...")
        # text
        config_opt = AutoConfig.from_pretrained(opt_model_name)
        
        config_opt.use_vis_prefix = use_vis_prefix
        config_opt.start_layer_idx = start_layer_idx
        config_opt.end_layer_idx = end_layer_idx
            
        print(config_opt)
        print("Loading: ", opt_model_name)
        self.model_text = OPTForCausalLM.from_pretrained(opt_model_name, config=config_opt)
        
        # vision
        print("Loading: ", vision_model_name)
        vision_func = getattr(models.vit, vision_model_name)
        self.model_vision = vision_func(pretrained=True, return_hidden_state=return_hidden_state_vision)
        
        # connector
        self.injected_hidden_states = injected_hidden_states
        vis_dim = self.model_vision.embed_dim
        text_dim = config_opt.hidden_size
        self.connector = nn.ModuleList([nn.Linear(vis_dim, text_dim) for i in range(injected_hidden_states)])
        
        
    def forward(self, image=None, text=None, mode='generate', return_dict=True, labels=None, reduction='mean', **generation_kwargs):
        
        if image is not None:
            image_embed, image_feat = self.model_vision(image, external_features=None)

            image_feat = list(image_feat)
            image_feat = image_feat[-self.injected_hidden_states:]

            ## only cls token,  we can think of somthing else
            for i in range(1, self.injected_hidden_states + 1):
                image_feat[-i] = self.connector[-i](image_feat[-i][:, 0, :].unsqueeze(1))
        else:
            image_feat = None

            
        # image_feat = None
        if mode == 'train' or mode == 'evaluate':
            text_output = self.model_text(input_ids=text.input_ids, attention_mask=text.attention_mask, return_dict=return_dict, vis_prefix=image_feat, labels = labels, reduction=reduction)
            return text_output
        elif mode == 'generate':
            print('generation')
            gen = self.model_text.generate(input_ids=text.input_ids, vis_prefix=image_feat, **generation_kwargs)        
            return gen