import os
import random
import shutil
import tempfile

import mlflow
import numpy as np

from iquaflow.datasets import DSModifier, DSWrapper
from iquaflow.datasets.modifier_rer import DSModifier_rer
from iquaflow.experiments import ExperimentInfo, ExperimentSetup
from iquaflow.experiments.task_execution import PythonScriptTaskExecution
from iquaflow.metrics import BBDetectionMetrics, SharpnessMetric, SNRMetric

results_1 = [
    1.0,
    1.0,
    1.0,
    1.0,
    1.0,
    1.0,
    0.7012303556658395,
    0.9950372208436726,
    1.0,
    1.0,
    1.0,
    1.0,
]

results_2 = [0.010293033090993803, 0.17317447893884003, 0.08905812427554621]

current_path = os.path.dirname(os.path.realpath(__file__))
ml_models_path = os.path.join(current_path, "test_ml_models")
python_ml_script_path = os.path.join(ml_models_path, "train.py")
mlruns_path = mlflow.get_tracking_uri().replace("file://", "")

base_ds = os.path.join(current_path, "test_datasets")
data_path = os.path.join(base_ds, "ds_coco_dataset")

experiment_name = "test"


def remove_mlruns() -> None:
    if os.path.exists(mlruns_path) and os.path.isdir(mlruns_path):
        shutil.rmtree(mlruns_path)
    os.mkdir(mlruns_path)
    trash_path = os.path.join(mlruns_path, ".trash")
    if os.path.exists(trash_path) and os.path.isdir(trash_path):
        shutil.rmtree(trash_path)
    os.mkdir(trash_path)


class TestExperimentInfo:
    def test_experiment_info_class(self):
        remove_mlruns()
        ds_wrapper = DSWrapper(data_path=data_path)
        ds_modifiers_list = [DSModifier(), DSModifier(DSModifier())]
        task = PythonScriptTaskExecution(model_script_path=python_ml_script_path)
        experiment = ExperimentSetup(
            experiment_name=experiment_name,
            task_instance=task,
            ref_dsw_train=ds_wrapper,
            ds_modifiers_list=ds_modifiers_list,
        )
        experiment.execute()
        experiment_info = ExperimentInfo(experiment_name)
        # metrics
        metric_1 = BBDetectionMetrics()
        metric_2 = SharpnessMetric(experiment_info, ext="jpg")
        # apply metrics
        experiment_info.apply_metric_per_run(metric_1, str(ds_wrapper.json_annotations))
        experiment_info.apply_metric_per_run(metric_2, str(ds_wrapper.json_annotations))
        #
        run_name = "ds_coco_dataset#base_modifier#base_modifier"
        run = experiment_info.runs[run_name]
        # assert metric results
        for metric_name, val in zip(metric_1.metric_names, results_1):
            assert (
                run["metrics_dict"][metric_name] == val
            ), f"Unexpected result for {metric_name}"
        for metric_name in metric_2.metric_names:
            assert (
                0.0 < run["metrics_dict"][metric_name] < 3.0
            ), f"Unexpected result for {metric_name}"
        remove_mlruns()


rer_buffer = 0.05
image_val = 100
image_size = 40
kernel_size = 15
image_edge = (np.ones((image_size, image_size)) * image_val).astype(np.uint8)
image_edge[:, : np.int(image_size / 2)] = np.int(image_val / 2)
desired_rer = random.uniform(0.3, 0.5)


class TestSharpness:
    """Applies a known, random amount of blur, then checks that
    resulting RER value is close to expected value.
    Does not test robustness in accounting for image content.
    """

    def test_sharpness_blurred_image(self):

        expected_results_in_coco_ds = {
            "RER_X": 0.55,
            "RER_Y": 0.55,
            "RER_other": 0.56,
            "FWHM_X": 1.58,
            "FWHM_Y": 1.58,
            "FWHM_other": 1.57,
            "MTF_NYQ_X": 0.181,
            "MTF_NYQ_Y": 0.177,
            "MTF_NYQ_other": 0.192,
            "MTF_halfNYQ_X": 0.53,
            "MTF_halfNYQ_Y": 0.53,
            "MTF_halfNYQ_other": 0.55,
        }

        expected_results_in_coco_ds_after_blur = {
            "RER_X": 0.354,
            "RER_Y": 0.351,
            "RER_other": 0.348,
            "FWHM_X": 2.46,
            "FWHM_Y": 2.49,
            "FWHM_other": 2.52,
            "MTF_NYQ_X": 0.022,
            "MTF_NYQ_Y": 0.023,
            "MTF_NYQ_other": 0.022,
            "MTF_halfNYQ_X": 0.24,
            "MTF_halfNYQ_Y": 0.23,
            "MTF_halfNYQ_other": 0.23,
        }

        with tempfile.TemporaryDirectory() as tmp_dir:
            new_data_path = os.path.join(tmp_dir, "ds")
            shutil.copytree(data_path, new_data_path)

            ds_wrapper = DSWrapper(data_path=new_data_path)
            ds_modifiers_list = [
                DSModifier(),
                DSModifier_rer(params={"initial_rer": 0.55, "rer": 0.3}),
            ]
            task = PythonScriptTaskExecution(model_script_path=python_ml_script_path)
            experiment = ExperimentSetup(
                experiment_name="metric_test",
                task_instance=task,
                ref_dsw_train=ds_wrapper,
                ds_modifiers_list=ds_modifiers_list,
            )
            experiment.execute()
            experiment_info = ExperimentInfo("metric_test")
            np.random.seed(42)
            metric_sharpness = SharpnessMetric(
                experiment_info,
                ext="jpg",
                metrics=["RER", "FWHM", "MTF"],
                parallel=True,
            )

            assert (
                metric_sharpness.ext == "jpg"
            ), f"Unexpected ext for metric snr. It was set to jpg and it is {metric_sharpness.ext}"

            assert all(
                [
                    n in metric_sharpness.metric_names
                    for n in list(expected_results_in_coco_ds.keys())
                ]
            ), f"Unexpected metric_sharpness.metric_names: {metric_sharpness.metric_names}"
            assert hasattr(
                metric_sharpness, "apply"
            ), 'Missing method "apply" required in IQF-Metric'
            experiment_info.apply_metric_per_run(
                metric_sharpness, str(ds_wrapper.json_annotations)
            )
            run_name = "ds#base_modifier"
            myrun = experiment_info.runs[run_name]
            for key in expected_results_in_coco_ds:
                assert key in metric_sharpness.metric_names, f"Missing name {key}"
                assert (
                    key in myrun["metrics_dict"]
                ), f"Missing metric {key} in the results"
                assert (
                    abs(expected_results_in_coco_ds[key] - myrun["metrics_dict"][key])
                    / expected_results_in_coco_ds[key]
                    < 0.15
                ), f"Unexpected result for  {key} ({myrun['metrics_dict'][key]})"

            run_name = "ds#rer0.3_modifier"
            myrun = experiment_info.runs[run_name]

            for key in expected_results_in_coco_ds_after_blur:
                assert key in metric_sharpness.metric_names, f"Missing name {key}"
                assert (
                    key in myrun["metrics_dict"]
                ), f"Missing metric {key} in the results"
                assert (
                    abs(
                        expected_results_in_coco_ds_after_blur[key]
                        - myrun["metrics_dict"][key]
                    )
                    / expected_results_in_coco_ds_after_blur[key]
                    < 0.15
                ), f"Unexpected result for  {key} ({myrun['metrics_dict'][key]})"


class TestSNR:
    """
    Test Signal to Noise ratio metric
    """

    def test_snr_metric_hb_option(self):

        expected_results_in_coco_ds = {
            "snr_mean": 35.55032353607093,
            "snr_median": 28.91166092794036,
            "snr_std": 21.353973062148025,
        }

        with tempfile.TemporaryDirectory() as tmp_dir:

            new_data_path = os.path.join(tmp_dir, "ds")
            shutil.copytree(data_path, new_data_path)

            ds_wrapper = DSWrapper(data_path=new_data_path)
            ds_modifiers_list = [DSModifier()]
            task = PythonScriptTaskExecution(model_script_path=python_ml_script_path)
            experiment = ExperimentSetup(
                experiment_name="metric_test",
                task_instance=task,
                ref_dsw_train=ds_wrapper,
                ds_modifiers_list=ds_modifiers_list,
            )
            experiment.execute()
            experiment_info = ExperimentInfo("metric_test")
            metric_snr = SNRMetric(experiment_info, ext="jpg")
            assert (
                metric_snr.ext == "jpg"
            ), f"Unexpected ext for metric snr. It was set to jpg and it is {metric_snr.ext}"
            assert all(
                [
                    n in metric_snr.metric_names
                    for n in ["snr_median", "snr_mean", "snr_std"]
                ]
            ), "Unexpected metric_snr.metric_names"
            assert hasattr(
                metric_snr, "apply"
            ), 'Missing method "apply" required in IQF-Metric'
            experiment_info.apply_metric_per_run(
                metric_snr, str(ds_wrapper.json_annotations)
            )
            run_name = "ds#base_modifier"
            myrun = experiment_info.runs[run_name]
            for key in expected_results_in_coco_ds:
                assert key in metric_snr.metric_names, f"Missing name {key}"
                assert (
                    key in myrun["metrics_dict"]
                ), f"Missing metric {key} in the results"
                assert (
                    expected_results_in_coco_ds[key] == myrun["metrics_dict"][key]
                ), f"Unexpected result for SNR {key}"

    def test_snr_metric_ha_option(self):

        expected_results_in_coco_ds = {
            "snr_median": 32.34672691093317,
            "snr_mean": 32.24198884869302,
            "snr_std": 16.574449958484983,
        }

        with tempfile.TemporaryDirectory() as tmp_dir:

            new_data_path = os.path.join(tmp_dir, "ds")
            shutil.copytree(data_path, new_data_path)

            ds_wrapper = DSWrapper(data_path=new_data_path)
            ds_modifiers_list = [DSModifier()]
            task = PythonScriptTaskExecution(model_script_path=python_ml_script_path)
            experiment = ExperimentSetup(
                experiment_name="metric_test",
                task_instance=task,
                ref_dsw_train=ds_wrapper,
                ds_modifiers_list=ds_modifiers_list,
            )
            experiment.execute()
            experiment_info = ExperimentInfo("metric_test")
            metric_snr = SNRMetric(
                experiment_info,
                ext="jpg",
                method="HA",
                params={"patch_size": 4, "lbp_threshold": 0.6},
            )
            assert (
                metric_snr.ext == "jpg"
            ), f"Unexpected ext for metric snr. It was set to jpg and it is {metric_snr.ext}"
            assert all(
                [
                    n in metric_snr.metric_names
                    for n in ["snr_median", "snr_mean", "snr_std"]
                ]
            ), "Unexpected metric_snr.metric_names"
            assert hasattr(
                metric_snr, "apply"
            ), 'Missing method "apply" required in IQF-Metric'
            experiment_info.apply_metric_per_run(
                metric_snr, str(ds_wrapper.json_annotations)
            )
            run_name = "ds#base_modifier"
            myrun = experiment_info.runs[run_name]
            for key in expected_results_in_coco_ds:
                assert key in metric_snr.metric_names, f"Missing name {key}"
                assert (
                    key in myrun["metrics_dict"]
                ), f"Missing metric {key} in the results"
                assert (
                    expected_results_in_coco_ds[key] == myrun["metrics_dict"][key]
                ), f"Unexpected result for SNR {key}"
