class Test:
    """
    This class is designed for testing and visualizing synthetic images generated by the Generator model.

    Attributes:
        image_size (int): The size of the generated images (height and width).
        num_samples (int): The number of synthetic images to generate.
        latent_space (int): The dimensionality of the latent space for the generator.
        device (torch.device): The device (CPU, GPU, etc.) on which the generator model will be loaded and run.
        generator (Generator): The generator model for synthetic image generation.

    Methods:
        get_the_best_model: Loads the best performing generator model from a predefined directory.
        test: Generates and visualizes synthetic images using the generator model.
    """

    def __init__(self, num_samples=20, image_size=64, latent_space=100, device="cpu"):
        """
        Initializes the Test class with the specified configuration.

        Args:
            num_samples (int): Number of synthetic images to generate for testing.
            image_size (int): The height and width of the images to generate (assumed square).
            latent_space (int): Dimensionality of the latent space for the generator.
            device (str): The device to run the generator model on ('cpu', 'cuda', 'mps', etc.).
        """
        self.image_size = image_size
        self.num_samples = num_samples
        self.latent_space = latent_space
        self.device = device_init(device=device)
        self.generator = Generator(
            image_size=self.image_size, latent_space=self.latent_space
        ).to(self.device)

    def get_the_best_model(self):
        """
        Searches for and loads the best performing generator model from a specified directory.

        Raises:
            Exception: If no model is found in the specified directory.
        """
        model_path = "./models/best_model/"

        if os.path.exists(model_path):
            logging.info("Best model found at {}".format(model_path).title())

            best_model = os.listdir(model_path)[0]
            load_state_dict = torch.load(os.path.join(model_path, best_model))
            self.generator.load_state_dict(load_state_dict)

        else:
            raise Exception("No Best Model Found".title())

    def test(self):
        """
        Generates synthetic images using the best generator model and visualizes them.

        This method also attempts to save the generated images to a file and logs any exceptions encountered.
        """
        try:
            self.get_the_best_model()
            random_noise = torch.randn(self.num_samples, self.latent_space, 1, 1).to(
                self.device
            )
        except Exception as e:
            logging.exception("Exception caught: {}".format(e))
            print("Exception caught: {}".format(e))
        else:
            plt.figure(figsize=(10, 10))
            synthetic_images = self.generator(random_noise)

            for index, image in enumerate(synthetic_images):
                plt.subplot(4, 5, index + 1)
                image_to_plot = image.cpu().detach().permute(1, 2, 0).numpy()
                image_to_plot = (image_to_plot - image_to_plot.min()) / (
                    image_to_plot.max() - image_to_plot.min()
                )
                plt.imshow(image_to_plot)
                plt.axis("off")

            try:
                plt.savefig("./outputs/fake_image.png")
            except Exception as e:
                logging.exception("Exception caught: {}".format(e))

            plt.show()