import torch
import torch.nn as nn
import numpy as np
import pandas as pd
from torch.utils.data import Dataset

class WakeWordDataset(Dataset):
    def __init__(self, features, labels):
        self.features = features
        self.labels = labels

    def __len__(self):
        return len(self.features)

    def __getitem__(self, idx):
        feature = torch.tensor(self.features[idx], dtype=torch.float32)
        label = torch.tensor(self.labels[idx], dtype=torch.long)
        return feature, label

class CNNNetwork(nn.Module):
    def __init__(self):
        super(CNNNetwork, self).__init__()
        self.network = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Flatten(),
            nn.Linear(128 * 4 * 4, 128),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(128, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.network(x)

def preprocess_training_data(training_directories, background_directory, model_path, update_progress_callback):
    import librosa
    data_path_dict = {}
    data_path_dict[0] = list_files_in_directory(background_directory, extension='.wav')
    for idx, directory in enumerate(training_directories):
        data_path_dict[idx+1] = list_files_in_directory(directory, extension='.wav')

    all_data = []
    total_files = sum(len(files) for files in data_path_dict.values())
    processed_files = 0

    for class_label, list_of_files in data_path_dict.items():
        for single_file in list_of_files:
            try:
                audio, sample_rate = librosa.load(single_file)
                mfcc = librosa.feature.mfcc(y=audio, sr=sample_rate, n_mfcc=40)
                mfcc_processed = np.mean(mfcc.T, axis=0)
                all_data.append([mfcc_processed, class_label])
                processed_files += 1
                update_progress_callback(processed_files / total_files * 100)
            except Exception as e:
                print(f"Exception: {str(e)}")

    df = pd.DataFrame(all_data, columns=["feature", "class_label"])
    df.to_pickle(model_path + "audio_data.csv")

def train_model(model, train_loader, optimizer, criterion, epoch, total_epochs):
    model.train()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target.float().unsqueeze(1))
        loss.backward()
        optimizer.step()

        if batch_idx % 10 == 0:
            print(f'Train Epoch: {epoch}/{total_epochs} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')

def optimize_graph(model, device, optimizer, criterion, dataloader):
    model.train()
    total_loss = 0
    correct = 0

    for data, target in dataloader:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target.float().unsqueeze(1))
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        pred = torch.round(output)
        correct += pred.eq(target.float().unsqueeze(1)).sum().item()

    accuracy = correct / len(dataloader.dataset)
    return total_loss / len(dataloader), accuracy