import streamlit as st
import streamlit.components.v1 as components
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForSeq2SeqLM
from catboost import CatBoostRegressor
import pandas as pd
device = 'cpu'

model = CatBoostRegressor()
model.load_model("./model.cbm")

@st.cache(allow_output_mutation=True)
def language_model():
    language_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ru-en")
    language_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-ru-en")
    return language_model, language_tokenizer

def prepare_data(text):
    text = text.lower()
    text = re.sub(r"[,.;#?!&$]+\ *", " ", text)
    return text

def preapare_language(text, language_tokenizer, language_model):
    tokenized_text = language_tokenizer.prepare_seq2seq_batch([text], return_tensors='pt')
    translation = language_model.generate(**tokenized_text)
    text = language_tokenizer.batch_decode(translation, skip_special_tokens=False)[0][5:-4]
    return text
    
def inference_sweets():
    HtmlFile = open("index.html", 'r', encoding='utf-8')
    source_code = HtmlFile.read()
    print(source_code)
    components.html(source_code)

inference_sweets()
st.write('Select a language/ default is English')
rus_buton = st.button('Russian')
if rus_buton:
    ingredients = st.text_area("Write the ingredients of the dish")
    recipe = st.text_area("Write the recipe")
    language_model, language_tokenizer = language_model()
    ingredients = preapare_language(ingredients, language_tokenizer, language_model)
    recipe = preapare_language(recipe, language_tokenizer, language_model)
    servings = int(st.text_area("Write number of servings"))
    data = pd.DataFrame({'TranslatedInstructions':[recipe],'TranslatedIngredients':[ingredients],'Servings': [servings]})
    result_button = st.button('Click')
    if result_button:
        time = model.predict(data)[0]
        st.write('The dish will be prepared')
        st.write(int(time))
else:
    ingredients = st.text_area("Write the ingredients of the dish")
    recipe = st.text_area("Write the recipe")
    servings = st.text_area("Write number of servings")
    if len(servings):
        data = pd.DataFrame({'TranslatedInstructions':[recipe],'TranslatedIngredients':[ingredients],'Servings': [int(servings)]})
        result_button = st.button('Click')
        if result_button:
            time = model.predict(data)[0]
            st.write('The dish will be prepared')
            st.write(int(time))