import streamlit as st
import pandas as pd
import utils
import time

from transformers import pipeline
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification

#####################

model_berto='hackathon-somos-nlp-2023/DiagTrast-Berto'
tokenizer_berto = AutoTokenizer.from_pretrained(model_berto)
classifier_berto = pipeline("text-classification", model=model_berto)

model_xml='hackathon-somos-nlp-2023/DiagTrast-xlm-roberta-base'
tokenizer_xml = AutoTokenizer.from_pretrained(model_xml)
classifier_xml = pipeline("text-classification", model=model_xml)

#####################

st.title('🏥 Diagnóstico de Trastornos Mentales')

DemoTab, AboutTab = st.tabs(["Demo", "Acerca de"])

with AboutTab:
    st.subheader("Motivación")
    st.markdown(
        "[Colocar aquí la motivación]."
    )

    st.subheader("Recursos")
    st.markdown("""
        Modelos usados:
        - [hackathon-somos-nlp-2023/DiagTrast-Berto](https://huggingface.co/hackathon-somos-nlp-2023/DiagTrast-Berto)
        - [hackathon-somos-nlp-2023/DiagTrast-xlm-roberta-base](https://huggingface.co/hackathon-somos-nlp-2023/DiagTrast-xlm-roberta-base)
        
        Dataset:
        - [hackathon-somos-nlp-2023/DiagTrast](https://huggingface.co/datasets/hackathon-somos-nlp-2023/DiagTrast)
    """)

    st.subheader("Equipo")
    st.markdown("""
        - [Alberto Martín Garrido](https://huggingface.co/Stremie)
        - [Edgar Mencia]()
        - [Miguel Ángel Solís Orozco](https://huggingface.co/homosapienssapiens)
        - [Jose Carlos Vílchez Villegas](https://huggingface.co/JCarlos)
    """)

with DemoTab:
    with st.form(key="diagtrast_form"):
        sintomas = st.text_input(label = 'Introduce síntomas:',
                         value = 'La gente cuando me ve por la calle me mira y halagan mi belleza')

        submit_button = st.form_submit_button(label="Clasificar")

    if submit_button and not sintomas:
        st.warning("⚠️ Debe introducir los síntomas.")
        
    elif submit_button:  
        with st.spinner('Clasificando...'):
            pred_berto = classifier_berto.predict(utils.clean_text(sintomas))
            pred_xml = classifier_xml.predict(utils.clean_text(sintomas))
            
            df = pd.DataFrame({
                'Modelo': ["DiagTrast-Berto", "DiagTrast-xlm-roberta-base"],
                'Diagnóstico': [pred_berto[0]['label'], pred_xml[0]['label']],
                'Score': [f"{pred_berto[0]['score']:.2%}", f"{pred_xml[0]['score']:.2%}"]
            })

        st.markdown("### Resultados:")
        st.caption("")

        st.dataframe(df, use_container_width=True)
        st.caption("")
        alert = st.success("✅ ¡Hecho!")