import os
import math
import gradio as gr
import json
import logging
from datasets import load_dataset, concatenate_datasets #list_datasets, load_from_disk
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    DataCollatorForLanguageModeling,
    TrainingArguments,
    Trainer,
    pipeline,
    AutoModelForQuestionAnswering, 
)
import evaluate # type: ignore
from huggingface_hub import login
from translate import Translator
from datasets import load_dataset, DownloadConfig



# Modelo base
MODEL_KEY = "EleutherAI/gpt-neo-125M"
tokenizer = AutoTokenizer.from_pretrained(MODEL_KEY)
model = AutoModelForCausalLM.from_pretrained(MODEL_KEY)
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Cargar modelos para Preguntas y Respuestas (QA)
qa_model = AutoModelForQuestionAnswering.from_pretrained("deepset/roberta-base-squad2")
qa_tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")
qa_pipeline = pipeline("question-answering", model=qa_model, tokenizer=qa_tokenizer)

# Agregar más datasets para mejorar las respuestas
context_map = {
    "imdb": "Dom: Cine | Estilo: Opinión",
    "daily_dialog": "Dom: Conversación | Estilo: Diálogo diario",
    "go_emotions": "Dom: Emociones | Estilo: Clasificación emocional",
    "wikitext": "Dom: Enciclopedia | Estilo: Conocimiento general",
    "math": "Dom: Matemáticas | Estilo: Problema matemático",  # Agregar problemas de matemáticas
    "empathetic_dialogues": "Dom: Psicología | Estilo: Apoyo emocional",  # Para el comportamiento emocional
    "math_advanced": "Dom: Matemáticas Avanzadas | Estilo: Trigonometría y Cálculo",  # Dataset de matemáticas avanzadas
    "programming": "Dom: Programación | Estilo: Resolución de código",  # Dataset para programación
    "social_sciences": "Dom: Ciencias Sociales | Estilo: Preguntas sobre historia y geografía",  # Dataset de sociales
    "economics": "Dom: Economía | Estilo: Problemas económicos y teorías",  # Dataset de economía
    "health": "Dom: Salud | Estilo: Información médica y de salud",  # Dataset de salud
}

# Detectar idioma automáticamente
def detect_language(text):
    try:
        return detect(text)  # Detecta el idioma de la entrada
    except:
        return "en"  # Si no se puede detectar, se asume inglés

# Generación de texto (si se solicita un cuento, por ejemplo)
def generate_text(prompt, max_length=100):
    output = generator(prompt, max_length=int(max_length), num_return_sequences=1)[0]["generated_text"]
    return output

# Preguntas y respuestas basadas en contexto
def answer_question(question, context):
    result = qa_pipeline(question=question, context=context)
    return result['answer']

# Traducción
def translate_text(text, lang):
    translator = Translator(to_lang=lang)
    try:
        return translator.translate(text)
    except Exception as e:
        return f"Error: {str(e)}"

# Selección del contexto y modelo según la solicitud
def process_input(user_input):
    # Detectar idioma
    detected_lang = detect_language(user_input)
    print(f"Idioma detectado: {detected_lang}")

    # Si la entrada es una pregunta, usaremos un modelo de Preguntas y Respuestas
    if '?' in user_input:
        context = "Este es un contexto general. Puedo responder preguntas específicas sobre cine, ciencia, o emociones."
        
        # Identificar tipo de pregunta (matemáticas, programación, etc.)
        if "matemática" in user_input.lower() or "cálculo" in user_input.lower() or "trigonometría" in user_input.lower():
            context = "Dom: Matemáticas Avanzadas | Estilo: Trigonometría y Cálculo. Ejemplo: ¿Cuál es la derivada de x^2?"
        elif "programación" in user_input.lower():
            context = "Dom: Programación | Estilo: Resolución de código. Ejemplo: ¿Cómo se resuelve un bucle en Python?"
        elif "historia" in user_input.lower() or "geografía" in user_input.lower():
            context = "Dom: Ciencias Sociales | Estilo: Preguntas sobre historia y geografía. Ejemplo: ¿Quién fue Sócrates?"
        elif "economía" in user_input.lower():
            context = "Dom: Economía | Estilo: Problemas económicos y teorías. Ejemplo: ¿Qué es la inflación?"
        elif "salud" in user_input.lower():
            context = "Dom: Salud | Estilo: Información médica y de salud. Ejemplo: ¿Qué es la diabetes?"
        answer = answer_question(user_input, context)
        return answer
    elif "math" in user_input.lower():
        # Si se detecta una pregunta matemática
        return "Resolviendo la operación matemática..."
    elif "cuento" in user_input.lower():
        # Si el usuario solicita un cuento
        prompt = "Había una vez, en un reino lejano..."
        return generate_text(prompt, max_length=200)
    elif "emoción" in user_input.lower():
        # Si es una solicitud de apoyo emocional
        context = "Dom: Psicología | Estilo: Apoyo emocional. ¿Cómo te sientes hoy?"
        return generate_text(context + " Estoy aquí para apoyarte", max_length=100)
    else:
        return "No entiendo la solicitud, por favor intenta preguntar algo más específico."

# Interfaz de Gradio
with gr.Blocks() as demo:
    gr.Markdown("# 🧠 MultiDomain Text Generator + Translator")

    with gr.Tab("Generar respuestas y contar historias"):
        user_input = gr.Textbox(label="Tu pregunta o solicitud", placeholder="Haz una pregunta o pide un cuento...")
        output_text = gr.Textbox(label="Respuesta generada")
        btn_generate = gr.Button("Generar respuesta o cuento")
        btn_generate.click(process_input, inputs=user_input, outputs=output_text)

    with gr.Tab("Traducir texto"):
        input_text = gr.Textbox(label="Texto a traducir")
        lang = gr.Textbox(label="Código de idioma destino", value="en")
        output_translation = gr.Textbox(label="Texto traducido")
        btn_translate = gr.Button("Traducir")
        btn_translate.click(translate_text, inputs=[input_text, lang], outputs=output_translation)

demo.launch()