from asyncio.base_tasks import _task_print_stack
import math
import argparse
import tqdm 
from tensorflow.nn import log_softmax
from tensorflow.nn import softmax
from transformers import PreTrainedTokenizerFast, TFMarianMTModel, MarianTokenizer
from transformers import TFMT5ForConditionalGeneration, T5Tokenizer
from scipy.stats import entropy
import torch
import numpy as np
import tensorflow as tf 
import matplotlib.pyplot as plt

def evaluate(model, encoder, tokenizer, task_prefix, samples): 
  std_word_level_log_probs = []
  avg_softmax_entropy = []
  top1prob, top2prob, top3prob = [], [], []
  for sample in tqdm.tqdm(samples): 
    cache = []
    input_ids = tokenizer([task_prefix + sample], return_tensors="tf").input_ids
    encoded_sequence = encoder(input_ids)

    top1, top2, top3 = [], [], []
    softmax_entropy = []
    log_softmax_outputs = [] 
    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(torch.from_numpy(softmax(lm_logits[:, -1:], axis=-1).numpy()), axis=-1)
    cache.append(next_decoder_input_ids)
    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())
        x = softmax(lm_logits[:, -1:], axis=-1)
        y = entropy(x, axis = -1)
        z = log_softmax(lm_logits[:, -1:], axis=-1)
        
        top_k_decoder_probs, _ = torch.topk(torch.from_numpy(x.numpy()), 3, axis=-1)
        top1.append(top_k_decoder_probs.numpy()[0][0][0]) #Probability of the likliest token
        top2.append(top_k_decoder_probs.numpy()[0][0][1]) #Probability of the next likliest token
        top3.append(top_k_decoder_probs.numpy()[0][0][2]) # .. 
        
        softmax_entropy.append(y)
        log_softmax_outputs.append(z)
        next_decoder_input_ids = torch.argmax(torch.from_numpy(x.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 following generic Hugging Face .generate() pipeline: {tokenizer.decode(decoder_input_ids[0], skip_special_tokens=False)}")
    # print(np.average(top1[1:-1]))
    # print(np.average(top2[1:-1]))
    # print(np.average(top3[1:-1]))
    top1prob.append(np.average(top1[1:-1]))
    top2prob.append(np.average(top2[1:-1]))
    top3prob.append(np.average(top3[1:-1]))

    word_level_log_probs = []
    for inp in log_softmax_outputs:
      idx = torch.argmax(torch.from_numpy(inp.numpy()))
      word_level_log_probs.append(inp[:, -1:].numpy()[0][0][idx])
    # print(np.std(word_level_log_probs))
    std_word_level_log_probs.append(np.std(word_level_log_probs))

    # print(np.average(softmax_entropy))
    avg_softmax_entropy.append(np.average(softmax_entropy))
    print(np.average(softmax_entropy))

  return [std_word_level_log_probs, top1prob, top2prob, top3prob, avg_softmax_entropy]


def load_model(arch, ENTIRE_MODEL_SAVE_PATH, vocab_path):
  if 'mt5' in arch: 
    tokenizer = T5Tokenizer.from_pretrained("google/mt5-small")
    custom_tf_model = TFMT5ForConditionalGeneration.from_pretrained(ENTIRE_MODEL_SAVE_PATH, from_pt=True)
    custom_tf_model.save_pretrained(ENTIRE_MODEL_SAVE_PATH)
    model = TFMT5ForConditionalGeneration.from_pretrained(ENTIRE_MODEL_SAVE_PATH)
    encoder = model.encoder
  else: 
    tokenizer = PreTrainedTokenizerFast(tokenizer_file=args.vocab_path, bos_token = "<s>", eos_token = "</s>", pad_token = "<pad>", unk_token = "<unk>")
    custom_tf_model = TFMarianMTModel.from_pretrained(args.model_path, from_pt=True)
    custom_tf_model.save_pretrained(args.model_path)
    model =  TFMarianMTModel.from_pretrained(pretrained_model_name_or_path = args.model_path, from_pt = True)    
    encoder = model.model.encoder
        
    # tokenizer = PreTrainedTokenizerFast(tokenizer_file=vocab_path, bos_token = "<s>", eos_token = "</s>", pad_token = "<pad>", unk_token = "<unk>")
    # model = TFMarianMTModel.from_pretrained(ENTIRE_MODEL_SAVE_PATH, from_pt = True) 
    # encoder = model.model.encoder
    print('Model Loaded!')
  return model, encoder, tokenizer

if __name__ == '__main__':

  parser = argparse.ArgumentParser()
  parser.add_argument("--src_lang", type=str, default = None)
  parser.add_argument("--tgt_lang", type=str, default = None )
  parser.add_argument("--benchmark_path", type=str, default = None)
  parser.add_argument("--model_arch", type=str, default = 'marian')
  parser.add_argument("--model_path", type=str, default = None)
  parser.add_argument("--vocab_path", type=str, default = None)
  parser.add_argument("--task_prefix", type = str, default = "")
  parser.add_argument("--callibration_file", type=str)
  parser.add_argument("--callibration_samples", type=int, default = 10)

       
  args = parser.parse_args()

  model, encoder, tokenizer = load_model(args.model_arch, args.model_path, args.vocab_path)
  with open(args.callibration_file, 'r') as file: 
    samples = file.read().split('\n')[:args.callibration_samples]

  print(f'Computing for {args.src_lang}-{args.tgt_lang}')
  confidence_metrics = evaluate(model, encoder, tokenizer, args.task_prefix, samples)

  if args.benchmark_path is None:
    benchmark_path = f'{args.src_lang}_{args.tgt_lang}_{args.model_arch}_confidence_scores.txt'
  else: 
    benchmark_path = args.benchmark_path
  with open(benchmark_path, 'w') as file: 
    for val in confidence_metrics:
      file.write(str(val) + '\n')
