import torch
import torch.nn as nn
from torch.utils.data import Dataset
from utils.file_utils import list_files_in_directory
import numpy as np
import wave
import logging
from pydub import AudioSegment
import matplotlib.pyplot as plt
import pandas as pd
import librosa.display

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):
        return self.features[idx], self.labels[idx]

class CNNNetwork(nn.Module):
    def __init__(self):
        super(CNNNetwork, self).__init__()
        self.network = nn.Sequential(
            nn.Conv1d(in_channels=1, out_channels=32, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2, stride=2),
            nn.Conv1d(in_channels=32, out_channels=64, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2, stride=2),
            nn.Flatten(),
            nn.Linear(64 * 10, 128),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(128, 1)
        )

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

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

    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, device):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

    for batch_X, batch_y in train_loader:
        batch_X, batch_y = batch_X.to(device), batch_y.to(device)
        optimizer.zero_grad()
        outputs = model(batch_X.float())
        loss = criterion(outputs, batch_y.unsqueeze(1).float())
        loss.backward()
        optimizer.step()
        global current_loss, current_accuracy, total_batches
        current_loss = loss.item()
        current_accuracy = (outputs.argmax(dim=1) == batch_y).float().mean().item()

        running_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += batch_y.size(0)
        correct += (predicted == batch_y).sum().item()

    epoch_loss = running_loss / len(train_loader)
    epoch_acc = 100 * correct / total

    print(f'Epoch [{epoch + 1}/{total_epochs}], Loss: {epoch_loss:.4f}, Accuracy: {epoch_acc:.2f}%')

def optimize_graph(model, device, optimizer, criterion, dataloader):
    pass