import os
import time
import json
import torch
import torch.optim as optim
from torch.utils.data import DataLoader
from utils.training_interface import TrainerInterface
from utils.file_utils import list_files_in_directory
from utils.pytorch_utils import WakeWordDataset, CNNNetwork, preprocess_training_data, train_model, optimize_graph

# Global variables for tracking progress
current_progress = 0
current_loss = 0.0
current_accuracy = 0.0
total_epochs = 0
total_batches = 0

class pytorch_adv_trainer(TrainerInterface):
    def __init__(self, learning_rate=0.001):
        self.model_path = None
        self.training_directories = []
        self.sample_directory = None
        self.background_directory = None
        self.dropout_rate = 0.5
        self.epochs = 10
        self.batch_size = 32
        self.learning_rate = learning_rate
        self.model = None
        self.optimizer = None
        self.criterion = torch.nn.BCEWithLogitsLoss()
        self.current_epoch = 0
        self.stop_training = False
        self.training_status = "Not started"
        self.train_loader = None
        self.val_loader = None

    def set_model_path(self, model_path):
        self.model_path = model_path
        self.model = self.load_model()
        self.optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate)

    def set_training_directories(self, directories):
        self.training_directories = directories

    def set_sample_directory(self, directory):
        self.sample_directory = directory

    def set_background_directory(self, directory):
        self.background_directory = directory

    def set_dropout_rate(self, rate):
        self.dropout_rate = rate

    def set_epochs(self, epochs):
        self.epochs = epochs

    def set_batch_size(self, batch_size):
        self.batch_size = batch_size

    def get_training_data(self):
        return self.training_directories

    def get_model_path(self):
        return self.model_path

    def get_epochs(self):
        return self.epochs

    def get_batch_size(self):
        return self.batch_size

    def get_dropout_rate(self):
        return self.dropout_rate

    def get_training_directories(self):
        return self.training_directories

    def get_sample_directory(self):
        return self.sample_directory

    def get_background_directory(self):
        return self.background_directory

    def set_training_data(self, features, labels):
        X_train, X_val, y_train, y_val = train_test_split(features, labels, test_size=0.2, random_state=42)
        self.train_loader = DataLoader(WakeWordDataset(X_train, y_train), batch_size=self.batch_size, shuffle=True)
        self.val_loader = DataLoader(WakeWordDataset(X_val, y_val), batch_size=self.batch_size, shuffle=False)

    def load_model(self):
        model = CNNNetwork()
        if os.path.exists(self.model_path + "hotword_model.pth"):
            model.load_state_dict(torch.load(self.model_path + "hotword_model.pth"))
        return model

    def start(self):
        self.training_status = "Training"
        self.stop_training = False
        self.run_training()

    def start_new(self):
        self.training_status = "Training"
        self.stop_training = False
        self.model = self.load_model()
        self.current_epoch = 0
        self.run_training()

    def pause(self):
        self.training_status = "Paused"
        self.stop_training = True

    def resume(self):
        self.training_status = "Training"
        self.stop_training = False
        self.run_training()

    def stop(self):
        self.training_status = "Stopped"
        self.stop_training = True

    def preprocess_training_data(self, update_progress_callback):
        preprocess_training_data(self.training_directories, self.background_directory, self.model_path, update_progress_callback)

    def run_training(self):
        global current_progress, current_loss, current_accuracy, total_epochs, total_batches
        self.preprocess_training_data(lambda progress: setattr(self, 'current_progress', progress))

        df = pd.read_pickle(self.model_path + "audio_data.csv")
        X = df["feature"].values
        X = np.concatenate(X, axis=0).reshape(len(X), 40)
        y = np.array(df["class_label"].tolist())
        self.set_training_data(X, y)
        
        total_batches = int(np.ceil(len(self.train_loader.dataset) / self.batch_size))
        total_epochs = self.epochs

        for epoch in range(self.current_epoch, self.epochs):
            if self.stop_training:
                break
            self.current_epoch = epoch
            train_model(self.model, self.train_loader, self.optimizer, self.criterion, epoch, self.epochs)

            current_progress = (epoch + 1) / self.epochs * 100
            print(f"Epoch {epoch + 1}/{self.epochs}")

        if not self.stop_training:
            torch.save(self.model.state_dict(), self.model_path + "hotword_model.pth")
            self.evaluate_model(X, y)
            self.training_status = "Finished"

    def evaluate_model(self, X_test, y_test):
        self.model.eval()
        with torch.no_grad():
            X_test_tensor = torch.FloatTensor(X_test)
            y_test_tensor = torch.LongTensor(y_test)
            outputs = self.model(X_test_tensor)
            _, predicted = torch.max(outputs, 1)
            accuracy = (predicted == y_test_tensor).sum().item() / len(y_test_tensor)
            print(f"Test Accuracy: {accuracy * 100:.2f}%")

    def get_training_info(self):
        return {
            "status": self.training_status,
            "current_epoch": self.current_epoch,
            "total_epochs": self.epochs,
            "loss": current_loss,
            "accuracy": current_accuracy
        }