import os
import shutil
import gradio as gr
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 matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from sklearn.metrics import classification_report, confusion_matrix
import tempfile
import warnings
warnings.filterwarnings("ignore")

print("🖥️ Iniciando sistema...")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}")

# Modelos disponíveis
MODELS = {
    'ResNet18': models.resnet18,
    'ResNet34': models.resnet34, 
    'MobileNetV2': models.mobilenet_v2
}

# Estado global
class AppState:
    def __init__(self):
        self.model = None
        self.train_loader = None
        self.val_loader = None
        self.test_loader = None
        self.dataset_path = None
        self.class_dirs = []
        self.class_labels = ['classe_0', 'classe_1']
        self.num_classes = 2
        self.image_queue = []  # Para armazenar imagens uploaded

state = AppState()

def setup_classes(num_classes_value):
    """Configura número de classes"""
    try:
        state.num_classes = int(num_classes_value)
        state.dataset_path = tempfile.mkdtemp()
        state.class_labels = [f'classe_{i}' for i in range(state.num_classes)]
        
        # Criar diretórios
        state.class_dirs = []
        for i in range(state.num_classes):
            class_dir = os.path.join(state.dataset_path, f'classe_{i}')
            os.makedirs(class_dir, exist_ok=True)
            state.class_dirs.append(class_dir)
        
        return f"✅ Sistema configurado para {state.num_classes} classes"
    except Exception as e:
        return f"❌ Erro: {str(e)}"

def set_class_labels(labels_text):
    """Define rótulos das classes"""
    try:
        labels = [label.strip() for label in labels_text.split(',')]
        if len(labels) != state.num_classes:
            return f"❌ Forneça {state.num_classes} rótulos separados por vírgula"
        
        state.class_labels = labels
        return f"✅ Rótulos definidos: {', '.join(state.class_labels)}"
    except Exception as e:
        return f"❌ Erro: {str(e)}"

def add_image_to_queue(image):
    """Adiciona imagem à fila"""
    if image is None:
        return "❌ Selecione uma imagem", 0
    
    state.image_queue.append(image)
    return f"✅ Imagem adicionada à fila. Total: {len(state.image_queue)}", len(state.image_queue)

def save_images_to_class(class_id, clear_queue=True):
    """Salva todas as imagens da fila para uma classe"""
    try:
        if not state.image_queue:
            return "❌ Nenhuma imagem na fila"
        
        if not state.class_dirs:
            return "❌ Configure as classes primeiro"
        
        class_idx = int(class_id)
        if class_idx >= len(state.class_dirs):
            return "❌ Classe inválida"
        
        class_dir = state.class_dirs[class_idx]
        count = 0
        
        for i, image in enumerate(state.image_queue):
            try:
                import time
                filename = f"img_{int(time.time())}_{i}.jpg"
                filepath = os.path.join(class_dir, filename)
                image.save(filepath)
                count += 1
            except Exception as e:
                print(f"Erro salvando imagem {i}: {e}")
        
        if clear_queue:
            state.image_queue = []
        
        class_name = state.class_labels[class_idx]
        return f"✅ {count} imagens salvas em '{class_name}'"
    except Exception as e:
        return f"❌ Erro: {str(e)}"

def clear_image_queue():
    """Limpa a fila de imagens"""
    state.image_queue = []
    return "✅ Fila limpa", 0

def prepare_data(batch_size):
    """Prepara dados para treinamento"""
    try:
        if not state.dataset_path:
            return "❌ Configure as classes primeiro"
        
        transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        
        dataset = datasets.ImageFolder(state.dataset_path, transform=transform)
        
        if len(dataset) < 6:
            return f"❌ Poucas imagens ({len(dataset)}). Mínimo: 6"
        
        # Divisão: 70% treino, 20% val, 10% 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],
            generator=torch.Generator().manual_seed(42)
        )
        
        batch_size = max(1, min(int(batch_size), 32))
        state.train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
        state.val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
        state.test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
        
        return f"✅ Dados preparados:\n• Treino: {train_size}\n• Validação: {val_size}\n• Teste: {test_size}\n• Batch size: {batch_size}"
    except Exception as e:
        return f"❌ Erro: {str(e)}"

def train_model(model_name, epochs, lr):
    """Treina o modelo"""
    try:
        if state.train_loader is None:
            return "❌ Prepare os dados primeiro"
        
        # Carregar modelo
        state.model = MODELS[model_name](pretrained=True)
        
        # Adaptar última camada
        if hasattr(state.model, 'fc'):
            state.model.fc = nn.Linear(state.model.fc.in_features, state.num_classes)
        elif hasattr(state.model, 'classifier'):
            if isinstance(state.model.classifier, nn.Sequential):
                state.model.classifier[-1] = nn.Linear(state.model.classifier[-1].in_features, state.num_classes)
        
        state.model = state.model.to(device)
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.Adam(state.model.parameters(), lr=float(lr))
        
        results = [f"🚀 Treinando {model_name}"]
        state.model.train()
        
        for epoch in range(int(epochs)):
            running_loss = 0.0
            correct = 0
            total = 0
            
            for inputs, labels in state.train_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                
                optimizer.zero_grad()
                outputs = state.model(inputs)
                loss = criterion(outputs, labels)
                loss.backward()
                optimizer.step()
                
                running_loss += loss.item()
                _, predicted = torch.max(outputs, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
            
            epoch_loss = running_loss / len(state.train_loader)
            epoch_acc = 100. * correct / total
            results.append(f"Época {epoch+1}: Loss={epoch_loss:.4f}, Acc={epoch_acc:.2f}%")
        
        results.append("✅ Treinamento concluído!")
        return "\n".join(results)
    except Exception as e:
        return f"❌ Erro: {str(e)}"

def evaluate_model():
    """Avalia o modelo"""
    try:
        if state.model is None or state.test_loader is None:
            return "❌ Modelo/dados não disponíveis"
        
        state.model.eval()
        all_preds = []
        all_labels = []
        
        with torch.no_grad():
            for inputs, labels in state.test_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = state.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, target_names=state.class_labels, zero_division=0)
        return f"📊 RELATÓRIO DE AVALIAÇÃO:\n\n{report}"
    except Exception as e:
        return f"❌ Erro: {str(e)}"

def generate_confusion_matrix():
    """Gera matriz de confusão"""
    try:
        if state.model is None or state.test_loader is None:
            return None
        
        state.model.eval()
        all_preds = []
        all_labels = []
        
        with torch.no_grad():
            for inputs, labels in state.test_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = state.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)
        
        plt.figure(figsize=(8, 6))
        sns.heatmap(cm, annot=True, fmt="d", cmap="Blues",
                    xticklabels=state.class_labels,
                    yticklabels=state.class_labels)
        plt.xlabel('Predições')
        plt.ylabel('Valores Reais')
        plt.title('Matriz de Confusão')
        plt.tight_layout()
        
        temp_path = tempfile.NamedTemporaryFile(suffix='.png', delete=False).name
        plt.savefig(temp_path, dpi=150, bbox_inches='tight')
        plt.close()
        
        return temp_path
    except Exception as e:
        return None

def predict_image(image):
    """Prediz uma única imagem"""
    try:
        if state.model is None:
            return "❌ Treine o modelo primeiro"
        
        if image is None:
            return "❌ Selecione uma imagem"
        
        transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        
        img_tensor = transform(image).unsqueeze(0).to(device)
        
        state.model.eval()
        with torch.no_grad():
            outputs = state.model(img_tensor)
            probs = torch.nn.functional.softmax(outputs[0], dim=0)
            _, predicted = torch.max(outputs, 1)
            
            class_id = predicted.item()
            confidence = probs[class_id].item() * 100
            class_name = state.class_labels[class_id]
            
            return f"🎯 Predição: {class_name}\n📊 Confiança: {confidence:.2f}%"
    except Exception as e:
        return f"❌ Erro: {str(e)}"

# Interface usando componentes mais antigos/estáveis
def create_interface():
    with gr.Blocks(title="🖼️ Classificador Completo") as demo:
        
        gr.Markdown("# 🖼️ Sistema de Classificação de Imagens Completo")
        
        # Configuração
        with gr.Group():
            gr.Markdown("## 1️⃣ Configuração")
            with gr.Row():
                num_classes = gr.Number(label="Número de Classes (2-5)", value=2, precision=0)
                setup_btn = gr.Button("🔧 Configurar")
            setup_status = gr.Textbox(label="Status")
            
            labels_input = gr.Textbox(label="Rótulos (separados por vírgula)", value="gato,cachorro")
            labels_btn = gr.Button("🏷️ Definir Rótulos")
            labels_status = gr.Textbox(label="Status dos Rótulos")
        
        # Upload de Imagens
        with gr.Group():
            gr.Markdown("## 2️⃣ Upload de Imagens")
            with gr.Row():
                upload_image = gr.Image(type="pil", label="Upload de Imagem")
                with gr.Column():
                    add_btn = gr.Button("➕ Adicionar à Fila")
                    queue_status = gr.Textbox(label="Fila de Imagens")
                    queue_count = gr.Number(label="Total na Fila", value=0)
            
            with gr.Row():
                class_id = gr.Number(label="Classe (0, 1, 2...)", value=0, precision=0)
                save_btn = gr.Button("💾 Salvar Fila na Classe", variant="primary")
                clear_btn = gr.Button("🗑️ Limpar Fila")
            save_status = gr.Textbox(label="Status do Upload")
        
        # Treinamento
        with gr.Group():
            gr.Markdown("## 3️⃣ Preparação e Treinamento")
            batch_size = gr.Number(label="Batch Size", value=8, precision=0)
            prepare_btn = gr.Button("⚙️ Preparar Dados", variant="primary")
            prepare_status = gr.Textbox(label="Status da Preparação", lines=4)
            
            with gr.Row():
                model_choice = gr.Dropdown(choices=list(MODELS.keys()), value="MobileNetV2", label="Modelo")
                epochs = gr.Number(label="Épocas", value=5, precision=0)
                learning_rate = gr.Number(label="Learning Rate", value=0.001)
            
            train_btn = gr.Button("🚀 Treinar Modelo", variant="primary")
            train_status = gr.Textbox(label="Status do Treinamento", lines=8)
        
        # Avaliação
        with gr.Group():
            gr.Markdown("## 4️⃣ Avaliação")
            with gr.Row():
                eval_btn = gr.Button("📊 Avaliar Modelo", variant="primary")
                matrix_btn = gr.Button("📈 Matriz de Confusão")
            
            eval_results = gr.Textbox(label="Relatório de Avaliação", lines=12)
            confusion_plot = gr.Image(label="Matriz de Confusão")
        
        # Predição
        with gr.Group():
            gr.Markdown("## 5️⃣ Predição")
            predict_img = gr.Image(type="pil", label="Imagem para Predição")
            predict_btn = gr.Button("🔮 Predizer", variant="primary")
            predict_result = gr.Textbox(label="Resultado da Predição", lines=3)
        
        # Conectar eventos
        setup_btn.click(setup_classes, [num_classes], [setup_status])
        labels_btn.click(set_class_labels, [labels_input], [labels_status])
        
        add_btn.click(add_image_to_queue, [upload_image], [queue_status, queue_count])
        save_btn.click(save_images_to_class, [class_id], [save_status])
        clear_btn.click(clear_image_queue, outputs=[queue_status, queue_count])
        
        prepare_btn.click(prepare_data, [batch_size], [prepare_status])
        train_btn.click(train_model, [model_choice, epochs, learning_rate], [train_status])
        
        eval_btn.click(evaluate_model, outputs=[eval_results])
        matrix_btn.click(generate_confusion_matrix, outputs=[confusion_plot])
        
        predict_btn.click(predict_image, [predict_img], [predict_result])
    
    return demo

if __name__ == "__main__":
    print("🎯 Criando interface...")
    demo = create_interface()
    print("🚀 Iniciando aplicação...")
    demo.launch(server_name="0.0.0.0", server_port=7860, show_api=False)