#!/usr/bin/env python3
"""
Prueba Interactiva: Reconocimiento de Dígitos
Dibuja un dígito y entrena la red en tiempo real

Copyright (C) 2024 GR Code Digital Solutions
Licenciado bajo GPL v3
"""

import sys

from grnexus import NeuralNetwork, KNeighborsClassifier
import random

try:
    import flet as ft
    FLET_DISPONIBLE = True
except ImportError:
    FLET_DISPONIBLE = False
    print("❌ Flet no está instalado")
    print("   Instala con: pip install flet")



class ClasificadorDigitos:
    """
    Clasificador de dígitos con capacidad de aprendizaje
    
    FLUJO DE DATOS:
    1. Imagen 28x28 píxeles → Vector 784D (aplanado)
    2. Cada píxel tiene valor 0.0 (blanco) a 1.0 (negro)
    3. La red procesa estos 784 números y predice qué dígito (0-9) es
    
    ARQUITECTURA:
    - Capa 1: 784 entradas → 128 neuronas (ReLU)
    - Dropout: Desactiva 20% de neuronas para evitar overfitting
    - Capa 2: 128 → 64 neuronas (ReLU)
    - Dropout: Desactiva 20% de neuronas
    - Capa 3: 64 → 10 neuronas (salida, una por dígito)
    - Softmax: Convierte salidas en probabilidades que suman 1.0
    """
    
    def __init__(self, use_ml=True):
        # Crear red neuronal
        # Importar desde el paquete instalado
        from grnexus import DenseLayer, DropoutLayer, SoftmaxLayer
        from grnexus import ReLU, Tanh
        
        self.nn = NeuralNetwork(learning_rate=0.01, optimizer='adam', loss='cross_entropy')
        self.nn.add(DenseLayer(units=128, input_dim=784, activation=ReLU()))
        self.nn.add(DropoutLayer(rate=0.2))
        self.nn.add(DenseLayer(units=64, input_dim=128, activation=ReLU()))
        self.nn.add(DropoutLayer(rate=0.2))
        self.nn.add(DenseLayer(units=10, input_dim=64))
        self.nn.add(SoftmaxLayer())
        
        print("✅ Red neuronal creada (784 → 128 → 64 → 10)")
        print("   Optimizer: Adam, Learning Rate: 0.01")
        
        # Crear clasificador KNN si está disponible
        self.knn = None
        self.knn_data = []  # Almacenar datos de entrenamiento
        self.knn_labels = []  # Almacenar etiquetas
        
        if use_ml:
            try:
                self.knn = KNeighborsClassifier(n_neighbors=5)
                print("✅ KNN classifier creado (k=5)")
                print("   Modo HÍBRIDO activado: NN + KNN")
            except:
                print("ℹ️  Modo solo Neural Network (KNN no disponible)")
        else:
            print("ℹ️  Modo solo Neural Network")
    
    def clasificar(self, imagen_28x28, usar_knn=False):
        """
        Clasifica una imagen usando NN o KNN
        
        Args:
            imagen_28x28: Vector de 784 valores (imagen aplanada)
            usar_knn: Si True, usa KNN; si False, usa Neural Network
        
        Returns:
            (digito, probs): Dígito predicho y probabilidades
        """
        if usar_knn and self.knn and len(self.knn_data) > 0:
            # Usar KNN
            pred = self.knn.predict([imagen_28x28])[0]
            digito = int(round(pred))
            # Crear pseudo-probabilidades (100% para el dígito predicho)
            probs = [0.0] * 10
            probs[digito] = 1.0
            return digito, probs
        else:
            # Usar Neural Network
            probs = self.nn.predict([imagen_28x28])[0]
            digito = probs.index(max(probs))
            return digito, probs
    
    def entrenar(self, imagen_28x28, etiqueta_real, entrenar_knn=True):
        """
        Entrena ambos modelos (NN y KNN)
        
        Args:
            imagen_28x28: Vector de 784 valores
            etiqueta_real: Dígito correcto (0-9)
            entrenar_knn: Si True, también entrena KNN
        
        Returns:
            loss: Pérdida de la red neuronal
        """
        # Entrenar Neural Network
        target = [0.0] * 10
        target[etiqueta_real] = 1.0
        loss = self.nn.train([imagen_28x28], [target], epochs=1, batch_size=1, verbose=False)['loss'][0]
        
        # Agregar datos para KNN
        if entrenar_knn and self.knn is not None:
            self.knn_data.append(imagen_28x28)
            self.knn_labels.append(float(etiqueta_real))
            # Re-entrenar KNN con todos los datos
            self.knn.fit(self.knn_data, self.knn_labels)
        
        return loss


def main(page: ft.Page):
    """Aplicación Flet para dibujar y entrenar"""
    
    if not FLET_DISPONIBLE:
        return
    
    page.title = "🎨 Reconocimiento de Dígitos - GRNexus"
    page.window_width = 950
    page.window_height = 750
    page.padding = 30
    page.bgcolor = "#0a0e27"
    page.theme_mode = ft.ThemeMode.DARK
    
    # Crear clasificador
    clasificador = ClasificadorDigitos()
    
    # Configuración del Canvas (más pequeño para que quepa)
    GRID_SIZE = 28
    CELL_SIZE = 12  # Reducido de 16 a 12
    CANVAS_SIZE = GRID_SIZE * CELL_SIZE  # 336px
    canvas_data = [[0.0 for _ in range(GRID_SIZE)] for _ in range(GRID_SIZE)]
    
    # Estado
    is_drawing = False
    ultimo_digito_predicho = -1
    iteraciones_entrenamiento = 0
    usar_knn_mode = False  # Toggle entre NN y KNN
    
    # Elementos UI (más compactos)
    resultado_text = ft.Text("Dibuja un dígito", size=24, weight=ft.FontWeight.BOLD)
    confianza_text = ft.Text("", size=14)
    modelo_actual_text = ft.Text("Modelo: Neural Network", size=12, color=ft.Colors.BLUE_300)
    
    # Barras de probabilidad
    barras_container = ft.Column(spacing=5)
    
    def crear_barra(digito, valor):
        color = ft.Colors.GREEN_400 if valor > 0.5 else (ft.Colors.BLUE_400 if valor > 0.2 else ft.Colors.GREY_700)
        return ft.Row([
            ft.Text(f"{digito}", size=12, weight=ft.FontWeight.BOLD, width=15),
            ft.Container(
                content=ft.ProgressBar(value=valor, color=color, bgcolor=ft.Colors.GREY_800, height=8),
                expand=True
            ),
            ft.Text(f"{valor:.0%}", size=10, width=35, text_align=ft.TextAlign.RIGHT)
        ], spacing=5)
    
    def actualizar_barras(probs):
        barras_container.controls.clear()
        for i, prob in enumerate(probs):
            barras_container.controls.append(crear_barra(i, prob))
        page.update()
    
    def obtener_imagen_plana():
        imagen = []
        for fila in canvas_data:
            imagen.extend(fila)
        return imagen
    
    def clasificar_dibujo(e=None):
        nonlocal usar_knn_mode
        imagen = obtener_imagen_plana()
        if sum(imagen) < 1.0:
            resultado_text.value = "⚠️ Dibuja algo primero"
            resultado_text.color = ft.Colors.ORANGE_400
            page.update()
            return
        
        # Clasificar con el modelo seleccionado
        digito, probs = clasificador.clasificar(imagen, usar_knn=usar_knn_mode)
        nonlocal ultimo_digito_predicho
        ultimo_digito_predicho = digito
        
        # Mostrar resultado
        modelo_nombre = "KNN" if usar_knn_mode else "Neural Network"
        resultado_text.value = f"Predicción ({modelo_nombre}): {digito}"
        resultado_text.color = ft.Colors.GREEN_400 if probs[digito] > 0.5 else ft.Colors.ORANGE_400
        confianza_text.value = f"Confianza: {probs[digito]:.1%}"
        
        actualizar_barras(probs)
    
    def comparar_modelos(e):
        """Compara predicciones de ambos modelos"""
        imagen = obtener_imagen_plana()
        if sum(imagen) < 1.0:
            page.snack_bar = ft.SnackBar(content=ft.Text("Dibuja algo primero"))
            page.snack_bar.open = True
            page.update()
            return
        
        if not clasificador.knn or len(clasificador.knn_data) == 0:
            page.snack_bar = ft.SnackBar(content=ft.Text("KNN necesita datos de entrenamiento primero"))
            page.snack_bar.open = True
            page.update()
            return
        
        # Predecir con ambos modelos
        digito_nn, probs_nn = clasificador.clasificar(imagen, usar_knn=False)
        digito_knn, probs_knn = clasificador.clasificar(imagen, usar_knn=True)
        
        # Mostrar comparación
        coinciden = "✅" if digito_nn == digito_knn else "❌"
        mensaje = f"{coinciden} NN: {digito_nn} ({probs_nn[digito_nn]:.0%}) | KNN: {digito_knn} ({probs_knn[digito_knn]:.0%})"
        
        page.snack_bar = ft.SnackBar(
            content=ft.Text(mensaje),
            bgcolor=ft.Colors.GREEN_700 if digito_nn == digito_knn else ft.Colors.ORANGE_700
        )
        page.snack_bar.open = True
        page.update()
    
    def toggle_modelo(e):
        """Cambia entre Neural Network y KNN"""
        nonlocal usar_knn_mode
        usar_knn_mode = e.control.value
        modelo_nombre = "KNN" if usar_knn_mode else "Neural Network"
        modelo_actual_text.value = f"Modelo: {modelo_nombre}"
        modelo_actual_text.color = ft.Colors.PURPLE_300 if usar_knn_mode else ft.Colors.BLUE_300
        
        if usar_knn_mode and (not clasificador.knn or len(clasificador.knn_data) == 0):
            page.snack_bar = ft.SnackBar(
                content=ft.Text("⚠️ KNN necesita datos de entrenamiento. Entrena algunos ejemplos primero."),
                bgcolor=ft.Colors.ORANGE_700
            )
            page.snack_bar.open = True
        
        page.update()
        
        # Re-clasificar si hay algo dibujado
        imagen = obtener_imagen_plana()
        if sum(imagen) >= 1.0:
            clasificar_dibujo()
    
    def entrenar_red(e):
        """Entrena la red con el dibujo actual y la etiqueta seleccionada"""
        try:
            etiqueta = int(dropdown_entrenamiento.value)
        except (ValueError, TypeError):
            page.snack_bar = ft.SnackBar(content=ft.Text("Selecciona el número correcto primero"))
            page.snack_bar.open = True
            page.update()
            return
            
        imagen = obtener_imagen_plana()
        if sum(imagen) < 1.0:
            page.snack_bar = ft.SnackBar(content=ft.Text("Dibuja algo para entrenar"))
            page.snack_bar.open = True
            page.update()
            return
            
        # Entrenar
        loss = clasificador.entrenar(imagen, etiqueta)
        nonlocal iteraciones_entrenamiento
        iteraciones_entrenamiento += 1
        contador_entrenamientos.value = f"Entrenamientos: {iteraciones_entrenamiento}"
        
        snackbar = ft.SnackBar(
            content=ft.Text(f"✅ ¡Aprendido! (Loss: {loss:.4f})"),
            bgcolor=ft.Colors.GREEN_700,
        )
        # Feedback visual
        snackbar.open = True
        page.overlay.append(snackbar)

        page.update()
        
        # Re-clasificar para ver la mejora
        clasificar_dibujo()
    
    def limpiar_canvas(e=None):
        for i in range(GRID_SIZE):
            for j in range(GRID_SIZE):
                canvas_data[i][j] = 0.0
        
        # Limpiar visualmente (reconstruir grid es costoso, mejor actualizar colores)
        # Pero con GestureDetector único, necesitamos repintar los contenedores
        for i in range(GRID_SIZE):
            for j in range(GRID_SIZE):
                idx = i * GRID_SIZE + j
                grid_container.controls[idx].bgcolor = ft.Colors.WHITE
        
        resultado_text.value = "Dibuja un dígito"
        resultado_text.color = ft.Colors.WHITE
        confianza_text.value = ""
        barras_container.controls.clear()
        page.update()
    
    # --- Lógica de Dibujo Mejorada (Single GestureDetector) ---
    
    def pintar_en_posicion(local_x, local_y):
        # Calcular índices de celda
        j = int(local_x // CELL_SIZE)
        i = int(local_y // CELL_SIZE)
        
        if 0 <= i < GRID_SIZE and 0 <= j < GRID_SIZE:
            # Pintar celda y vecinas (pincel suave)
            for di in range(-1, 2):
                for dj in range(-1, 2):
                    ni, nj = i + di, j + dj
                    if 0 <= ni < GRID_SIZE and 0 <= nj < GRID_SIZE:
                        distancia = abs(di) + abs(dj)
                        factor = 1.0 if distancia == 0 else 0.5 if distancia == 1 else 0.2
                        
                        canvas_data[ni][nj] = min(1.0, canvas_data[ni][nj] + factor)
                        
                        # Actualizar visualmente
                        idx = ni * GRID_SIZE + nj
                        gray = int((1.0 - canvas_data[ni][nj]) * 255)
                        color = f"#{gray:02x}{gray:02x}{gray:02x}"
                        grid_container.controls[idx].bgcolor = color
            page.update()

    def on_pan_start(e: ft.DragStartEvent):
        pintar_en_posicion(e.local_x, e.local_y)

    def on_pan_update(e: ft.DragUpdateEvent):
        pintar_en_posicion(e.local_x, e.local_y)
        
    def on_tap_down(e: ft.ContainerTapEvent):
        pintar_en_posicion(e.local_x, e.local_y)

    # Grid visual (celdas estáticas)
    grid_controls = []
    for i in range(GRID_SIZE):
        for j in range(GRID_SIZE):
            grid_controls.append(
                ft.Container(
                    width=CELL_SIZE,
                    height=CELL_SIZE,
                    bgcolor=ft.Colors.WHITE,
                    border=ft.border.all(0.5, ft.Colors.GREY_300)
                )
            )
            
    grid_container = ft.Row(
        controls=grid_controls,
        wrap=True,
        width=CANVAS_SIZE,
        spacing=0,
        run_spacing=0,
    )
    
    # Contenedor interactivo único
    canvas_stack = ft.Stack([
        grid_container,
        ft.GestureDetector(
            on_pan_start=on_pan_start,
            on_pan_update=on_pan_update,
            on_tap_down=on_tap_down,
            width=CANVAS_SIZE,
            height=CANVAS_SIZE,
        )
    ])

    # --- Botones de acciones ---
    btn_limpiar = ft.ElevatedButton(
        "Limpiar Canvas",
        icon=ft.Icons.CLEAR,
        on_click=limpiar_canvas,
        bgcolor="#e63946",
        color=ft.Colors.WHITE,
    )

    btn_predecir = ft.ElevatedButton(
        "🤖 Clasificar",
        icon=ft.Icons.PSYCHOLOGY,
        on_click=lambda _: clasificar_dibujo(), # Changed from clasificar_btn to clasificar_dibujo
        bgcolor="#4361ee",
        color=ft.Colors.WHITE,
    )

    # --- Panel de Entrenamiento ---
    dropdown_entrenamiento = ft.Dropdown(
        label="Es en realidad un...",
        options=[ft.dropdown.Option(str(i)) for i in range(10)],
        width=150,
    )

    btn_entrenar = ft.ElevatedButton(
        "🎓 ¡Aprender!",
        icon=ft.Icons.SCHOOL,
        on_click=entrenar_red,
        bgcolor="#7209b7",
        color=ft.Colors.WHITE,
    )
    
    contador_entrenamientos = ft.Text("Entrenamientos: 0", size=11, color=ft.Colors.GREY_400)
    
    # --- Campo de texto para nombre de archivo ---
    filename_field = ft.TextField(
        label="Nombre del archivo",
        value="modelo_digitos.nexus",
        hint_text="Escribe el nombre del archivo .nexus",
        width=250,
        dense=True,
        text_size=12
    )
    
    def guardar_modelo_simple(e):
        """Guarda el modelo con el nombre especificado"""
        filename = filename_field.value.strip()
        if not filename:
            filename = "modelo.nexus"
        if not filename.endswith('.nexus'):
            filename += '.nexus'
        
        try:
            clasificador.nn.save(filename)
            page.snack_bar = ft.SnackBar(
                content=ft.Text(f"💾 Guardado: {filename}"),
                bgcolor=ft.Colors.GREEN_700
            )
            page.snack_bar.open = True
            page.update()
        except Exception as ex:
            page.snack_bar = ft.SnackBar(
                content=ft.Text(f"❌ Error: {str(ex)}"),
                bgcolor=ft.Colors.RED_700
            )
            page.snack_bar.open = True
            page.update()
    
    def cargar_modelo_simple(e):
        """Carga el modelo con el nombre especificado"""
        filename = filename_field.value.strip()
        if not filename:
            page.snack_bar = ft.SnackBar(
                content=ft.Text("⚠️ Ingresa un nombre de archivo"),
                bgcolor=ft.Colors.ORANGE_700
            )
            page.snack_bar.open = True
            page.update()
            return
        
        if not os.path.exists(filename):
            page.snack_bar = ft.SnackBar(
                content=ft.Text(f"❌ Archivo no encontrado: {filename}"),
                bgcolor=ft.Colors.RED_700
            )
            page.snack_bar.open = True
            page.update()
            return
        
        try:
            nonlocal clasificador, iteraciones_entrenamiento
            clasificador.nn = GRNexus.load(filename)
            iteraciones_entrenamiento = 0
            contador_entrenamientos.value = "Entrenamientos: 0 (modelo cargado)"
            
            page.snack_bar = ft.SnackBar(
                content=ft.Text(f"📂 Cargado: {filename}"),
                bgcolor=ft.Colors.GREEN_700
            )
            page.snack_bar.open = True
            
            # Re-clasificar si hay algo dibujado
            imagen = obtener_imagen_plana()
            if sum(imagen) >= 1.0:
                clasificar_dibujo()
            
            page.update()
        except Exception as ex:
            page.snack_bar = ft.SnackBar(
                content=ft.Text(f"❌ Error: {str(ex)}"),
                bgcolor=ft.Colors.RED_700
            )
            page.snack_bar.open = True
            page.update()
    
    def resetear_red(e):
        """Reinicia la red neuronal a su estado inicial"""
        nonlocal clasificador, iteraciones_entrenamiento
        clasificador = ClasificadorDigitos()
        iteraciones_entrenamiento = 0
        contador_entrenamientos.value = "Entrenamientos: 0"
        limpiar_canvas()
        page.snack_bar = ft.SnackBar(
            content=ft.Text("🔄 Red neuronal reiniciada"),
            bgcolor=ft.Colors.BLUE_700
        )
        page.snack_bar.open = True
        page.update()
    
    # Layout Principal
    page.add(
        ft.Row([
            # Izquierda: Dibujo
            ft.Container(
                content=ft.Column([
                    ft.Text("✏️ Zona de Dibujo", size=20, weight=ft.FontWeight.BOLD),
                    ft.Container(
                        content=canvas_stack,
                        border=ft.border.all(2, ft.Colors.BLUE_500),
                        padding=2,
                        bgcolor=ft.Colors.GREY_900
                    ),
                    ft.Row([
                        ft.ElevatedButton("Clasificar", icon=ft.Icons.CHECK, on_click=lambda _: clasificar_dibujo()),
                        btn_limpiar,
                    ], alignment=ft.MainAxisAlignment.CENTER),
                ], horizontal_alignment=ft.CrossAxisAlignment.CENTER),
                padding=20,
                bgcolor=ft.Colors.GREY_900,
                border_radius=10
            ),
            
            # Derecha: Resultados y Entrenamiento (con scroll)
            ft.Container(
                content=ft.Column([
                    ft.Text("🧠 Cerebro Digital", size=18, weight=ft.FontWeight.BOLD),
                    ft.Divider(height=1),
                    resultado_text,
                    confianza_text,
                    modelo_actual_text,
                    
                    # Toggle de modelo
                    ft.Row([
                        ft.Text("Modelo:", size=12),
                        ft.Switch(
                            label="KNN",
                            value=False,
                            disabled=clasificador.knn is None,
                            on_change=lambda e: toggle_modelo(e)
                        ),
                    ], spacing=10),
                    
                    # Botón de comparación
                    ft.ElevatedButton(
                        "🎯 Comparar NN vs KNN",
                        icon=ft.Icons.COMPARE_ARROWS,
                        on_click=comparar_modelos,
                        style=ft.ButtonStyle(bgcolor=ft.Colors.PURPLE_700, color=ft.Colors.WHITE),
                        disabled=clasificador.knn is None
                    ),
                    
                    ft.Container(height=5),
                    ft.Text("Probabilidades:", weight=ft.FontWeight.BOLD, size=12),
                    barras_container,
                    ft.Divider(height=1),
                    ft.Text("🎓 Entrenamiento", size=14, weight=ft.FontWeight.BOLD, color=ft.Colors.PURPLE_300),
                    ft.Text("Corrige si se equivoca:", size=11),
                    ft.Row([dropdown_entrenamiento, btn_entrenar], spacing=5),
                    contador_entrenamientos,
                    ft.Divider(height=1),
                    ft.Text("💾 Modelos", size=14, weight=ft.FontWeight.BOLD, color=ft.Colors.BLUE_300),
                    filename_field,
                    ft.Row([
                        ft.ElevatedButton(
                            "💾 Guardar",
                            on_click=guardar_modelo_simple,
                            style=ft.ButtonStyle(bgcolor=ft.Colors.GREEN_700, color=ft.Colors.WHITE)
                        ),
                        ft.ElevatedButton(
                            "📂 Cargar",
                            on_click=cargar_modelo_simple,
                            style=ft.ButtonStyle(bgcolor=ft.Colors.BLUE_700, color=ft.Colors.WHITE)
                        ),
                    ], spacing=5),
                    ft.Text("(Archivos en carpeta actual)", size=10, color=ft.Colors.GREY_500),
                    ft.ElevatedButton(
                        "Resetear Red",
                        icon=ft.Icons.REFRESH,
                        on_click=resetear_red,
                        style=ft.ButtonStyle(bgcolor=ft.Colors.ORANGE_700, color=ft.Colors.WHITE)
                    ),
                ], 
                scroll=ft.ScrollMode.ALWAYS,  # Siempre mostrar scroll
                spacing=8),
                padding=15,
                expand=True,
                height=700  # Altura fija para forzar scroll
            )
        ], vertical_alignment=ft.CrossAxisAlignment.START)
    )

if __name__ == "__main__":
    if FLET_DISPONIBLE:
        print("\n🎨 Reconocimiento de Dígitos - GRNexus")
        print("   Modo Interactivo con Aprendizaje en Tiempo Real\n")
        ft.app(target=main)
    else:
        print("\n❌ Error: Flet no instalado")
