import streamlit as st
import tensorflow as tf
from transformers import pipeline
import os

# Define las rutas absolutas a tus modelos
MODEL_TRANSFORMER_ENCODER_PATH = 'stacked_transformer_encoder.keras'
MODEL_TRANSFORMER_POS_PATH = 'transformer_encoder_pos.keras'
MODEL_SIMPLE_RNN_PATH = 'model_simple_rnn.keras'
MODEL_LSTM_PATH = 'model_lstm.keras'

# Cargar tus modelos de clasificación forzando la ruta absoluta
try:
    st.info(f"Intentando cargar modelo desde: {MODEL_TRANSFORMER_ENCODER_PATH}")
    model_transformer_encoder = tf.keras.models.load_model(MODEL_TRANSFORMER_ENCODER_PATH)

    st.info(f"Intentando cargar modelo desde: {MODEL_TRANSFORMER_POS_PATH}")
    model_transformer_positional_encoding = tf.keras.models.load_model(MODEL_TRANSFORMER_POS_PATH)

    st.info(f"Intentando cargar modelo desde: {MODEL_SIMPLE_RNN_PATH}")
    model_simple_rnn = tf.keras.models.load_model(MODEL_SIMPLE_RNN_PATH)

    st.info(f"Intentando cargar modelo desde: {MODEL_LSTM_PATH}")
    model_lstm = tf.keras.models.load_model(MODEL_LSTM_PATH)

    st.info("Modelos de clasificación cargados correctamente.")

except Exception as e:
    st.error(f"Error al cargar los modelos de clasificación: {e}")
    model_transformer_encoder = None
    model_transformer_positional_encoding = None
    model_simple_rnn = None
    model_lstm = None

# Cargar el pipeline de traducción de Hugging Face (inglés a español)
try:
    hf_token = os.environ.get("HF_TOKEN")
    modelo_traduccion_en_es = pipeline("translation", model="Helsinki-NLP/opus-mt-en-es",token = hf_token)
    st.info("Modelo de traducción (inglés a español) cargado correctamente con token.")
except Exception as e:
    st.error(f"Error al cargar el modelo de traducción (inglés a español): {e}")
    translator_en_es = None

def clasificar_noticia(texto, modelo_seleccionado):
    """Función para clasificar el texto usando el modelo seleccionado."""
    modelo_a_usar = None
    if modelo_seleccionado == "Transformer Encoder Apilado" and model_transformer_encoder is not None:
        modelo_a_usar = model_transformer_encoder
    elif modelo_seleccionado == "Transformer Positional Encoding" and model_transformer_positional_encoding is not None:
        modelo_a_usar = model_transformer_positional_encoding
    elif modelo_seleccionado == "Simple RNN" and model_simple_rnn is not None:
        modelo_a_usar = model_simple_rnn
    elif modelo_seleccionado == "LSTM" and model_lstm is not None:
        modelo_a_usar = model_lstm

    if modelo_a_usar:
        # *** IMPORTANTE: Implementa aquí la lógica de preprocesamiento del texto ***
        # para que coincida con la entrada esperada por el modelo.
        # Esto incluye tokenización, padding, etc.
        prediction = modelo_a_usar.predict([texto])
        # *** IMPORTANTE: Interpreta la salida del modelo para obtener la clase predicha. ***
        # Esto dependerá de cómo entrenaste tu modelo.
        return f"Clase predicha ({modelo_seleccionado}): {prediction}"
    else:
        return f"El modelo '{modelo_seleccionado}' no está disponible."

def traducir_texto_en_es(texto_en):
    """Función para traducir texto de inglés a español."""
    if translator_en_es:
        result = translator_en_es(texto_en)[0]['translation_text']
        return result
    else:
        return "El modelo de traducción (inglés a español) no está disponible."

def main():
    st.title("Aplicación de Clasificación y Traducción")

    tab1, tab2 = st.tabs(["Clasificación de Noticias", "Traducción (Inglés a Español)"])

    with tab1:
        st.header("Clasificación de Noticias")
        input_texto_clasificacion = st.text_area("Ingresa la noticia aquí:")
        modelos = ["Transformer Encoder Apilado", "Transformer Positional Encoding", "Simple RNN", "LSTM"]
        modelo_seleccion = st.selectbox("Selecciona el modelo:", modelos)
        boton_clasificar = st.button("Clasificar")
        if boton_clasificar:
            if input_texto_clasificacion:
                resultado_clasificacion = clasificar_noticia(input_texto_clasificacion, modelo_seleccion)
                st.subheader("Resultado de la clasificación:")
                st.write(resultado_clasificacion)
            else:
                st.warning("Por favor, ingresa el texto de la noticia.")

    with tab2:
        st.header("Traducción de Inglés a Español")
        input_texto_traduccion = st.text_area("Ingresa el texto en inglés:")
        boton_traducir = st.button("Traducir")
        if boton_traducir:
            if input_texto_traduccion:
                texto_traducido = traducir_texto_en_es(input_texto_traduccion)
                st.subheader("Texto traducido al español:")
                st.write(texto_traducido)
            else:
                st.warning("Por favor, ingresa el texto en inglés.")

if __name__ == "__main__":
    main()