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

# ! /usr/bin/env python

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

# run with python -m tests.test_model_selector

import logging

import sys

import pytest

# 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 = [

    (

        "00001_DCGAN_MMG_CALC_ROI",

        {},

        100,

    ),

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

    ("00003_CYCLEGAN_MMG_DENSITY_FULL", {"translate_all_images": False}, 2),

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

    # Further models can be added here if/when needed.

]

# class TestMediganSelectorMethods(unittest.TestCase):

class TestMediganSelectorMethods:

    def setup_method(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_init_generators()

    def test_init_generators(self):

        from src.medigan.generators import Generators

        self.generators = Generators()

    @pytest.mark.parametrize(

        "values_list",

        [

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

            (["DCGAN", "Mammography"]),

        ],

    )

    def test_search_for_models_method(self, values_list):

        found_models = self.generators.find_matching_models_by_values(

            values=values_list,

            target_values_operator="AND",

            are_keys_also_matched=True,

            is_case_sensitive=False,

        )

        self.logger.debug(

            f"For value {values_list}, these models were found: {found_models}"

        )

        assert len(found_models) > 0

    @pytest.mark.parametrize(

        "models, values_list, metric",

        [

            (

                models,

                ["dcgan", "MMG"],

                "CLF.trained_on_real_and_fake.f1",

            ),

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

        ],

    )

    def test_find_and_rank_models_by_performance(self, models, values_list, metric):

        # These values would need to find at least two models. See metrics and values in the config/global.json file.

        found_ranked_models = self.generators.find_models_and_rank(

            values=values_list,

            target_values_operator="AND",

            are_keys_also_matched=True,

            is_case_sensitive=False,

            metric=metric,

            order="desc",

        )

        assert (

            len(found_ranked_models) > 0  # some models were found as is expected

            and found_ranked_models[0]["model_id"] is not None  # has a model id

            and (

                len(found_ranked_models) < 2

                or found_ranked_models[0][metric] > found_ranked_models[1][metric]

            )  # descending order (the higher a model's value, the lower its index in the list) is working

        )

    @pytest.mark.parametrize(

        "models, metric, order",

        [

            (

                models,

                "FID",

                "asc",

            ),  # Note: normally a lower FID is better, therefore asc (model with lowest FID has lowest result list index).

            (

                models,

                "FID_RADIMAGENET_ratio",

                "desc",  # descending, as the higher the FID ratio the better.

            ),

            # Note: normally a lower FID is better, therefore asc (model with lowest FID has lowest result list index).

            (models, "CLF.trained_on_real_and_fake.f1", "desc"),

            (models, "turing_test.AUC", "desc"),

        ],

    )

    def test_rank_models_by_performance(self, models, metric, order):

        """Ranking according to metrics in the config/global.json file."""

        ranked_models = self.generators.rank_models_by_performance(

            model_ids=None,

            metric=metric,

            order=order,

        )

        assert (

            len(ranked_models) > 0  # at least one model was found

            and (

                len(ranked_models) >= 21 or metric != "FID"

            )  # we should find at least 21 models with FID in medigan

            and ranked_models[0]["model_id"]

            is not None  # found model has a model id (i.e. correctly formatted results)

            and (

                len(ranked_models) == 1

                or (

                    ranked_models[0][metric] > ranked_models[1][metric]

                    or metric == "FID"

                )

            )  # descending order (the higher a model's value, the lower its index in the list) is working. In case of FID it is the other way around (ascending order is better).

        )

    @pytest.mark.parametrize(

        "models, metric, order",

        [

            (models, "CLF.trained_on_real_and_fake.f1", "desc"),

            (models, "turing_test.AUC", "desc"),

        ],

    )

    def test_rank_models_by_performance_with_given_ids(self, models, metric, order):

        """Ranking a specified set of models according to metrics in the config/global.json file."""

        ranked_models = self.generators.rank_models_by_performance(

            model_ids=[models[1][0], models[2][0]],

            metric=metric,

            order=order,

        )

        assert 0 < len(ranked_models) <= 2 and (

            len(ranked_models) < 2

            or (ranked_models[0][metric] > ranked_models[1][metric])

        )  # checking if descending order (the higher a model's value, the lower its index in the list) is working.

    @pytest.mark.parametrize(

        "key1, value1, expected",

        [

            ("modality", "Full-Field Mammography", 2),

            ("license", "BSD", 2),

            ("performance.CLF.trained_on_real_and_fake.f1", "0.96", 0),

            ("performance.turing_test.AUC", "0.56", 0),

        ],

    )

    def test_get_models_by_key_value_pair(self, key1, value1, expected):

        found_models = self.generators.get_models_by_key_value_pair(

            key1=key1, value1=value1, is_case_sensitive=False

        )

        assert len(found_models) >= expected