import os
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, models, transforms
from torch.utils.data import DataLoader

def train_audio_model():
    print("Configuration du Device...")
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(f"Utilisation de l'appareil: {device}")

    # Transformations
    data_transforms = {
        'train': transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
        'test': transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
    }

    data_dir = 'data/audio/spectrograms'
    
    class SafeImageFolder(datasets.ImageFolder):
        def find_classes(self, directory):
            classes = sorted(entry.name for entry in os.scandir(directory) if entry.is_dir() and not entry.name.startswith('.'))
            if not classes:
                raise FileNotFoundError(f"Couldn't find any class folder in {directory}.")
            class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
            return classes, class_to_idx

    print("Préparation des DataLoader pour le modèle Audio...")
    image_datasets = {
        x: SafeImageFolder(os.path.join(data_dir, x), data_transforms[x])
        for x in ['train', 'test']
    }
    
    dataloaders = {
        x: DataLoader(image_datasets[x], batch_size=16, shuffle=(True if x == 'train' else False), num_workers=0)
        for x in ['train', 'test']
    }
    
    dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'test']}
    
    print("Initialisation du modèle torchvision.models.resnet18(pretrained=True)...")
    model = models.resnet18(pretrained=True)

    for param in model.parameters():
        param.requires_grad = False

    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, 2)
    model = model.to(device)

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

    num_epochs = 3
    print(f"\nDémarrage de l'entraînement sur {num_epochs} epochs...")

    for epoch in range(num_epochs):
        print(f"Epoch {epoch+1}/{num_epochs}")
        print("-" * 10)

        for phase in ['train', 'test']:
            if phase == 'train':
                model.train()
            else:
                model.eval()

            running_loss = 0.0
            running_corrects = 0

            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                optimizer.zero_grad()

                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]

            print(f"{phase.capitalize()} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}")

    print("\nÉvaluation finale terminée.")
    model.eval()
    final_corrects = 0
    with torch.no_grad():
        for inputs, labels in dataloaders['test']:
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            final_corrects += torch.sum(preds == labels.data)
            
    final_acc = final_corrects.double() / dataset_sizes['test']
    print("=" * 40)
    print(f"RÉSULTATS DE L'ÉVALUATION AUDIO (Test)")
    print("=" * 40)
    print(f"Accuracy finale sur le jet de test : {final_acc * 100:.2f} %")
    print("=" * 40)

    print("\nSauvegarde du modèle...\n")
    os.makedirs('models', exist_ok=True)
    save_path = 'models/audio_model.pth'
    
    torch.save(model.state_dict(), save_path)
    print(f"Modèle sauvegardé dans : {save_path}")
    print("Terminé avec succès !")

if __name__ == '__main__':
    train_audio_model()