# -*- coding: utf-8 -*-

# ! /usr/bin/env python

""" main test script to test the primary functions/classes/methods. """

# run with python -m tests.test_generator

import glob

import logging

import os

import shutil

import sys

import pytest

import torch

# import unittest

# Set the logging level depending on the level of detail you would like to have in the logs while running the tests.

LOGGING_LEVEL = logging.INFO  # WARNING  # logging.INFO

models_with_args = [

    (

        "00001_DCGAN_MMG_CALC_ROI",

        {},

        100,

    ),  # 100 samples to test automatic batch-wise image generation in model_executor

    (

        "00002",

        {},

        3,

    ),  # "00002" instead of "00002_DCGAN_MMG_MASS_ROI" to test shortcut model_ids

    (

        "03",

        {"translate_all_images": False},

        2,

    ),  # "03" instead of "00003_CYCLEGAN_MMG_DENSITY_FULL" to test shortcut model_ids

    (

        4,  # 4 instead of "00004_PIX2PIX_MMG_MASSES_W_MASKS" to test shortcut model_ids

        {

            "shapes": ["oval"],

            "ssim_threshold": 0.18,

            "image_size": [128, 128],

            "patch_size": [30, 30],

        },

        3,

    ),

    ("00005_DCGAN_MMG_MASS_ROI", {}, 3),

    ("00006_WGANGP_MMG_MASS_ROI", {}, 3),

    (

        "00007_INPAINT_BRAIN_MRI",

        {

            "image_size": (256, 256),

            "num_inpaints_per_sample": 2,

            "randomize_input_image_order": False,

            "add_variations_to_mask": False,

            "x_center": 120,

            "y_center": 140,

            "radius_1": 8,

            "radius_2": 12,

            "radius_3": 24,

        },

        3,

    ),

    (

        "00008_C-DCGAN_MMG_MASSES",

        {"condition": 0, "is_cbisddsm_training_data": False},

        3,

    ),

    ("00009_PGGAN_POLYP_PATCHES_W_MASKS", {"save_option": "image_only"}, 3),

    ("00010_FASTGAN_POLYP_PATCHES_W_MASKS", {"save_option": "image_only"}, 3),

    # ("00011_SINGAN_POLYP_PATCHES_W_MASKS", {"checkpoint_ids": [999]}, 3), # removed after successful testing due to limited CI pipeline capacity

    # ("00012_C-DCGAN_MMG_MASSES", {"condition": 0}, 3), # removed after successful testing due to limited CI pipeline capacity

    # ("00013_CYCLEGAN_MMG_DENSITY_OPTIMAM_MLO", {"translate_all_images": False}, 2), # removed after successful testing due to limited CI pipeline capacity

    # ("00014_CYCLEGAN_MMG_DENSITY_OPTIMAM_CC", {"translate_all_images": False}, 2), # removed after successful testing due to limited CI pipeline capacity

    # ("00015_CYCLEGAN_MMG_DENSITY_CSAW_MLO", {"translate_all_images": False}, 2), # removed after successful testing due to limited CI pipeline capacity

    # ("00016_CYCLEGAN_MMG_DENSITY_CSAW_CC", {"translate_all_images": False}, 2), # removed after successful testing due to limited CI pipeline capacity

    ("00017_DCGAN_XRAY_LUNG_NODULES", {}, 3),

    ("00018_WGANGP_XRAY_LUNG_NODULES", {}, 3),

    ("00019_PGGAN_CHEST_XRAY", {}, 3),

    ("00020_PGGAN_CHEST_XRAY", {"resize_pixel_dim": 512, "image_size": 256}, 3),

    (

        "00021_CYCLEGAN_BRAIN_MRI_T1_T2",

        {

            "input_path": "models/00021_CYCLEGAN_Brain_MRI_T1_T2/inputs/T2",

            "gpu_id": 0,

            "T1_to_T2": False,

        },

        3,

    ),

    ("00022_WGAN_CARDIAC_AGING", {}, 3),

    (

        "00023_PIX2PIXHD_BREAST_DCEMRI",

        {

            "input_path": "input",

            "gpu_id": 0,

            "image_size": 448,

        },

        3,

    ),

]

# class TestMediganExecutorMethods(unittest.TestCase):

class TestMediganExecutorMethods:

    def setup_class(self):

        ## unittest logger config

        # This logger on root level initialized via logging.getLogger() will also log all log events

        # from the medigan library. Pass a logger name (e.g. __name__) instead if you only want logs from tests.py

        self.logger = logging.getLogger()  # (__name__)

        self.logger.setLevel(LOGGING_LEVEL)

        stream_handler = logging.StreamHandler(sys.stdout)

        stream_handler.setLevel(LOGGING_LEVEL)

        formatter = logging.Formatter(

            "%(asctime)s - %(name)s - %(levelname)s - %(message)s"

        )

        stream_handler.setFormatter(formatter)

        self.logger.addHandler(stream_handler)

        self.test_output_path = "test_output_path"

        self.num_samples = 2

        self.test_imports_and_init_generators(self)

        self._remove_dir_and_contents(self)  # in case something is left there.

        self.model_ids = self.generators.config_manager.model_ids

    def test_imports_and_init_generators(self):

        from src.medigan.constants import (

            CONFIG_FILE_KEY_EXECUTION,

            CONFIG_FILE_KEY_GENERATE,

            CONFIG_FILE_KEY_GENERATE_ARGS_INPUT_LATENT_VECTOR_SIZE,

        )

        from src.medigan.generators import Generators

        self.generators = Generators()

        self.CONFIG_FILE_KEY_EXECUTION = CONFIG_FILE_KEY_EXECUTION

        self.CONFIG_FILE_KEY_GENERATE = CONFIG_FILE_KEY_GENERATE

        self.CONFIG_FILE_KEY_GENERATE_ARGS_INPUT_LATENT_VECTOR_SIZE = (

            CONFIG_FILE_KEY_GENERATE_ARGS_INPUT_LATENT_VECTOR_SIZE

        )

    @pytest.mark.parametrize("models_with_args", [models_with_args])

    def test_sample_generation_methods(self, models_with_args: list):

        self.logger.debug(f"models: {models_with_args}")

        for i, model_id in enumerate(self.model_ids):

            # if (

            #    model_id != "00011_SINGAN_POLYP_PATCHES_W_MASKS"

            # ):

            ## avoiding full memory on Windows ci test server

            # continue

            self.logger.debug(f"Now testing model {model_id}")

            self._remove_dir_and_contents()  # Already done in each test independently, but to be sure, here again.

            self.test_generate_method(model_id=model_id)

            # Check if args available fo model_id. Note: The models list may not include the latest medigan models

            for model in models_with_args:

                if model_id == model[0]:

                    self.test_generate_method_with_additional_args(

                        model_id=model[0], args=model[1], expected_num_samples=model[2]

                    )

            self.test_get_generate_method(model_id=model_id)

            self.test_get_dataloader_method(model_id=model_id)

            # if i == 16:  # just for local testing

            # self._remove_model_dir_and_zip(

            #    model_ids=[model_id], are_all_models_deleted=False

            # )

    @pytest.mark.parametrize(

        "values_list, should_sample_be_generated",

        [

            (["dcgan", "mMg", "ClF", "modality", "inbreast"], True),

            (["dcgan", "mMg", "ClF", "modality", "optimam"], True),

            (["dcgan", "mMg", "ClF", "modalities"], False),

        ],

    )

    def test_find_model_and_generate_method(

        self, values_list, should_sample_be_generated

    ):

        self._remove_dir_and_contents()

        self.generators.find_model_and_generate(

            values=values_list,

            target_values_operator="AND",

            are_keys_also_matched=True,

            is_case_sensitive=False,

            num_samples=self.num_samples,

            output_path=self.test_output_path,

        )

        self._check_if_samples_were_generated(

            should_sample_be_generated=should_sample_be_generated

        )

    @pytest.mark.parametrize(

        "values_list, metric",

        [

            (["dcgan", "MMG"], "CLF.trained_on_real_and_fake.f1"),

            (["dcgan", "MMG"], "turing_test.AUC"),

        ],

    )

    def test_find_and_rank_models_then_generate_method(self, values_list, metric):

        self._remove_dir_and_contents()

        # TODO This test needs the respective metrics for any of these models to be available in config/global.json.

        # These values would need to find at least two models.

        self.generators.find_models_rank_and_generate(

            values=values_list,

            target_values_operator="AND",

            are_keys_also_matched=True,

            is_case_sensitive=False,

            metric=metric,

            order="asc",

            num_samples=self.num_samples,

            output_path=self.test_output_path,

        )

        self._check_if_samples_were_generated()

    # @pytest.mark.parametrize("model_id", [model[0] for model in models_with_args])

    @pytest.mark.skip

    def test_generate_method(self, model_id):

        self._remove_dir_and_contents()

        self.generators.generate(

            model_id=model_id,

            num_samples=self.num_samples,

            output_path=self.test_output_path,

            install_dependencies=True,

        )

        self._check_if_samples_were_generated(model_id=model_id)

    # @pytest.mark.parametrize("model_id, args, expected_num_samples", models_with_args)

    @pytest.mark.skip

    def test_generate_method_with_additional_args(

        self, model_id, args, expected_num_samples

    ):

        self._remove_dir_and_contents()

        self.generators.generate(

            model_id=model_id,

            num_samples=expected_num_samples,

            output_path=self.test_output_path,

            **args,

        )

        self._check_if_samples_were_generated(

            model_id=model_id, num_samples=expected_num_samples

        )

    # @pytest.mark.parametrize("model_id", [model[0] for model in models_with_args])

    @pytest.mark.skip

    def test_get_generate_method(self, model_id):

        self._remove_dir_and_contents()

        gen_function = self.generators.get_generate_function(

            model_id=model_id,

            num_samples=self.num_samples,

            output_path=self.test_output_path,

        )

        gen_function()

        self._check_if_samples_were_generated(model_id=model_id)

        del gen_function

    # @pytest.mark.parametrize("model_id", [model[0] for model in models_with_args])

    @pytest.mark.skip

    def test_get_dataloader_method(self, model_id):

        self._remove_dir_and_contents()

        data_loader = self.generators.get_as_torch_dataloader(

            model_id=model_id, num_samples=self.num_samples

        )

        self.logger.debug(f"{model_id}: len(data_loader): {len(data_loader)}")

        if len(data_loader) != self.num_samples:

            logging.warning(

                f"{model_id}: The number of samples in the dataloader (={len(data_loader)}) is not equal the number of samples requested (={self.num_samples}). "

                f"Hint: Revise if the model's internal generate() function returned tuples as required in get_as_torch_dataloader()."

            )

        #### Get the object at index 0 from the dataloader

        data_dict = next(iter(data_loader))

        # Test if the items at index [0] of the aforementioned object is of type torch tensor (e.g. torch.uint8) and not None, as expected by data structure design decision.

        assert torch.is_tensor(data_dict.get("sample"))

        # Test if the items at index [1], [2] of the aforementioned object are None and, if not, whether they are of type torch tensor, as expected

        assert data_dict.get("mask") is None or torch.is_tensor(data_dict.get("mask"))

        assert data_dict.get("other_imaging_output") is None or torch.is_tensor(

            data_dict.get("other_imaging_output")

        )

        # Test if the items at index [3] of the aforementioned object is None and, if not, whether it is of type list of strings, as expected.

        assert data_dict.get("label") is None or (

            isinstance(data_dict.get("label"), list)

            and isinstance(data_dict.get("label")[0], str)

        )

        del data_dict

        del data_loader

    # @pytest.mark.parametrize("model_id", [model[0] for model in models_with_args])

    @pytest.mark.skip

    def test_visualize_method(self, model_id):

        if (

            self.CONFIG_FILE_KEY_GENERATE_ARGS_INPUT_LATENT_VECTOR_SIZE

            in self.generators.config_manager.config_dict[model_id][

                self.CONFIG_FILE_KEY_EXECUTION

            ][self.CONFIG_FILE_KEY_GENERATE]

        ):

            self.generators.visualize(model_id, auto_close=True)

        else:

            with pytest.raises(Exception) as e:

                self.generators.visualize(model_id, auto_close=True)

                assert e.type == ValueError

    @pytest.mark.skip

    def _check_if_samples_were_generated(

        self, model_id=None, num_samples=None, should_sample_be_generated: bool = True

    ):

        # check if the number of generated samples of model_id_1 is as expected.

        file_list = glob.glob(self.test_output_path + "/*")

        self.logger.debug(f"{model_id}: {len(file_list)} == {self.num_samples} ?")

        if num_samples is None:

            num_samples = self.num_samples

        if should_sample_be_generated:

            assert (

                len(file_list) == num_samples

                or len(file_list)

                == num_samples

                * 2

                * 6  # 00007_INPAINT_BRAIN_MRI: 2 inpaints per sample, 6 outputs per sample

                or len(file_list)

                == num_samples * 2  # Temporary fix for different outputs per model.

                or len(file_list) == num_samples + 1

            ), f"Model {model_id} generated {len(file_list)} samples instead of the expected {num_samples}, {num_samples*2*6}, or {num_samples + 1}."

            # Some models are balanced per label by default: If num_samples is odd, then len(file_list)==num_samples +1

        else:

            assert len(file_list) == 0

    # @pytest.mark.skip

    def _remove_dir_and_contents(self):

        """After each test, empty the created folders and files to avoid corrupting a new test."""

        try:

            shutil.rmtree(self.test_output_path)

        except OSError as e:

            # This may give an error if the folders are not created.

            self.logger.debug(

                f"Exception while trying to delete folder. Likely it simply had not yet been created: {e}"

            )

        except Exception as e2:

            self.logger.error(f"Error while trying to delete folder: {e2}")

    @pytest.mark.skip

    def _remove_model_dir_and_zip(

        self, model_ids=[], are_all_models_deleted: bool = False

    ):

        """After a specific model folders, model_executor, and model zip file to avoid running out-of-disk space."""

        try:

            for i, model_executor in enumerate(self.generators.model_executors):

                if are_all_models_deleted or (

                    model_ids is not None and model_executor.model_id in model_ids

                ):

                    try:

                        # Delete the folder containing the model

                        model_path = os.path.dirname(

                            model_executor.deserialized_model_as_lib.__file__

                        )

                        shutil.rmtree(model_path)

                        self.logger.info(

                            f"Deleted directory of model {model_executor.model_id}. ({model_path})"

                        )

                    except OSError as e:

                        # This may give an error if the FOLDER is not present

                        self.logger.warning(

                            f"Exception while trying to delete the model folder of model {model_executor.model_id}: {e}"

                        )

                    try:

                        # If the downloaded zip package of the model was not deleted inside the model_path, we explicitely delete it now.

                        if model_executor.package_path.is_file():

                            os.remove(model_executor.package_path)

                            self.logger.info(

                                f"Deleted zip file of model {model_executor.model_id}. ({model_executor.package_path})"

                            )

                    except Exception as e:

                        self.logger.warning(

                            f"Exception while trying to delete the ZIP file ({model_executor.package_path}) of model {model_executor.model_id}: {e}"

                        )

            # Deleting the stateful model_executors instantiated by the generators module, after deleting folders and zips

            if are_all_models_deleted:

                self.generators.model_executors.clear()

            else:

                if model_ids is not None:

                    for model_id in model_ids:

                        model_executor = self.generators.find_model_executor_by_id(

                            model_id

                        )

                        if model_executor is not None:

                            self.generators.model_executors.remove(model_executor)

                        del model_executor

        except Exception as e2:

            self.logger.error(

                f"Error while trying to delete model folders and zips: {e2}"

            )

    # @pytest.fixture(scope="session", autouse=True)

    def teardown_class(self):

        """After all tests, empty the large model folders, model_executors, and zip files to avoid running out-of-disk space."""

        # yield is at test-time, signaling that things after yield are run after the execution of the last test has terminated

        # https://docs.pytest.org/en/7.1.x/reference/reference.html?highlight=fixture#pytest.fixture

        # yield None

        # Remove all test outputs in test_output_path

        self._remove_dir_and_contents(self)

        # Remove all model folders, zip files and model executors

        # self._remove_model_dir_and_zip(

        #    self, model_ids=["00006_WGANGP_MMG_MASS_ROI"], are_all_models_deleted=False

        # )  # just for local testing

        # self._remove_model_dir_and_zip(

        #    self, model_ids=None, are_all_models_deleted=True

        # )