import gradio as gr
from time import time

from bertviz import model_view, head_view
from bertviz_gradio import head_view_mod

import faiss
import torch
import os
# import nltk
import argparse
import random
import numpy as np
import pandas as pd

from argparse import Namespace
from tqdm.notebook import tqdm
from torch.utils.data import DataLoader
from functools import partial

from transformers import AutoTokenizer, MarianTokenizer, AutoModel, AutoModelForSeq2SeqLM, MarianMTModel

model_es = "Helsinki-NLP/opus-mt-fr-es"
model_en = "Helsinki-NLP/opus-mt-fr-en"
model_de = "Helsinki-NLP/opus-mt-fr-de"
# model_ar = "Helsinki-NLP/opus-mt-fr-ar"

tokenizer_es = AutoTokenizer.from_pretrained(model_es)
tokenizer_en = AutoTokenizer.from_pretrained(model_en)
tokenizer_de = AutoTokenizer.from_pretrained(model_de)
# tokenizer_ar = AutoTokenizer.from_pretrained(model_ar)

model_tr_es = MarianMTModel.from_pretrained(model_es)
model_tr_en = MarianMTModel.from_pretrained(model_en)
model_tr_de = MarianMTModel.from_pretrained(model_de)
# model_tr_ar = MarianMTModel.from_pretrained(model_ar)

from faiss import write_index, read_index
import pickle 

def load_index(model): 
	with open('index/'+ model + '_metadata_ref.pkl', 'rb') as f:
	# with open('local_index/'+ model + '_metadata_ref.pkl', 'rb') as f:
		loaded_dict = pickle.load(f)
	for type in ['tokens','words']:
		for kind in ['input', 'output']: 
			## save index file
			name = 'index/'+ model + "_" + kind + "_"+ type + ".index" 
			loaded_dict[kind][type][1] = read_index(name)
			# write_index(metadata_all[kind][type][1], name)
	return loaded_dict


dict_models = {
	'fr-es': model_es,
	'fr-en': model_en,
	'fr-de': model_de,
	# 'en-ar': model_ar,
}

dict_models_tr = {
	'fr-es': model_tr_es,
	'fr-en': model_tr_en,
	'fr-de': model_tr_de,
	# 'en-ar': model_tr_ar,
}

dict_tokenizer_tr = {
	'fr-es': tokenizer_es,
	'fr-en': tokenizer_en,
	'fr-de': tokenizer_de,
	# 'en-ar': tokenizer_ar,
}
# dict_reference_faiss = {'en-es':[]}
dict_reference_faiss = {
	'fr-es': [],  #load_index('en-es'),
	'fr-en': [], #load_index('en-ar'),
	'fr-de': [], #load_index('en-fr'),
	'fr-zh': [], #load_index('en-zh'),
}

# print("dict", dict_reference_faiss['en-es']['input']['tokens'][1])

saliency_examples = [
	"Peace of Mind: Protection for consumers.",
	"The sustainable development goals report: towards a rescue plan for people and planet",
	"We will leave no stone unturned to hold those responsible to account.",
	"The clock is now ticking on our work to finalise the remaining key legislative proposals presented by this Commission to ensure that citizens and businesses can reap the benefits of our policy actions.",
	"Pumpkins, squash and gourds, fresh or chilled, excluding courgettes",
	"The labour market participation of mothers with infants has even deteriorated over the past two decades, often impacting their career and incomes for years.",
]

contrastive_examples = [
["Peace of Mind: Protection for consumers.",
"Paz mental: protección de los consumidores",
"Paz de la mente: protección de los consumidores"],
["the slaughterer has finished his work.",
"l'abatteur a terminé son travail.",
"l'abatteuse a terminé son travail."],
['A fundamental shift is needed - in commitment, solidarity, financing and action - to put the world on a better path.',
 '需要在承诺、团结、筹资和行动方面进行根本转变,使世界走上更美好的道路。',
 '我们需要从根本上转变承诺、团结、资助和行动，使世界走上更美好的道路。',]
	]


french_examples = [
    ["","","Les appels répétés de sa mère [auraient] dû nous alerter.", "", ""],
    ["","","Elle a promis à son médecin de demeurer [active] après s’être retirée.", "", ""],
    ["","","Nous [avons] lancé une insulte et nous [sommes] partis brusquement.", "",""],
    ["","","La vache et la poule [doivent] être [nourries].","",""],
    ["","","Le bruit soudain dans les chambres supérieures [aurait] dû nous alerter.","",""],
    ["","","le bailleur de fonds a terminé son travail.","",""],
    ["","","la bailleuse de fonds a terminé son travail.","",""],
    ["","","l'esthéticien a terminé son travail.","",""],
    ["","","l'esthéticienne a terminé son travail.","",""],
    ["","","l'esthéticienne a terminé son travail.","",""],
    ["","","chaque ingénieure a terminé son travail.","",""],
    ["","","chaque ingénieur a fini son travail.","",""],
    

    
    
]


#Load challenge set examples
# df_challenge_set = pd.read_csv("challenge_sets.csv")
# arr_challenge_set = df_challenge_set.values
# arr_challenge_set = [[x[2], x[3], x[4], x[5]] for x in arr_challenge_set]
arr_challenge_set = [[x[2]] for x in french_examples]




def get_k_prob_tokens(transition_scores, result, model, k_values=5):
	tokenizer_tr = dict_tokenizer_tr[model]
	gen_sequences = result.sequences[:, 1:]

	result_output = []

	# First beam only... 
	bs = 0
	text = ' '
	for tok, score, i_step in zip(gen_sequences[bs], transition_scores[bs],range(len(gen_sequences[bs]))):

		beam_i = result.beam_indices[0][i_step]
		if beam_i < 0:
			beam_i = bs
		bs_alt = [tokenizer_tr.decode(tok) for tok in result.scores[i_step][beam_i].topk(k_values).indices ]
		bs_alt_scores = np.exp(result.scores[i_step][beam_i].topk(k_values).values)
		result_output.append([np.array(result.scores[i_step][beam_i].topk(k_values).indices), np.array(bs_alt_scores),bs_alt])

	return result_output


def split_token_from_sequences(sequences, model) -> dict :
	n_sentences = len(sequences)

	gen_sequences_texts = []
	for bs in range(n_sentences): 
		# gen_sequences_texts.append(dict_tokenizer_tr[model].decode(sequences[:, 1:][bs],  skip_special_tokens=True).split(' '))
		#### decoder per token.
		seq_bs = []
		
		for token in sequences[:, 1:][bs]:
			seq_bs.append(dict_tokenizer_tr[model].decode(token,  skip_special_tokens=True))
		gen_sequences_texts.append(seq_bs)

	score = 0
	#raw dict is bos
	text = 'bos'
	new_id = text +'--1'
	dict_parent = [{'id': new_id, 'parentId': None , 'text': text, 'name': 'bos', 'prob': score }]
	id_dict_pos = {}
	step_i = 0
	cont = True
	words_by_step = [] #[['bos' for i in range(n_sentences)]]

	while cont: 
		# append to dict_parent for all beams of step_i
		cont = False
		step_words = []
		for beam in range(n_sentences):
			app_text = '<empty_word>'
			if step_i < len(gen_sequences_texts[beam]): 
				app_text = gen_sequences_texts[beam][step_i]
				cont = True
			step_words.append(app_text)
		words_by_step.append(step_words)
		print(words_by_step)

		for i_bs, step_w in enumerate(step_words):
			if not step_w in ['<empty_word>', '<pad>']:
				#new id if the same word is not in another beam (?) [beam[i] was a token id]
				#parent id = previous word and previous step.

				
				# new_parent_id = "-".join([str(beam[i]) for i in range(step_i)])
				
				new_id = "-".join([str(words_by_step[i][i_bs])+ '-' + str(i) for i in range(step_i+1)])
				parent_id = "-".join([words_by_step[i][i_bs] + '-' + str(i) for i in range(step_i) ])
				
				# new_id = step_w +'-' + str(step_i)
				# parent_id = words_by_step[step_i-1][i_bs] + '-' + str(step_i -1)
				next_word_flag = 1
				if step_i == 0 : 
					parent_id =  'bos--1'
				## if the dict already exists remove it, if it is not a root... 
				## root?? then next is ''
				else:
					next_word_flag = len(gen_sequences_texts[i_bs][step_i]) > step_i ## Not in step_i = 0;
				if next_word_flag:
					if not (new_id in id_dict_pos):
						dict_parent.append({'id': new_id, 'parentId': parent_id , 'text': step_w, 'name': step_w, 'prob' : score })
						id_dict_pos[new_id] = len(dict_parent) - 1
				else: 
					if not (new_id in id_dict_pos):
						dict_parent.append({'id': new_id, 'parentId': parent_id , 'text': step_w, 'name': step_w, 'prob' : score  })
						id_dict_pos[new_id] = len(dict_parent) - 1


		step_i += 1
	return dict_parent


## Tokenization 
def compute_tokenization(inputs, targets, w1, model): 
	colors = ['tok-first-color', 'tok-second-color', 'tok-third-color', 'tok-fourth-color']
	len_colors = len(colors);
	inputs = inputs.input_ids
	html_tokens = ""
	i = 0 
	for sentence in inputs: 
		html_tokens += "<p>"
		# print("TOKENS", inputs, targets)
		# print("input", [dict_tokenizer_tr[model].decode(tok) for tok in sentence])
		tokens = [dict_tokenizer_tr[model].decode(tok) for tok in sentence]
		for token in tokens:
			token = token.replace("<", "<'") # .substring(0, token.length - 2) 
			html_tokens += "<span class='" + colors[i % len_colors] + "'>" + token + " </span>"
			i +=1
	html_tokens += "</p>"
	i = 0 
	# for tgt_sentence in targets : 
	html_tokens_tgt = ""
	html_tokens_tgt += "<p>"
	# print("targets", [dict_tokenizer_tr[model].decode(tok) for tok in targets])
	# print("targets", dict_tokenizer_tr[model].decode(targets))
	tokens = [dict_tokenizer_tr[model].decode(tok) for tok in targets]
	for token in tokens:
		token = token.replace("<", "<'") # .substring(0, token.length - 2) 
		html_tokens_tgt += "<span class='" + colors[i % len_colors] + "'>" + token + " </span>"
		i +=1
	html_tokens_tgt += "</p>"
	# print("HTML", html_tokens, html_tokens_tgt)
	return html_tokens, html_tokens_tgt


def create_vocab_multiple(embeddings_list, model): 
	"""_summary_

	Args:
		embeddings_list (list): embedding array 

	Returns:
		Dict: vocabulary of tokens' embeddings
	"""
	print("START VOCAB CREATION MULTIPLE \n \n ")
	vocab = {} ## add embedds. 
	sentence_tokens_text_list = []
	for embeddings in embeddings_list: 
		tokens_id = embeddings['tokens'] # [[tokens_id]x n_sentences ]
		for sent_i, sentence in enumerate(tokens_id):
			sentence_tokens = []
			for tok_i, token in enumerate(sentence): 
				sentence_tokens.append(token)
				if not (token in vocab):
					vocab[token] = {
						'token' : token,
						'count': 1, 
						# 'text': embeddings['texts'][sent_i][tok_i],
						'text': dict_tokenizer_tr[model].decode([token]),
						# 'text': src_token_lists[sent_i][tok_i], 
						'embed': embeddings['embeddings'][sent_i][tok_i]}
				else: 
					vocab[token]['count'] = vocab[token]['count'] + 1  
		# print(vocab)
			sentence_tokens_text_list.append(sentence_tokens)
	print("END VOCAB CREATION MULTIPLE \n \n ")
	return vocab, sentence_tokens_text_list

def vocab_words_all_prefix(token_embeddings, model, sufix="@@",prefix = '▁' ):
	vocab = {} 
	# inf_model = dict_models_tr[model]
	sentence_words_text_list = []
	if prefix : 
		n_prefix = len(prefix)
		for input_sentences in token_embeddings: 
			# n_tokens_in_word 
			for sent_i, sentence in enumerate(input_sentences['tokens']):
				words_text_list = []
				# embedding = input_sentences['embed'][sent_i]
				word = '' 
				tokens_ids = []
				embeddings = []
				ids_to_tokens = dict_tokenizer_tr[model].convert_ids_to_tokens(sentence)
				# print("validate same len", len(sentence) == len(ids_to_tokens), len(sentence), len(ids_to_tokens), ids_to_tokens)

				to_save= False
				for tok_i, token_text in enumerate(ids_to_tokens): 
					token_id = sentence[tok_i]
					if token_text[:n_prefix] == prefix : 
						#first we save the previous word 
						if to_save: 
							vocab[word] = {
									'word' : word,
									'text': word,
									'count': 1, 
									'tokens_ids' : tokens_ids, 
									'embed': np.mean(np.array(embeddings), 0).tolist()
								}
							words_text_list.append(word)
						#word is starting if prefix
						tokens_ids = [token_id]
						embeddings = [input_sentences['embeddings'][sent_i][tok_i]]
						word = token_text[n_prefix:]
						## if word 
						to_save = True 
						
					else : 
						if (token_text in dict_tokenizer_tr[model].special_tokens_map.values()):
							# print('final or save', token_text, token_id, to_save, word)
							if to_save: 
								# vocab[word] = ids
								vocab[word] = {
									'word' : word,
									'text': word,
									'count': 1, 
									'tokens_ids' : tokens_ids, 
									'embed': np.mean(np.array(embeddings), 0).tolist()
								}
								words_text_list.append(word)
							#special token is one token element, no continuation 
							# vocab[token_text] = [token_id]
							tokens_ids = [token_id]
							embeddings = [input_sentences['embeddings'][sent_i][tok_i]]
							vocab[token_text] = {
									'word' : token_text,
									'count': 1, 
									'text': word,
									'tokens_ids' : tokens_ids, 
									'embed': np.mean(np.array(embeddings), 0).tolist()
								}
							words_text_list.append(token_text)
							to_save = False
						else: 
							# is a continuation; we do not know if it is final; we don't save here.
							to_save = True 
							word += token_text 
							tokens_ids.append(token_id)
							embeddings.append(input_sentences['embeddings'][sent_i][tok_i])
				if to_save: 
					# print('final save', token_text, token_id, to_save, word)
					vocab[word] = tokens_ids
					if not (word in vocab):
						vocab[word] = {
							'word' : word,
							'count': 1, 
							'text': word,
							'tokens_ids' : tokens_ids, 
							'embed': np.mean(np.array(embeddings), 0).tolist()
							}
						words_text_list.append(word)
					else: 
						vocab[word]['count'] = vocab[word]['count'] + 1 
				sentence_words_text_list.append(words_text_list)

	return vocab, sentence_words_text_list
def search_query_vocab(index, vocab_queries,  topk = 10, limited_search = []):
	""" the embed queries are a vocabulary of words : embds_input_voc

	Args:
		index (_type_): faiss index
		embed_queries (_type_): vocab format.
			{   'token' : token,
				'count': 1, 
				'text': src_token_lists[sent_i][tok_i], 
				'embed': embeddings[0]['embeddings'][sent_i][tok_i] }
		nb_ids (_type_): hash to find the token_id w.r.t the faiss index id. 
		topk (int, optional): nb of similar tokens. Defaults to 10.

	Returns:
		_type_: Distance matrix D, indices matrix I and tokens ids (using nb_ids)
	"""
	# nb_qi_ids = [] ##ordered ids list
	nb_q_embds = [] ##ordered embeddings list
	metadata = {}
	qi_pos = 0
	for key , token_values in vocab_queries.items():
		# nb_qi_ids.append(token_values['token']) # for x in vocab_tokens]
		metadata[qi_pos] = {'word': token_values['word'], 'tokens': token_values['tokens_ids'], 'text': token_values['text']}
		qi_pos += 1
		nb_q_embds.append(token_values['embed']) # for x in vocab_tokens]
	
	xq = np.array(nb_q_embds).astype('float32') #elements to query

	D,I = index.search(xq, topk)

	return D,I, metadata

def search_query_vocab_token(index, vocab_queries,  topk = 10, limited_search = []):
	""" the embed queries are a vocabulary of words : embds_input_vov
	Returns:
		_type_: Distance matrix D, indices matrix I and tokens ids (using nb_ids)
	"""
	# nb_qi_ids = [] ##ordered ids list
	nb_q_embds = [] ##ordered embeddings list
	metadata = {}
	qi_pos = 0
	for key , token_values in vocab_queries.items():
		# nb_qi_ids.append(token_values['token']) # for x in vocab_tokens]
		metadata[qi_pos] = {'token': token_values['token'], 'text': token_values['text']}
		qi_pos += 1
		nb_q_embds.append(token_values['embed']) # for x in vocab_tokens]
	
	xq = np.array(nb_q_embds).astype('float32') #elements to query

	D,I = index.search(xq, topk)

	return D,I, metadata	


def build_search(query_embeddings, model,type="input"):
	metadata_all = dict_reference_faiss[model]

	# ## biuld vocab for index 
	vocab_queries, sentence_tokens_list = create_vocab_multiple(query_embeddings, model)
	words_vocab_queries, sentence_words_list = vocab_words_all_prefix(query_embeddings, model, sufix="@@",prefix="▁")
	
	index_vor_tokens = metadata_all[type]['tokens'][1]
	md_tokens = metadata_all[type]['tokens'][2]
	D, I, meta = search_query_vocab_token(index_vor_tokens, vocab_queries)

	qi_pos = 0 
	similar_tokens = {}
	# similar_tokens = []
	for dist, ind in zip(D,I):
		try: 
			# similar_tokens.append({
			similar_tokens[str(meta[qi_pos]['token'])] = {
				'token': meta[qi_pos]['token'], 
				'text': meta[qi_pos]['text'], 
				# 'text': dict_tokenizer_tr[model].decode(meta[qi_pos]['token'])
				# 'text': meta[qi_pos]['text'], 
				"similar_topk": [md_tokens[i_index]['token'] for i_index in ind if (i_index != -1) ], 
				"distance": [dist[i] for (i, i_index) in enumerate(ind) if (i_index != -1)], 
				}
			# )
		except: 
			print("\n ERROR ", qi_pos, dist, ind)
		qi_pos += 1


	index_vor_words = metadata_all[type]['words'][1]
	md_words = metadata_all[type]['words'][2]

	Dw, Iw, metaw = search_query_vocab(index_vor_words, words_vocab_queries)
	# D, I, meta, vocab_words, sentence_words_list = result_input['words']# [2] # D ; I ; meta
	qi_pos = 0 
	# similar_words = []
	similar_words = {}
	for dist, ind in zip(Dw,Iw):
		try: 
			# similar_words.append({
			similar_words[str(metaw[qi_pos]['word']) ] = {
				'word': metaw[qi_pos]['word'], 
				'text': metaw[qi_pos]['word'], 
				"similar_topk": [md_words[i_index]['word'] for i_index in ind if (i_index != -1) ], 
				"distance": [dist[i] for (i, i_index) in enumerate(ind) if (i_index != -1)], 
				}
			# )
		except: 
			print("\n ERROR ", qi_pos, dist, ind)
		qi_pos += 1


	return {'tokens': {'D': D, 'I': I, 'meta': meta, 'vocab_queries': vocab_queries, 'similar':similar_tokens, 'sentence_key_list': sentence_tokens_list}, 
			'words': {'D':Dw,'I': Iw, 'meta': metaw, 'vocab_queries':words_vocab_queries, 'sentence_key_list': sentence_words_list, 'similar': similar_words}
			}
from sklearn.manifold import TSNE
def embds_input_projection_vocab(vocab, key="token"): 
	t0 = time()
	
	nb_ids = [] ##ordered ids list
	nb_embds = [] ##ordered embeddings list
	nb_text = [] ##ordered embeddings list
	tnse_error = []
	for _ , token_values in vocab.items():
		tnse_error.append([0,0])
		nb_ids.append(token_values[key]) # for x in vocab_tokens]
		nb_text.append(token_values['text']) # for x in vocab_tokens]
		nb_embds.append(token_values['embed']) # for x in vocab_tokens]

	X = np.array(nb_embds).astype('float32') #elements to project 
	try:
		tsne = TSNE(random_state=0, n_iter=1000)
		tsne_results = tsne.fit_transform(X)
	
		tsne_results = np.c_[tsne_results, nb_ids, nb_text, range(len(nb_ids))] ## creates a zip array : [[TNSE[X,Y], tokenid, token_text], ...]
	except: 
		tsne_results = np.c_[tnse_error, nb_ids, nb_text, range(len(nb_ids))] ## creates a zip array : [[TNSE[X,Y], tokenid, token_text], ...]

	t1 = time()
	print("t-SNE: %.2g sec" % (t1 - t0))
	# print(tsne_results)
	
	return tsne_results.tolist()



def filtered_projection(similar_key, vocab, model, type="input", key="word"): 
	metadata_all = dict_reference_faiss[model]
	vocab_proj = vocab.copy()
	## tnse projection Input words
	source_words_voc_similar = set()
	# for words_set in similar_key:
	for key_i in similar_key:
		words_set = similar_key[key_i]
		source_words_voc_similar.update(words_set['similar_topk'])

	# print(len(source_words_voc_similar))
	# source_embeddings_filtered = {key:  metadata_all['input']['words'][0][key] for key in source_words_voc_similar}
	source_embeddings_filtered = {key_value:  metadata_all[type][key][0][key_value] for key_value in source_words_voc_similar}
	vocab_proj.update(source_embeddings_filtered)
	## 	vocab_proj add 
	try:
		result_TSNE = embds_input_projection_vocab(vocab_proj, key=key[:-1]) ## singular => without 's'
		dict_projected_embds_all = {str(embds[2]): [embds[0], embds[1], embds[2], embds[3], embds[4]] for embds in result_TSNE}
	except: 
		print('TSNE error', type, key)
		dict_projected_embds_all = {}

	

	# print(result_TSNE)
	return dict_projected_embds_all 

def get_bertvis_data(input_text, lg_model):
	tokenizer_tr = dict_tokenizer_tr[lg_model]
	model_tr = dict_models_tr[lg_model]

	# input_ids = tokenizer_tr(input_text, return_tensors="pt", padding=True)
	input_ids = tokenizer_tr(input_text, return_tensors="pt", padding=False)
	result_att = model_tr.generate(**input_ids,
		num_beams=4,
        num_return_sequences=4,
		return_dict_in_generate=True,
		output_attentions =True,
		output_scores=True,
	)

	# tokenizer_tr.convert_ids_to_tokens(result_att.sequences[0])
	# tokenizer_tr.convert_ids_to_tokens(input_ids.input_ids[0])

	tgt_text = tokenizer_tr.decode(result_att.sequences[0], skip_special_tokens=True)


	outputs = model_tr(input_ids=input_ids.input_ids,
					decoder_input_ids=result_att.sequences[:1],
					output_attentions =True,
					)
	html_attentions = head_view_mod(
		encoder_attention = outputs.encoder_attentions,
		cross_attention = outputs.cross_attentions,
		decoder_attention = outputs.decoder_attentions,
		encoder_tokens = tokenizer_tr.convert_ids_to_tokens(input_ids.input_ids[0]),
		decoder_tokens = tokenizer_tr.convert_ids_to_tokens(result_att.sequences[0]),
		html_action='gradio'
		)
	return html_attentions, tgt_text, result_att, outputs


def translation_model(w1, model):
  #translate and get internal values and visualizations; 
  # src_text = saliency_examples[0]
  inputs = dict_tokenizer_tr[model](w1, return_tensors="pt", padding=True)

  num_ret_seq = 4
  translated  = dict_models_tr[model].generate(**inputs, 
                  num_beams=4,
                  num_return_sequences=num_ret_seq,
                  return_dict_in_generate=True, 
                  output_attentions =True,  
                  output_hidden_states = True, 
                  output_scores=True,)


  beam_dict = split_token_from_sequences(translated.sequences,model )

  tgt_text = dict_tokenizer_tr[model].decode(translated.sequences[0], skip_special_tokens=True)
  ## Attentions
  outputs = dict_models_tr[model](input_ids=inputs.input_ids,
    decoder_input_ids=translated.sequences[:1],
    output_attentions =True,
  )
  encoder_tokens = dict_tokenizer_tr[model].convert_ids_to_tokens(inputs.input_ids[0])
  decoder_tokens = dict_tokenizer_tr[model].convert_ids_to_tokens(translated.sequences[0]) 
#   decoder_tokens = [tok for tok in decoder_tokens if tok != '<pad>']
#   decoder_tokens = [tok for tok in decoder_tokens if tok != '<pad>']

  # html_attentions = head_view_mod(
	# 	encoder_attention = outputs.encoder_attentions,
	# 	cross_attention = outputs.cross_attentions,
	# 	decoder_attention = outputs.decoder_attentions,
	# 	encoder_tokens = encoder_tokens,
	# 	decoder_tokens = decoder_tokens,
	# 	html_action='gradio'
	# 	)
  
  html_attentions_enc = head_view_mod(
		encoder_attention = outputs.encoder_attentions,
		encoder_tokens = encoder_tokens,
		decoder_tokens = decoder_tokens,
		html_action='gradio'
		)
  
  html_attentions_dec = head_view_mod(
		# encoder_attention = outputs.encoder_attentions,
    decoder_attention = outputs.decoder_attentions,
		encoder_tokens = encoder_tokens,
		decoder_tokens = decoder_tokens,
		html_action='gradio'
		)
  
  html_attentions_cross = head_view_mod(
    cross_attention = outputs.cross_attentions,
    encoder_tokens = encoder_tokens,
    decoder_tokens = decoder_tokens,
    html_action='gradio'
    )

  # tokenization 
  html_in, html_out = compute_tokenization(inputs, translated.sequences[0],w1, model)

  transition_scores = dict_models_tr[model].compute_transition_scores(
	translated.sequences, translated.scores, translated.beam_indices  , normalize_logits=True
	)
  prob_tokens = get_k_prob_tokens(transition_scores, translated, model, k_values=10)

  input_embeddings = dict_models_tr[model].get_encoder().embed_tokens(inputs.input_ids)
  target_embeddings = dict_models_tr[model].get_decoder().embed_tokens(translated.sequences)


  return [tgt_text,
          [beam_dict,prob_tokens, html_in, html_out, translated, inputs.input_ids,input_embeddings,target_embeddings],
          [html_attentions_enc['params'], html_attentions_enc['html2'].data], 
          [html_attentions_dec['params'], html_attentions_dec['html2'].data], 
          [html_attentions_cross['params'], html_attentions_cross['html2'].data] ]



html = """
<html>
<script async src="https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.6/require.min.js"></script>

  <style>
  .tok-first-color {
    background: #e0ffcd;
  }

  .tok-second-color {
    background: #fdffcd;
  }

  .tok-third-color {
    background: #ffebbb;
  }

  .tok-fourth-color {
    background: #ffcab0;
  }
</style>    
  <body>
  
    <p id="demo"></p>
    <p id="viz"></p>

    <p id="demo2"></p>
    <h4> Exploring top-k probable tokens </h4>
    <div id="d3_text_grid">... top 10 tokens generated at each step ...</div>

    <h4> Exploring the Beam Search sequence generation</h4>
    <div id="d3_beam_search">... top 4 generated sequences using Beam Search...</div>

  </body>
</html>
"""

html_tok = """
 <div id="d3_tok">... tokenization visualization ...</div> 
 """

html_embd = """
 <div id="d3_embd">... token embeddings visualization ...</div>
	<div id="select_div">
    	<select id="select_type" class="form-select" aria-label="select example" hidden>
      	<option selected value="words">Words</option>
      	<option value="tokens">Tokens</option>
    	</select>
  	</div>
	  <div class="row">
		<div class="col-9"> 
		<div id="d3_graph_input_words" class="d3_graph words"></div>
		</div>
		<div class="col-3">
		<div id="similar_input_words" class=""></div>
		</div>
		</div>
	<div id="d3_graph_input_tokens" class="d3_graph tokens"></div>
	<div id="similar_input_tokens" class=" "></div>

 """

html_tok_target ="""
 <div id="d3_tok_target"> </div> 
 """

html_embd_target= """
 <div id="d3_embd_target"> </div> 
 <div id="d3_graph_output_words" class="d3_graph words"></div>
	<div id="d3_graph_output_tokens" class="d3_graph tokens"></div>
  <div id="similar_output_words" class=""></div>
  <div id="similar_output_tokens" class=" "></div>
 """

html_att_enc =  """
 <div id="d3_att_enc">... Encoder self attention only ... Always read from left to right</div> 
 <div id="bertviz_enc"></div>
 """

html_att_cross =  """
 <div id="d3_att_cross">... Encoder-decoder cross attention only ...</div> 
 """

html_att_dec =  """
 <div id="d3_att_dec">... Decoder self attention only ...</div> 
 """



def sentence_maker2(w1,j2):
   print(w1,j2)
   return "in sentence22..."


def first_function(w1, model):
	global metadata_all
	#translate and get internal values
	sentences = w1.split("\n")
	all_sentences = []
	translated_text = ''
	input_embeddings = []
	output_embeddings = []
	for sentence in sentences :
		# print(sentence, end=";") 
		params = translation_model(sentence, model)
		all_sentences.append(params)
		# print(len(params))
		translated_text +=  params[0] + ' \n'
		# input_embeddings.append({	
		# 	'embeddings': params[1][6].detach(), ## create a vocabulary with the set of embeddings 
		# 	'tokens': params[1][3+2].tolist(), # one translation = one sentence
		# 	# 'texts' : 	dict_tokenizer_tr[model].decode(params[2].tolist())

		# }) 
		# output_embeddings.append({
		# 	'embeddings' : params[1][7].detach(),
		# 	'tokens': params[1][3+1].sequences.tolist(),
		# 	# 'texts' : 	dict_tokenizer_tr[model].decode(params[1].sequences.tolist())
		# })
	
	## load_reference; ERROR 
	## Build FAISS index 
	# ---> preload faiss using the respective model with a initial dataset. 
	## dict_reference_faiss[model] = metadata_all [per language]
	# result_input = build_reference(input_embeddings,model)
	# result_output = build_reference(output_embeddings,model)
	# metadata_all = {'input': result_input, 'output': result_output}

	## Build FAISS index 
	# ---> preload faiss using the respective model with a initial dataset. 

    ### to uncomment gg1 ###
	# result_search = {}
	# result_search['input'] = build_search(input_embeddings, model, type='input')
	# result_search['output'] = build_search(output_embeddings, model, type='output')

	# json_out = {'input': {'tokens': {}, 'words': {}}, 'output': {'tokens': {}, 'words': {}}}
	# dict_projected = {}
	# for type in ['input', 'output']:
	# 	dict_projected[type] = {} 
	# 	for key in ['tokens', 'words']: 
	# 		similar_key = result_search[type][key]['similar']
	# 		vocab = result_search[type][key]['vocab_queries']
	# 		dict_projected[type][key] =  filtered_projection(similar_key, vocab, model, type=type, key=key)
	# 		json_out[type][key]['similar_queries'] = similar_key
	# 		json_out[type][key]['tnse'] = dict_projected[type][key]
	# 		json_out[type][key]['key_text_list'] = result_search[type][key]['sentence_key_list']
    ## to uncomment gg1 ### 
    
	## bertviz
	# paramsbv, tgtbv = get_bertvis_data(w1, model)
	
	# params.append(json_out)
	html_att_enc = params[2][1]#.root_div_id = "bertviz_enc"
	html_att_dec = params[3][1]
	html_att_cross = params[4][1]
	
    ### to uncomment gg1 ###
	# params = [params[0], params[1], json_out, params[2][0], params[3][0], params[4][0]]
    ### to uncomment gg1 
	params = [params[0], params[1], [], params[2][0], params[3][0], params[4][0]]
	### to comment gg1 ###
    
    # params.append([tgt, params['params'], params['html2'].data]

	return [translated_text, params, html_att_enc, html_att_dec, html_att_cross]

def second_function(w1,j2):
	#  json_value = {'one':1}#  return f"{w1['two']} in sentence22..."
	# to transfer the data to json.
	print("second_function -- after the js", w1,j2)
	return "transition to second js function finished."


with gr.Blocks(js="plotsjs.js") as demo:
	gr.Markdown(
	"""
	# MAKE NMT Viz \t `Literacy task` 
	""")

	gr.Markdown(
		"""
		### Translation 
		""")
	
	gr.Markdown(
		"""
		1. Select the language pair for the translation 
		""")
	radio_c = gr.Radio(choices=[['french to english','fr-en'], ['french to spanish','fr-es'], ['french to german', 'fr-de']], value="fr-en", label= ['French to English', "French to Spanish", "French to German"], container=False)
	gr.Markdown(
		"""
		2. Source text to translate
		""")
	in_text = gr.Textbox(label="source text")
	with gr.Accordion("Optional: Challenge selection:", open=False):	
		gr.Markdown(
			"""
			### select an example from the challenge set listed bellow
			""")
		# challenge_ex  = gr.Textbox(label="Challenge", interactive=False)
		# category_minor  = gr.Textbox(label="category_minor", interactive=False)
		# category_major  = gr.Textbox(label="category_major", interactive=False)
		
		with gr.Accordion("Examples:"):
			gr.Examples(arr_challenge_set,[in_text], label="")
			# gr.Examples(arr_challenge_set,[in_text, challenge_ex,category_minor,category_major], label="")

	btn = gr.Button("Translate")
	
	
	with gr.Accordion("3. Review the source tokenization:", open=False):
		input_tokenisation = gr.HTML(html_tok)

	# with gr.Accordion("4. Review similar source tokens in the embedding space:", open=False):
	# 	input_embd= gr.HTML(html_embd)

	with gr.Accordion("4. Review the attention between the source tokens:", open=False):
		gr.Markdown(
			"""
			`Bertviz ` 
			""")
		input_embd= gr.HTML(html_att_enc)
		enc_html = gr.HTML()

	gr.Markdown(
		"""
		### Text is translated into Target Language 
		""")
	out_text  = gr.Textbox(label="target text")

	with gr.Accordion("1. Review the target tokenization:", open=False):
		target_tokenisation = gr.HTML(html_tok_target)

	# with gr.Accordion("2. Review similar target tokens in the embedding space:", open=False):
	# 	target_embd= gr.HTML(html_embd_target)

	with gr.Accordion("2. Review the attention between the target and source tokens:", open=False):
		gr.Markdown(
			"""
			`Bertviz -cross attention` 
			""")
		input_embd= gr.HTML(html_att_cross)
		cross_html = gr.HTML()

	with gr.Accordion("3. Review the attention between the target tokens:", open=False):
		gr.Markdown(
			"""
			`Bertviz -dec attention` 
			""")
		input_embd= gr.HTML(html_att_dec)
		dec_html = gr.HTML()

	with gr.Accordion("4. Review the alternative translations tokens:", open=False):
		gr.Markdown(
			"""
			Generation process : `topk - beam search ` 
			""")
		input_mic = gr.HTML(html)


	out_text2  = gr.Textbox(visible=False)
	var2 = gr.JSON(visible=False)
		

	btn.click(first_function, [in_text, radio_c], [out_text,var2,enc_html, dec_html, cross_html], js="(in_text,radio_c) => testFn_out(in_text,radio_c)") #should return an output comp.
	out_text.change(second_function, [out_text, var2], out_text2, js="(out_text,var2) => testFn_out_json(var2)") #

	# run script function on load,
	# demo.load(None,None,None,js="plotsjs.js")

if __name__ == "__main__":   
    demo.launch()