import gradio as gr
from datasets import load_dataset
import requests

# Diccionario de tareas por modalidad
MODALIDAD_TAREAS = {
    "NLP": ["Clasificación de texto", "Generación de texto", "Traducción", "Resumen"],
    "Audio": ["Reconocimiento de voz", "Clasificación de audio", "Separación de fuentes"],
    "Vision": ["Clasificación de imágenes", "Detección de objetos", "Segmentación"],
    "RAG": ["Respuesta a preguntas", "Recuperación de información"],
    "Code": ["Generación de código", "Completado de código"],
    "Tabular": ["Clasificación tabular", "Regresión tabular"],
    "Time Series": ["Predicción de series temporales"]
}

# Función para actualizar las tareas según las modalidades seleccionadas
def actualizar_tareas(modalidades_seleccionadas):
    tareas_disponibles = []
    for modalidad in modalidades_seleccionadas:
        if modalidad in MODALIDAD_TAREAS:
            tareas_disponibles.extend(MODALIDAD_TAREAS[modalidad])
    # Eliminar duplicados y devolver un diccionario con valores predeterminados (no seleccionados)
    return {tarea: False for tarea in set(tareas_disponibles)}

# Función para generar el encabezado del CSV (ahora incluye 'id')
def generar_encabezado(modalidades_seleccionadas, tareas_seleccionadas):
    if not modalidades_seleccionadas or not tareas_seleccionadas:
        raise ValueError("Debes seleccionar al menos una modalidad y una tarea.")
    columnas = ["id"]  # Añadimos 'id' como primer elemento
    for modalidad in modalidades_seleccionadas:
        columnas.append(f"{modalidad.lower()}_input")
        columnas.append(f"{modalidad.lower()}_output")
    for tarea in tareas_seleccionadas:
        columnas.append(f"{tarea.lower().replace(' ', '_')}_label")
    return ",".join(columnas)

# Función para buscar datasets compatibles en HuggingFace
def buscar_datasets(modalidades_seleccionadas, tareas_seleccionadas, filtro_tamaño=None, filtro_licencia=None):
    query = "+".join(modalidades_seleccionadas + tareas_seleccionadas)
    url = f"https://huggingface.co/api/datasets?search={query}"
    response = requests.get(url)
    datasets = response.json()
    
    resultados = []
    for dataset in datasets:
        # Aplicar filtros adicionales
        if filtro_tamaño and dataset.get("size_categories") != filtro_tamaño:
            continue
        if filtro_licencia and dataset.get("license") != filtro_licencia:
            continue
        resultados.append(f"- {dataset['id']}: {dataset['description']}")
    return "\n".join(resultados)

# Función para generar el dataset
def generar_dataset(encabezado, datasets_seleccionados, pagina_actual=1, filas_por_pagina=5):
    if not datasets_seleccionados:
        raise ValueError("Debes seleccionar al menos un dataset.")
    
    columnas = encabezado.split(",")
    filas = []
    
    # Cargar datos reales desde los datasets seleccionados
    for dataset_id in datasets_seleccionados.split("\n"):
        dataset_id = dataset_id.strip("- ").split(":")[0]  # Extraer ID del dataset
        try:
            dataset = load_dataset(dataset_id, split="train")
            inicio = (pagina_actual - 1) * filas_por_pagina
            fin = pagina_actual * filas_por_pagina
            for i, fila in enumerate(dataset[inicio:fin]):
                valores = [str(fila.get(col, "valor_default")) for col in columnas[1:]]  # Ignorar 'id'
                filas.append(f"id_{inicio + i}," + ",".join(valores))
        except Exception as e:
            filas.append(f"Error cargando dataset {dataset_id}: {str(e)}")
    
    contenido_csv = "\n".join([encabezado] + filas)
    return contenido_csv

# Interfaz Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Diseñador de Modelos de Redes Neuronales y Generador de Datasets Multimodales")

    # Panel 1: Selección de Modalidades
    with gr.Row():
        modalidades = list(MODALIDAD_TAREAS.keys())
        modalidades_seleccionadas = gr.CheckboxGroup(choices=modalidades, label="Selecciona Modalidades")
    
    # Panel 2: Configuración de Tareas (Refactorizado)
    with gr.Row():
        gr.Markdown("### Selecciona Tareas")
        tareas_disponibles = gr.CheckboxGroup(choices=[], label="Tareas Disponibles", interactive=True, elem_classes=["expanded-checkbox"])
        modalidades_seleccionadas.change(
            actualizar_tareas,
            inputs=modalidades_seleccionadas,
            outputs=tareas_disponibles
        )

    # Panel 3: Generar Encabezado del CSV
    with gr.Row():
        btn_generar_encabezado = gr.Button("Generar Encabezado del CSV")
        encabezado_generado = gr.Textbox(label="Encabezado Generado", interactive=False)
        btn_generar_encabezado.click(generar_encabezado, inputs=[modalidades_seleccionadas, tareas_disponibles], outputs=encabezado_generado)

    # Panel 4: Buscar Datasets Compatibles con Filtros Adicionales
    with gr.Row():
        with gr.Column():
            filtro_tamaño = gr.Dropdown(choices=["small", "medium", "large"], label="Filtrar por Tamaño")
            filtro_licencia = gr.Dropdown(choices=["mit", "apache-2.0", "cc-by-4.0"], label="Filtrar por Licencia")
        btn_buscar_datasets = gr.Button("Buscar Datasets Compatibles")
        datasets_encontrados = gr.Textbox(label="Datasets Encontrados", interactive=False)
        btn_buscar_datasets.click(
            buscar_datasets,
            inputs=[modalidades_seleccionadas, tareas_disponibles, filtro_tamaño, filtro_licencia],
            outputs=datasets_encontrados
        )

    # Panel 5: Generar Dataset con Paginación
    with gr.Row():
        pagina_actual = gr.Number(value=1, label="Página Actual")
        filas_por_pagina = gr.Number(value=5, label="Filas por Página")
        btn_generar_dataset = gr.Button("Generar Dataset")
        vista_previa_csv = gr.Textbox(label="Vista Previa del CSV", interactive=False)
        btn_descargar_csv = gr.File(label="Descargar CSV")
        btn_generar_dataset.click(
            generar_dataset,
            inputs=[encabezado_generado, datasets_encontrados, pagina_actual, filas_por_pagina],
            outputs=vista_previa_csv
        )

# Estilo CSS para expandir el área de selección de tareas
demo.css = """
.expanded-checkbox {
    max-height: 200px; /* Altura máxima */
    overflow-y: auto;  /* Scroll vertical si es necesario */
    border: 1px solid #ccc; /* Borde para mejorar la visibilidad */
    padding: 10px; /* Espaciado interno */
}
"""

# Ejecutar la aplicación
demo.launch()