class Trainer:
    """
    The `Trainer` class encapsulates the training process for a Wasserstein Generative Adversarial Network (WGAN) composed of a generator and a critic. It manages the training loop, loss computations, parameter updates, and enforces the Lipschitz constraint through weight clipping.

    Attributes:
        latent_space (int): Dimensionality of the generator's input latent space.
        epochs (int): Total number of training epochs.
        learning_rate (float): Learning rate for the Adam optimizers.
        beta1 (float): Beta1 hyperparameter for the Adam optimizer.
        beta2 (float): Beta2 hyperparameter for the Adam optimizer.
        generator (Generator): The generator model of the WGAN.
        critic (Critic): The critic model of the WGAN.
        dataloader (DataLoader): DataLoader providing the training data.
        optimizer_generator (optim.Optimizer): Optimizer for updating the generator's weights.
        optimizer_critic (optim.Optimizer): Optimizer for updating the critic's weights.
        critic_loss (list): List to record the critic's loss after each epoch.
        generator_loss (list): List to record the generator's loss after each epoch.
        device (str): The device ('cuda', 'mps', or 'cpu') on which the models will run.
        n_critic_step (int): Number of critic updates per generator update.
        clamp_value (float): The clamp value for the critic's weight clipping.

    Methods:
        connect_gpu(generator, critic, device): Assigns the generator and critic models to the specified device.
        saved_checkpoints(model, epoch): Saves a checkpoint of the model at the specified epoch.
        train_critic(real_samples, fake_samples): Trains the critic model for one batch of data.
        train_generator(generated_samples): Trains the generator model for one batch of data.
        train_WGAN(): Conducts the training loop for the WGAN.
    """

    def __init__(
        self,
        latent_space=100,
        epochs=100,
        lr=0.00005,
        beta1=0.5,
        beta2=0.999,
        device="cpu",
        n_critic_step=5,
        display=True,
    ):
        """
        Initializes the Trainer object with the specified configuration and sets up the neural network models, dataloader, loss function, and optimizers.

        Args:
            latent_space (int): Size of the latent space (input vector for the generator).
            epochs (int): Number of epochs for training the models.
            lr (float): Learning rate for the Adam optimizers.
            beta1 (float): Beta1 hyperparameter for the Adam optimizer.
            beta2 (float): Beta2 hyperparameter for the Adam optimizer.
            device (str): The device ('cuda', 'mps', or 'cpu') on which the models will run.
            n_critic_step (int): Number of critic updates per generator update.
        """
        self.latent_space = latent_space
        self.epochs = epochs
        self.learning_rate = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.device = device
        self.clamp_value = 0.01
        self.n_critic_step = n_critic_step
        self.display = display

        self.device_cuda = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.device_mps = torch.device(
            "mps" if torch.backends.mps.is_available() else "cpu"
        )
        self.device_cpu = torch.device("cpu")

        self.generator = Generator()
        self.critic = Critic()

        self.generator, self.critic, self.device = self.connect_gpu(
            generator=self.generator, critic=self.critic, device=self.device
        )

        self.generator.apply(weight_initialization)
        self.critic.apply(weight_initialization)

        try:
            self.dataloader = joblib.load(filename="./data/processed/dataloader.pkl")
        except Exception as e:
            logging.exception("Dataloader is not transformed from pickle".capitalize())

        self.optimizer_generator = optim.RMSprop(
            params=self.generator.parameters(),
            lr=self.learning_rate,
        )
        self.optimizer_critic = optim.RMSprop(
            params=self.critic.parameters(),
            lr=self.learning_rate,
        )

        self.critic_loss = list()
        self.generator_loss = list()

    def connect_gpu(self, generator, critic, device):
        """
        Connects the generator and critic models to the specified computing device.

        Args:
            generator (Generator): The generator model.
            critic (Critic): The critic model.
            device (str): The target device ('cuda', 'mps', or 'cpu').

        Returns:
            tuple: The generator, critic, and device after assignment to the target device.
        """
        if device == "cuda":
            generator = generator.to(self.device_cuda)
            critic = critic.to(self.device_cuda)
        elif device == "mps":
            generator = generator.to(self.device_mps)
            critic = critic.to(self.device_mps)
        else:  # Default to CPU if neither 'cuda' nor 'mps'
            generator = generator.to(self.device_cpu)
            critic = critic.to(self.device_cpu)

        return generator, critic, device

    def saved_checkpoints(self, model=None, epoch=None):
        """
        Saves a checkpoint of the given model at the specified epoch.

        Args:
            model (nn.Module): The model to be saved.
            epoch (int): The current epoch number for naming the saved file.

        Side Effects:
            Saves the model's state dictionary to the file system.
        """
        try:
            torch.save(
                model.state_dict(), f"./models/checkpoints/generator_{epoch}.pth"
            )
            logging.info(f"Checkpoint saved for epoch {epoch}")
        except Exception as e:
            logging.error(f"Error saving checkpoint at epoch {epoch}: {e}")
            raise e

    def train_critic(self, real_samples, fake_samples):
        """
        Trains the critic model for one batch of data.

        Args:
            real_samples (Tensor): Real samples from the dataset.
            fake_samples (Tensor): Fake samples generated by the generator.

        Returns:
            float: The total loss for the critic for the current batch.

        Side Effects:
            Updates the weights of the critic model.
        """
        real_predict = self.critic(real_samples)
        fake_predict = self.critic(fake_samples.detach())

        total_loss = -torch.mean(real_predict) + torch.mean(fake_predict)

        self.optimizer_critic.zero_grad()
        total_loss.backward()
        self.optimizer_critic.step()

        return total_loss.item()

    def train_generator(self, generated_samples):
        """
        Trains the generator model for one batch of data.

        Args:
            generated_samples (Tensor): Samples generated by the generator.

        Returns:
            float: The loss for the generator for the current batch.

        Side Effects:
            Updates the weights of the generator model.
        """
        generated_loss = -torch.mean(self.critic(generated_samples))

        self.optimizer_generator.zero_grad()
        generated_loss.backward()
        self.optimizer_generator.step()

        return generated_loss.item()

    def display_performance(self, **kwargs):
        print(f"Epoch [{kwargs['epoch'] + 1}/{kwargs['epochs']}] Completed")
        print(
            f"Average Critic Loss: {kwargs['avg_critic_loss']:.4f},\
            Average Generator Loss: {kwargs['avg_generator_loss']:.4f}"
        )

    def train_WGAN(self):
        """
        Conducts the training loop for the Wasserstein Generative Adversarial Network (WGAN).
        The loop iterates over the dataset, trains the critic and generator in alternation,
        and records the loss for each epoch.

        Process:
        - For each epoch:
            - For each batch in the dataloader:
                - Train the critic using both real and fake data.
                - Generate new fake samples and train the generator.
                - Record and accumulate the loss for both the critic and generator.
        - After each epoch, print the average losses and save the generator's state as a checkpoint.

        Side Effects:
        - Updates the weights of both the critic and generator models.
        - Appends the average loss of each epoch to the respective loss lists (`critic_loss`, `generator_loss`).
        - Saves the generator's state after each epoch.
        - Prints the progress and average losses to the console.

        Error Handling:
        - If the model checkpoint cannot be saved, an exception is raised.
        """
        for epoch in range(self.epochs):
            c_loss = []
            g_loss = []
            for index, (real_samples, _) in enumerate(self.dataloader):
                real_samples = real_samples.to(self.device)
                batch_size = real_samples.shape[0]

                noise_samples = torch.randn(batch_size, self.latent_space).to(
                    self.device
                )
                fake_samples = self.generator(noise_samples)

                D_loss = self.train_critic(
                    real_samples=real_samples, fake_samples=fake_samples
                )

                # Clamp the weights of the critic
                for params in self.critic.parameters():
                    params.data.clamp_(-self.clamp_value, self.clamp_value)

                if (index + 1) % self.n_critic_step == 0:
                    generated_samples = self.generator(noise_samples)
                    G_loss = self.train_generator(generated_samples=generated_samples)

                    c_loss.append(D_loss)
                    g_loss.append(G_loss)

            avg_critic_loss = np.mean(c_loss)
            avg_generator_loss = np.mean(g_loss)
            self.critic_loss.append(avg_critic_loss)
            self.generator_loss.append(avg_generator_loss)

            if self.display:
                self.display_performance(
                    epochs=self.epochs,
                    epoch=epoch,
                    avg_critic_loss=avg_critic_loss,
                    avg_generator_loss=avg_generator_loss,
                )
            else:
                raise Exception("Please set display to True".capitalize)

            logging.info(f"Epoch [{epoch + 1}/{self.epochs}] Completed")
            logging.info(
                f"Average Critic Loss: {avg_critic_loss:.4f}, Average Generator Loss: {avg_generator_loss:.4f}"
            )

            self.saved_checkpoints(model=self.generator, epoch=epoch + 1)