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')  # Use non-interactive backend
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")

# Configuração
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"🖥️ Device: {device}")

# Modelos disponíveis
MODELS = {
    'ResNet18': models.resnet18,
    '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 = []
        self.num_classes = 2

state = AppState()

def setup_classes(num_classes_value):
    """Configura 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)]
        
        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"✅ Criados {state.num_classes} diretórios"
    except Exception as e:
        return f"❌ Erro: {str(e)}"

def set_class_labels(labels_text):
    """Define rótulos das classes (separados por vírgula)"""
    try:
        labels = [label.strip() for label in labels_text.split(',') if label.strip()]
        
        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: {', '.join(state.class_labels)}"
    except Exception as e:
        return f"❌ Erro: {str(e)}"

def upload_images(class_id, images):
    """Upload de imagens"""
    try:
        if not images:
            return "❌ Selecione imagens"
        
        class_idx = int(class_id)
        if class_idx >= len(state.class_dirs):
            return f"❌ Classe inválida"
        
        class_dir = state.class_dirs[class_idx]
        count = 0
        
        for image in images:
            if image is not None:
                shutil.copy2(image, class_dir)
                count += 1
        
        class_name = state.class_labels[class_idx]
        return f"✅ {count} imagens → {class_name}"
    except Exception as e:
        return f"❌ Erro: {str(e)}"

def prepare_data(batch_size):
    """Prepara dados"""
    try:
        if not state.dataset_path:
            return "❌ Configure 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)
        )
        
        state.train_loader = DataLoader(train_dataset, batch_size=int(batch_size), shuffle=True)
        state.val_loader = DataLoader(val_dataset, batch_size=int(batch_size), shuffle=False)
        state.test_loader = DataLoader(test_dataset, batch_size=int(batch_size), shuffle=False)
        
        return f"✅ Dados preparados:\n• Treino: {train_size}\n• Validação: {val_size}\n• Teste: {test_size}"
    except Exception as e:
        return f"❌ Erro: {str(e)}"

def train_model(model_name, epochs, lr):
    """Treina modelo"""
    try:
        if state.train_loader is None:
            return "❌ Prepare os dados primeiro"
        
        # Carregar modelo
        state.model = MODELS[model_name](pretrained=True)
        
        # Adaptar camada final
        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))
        
        state.model.train()
        results = [f"🚀 Treinando {model_name}"]
        
        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 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:\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:
        print(f"Erro matriz confusão: {e}")
        return None

def predict_images(images):
    """Prediz imagens"""
    try:
        if state.model is None:
            return "❌ Treine o modelo primeiro"
        
        if not images:
            return "❌ Selecione imagens"
        
        transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        
        state.model.eval()
        results = []
        
        for image_path in images:
            if image_path:
                image = Image.open(image_path).convert('RGB')
                img_tensor = transform(image).unsqueeze(0).to(device)
                
                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]
                    
                    results.append(f"📸 {os.path.basename(image_path)}")
                    results.append(f"   🎯 {class_name}")
                    results.append(f"   📊 {confidence:.2f}%")
                    results.append("-" * 30)
        
        return "\n".join(results) if results else "❌ Nenhuma predição"
    except Exception as e:
        return f"❌ Erro: {str(e)}"

# Interface
with gr.Blocks(title="🖼️ Classificador", theme=gr.themes.Soft()) as demo:
    
    gr.Markdown("""
    # 🖼️ Sistema de Classificação de Imagens
    **Instruções:** Configure → Upload → Treine → Avalie → Prediga
    """)
    
    with gr.Tab("1️⃣ Configuração"):
        gr.Markdown("### 🎯 Configurar Classes")
        
        num_classes = gr.Slider(
            minimum=2, maximum=5, value=2, step=1,
            label="Número de Classes"
        )
        
        setup_btn = gr.Button("🔧 Configurar", variant="primary")
        setup_status = gr.Textbox(label="Status", lines=2)
        
        gr.Markdown("### 🏷️ Definir Rótulos")
        
        labels_input = gr.Textbox(
            label="Rótulos (separados por vírgula)",
            placeholder="gato, cachorro",
            value="gato, cachorro"
        )
        
        labels_btn = gr.Button("🏷️ Definir Rótulos")
        labels_status = gr.Textbox(label="Status Rótulos")
        
    with gr.Tab("2️⃣ Upload"):
        gr.Markdown("### 📤 Upload de Imagens")
        
        class_selector = gr.Slider(
            minimum=0, maximum=1, value=0, step=1,
            label="Classe (0, 1, 2...)"
        )
        
        images_upload = gr.File(
            label="Imagens",
            file_count="multiple",
            file_types=["image"]
        )
        
        upload_btn = gr.Button("📤 Upload", variant="primary")
        upload_status = gr.Textbox(label="Status")
        
    with gr.Tab("3️⃣ Treinamento"):
        gr.Markdown("### ⚙️ Preparar Dados")
        
        batch_size = gr.Slider(1, 32, 8, step=1, label="Batch Size")
        prepare_btn = gr.Button("⚙️ Preparar", variant="primary")
        prepare_status = gr.Textbox(label="Status", lines=4)
        
        gr.Markdown("### 🚀 Treinar Modelo")
        
        with gr.Row():
            model_choice = gr.Radio(
                choices=list(MODELS.keys()),
                value="MobileNetV2",
                label="Modelo"
            )
            epochs = gr.Slider(1, 10, 3, step=1, label="Épocas")
            learning_rate = gr.Slider(0.0001, 0.01, 0.001, label="Learning Rate")
        
        train_btn = gr.Button("🚀 Treinar", variant="primary")
        train_status = gr.Textbox(label="Status Treinamento", lines=8)
        
    with gr.Tab("4️⃣ Avaliação"):
        gr.Markdown("### 📊 Avaliar Modelo")
        
        with gr.Row():
            eval_btn = gr.Button("📊 Avaliar", variant="primary")
            matrix_btn = gr.Button("📈 Matriz Confusão")
        
        eval_results = gr.Textbox(label="Relatório", lines=12)
        confusion_matrix_plot = gr.Image(label="Matriz de Confusão")
        
    with gr.Tab("5️⃣ Predição"):
        gr.Markdown("### 🔮 Predizer Novas Imagens")
        
        predict_images_input = gr.File(
            label="Imagens para Predição",
            file_count="multiple",
            file_types=["image"]
        )
        
        predict_btn = gr.Button("🔮 Predizer", variant="primary")
        predict_results = gr.Textbox(label="Resultados", lines=10)
    
    # Conectar eventos
    setup_btn.click(setup_classes, [num_classes], [setup_status])
    labels_btn.click(set_class_labels, [labels_input], [labels_status])
    upload_btn.click(upload_images, [class_selector, images_upload], [upload_status])
    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, [], [eval_results])
    matrix_btn.click(generate_confusion_matrix, [], [confusion_matrix_plot])
    predict_btn.click(predict_images, [predict_images_input], [predict_results])

if __name__ == "__main__":
    demo.launch()