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)

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

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

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

with DemoTab:
    with st.form(key="diagtrast_form"):
        sintomas = st.text_input(label = 'Introduce texto:',
                                 value = 'Soy el más guapo de mi clase')

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

    with st.form(key="examples"):
        ex_1 = st.form_submit_button(label="Yo también")

    if ex_1:
        sintomas.value = ex_1

    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))
            
            df = pd.DataFrame({
                'Modelo': ["DiagTrast-Berto"],
                'Diagnóstico': [pred_berto[0]['label']]
            })

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

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

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)
        
        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](https://huggingface.co/edmenciab)
        - [Miguel Ángel Solís Orozco](https://huggingface.co/homosapienssapiens)
        - [Jose Carlos Vílchez Villegas](https://huggingface.co/JCarlos)
    """)