# Essential libraries
import torch
import gradio as gr
import os
# load and split data in train & test
from sklearn.model_selection import train_test_split
import pandas as pd

# Create Dataset
import os
import pickle
import pandas as pd
from torch.utils.data import Dataset

# pad sequences
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader

# Libraries for find checkpoint
import re
import glob
import shutil

# animation for iteration
from tqdm.notebook import tqdm, trange

# Parameters 
import config as cfg

from transformers import (
    WEIGHTS_NAME,
    AdamW,
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    get_linear_schedule_with_warmup,
)

import logging
# Instance logger
logger = logging.getLogger(__name__)
# Setup logging
logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
    datefmt="%m/%d/%Y %H:%M:%S",
    level=logging.INFO
)

import requests
import json
import hashlib

app_id = "20230324001613137"
app_key = "pFJkfuVq673YMAWsXNRo"

# Define function for Translate API call
def translate(text, from_lang, to_lang):
    url = "https://fanyi-api.baidu.com/api/trans/vip/translate"
    salt = "1234567890"
    sign = app_id + text + salt + app_key
    sign = hashlib.md5(sign.encode()).hexdigest()
    payload = {"q": text, "from": from_lang, "to": to_lang, "appid": app_id, "salt": salt, "sign": sign}
    response = requests.get(url, params=payload)
    result = json.loads(response.text)
    if "error_code" in result:
        raise Exception(result["error_msg"])
    return result["trans_result"][0]["dst"]

cfg.output_dir = "Bot_Rick"
tokenizer = AutoTokenizer.from_pretrained(cfg.output_dir)
model = AutoModelForCausalLM.from_pretrained(cfg.output_dir)

# Define chatbot function with translation
def chatbot(text):
    # Translate input text to English
    text_translated = translate(text, "it", "en")
    
    # Generate response in English
    new_user_input_ids = tokenizer.encode(text_translated + tokenizer.eos_token, return_tensors='pt')
    bot_input_ids = new_user_input_ids
    chat_history_ids = model.generate(
        bot_input_ids, max_length=200,
        pad_token_id=tokenizer.eos_token_id,
        no_repeat_ngram_size=3,
    )
    response_translated = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
    
    # Translate response back to Chinese
    response = translate(response_translated, "en", "it")
    
    return response
   
title = "chatbot"
description = "A chatbot by gpt2 "

gr.Interface(
    fn=chatbot, 
    inputs="text", 
    outputs="text",
    title = title,
    examples =[ 
    "Allora perche' non lo facciamo?"
    ],

    description="A chatbot by the  Rick and Morty series datasets ",
    theme="default",
    
    #allow_screenshot=False,
    allow_flagging="never",
).launch(debug=False)