import gradio as gr
import numpy as np
from PIL import Image

from fastai.vision.all import *
from huggingface_hub import from_pretrained_fastai

import pathlib

# Cargar el modelo preentrenado
repo_id = "ignaciobfp/moonboard_difficulty"
learner = from_pretrained_fastai(repo_id)
labels = learner.dls.vocab

# Función para realizar predicciones
def predict(img):
    pred,pred_idx,probs = learner.predict(img)
    return {labels[i]: float(probs[i]) for i in range(len(labels))}


def row_to_image(presas, ano, inclinacion):
    m, n = 20, 18
    image_matrix = np.zeros((n, m, 3), dtype=np.uint8)

    for presa in presas:
        letra, numero = presa[0], int(presa[1:])
        valor = 1  # Se asume que todas las presas marcadas son blancas
        image_matrix[numero-1, ord(letra) - ord('A')] = [255 * valor, 255 * valor, 255 * valor]

    ano_scaled = (ano - 2016) / (2019 - 2016)
    grey_intensity_ano = int(ano_scaled * (255 - 50) + 50)
    image_matrix[:, -7:-5] = [grey_intensity_ano, grey_intensity_ano, grey_intensity_ano]

    inclinacion_scaled = (inclinacion - 25) / (40 - 25)
    grey_intensity_inclinacion = int(inclinacion_scaled * (100) + 100)
    image_matrix[:, -3:-1] = [grey_intensity_inclinacion, grey_intensity_inclinacion, grey_intensity_inclinacion]

    image_matrix_padded = np.pad(image_matrix, ((3, 3), (2, 2), (0, 0)), mode='constant')
    image_pil = Image.fromarray(image_matrix_padded, 'RGB')

    new_size = (image_pil.width * 3, image_pil.height * 3)
    resized_image_pil = image_pil.resize(new_size, resample=Image.NEAREST)

    return resized_image_pil


def mostrar_parametros(ano, inclinacion, *seleccion):
    # Imprimir en el log la variable seleccion
    print("Seleccion:", seleccion)
    # Asegurarse de que la selección sea una lista plana
    seleccion_plana = [item for sublist in seleccion for item in sublist]
    imagen = row_to_image(seleccion_plana, int(ano), int(inclinacion))
    return (imagen, predict(imagen))

# Generar todas las combinaciones posibles de letras y números
letras = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K"]
numeros = list(map(str, range(1, 19)))
combinaciones = [f"{letra}{numero}" for numero in numeros for letra in letras]

# Dividir las combinaciones en listas de 11 elementos (A-K)
combinaciones_por_filas = [combinaciones[i:i+11] for i in range(0, len(combinaciones), 11)]

# Crear la interfaz
inputs = [
    gr.Dropdown(["2016", "2017", "2019"], label="Año"),
    gr.Dropdown(["25", "40"], label="Inclinación"),
]

# Agregar un CheckboxGroup para cada fila (del 1 al 18)
for i, fila in enumerate(combinaciones_por_filas, start=1):
    inputs.append(gr.CheckboxGroup(fila, label=f"Fila {i}", scale=4, min_width=2000))

iface = gr.Interface(
    fn=mostrar_parametros,
    inputs=inputs,
    outputs=[gr.Image(label="Output"), gr.Label("Predicción del Modelo")],
)

iface.launch(share=False)