import gradio as gr
import os
import shutil
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import io
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader, random_split
from PIL import Image
import joblib  # Para salvar como .pkl

# Modelos para seleção
model_dict = {
    'AlexNet': models.alexnet,
    'ResNet18': models.resnet18,
    'ResNet34': models.resnet34,
    'ResNet50': models.resnet50,
    'MobileNetV2': models.mobilenet_v2
}

# Variáveis globais
model = None
train_loader = None
val_loader = None
test_loader = None
dataset_path = 'dataset'
class_dirs = []
test_dataset_path = 'test_dataset'
test_class_dirs = []
num_classes = 2  # Valor padrão para o número de classes
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Função para configurar as classes
def setup_classes(num_classes_value):
    global class_dirs, dataset_path, num_classes

    num_classes = int(num_classes_value)  # Atualizar a variável global num_classes
    # Limpar diretório antigo se existir
    if os.path.exists(dataset_path):
        shutil.rmtree(dataset_path)
    os.makedirs(dataset_path)

    # Criar diretórios para cada classe
    class_dirs = [os.path.join(dataset_path, f'class_{i}') for i in range(num_classes)]
    for class_dir in class_dirs:
        os.makedirs(class_dir)

    return f"Criados {num_classes} diretórios para classes."

# Função para upload de imagens
def upload_images(class_id, images):
    class_dir = class_dirs[int(class_id)]
    for image in images:
        shutil.copy(image, class_dir)
    return f"Imagens salvas na classe {class_id}."

# Função para preparação dos dados
def prepare_data(batch_size=32, resize=(224, 224)):
    global train_loader, val_loader, test_loader, num_classes

    # Transformações para os dados de treinamento e validação
    transform = transforms.Compose([
        transforms.Resize(resize),
        transforms.ToTensor(),
    ])

    dataset = datasets.ImageFolder(dataset_path, transform=transform)

    if len(dataset.classes) != num_classes:
        return f"Erro: Número de classes detectadas ({len(dataset.classes)}) não corresponde ao número esperado ({num_classes}). Verifique suas imagens."

    # Divisão do conjunto de dados em treinamento, validação e teste
    train_size = int(0.7 * len(dataset))
    val_size = int(0.2 * len(dataset))
    test_size = len(dataset) - train_size - val_size
    train_dataset, val_dataset, test_dataset = random_split(dataset, [train_size, val_size, test_size])

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    return "Preparação dos dados concluída com sucesso."

# Função para iniciar o treinamento
def start_training(model_name, epochs, lr):
    global model, train_loader, val_loader, device

    if train_loader is None or val_loader is None:
        return "Erro: Dados não preparados."

    model = model_dict[model_name](pretrained=True)
    model.fc = nn.Linear(model.fc.in_features, num_classes)
    model = model.to(device)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=float(lr))

    for epoch in range(int(epochs)):
        model.train()
        running_loss = 0.0
        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

        print(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss/len(train_loader)}")

    torch.save(model.state_dict(), 'modelo.pth')
    return f"Treinamento concluído com sucesso. Modelo salvo."

# Função para avaliação do modelo com conjunto de teste
def evaluate_model(loader):
    global model, device, num_classes

    if model is None:
        return "Erro: Modelo não treinado."

    if loader is None:
        return "Erro: Conjunto de dados de teste não está preparado."

    model.eval()
    all_preds = []
    all_labels = []
    try:
        with torch.no_grad():
            for inputs, labels in loader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                _, preds = torch.max(outputs, 1)
                all_preds.extend(preds.cpu().numpy())
                all_labels.extend(labels.cpu().numpy())

        report = classification_report(all_labels, all_preds, labels=list(range(num_classes)), target_names=[f"class_{i}" for i in range(num_classes)], zero_division=0)
        return report
    except Exception as e:
        return f"Erro durante a avaliação: {str(e)}"

# Função para mostrar a matriz de confusão
def show_confusion_matrix(loader):
    global model, device, num_classes

    if model is None:
        return "Erro: Modelo não treinado."

    model.eval()
    all_preds = []
    all_labels = []
    with torch.no_grad():
        for inputs, labels in loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())

    cm = confusion_matrix(all_labels, all_preds, labels=list(range(num_classes)))

    plt.figure(figsize=(6, 4.8))
    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", xticklabels=[f"class_{i}" for i in range(num_classes)], yticklabels=[f"class_{i}" for i in range(num_classes)])
    plt.xlabel('Predictions')
    plt.ylabel('Actuals')
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    plt.close()
    buf.seek(0)
    return Image.open(buf)

# Função para predição de imagens desconhecidas
def predict_images(images):
    global model, device, num_classes

    if model is None:
        return "Erro: Modelo não treinado."

    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
    ])

    model.eval()
    results = []

    for image in images:
        try:
            img = transform(Image.open(image)).unsqueeze(0).to(device)
            with torch.no_grad():
                outputs = model(img)
                _, preds = torch.max(outputs, 1)
                predicted_class = preds.item()
                results.append(f"Imagem {os.path.basename(image)} - Classe prevista: class_{predicted_class}")
        except Exception as e:
            results.append(f"Erro ao processar a imagem {image}: {str(e)}")

    return results

# Função para exportar o modelo
def export_model(format):
    global model

    if model is None:
        return "Erro: Modelo não treinado."

    file_path = f"modelo_exportado.{format}"
    if format == "pth":
        torch.save(model.state_dict(), file_path)
    elif format == "onnx":
        try:
            dummy_input = torch.randn(1, 3, 224, 224).to(device)
            torch.onnx.export(model, dummy_input, file_path, export_params=True, opset_version=10, input_names=['input'], output_names=['output'])
        except Exception as e:
            return f"Erro ao exportar para ONNX: {str(e)}"
    elif format == "pkl":
        joblib.dump(model, file_path)
    else:
        return f"Formato {format} não suportado."

    return f"Modelo exportado com sucesso para {file_path}"

# Função para configurar diretórios de teste
def setup_test_classes():
    global test_class_dirs, test_dataset_path

    if os.path.exists(test_dataset_path):
        shutil.rmtree(test_dataset_path)
    os.makedirs(test_dataset_path)

    # Criar diretórios para cada classe
    test_class_dirs = [os.path.join(test_dataset_path, f'class_{i}') for i in range(num_classes)]
    for class_dir in test_class_dirs:
        os.makedirs(class_dir)

    return f"Criados {num_classes} diretórios para classes de teste."

# Função para upload de imagens de teste
def upload_test_images(class_id, images):
    class_dir = test_class_dirs[int(class_id)]
    for image in images:
        shutil.copy(image, class_dir)
    return f"Imagens de teste salvas na classe {class_id}."

# Função para preparar dados de teste
def prepare_test_data(batch_size=32, resize=(224, 224)):
    global test_loader, num_classes

    transform = transforms.Compose([
        transforms.Resize(resize),
        transforms.ToTensor(),
    ])

    test_dataset = datasets.ImageFolder(test_dataset_path, transform=transform)

    if len(test_dataset.classes) != num_classes:
        return f"Erro: Número de classes detectadas ({len(test_dataset.classes)}) não corresponde ao número esperado ({num_classes}). Verifique suas imagens."

    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    return "Preparação dos dados de teste concluída com sucesso."

# Interface Gradio
def main():
    with gr.Blocks() as demo:
        gr.Markdown("# Image Classification Training")

        with gr.Tab("Configurar Classes"):
            num_classes_input = gr.Number(label="Número de Classes", value=2, precision=0)
            setup_button = gr.Button("Configurar Classes")
            setup_output = gr.Textbox()
            setup_button.click(setup_classes, inputs=num_classes_input, outputs=setup_output)

        with gr.Tab("Upload de Imagens"):
            upload_inputs = []
            for i in range(num_classes):
                with gr.Column():
                    gr.Markdown(f"### Classe {i}")
                    class_id = gr.Number(label=f"ID da Classe {i}", value=i, precision=0)
                    images = gr.File(label="Upload de Imagens", file_count="multiple", type="filepath")
                    upload_button = gr.Button("Upload")
                    upload_output = gr.Textbox()

                    upload_inputs.append((class_id, images, upload_button, upload_output))
                    upload_button.click(upload_images, inputs=[class_id, images], outputs=upload_output)

        with gr.Tab("Preparação de Dados"):
            batch_size = gr.Number(label="Tamanho do Batch", value=32)
            resize = gr.Textbox(label="Resize (Ex: 224,224)", value="224,224")
            prepare_button = gr.Button("Preparar Dados")
            prepare_output = gr.Textbox()
            prepare_button.click(lambda batch_size, resize: prepare_data(batch_size=batch_size, resize=tuple(map(int, resize.split(',')))), inputs=[batch_size, resize], outputs=prepare_output)

        with gr.Tab("Treinamento"):
            model_name = gr.Dropdown(label="Modelo", choices=list(model_dict.keys()))
            epochs = gr.Number(label="Épocas", value=30)
            lr = gr.Number(label="Taxa de Aprendizado", value=0.001)
            train_button = gr.Button("Iniciar Treinamento")
            train_output = gr.Textbox()
            train_button.click(start_training, inputs=[model_name, epochs, lr], outputs=train_output)

        with gr.Tab("Avaliação do Modelo"):
            eval_button = gr.Button("Avaliar Modelo")
            eval_output = gr.Textbox()
            eval_button.click(lambda: evaluate_model(test_loader), outputs=eval_output)

            cm_button = gr.Button("Mostrar Matriz de Confusão")
            cm_output = gr.Image()
            cm_button.click(lambda: show_confusion_matrix(test_loader), outputs=cm_output)

        with gr.Tab("Predição e Avaliação"):
            predict_images_input = gr.File(label="Upload de Imagens para Predição", file_count="multiple", type="filepath")
            predict_button = gr.Button("Predizer")
            predict_output = gr.Textbox()
            predict_button.click(predict_images, inputs=predict_images_input, outputs=predict_output)

            gr.Markdown("### Upload de Imagens de Teste")
            setup_test_button = gr.Button("Configurar Diretórios de Teste")
            setup_test_output = gr.Textbox()
            setup_test_button.click(setup_test_classes, outputs=setup_test_output)

            upload_test_inputs = []
            for i in range(num_classes):
                with gr.Column():
                    gr.Markdown(f"### Classe de Teste {i}")
                    test_class_id = gr.Number(label=f"ID da Classe {i}", value=i, precision=0)
                    test_images = gr.File(label="Upload de Imagens de Teste", file_count="multiple", type="filepath")
                    upload_test_button = gr.Button("Upload Imagens de Teste")
                    upload_test_output = gr.Textbox()

                    upload_test_inputs.append((test_class_id, test_images, upload_test_button, upload_test_output))
                    upload_test_button.click(upload_test_images, inputs=[test_class_id, test_images], outputs=upload_test_output)

            prepare_test_button = gr.Button("Preparar Dados de Teste")
            prepare_test_output = gr.Textbox()
            prepare_test_button.click(lambda batch_size, resize: prepare_test_data(batch_size=batch_size, resize=tuple(map(int, resize.split(',')))), inputs=[batch_size, resize], outputs=prepare_test_output)

            eval_test_button = gr.Button("Avaliar Conjunto de Teste")
            eval_test_output = gr.Textbox()
            eval_test_button.click(lambda: evaluate_model(test_loader), outputs=eval_test_output)

            cm_test_button = gr.Button("Mostrar Matriz de Confusão do Conjunto de Teste")
            cm_test_output = gr.Image()
            cm_test_button.click(lambda: show_confusion_matrix(test_loader), outputs=cm_test_output)

        with gr.Tab("Exportação"):
            export_format = gr.Radio(label="Formato", choices=["pth", "onnx", "pkl"])
            export_button = gr.Button("Exportar Modelo")
            export_output = gr.Textbox()
            export_button.click(export_model, inputs=export_format, outputs=export_output)

    demo.launch()

if __name__ == "__main__":
    main()