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 Trastornos Mentales')

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

with DemoTab:
    with st.form(key="diagtrast_form"):
        sintomas = st.text_input(label = 'Introduce texto:',
                                 value = 'El paciente piensa que es la persona más bella, y se enfada cuando los demás no lo ven así.')

        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))
            
            df = pd.DataFrame({
                'Texto': [(sintomas[:50] + '...') if len(sintomas) > 50 else sintomas],
                'Diagnóstico': [pred_berto[0]['label']]
            })

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

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

    st.markdown("##### Ejemplos")
    st.markdown("Se muestra impasivo emocionalmente.")
    st.markdown("Irresponsable en su trabajo, suele saltarse las normas. No le importa la opinión de los demás.")
    st.markdown("El paciente piensa que es la persona más bella, y se enfada cuando los demás no lo ven así.")
    st.markdown("Él siempre se siente incómodo cuando no es el centro de atención. Es una persona muy necesitada de reconocimiento y se siente ansioso cuando no es el foco de atención de los demás. A menudo busca formas de atraer la atención de los demás y siente que su autoestima depende de ello.")
    st.markdown("El paciente tiene problemas con el alcohol. Normalmente toma decisiones importantes sin pensarlo profundamente. Tiene una idea pesimista de su persona y acude a sus familiares para sentirse mejor. No tiene la capacidad de controlar sus sentimientos, la mayoría de las veces los reprime.")

with ConclusionTab:
    st.subheader("Conclusiones")
    st.markdown("El presente proyecto muestra una herramienta que facilita al profesional la tarea de diagnosticar a pacientes con trastornos mentales. Aunque el proyecto se encuentra en la fase de prototipado, demuestra que los modelos de aprendizaje profundo basados en el lenguaje ayudan a identificar trastornos mentales con precisión, facilitando la tarea a los profesionales.")

    st.subheader("Trabajo futuro")
    st.markdown("- El modelo actual no tiene en cuenta la ausencia de trastornos mentales, ya que no se incluyó en el dataset dicha categoría.")
    st.markdown("- Incluir todo el conjunto de trastornos mentales del manual DSM-5.")
    st.markdown("- Implementación de un modelo de pregunta/respuesta donde el modelo realizará una pregunta al profesional en caso de no tener claro el diagnóstico a partir del texto inicial. A partir del texto inicial y las futuras respuestas, el modelo realizará un diagnóstico con mayor certeza.")
    st.markdown("- En el caso de implementar el último punto, sustituir la arquitectura por un LLM, de forma que el modelo pueda manejar una mayor cantidad de información con mayor precisión.")
    
with AboutTab:
    st.subheader("Motivación")
    st.markdown(
        "Actualmente el proceso de diagnóstico de enfermedades mentales enfrenta retos importantes de subjetividad que podría llevar a un diagnóstico erróneo en un paciente. Uno de los documentos más avalados como refuerzo de diagnóstico es el DSM-5. Este conjunto de guías de diagnóstico han procedido a ser fundamentales en casos de pacientes difíciles de identificar. Sin embargo, sumergirse en las más de 500 hojas del DSM-5 puede llegar a ser abrumador. El objetivo de este proyecto ha sido tener un modelo que, por medio del lenguaje natural, los especialistas de la salud mental puedan describir el caso de un un paciente en concreto, dando así una sugerencia de diagnóstico para facilitar y concretar de manera más exacta un diagnóstico de salud mental."
    )

    st.subheader("Recursos")
    st.markdown("""
        Modelo:
        - [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("ODS")
    st.markdown("El presente proyecto se engloba en el objetivo número 3, Salud y bienestar, de los Objetivos de Desarrollo Sostenible de la ONU.")
    
    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)
    """)