import gradio as gr
print('Imported gradio', gr.__version__)
import transformers
print('Imported sentence_transformers', transformers.__version__)
import peft
print('Imported torch', peft.__version__)
import torch
print('Imported torch', torch.__version__)
import re
import gc

intermediate_loras = [
    'PTHQL_level1_Germanic',
    'PTHQL_level1_Romance',
    'PTHQL_level2_High_German',
    'PTHQL_level2_North_Sea_Germanic',
    'PTHQL_level2_Weser_Rhine_Germanic',
    'PTHQL_level2_Gallo_Romance',
    'PTHQL_level2_Iberian_Romance',
    'PTHQL_level2_Italo_Romance'
]

language_mapping = {
    'Asturian (ast_Latn)': ['PTHQL_level0_Indo_European', 'PTHQL_level1_Romance', 'PTHQL_level2_Iberian_Romance', 'PTHQL_language_Asturian'],
    'Dutch (nld_Latn)': ['PTHQL_level0_Indo_European', 'PTHQL_level1_Germanic', 'PTHQL_level2_Weser_Rhine_Germanic', 'PTHQL_language_Dutch'],
    'English (eng_Latn)': ['PTHQL_level0_Indo_European', 'PTHQL_level1_Germanic', 'PTHQL_level2_North_Sea_Germanic', 'PTHQL_language_English'],
    'French (fra_Latn)': ['PTHQL_level0_Indo_European', 'PTHQL_level1_Romance', 'PTHQL_level2_Gallo_Romance', 'PTHQL_language_French'],
    'German (deu_Latn)': ['PTHQL_level0_Indo_European', 'PTHQL_level1_Germanic', 'PTHQL_level2_High_German', 'PTHQL_language_German'],
    'Haitian Creole (hat_Latn)': ['PTHQL_level0_Indo_European', 'PTHQL_level1_Romance', 'PTHQL_level2_Gallo_Romance', 'PTHQL_language_Haitian_Creole'],
    'Italian (ita_Latn)': ['PTHQL_level0_Indo_European', 'PTHQL_level1_Romance', 'PTHQL_level2_Italo_Romance', 'PTHQL_language_Italian'],
    'Limburgish (lim_Latn)': ['PTHQL_level0_Indo_European', 'PTHQL_level1_Germanic', 'PTHQL_level2_Weser_Rhine_Germanic', 'PTHQL_language_Limburgish'],
    'Luxembourgish (ltz_Latn)': ['PTHQL_level0_Indo_European', 'PTHQL_level1_Germanic', 'PTHQL_level2_High_German', 'PTHQL_language_Luxembourgish'],
    'Sicilian (scn_Latn)': ['PTHQL_level0_Indo_European', 'PTHQL_level1_Romance', 'PTHQL_level2_Italo_Romance', 'PTHQL_language_Sicilian'],
    'Spanish (spa_Latn)': ['PTHQL_level0_Indo_European', 'PTHQL_level1_Romance', 'PTHQL_level2_Iberian_Romance', 'PTHQL_language_Spanish'],
    'Tok Pisin (tpi_Latn)': ['PTHQL_level0_Indo_European', 'PTHQL_level1_Germanic', 'PTHQL_level2_North_Sea_Germanic', 'PTHQL_language_Tok_Pisin'],
}

last_language = 'English (eng_Latn)'

print('Loading base model...')
model = transformers.MT5ForConditionalGeneration.from_pretrained('google/mt5-large', torch_dtype=torch.bfloat16)
tokenizer = transformers.AutoTokenizer.from_pretrained('google/mt5-large')
print('Base model loaded!')

print('Loading base LoRA...')
model = peft.PeftModel.from_pretrained(model, 'WilliamSotoM/PTHQL_level0_Indo_European', 'PTHQL_level0_Indo_European')
print('Base LoRA loaded!')

print('Loading Intermediate LoRAs...')
for adapter in intermediate_loras:
    model.load_adapter(f'WilliamSotoM/{adapter}', adapter)
    print(f"{adapter} loaded!")
print('Intermediate LoRAs loaded!')

print('Loading English (eng_Latn) LoRA...')
model.load_adapter('WilliamSotoM/PTHQL_language_English', 'language')
print('English (eng_Latn) LoRA loaded!')

print('Merging English (eng_Latn) related LoRAs...')
for adapter in language_mapping['English (eng_Latn)'][:-1]:
    model.merge_adapter([adapter])
    print(f'{adapter} merged!')
model.merge_adapter(['language'])
print('English (eng_Latn) related LoRAs merged!')
gc.collect()

print('Defining evaluate function...')
def evaluate(language, amr_graph):
    global last_language
    global model
    
    if language != last_language:
        print('Unmerging LoRAs...')
        model.unmerge_adapter()
        print('LoRAs unmerged')
        
        print('Removing old language LoRA...')
        model.delete_adapter('language')
        gc.collect()
        print('Old language LoRA removed!')
        
        print(f'Loading {language} LoRA...')
        language_lora = language_mapping[language][-1]
        model.load_adapter(f'WilliamSotoM/{language_lora}', 'language')
        print(f'{language}LoRA loaded!')

        print(f'Merging {language} related LoRAs...')
        for adapter in language_mapping[language][:-1]:
            model.merge_adapter([adapter])
            print(f'{adapter} merged!')
        model.merge_adapter(['language'])
        print(f'{language} related LoRAs merged!')
        last_language = language
    
    tokenized_input = tokenizer(amr_graph, return_tensors='pt')
    
    with torch.inference_mode():
        prediction = model.generate(**tokenized_input, max_length=128)
        generated_text = tokenizer.batch_decode(prediction, skip_special_tokens=True)[0]
    
    print(f'AMR Graph:\n{amr_graph}')
    print('-----')
    print(f'Generated Text:\n{generated_text}')
    print('=====')
    return generated_text
print('Evaluate function defined!')

print('Instantiating gradio interface...')
demo = gr.Interface(
    fn=evaluate,
    inputs = [
        gr.Dropdown(label = 'Language', choices=list(language_mapping.keys()), value='English (eng_Latn)'),
        gr.Textbox(label='AMR Graph', lines=10),
    ],
    outputs = [
        gr.Textbox(
            label="Lexicalization",
            interactive=False,
            show_copy_button=True
        )
    ],
    title = 'Multilingual AMR-to-Text via PTHQL',
    description = '''Select a language and write the input AMR Graph to obtain a Lexicalization.

The first Generation after changing language might take longer while the last LoRA is changed.'''
)
print('Gradio interface instantiated...')

print('Launching server...')
demo.launch()