from transformers import MT5ForConditionalGeneration, TFMT5ForConditionalGeneration, T5Tokenizer, PreTrainedTokenizerFast
import torch
from torch import nn
import numpy as np
import tensorflow as tf 
import tqdm
import io
import time 
import os 
import argparse

os.environ["CUDA_VISIBLE_DEVICES"] = "-1" 


def shape_list(x):
    """Deal with dynamic shape in tensorflow cleanly."""
    static = x.shape.as_list()
    dynamic = tf.shape(x)
    return [dynamic[i] if s is None else s for i, s in enumerate(static)]

predictions = []
def online_generate(model, tokenizer, src_dev, task_prefix_a):
    for sample in tqdm.tqdm(src_dev):
        sample =  task_prefix_a + sample
        batch = tokenizer([sample], return_tensors="tf")
        output = model.generate(input_ids = batch['input_ids'], attention_mask = batch['attention_mask'] , decoder_start_token_id=tokenizer.pad_token_id)
        custom = tokenizer.decode(output[0], skip_special_tokens=True)
        predictions.append(custom)


def native_model():
     
    input_ids = tokenizer([task_prefix_a + "We are going to the market."], return_tensors="tf").input_ids
    encoded_sequence = encoder(input_ids)
    next_decoder_input_ids = 1 #Initializing to the bos token 
    decoder_input_ids = tokenizer("<pad>", add_special_tokens=False, return_tensors="tf").input_ids
    outputs = model(input_ids, decoder_input_ids=tf.convert_to_tensor(decoder_input_ids), return_dict=True)  
    lm_logits = torch.from_numpy(outputs.logits.numpy())
    next_decoder_input_ids = torch.argmax(lm_logits[:, -1:], axis=-1)
    decoder_input_ids = torch.cat([torch.from_numpy(decoder_input_ids.numpy()), next_decoder_input_ids], axis=-1)
  
    while next_decoder_input_ids!=tokenizer.eos_token_id:
        lm_logits = torch.from_numpy(model(None, encoder_outputs=encoded_sequence, decoder_input_ids=tf.convert_to_tensor(decoder_input_ids), return_dict=True).logits.numpy())
        
        next_decoder_input_ids = torch.argmax(lm_logits[:, -1:], axis=-1)
        
        decoder_input_ids = torch.cat([torch.from_numpy(decoder_input_ids.numpy()), next_decoder_input_ids], axis=-1)
    print(f"Generated by manually generic Hugging Face .generate() pipeline: {tokenizer.decode(decoder_input_ids[0], skip_special_tokens=True)}")

def native_decoder():
        
    input_ids = tokenizer([task_prefix_a + sample], return_tensors="tf").input_ids
    encoded_sequence = encoder(input_ids)
    next_decoder_input_ids = -1 # Initializing to the bos token 
    decoder_input_ids = tokenizer("<pad>", add_special_tokens=False, return_tensors="tf").input_ids
    lm_head = model.lm_head

    while next_decoder_input_ids!=tokenizer.eos_token_id:
        decoder_outputs = decoder(tf.convert_to_tensor(decoder_input_ids),encoder_hidden_states=encoded_sequence[0])
        lm_logits = model.lm_head(decoder_outputs[0].numpy()) 
     
        next_decoder_input_ids = torch.argmax(torch.from_numpy(lm_logits[:, -1:].numpy()), axis=-1)
        decoder_input_ids = torch.cat([torch.from_numpy(decoder_input_ids.numpy()), next_decoder_input_ids], axis=-1)
    print(f"Generated by manually decoding: {tokenizer.decode(decoder_input_ids[0], skip_special_tokens=True)}")

def offline_inference(encoder_interpreter_path, decoder_interpreter_path, ENCODER_MAX_LEN, DECODER_MAX_LEN, src_dev, tokenizer, task_prefix):
    encoder_interpreter = tf.lite.Interpreter(model_path = encoder_interpreter_path)
    encoder_input_details = encoder_interpreter.get_input_details()
    encoder_output_details = encoder_interpreter.get_output_details()
    print(encoder_input_details)

    encoder_interpreter.allocate_tensors()  

    decoder_interpreter = tf.lite.Interpreter(model_path = decoder_interpreter_path)
    decoder_input_details = decoder_interpreter.get_input_details()
    decoder_output_details = decoder_interpreter.get_output_details()
    print(decoder_output_details)

    decoder_interpreter.allocate_tensors()
     
    for sample in tqdm.tqdm(src_dev):
        batch = tokenizer(task_prefix + sample, return_tensors = 'tf',  truncation = True, padding='max_length', max_length = ENCODER_MAX_LEN)
        
        input_ids = batch['input_ids']
        attention_mask = batch['attention_mask']
          
        encoder_interpreter.set_tensor(encoder_input_details[0]['index'],input_ids) # attention mask is the shape of the input - but all ones 
        encoder_interpreter.set_tensor(encoder_input_details[1]['index'],attention_mask) # Have to ideally shift this by one - by adding the POS token to the array as the first token 
        
        encoder_interpreter.invoke()
        
        encoder_outputs = encoder_interpreter.get_tensor(encoder_output_details[0]['index']) 
    
        initial = ['<pad>']
           
        decoder_input_ids = tokenizer.encode(initial, add_special_tokens=True, return_tensors="tf", truncation = True, padding='max_length', max_length = ENCODER_MAX_LEN)
        decoder_interpreter.set_tensor(decoder_input_details[0]['index'],decoder_input_ids)  
        decoder_interpreter.set_tensor(decoder_input_details[1]['index'],encoder_outputs) 

        decoder_interpreter.invoke()
        decoder_outputs = decoder_interpreter.get_tensor(decoder_output_details[0]['index']) 
        

        decoder_interpreter.set_tensor(decoder_input_details[1]['index'],encoder_outputs) 
        decoder_input_ids = tokenizer.encode(initial, add_special_tokens=True, return_tensors="tf", truncation = True, padding='max_length', max_length = ENCODER_MAX_LEN)
        next_decoder_input_ids = 0 #Initializing to the bos token 
        
        cache = []
        while True:
            # stime = time.time()
            decoder_input_ids = decoder_input_ids.numpy().astype('int32') # If this is an eager tensor 
            decoder_interpreter.set_tensor(decoder_input_details[0]['index'],decoder_input_ids)  
            decoder_interpreter.invoke()  
            lm_logits = decoder_interpreter.get_tensor(decoder_output_details[0]['index'])     
            next_decoder_input_ids = torch.argmax(torch.from_numpy(lm_logits[:, -1:]), axis=-1)
            
            cache.append(next_decoder_input_ids)
            decoder_input_ids = np.array([decoder_input_ids[0][1:]])  # 1: is to leave space for the previous token that would be concatenated to the decoder_input_ids             
            decoder_input_ids = torch.cat([torch.from_numpy(decoder_input_ids), next_decoder_input_ids], axis=-1)
            
            if len(cache) > 2 and next_decoder_input_ids==tokenizer.eos_token_id:
                break
            if len(cache) > ENCODER_MAX_LEN:
                break
        print(tokenizer.decode(decoder_input_ids[0], skip_special_tokens=True))
        predictions.append(tokenizer.decode(decoder_input_ids[0], skip_special_tokens=True))

if __name__ == '__main__':
    
    parser = argparse.ArgumentParser()
    parser.add_argument("--benchmark_path", type=str, default = './test.txt')
    parser.add_argument("--model_path", type=str, default = None)
    parser.add_argument("--vocab_path", type=str, default = None)
    parser.add_argument("--mode", type=str, default = "online")
    parser.add_argument("--task_prefix", type = str)
    parser.add_argument("--encoder_interpreter_path", type=str)
    parser.add_argument("--src", type=str)
    parser.add_argument("--decoder_interpreter_path", type=str)
    parser.add_argument("--encoder_max_len", type=str, default = 28)
    parser.add_argument("--decoder_max_len", type=str, default = 28)
        
    args = parser.parse_args()
    tokenizer = T5Tokenizer.from_pretrained("google/mt5-small")
    src_dev = io.open(args.src, encoding='UTF-8').read().strip().split('\n')[:10]

    if args.mode == "online":       
        custom_tf_model = TFMT5ForConditionalGeneration.from_pretrained(args.model_path, from_pt = True)
        custom_tf_model.save_pretrained(args.model_path)
        model = TFMT5ForConditionalGeneration.from_pretrained(args.model_path)
        decoder = model.decoder 
        encoder = model.encoder
        predictions = online_generate(model, tokenizer, src_dev, args.task_prefix)

    elif args.mode == "offline":
        predictions = offline_inference(encoder_interpreter_path = args.encoder_interpreter_path, decoder_interpreter_path=args.decoder_interpreter_path, ENCODER_MAX_LEN=args.encoder_max_len, DECODER_MAX_LEN=args.decoder_max_len, src_dev = src_dev, tokenizer=tokenizer, task_prefix = args.task_prefix)
    
    with open(args.benchmark_path, 'w+', encoding='UTF-8' ) as file:
        for pred in predictions:
            file.write(pred)
            file.write('\n')
