import os
import sys
sys.path.insert(0, os.path.abspath('./'))

import torch
from tqdm.auto import tqdm
from torch.utils.data import DataLoader, random_split
from transformers import AutoTokenizer, AutoModelForMaskedLM
from event_detection_dataset import *
from event_detection_model import *

import gradio as gr
#print(f"Gradio version: {gr.__version__}")


def predict(data):
    data=[data]
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    #print(f"Device {device}")    
    
    
    """Load Tokenizer"""  
    tokenizer = AutoTokenizer.from_pretrained('distilbert-base-cased', use_fast=True)
    
    
    """Tokenized Inputs"""
    tokenized_inputs = tokenizer(
        data,                          
        add_special_tokens=True,
        max_length=512,
        padding='max_length',
        return_token_type_ids=True,
        truncation=True,
        is_split_into_words=True
    )
    
    
    """Load Model"""
    model_path = "./"
    #print("model_path:", model_path)

    #print("================ load model ===========================")
    model = DistillBERTClass('distilbert-base-cased')
    
    #print("================  model init ===========================")
    pretrained_model=torch.load(model_path + "/event_domain_final.pt",map_location=torch.device('cpu'))
    model.load_state_dict(pretrained_model['model_state_dict'])
    model.to(device)

    
    """Make Prediction"""
    model.eval()
    
    ids = torch.tensor(tokenized_inputs['input_ids']).to(device)
    mask = torch.tensor(tokenized_inputs['attention_mask']).to(device)

    with torch.no_grad():
        outputs = model(ids, mask)

    max_val, max_idx = torch.max(outputs.data, dim=1)
    
    #print("=============== inference result =================")
    #print(f"predicted class {max_idx}")
    id2tags={0: "Acquisition",1: "I-Positive Clinical Trial & FDA Approval",2: "Dividend Cut",3: "Dividend Increase",4: "Guidance Increase",5: "New Contract",6: "Dividend",7: "Reverse Stock Split",8: "Special Dividend ",9: "Stock Repurchase",10: "Stock Split",11: "Others"}
    return id2tags[max_idx.item()]


title="Financial Event Detection"
description="Predict Finacial Events."
article="modified the model in the following paper: Zhou, Z., Ma, L., & Liu, H. (2021)."
example_list=[["Investors who receive dividends can choose to take them as cash or as additional shares."]]

# Create the Gradio demo
demo = gr.Interface(fn=predict, # mapping function from input to output
                    inputs="text", # what are the inputs?
                    outputs="text", # our fn has two outputs, therefore we have two outputs
                    examples=example_list, 
                    title=title,
                    description=description,
                    article=article)

# Launch the demo!
demo.launch(debug=False, share=True)