import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import torch
import matplotlib.pyplot as plt
import numpy as np
import gradio as gr


class CNN(nn.Module):
    def __init__(self, in_channels=1, num_classes=4):
        super(CNN, self).__init__()

        self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)

        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

        self.batch_norm1 = nn.BatchNorm2d(32)
        self.batch_norm2 = nn.BatchNorm2d(64)
        self.batch_norm3 = nn.BatchNorm2d(128)

        self.dropout = nn.Dropout(0.5)

        # Calcular el tamaño de la entrada a la capa fully connected
        self.fc1 = nn.Linear(128 * (200 // 8) * (200 // 8), 256)
        self.fc2 = nn.Linear(256, num_classes)

    def forward(self, x):
        x = F.relu(self.batch_norm1(self.conv1(x)))
        x = self.pool(x)

        x = F.relu(self.batch_norm2(self.conv2(x)))
        x = self.pool(x)

        x = F.relu(self.batch_norm3(self.conv3(x)))
        x = self.pool(x)

        x = x.view(x.shape[0], -1)  # Aplanar
        x = self.dropout(F.relu(self.fc1(x)))
        x = self.fc2(x)

        return x


model = CNN()
model.load_state_dict(
    torch.load("gabriel_complex_modelo.pth", map_location=torch.device("cpu"))
)


def inference(model, imagen, device="cpu"):
    label_mapping = {0: "Círculo", 1: "Triángulo", 2: "Cuadrado", 3: "Estrella"}

    model.eval()  # Ponemos el modelo en modo evaluación

    # Realizar la inferencia
    with torch.no_grad():
        scores = model(imagen)  # Output: tensor con logits
        probabilities = torch.softmax(
            scores, dim=1
        )  # Convertir logits a probabilidades
        _, prediction = scores.max(1)  # Obtener la clase con mayor probabilidad
        label_predicho = prediction.item()

    # Diccionario con las probabilidades
    probabilities_dict = {
        label_mapping[i]: float(probabilities[0, i]) for i in range(4)
    }

    return label_mapping[label_predicho], probabilities_dict


def predict(img):
    image_array = img["composite"][:, :, 3]
    image_array = 255 - image_array

    image_tensor = torch.from_numpy(image_array).unsqueeze(0)

    transform_to_gray = transforms.Compose(
        [
            transforms.Resize((200, 200)),
            transforms.ConvertImageDtype(dtype=torch.float32),  # Convertir a flotante
        ]
    )

    image = transform_to_gray(image_tensor)
    image = image.unsqueeze(0)  # Agregar dimensión extra

    # Hacemos la inferencia
    label_predict, probabilities = inference(model, image, device="cpu")
    print(label_predict)
    print(probabilities)

    return probabilities  # Retorna el diccionario con las probabilidades


with gr.Blocks() as demo:
    with gr.Row():
        im = gr.Sketchpad(type="numpy", crop_size="1:1")
        out = gr.Label()

    im.change(predict, outputs=out, inputs=im, show_progress="hidden")

demo.launch(share=True, debug=False)