"""Unittests for the 'ModelFitting' class."""
#  Copyright (c) European Space Agency, 2017, 2018, 2019, 2020.
#
#  This file is subject to the terms and conditions defined in file 'LICENCE.txt', which
#  is part of this Pyxel package. No part of the package, including
#  this file, may be copied, modified, propagated, or distributed except according to
#  the terms contained in the file ‘LICENCE.txt’.

from pathlib import Path

import numpy as np
import pytest

from pyxel import inputs_outputs as io
from pyxel.calibration import Calibration
from pyxel.calibration.fitting import ModelFitting
from pyxel.calibration.util import CalibrationMode, ResultType
from pyxel.detectors import CCD
from pyxel.inputs_outputs import Configuration
from pyxel.pipelines import DetectionPipeline, Processor

try:
    import pygmo as pg

    WITH_PYGMO = True
except ImportError:
    WITH_PYGMO = False


def configure(mf: ModelFitting, sim: Configuration) -> None:
    """TBW."""
    assert sim.calibration is not None

    pg.set_global_rng_seed(sim.calibration.seed)
    np.random.seed(sim.calibration.seed)

    mf.configure(
        calibration_mode=sim.calibration.calibration_mode,
        generations=sim.calibration.algorithm.generations,
        population_size=sim.calibration.algorithm.population_size,
        simulation_output=sim.calibration.result_type,
        fitness_func=sim.calibration.fitness_function,
        target_output=sim.calibration.target_data_path,
        target_fit_range=sim.calibration.target_fit_range,
        out_fit_range=sim.calibration.result_fit_range,
        weighting=sim.calibration.weighting_path,
        file_path=None,
    )


@pytest.mark.skipif(not WITH_PYGMO, reason="Package 'pygmo' is not installed.")
@pytest.mark.parametrize("yaml_file", ["tests/data/calibrate.yaml"])
def test_configure_params(yaml_file):
    """Test """
    cfg = io.load(yaml_file)
    detector = cfg.ccd_detector
    pipeline = cfg.pipeline
    processor = Processor(detector, pipeline)
    calibration = cfg.calibration
    mf = ModelFitting(processor, calibration.parameters)

    if not isinstance(mf.processor, Processor):
        raise ValueError

    configure(mf, cfg)

    assert mf.calibration_mode == CalibrationMode.Pipeline
    assert mf.sim_fit_range == slice(2, 5, None)
    assert mf.targ_fit_range == slice(1, 4, None)
    assert mf.sim_output == ResultType.Image


@pytest.mark.skipif(not WITH_PYGMO, reason="Package 'pygmo' is not installed.")
@pytest.mark.parametrize("yaml", ["tests/data/calibrate_fits.yaml"])
def test_configure_fits_target(yaml):
    """Test """
    cfg = io.load(yaml)
    detector = cfg.ccd_detector
    pipeline = cfg.pipeline
    processor = Processor(detector, pipeline)
    calibration = cfg.calibration
    mf = ModelFitting(processor, calibration.parameters)
    configure(mf, cfg)
    assert mf.sim_fit_range == (slice(2, 5, None), slice(4, 7, None))
    assert mf.targ_fit_range == (slice(1, 4, None), slice(5, 8, None))
    assert mf.sim_output == ResultType.Image
    expected = np.array(
        [
            [3858.44799859, 3836.11204939, 3809.85008514],
            [4100.87410744, 4053.26348117, 4018.33656962],
            [4233.53215652, 4021.60164244, 3969.79740826],
        ]
    )
    np.testing.assert_array_equal(
        np.around(mf.all_target_data[0], decimals=4), np.around(expected, decimals=4)
    )


@pytest.mark.skipif(not WITH_PYGMO, reason="Package 'pygmo' is not installed.")
@pytest.mark.parametrize(
    "yaml", ["tests/data/calibrate.yaml", "tests/data/calibrate_fits.yaml"]
)
def test_boundaries(yaml):
    """Test """
    cfg = io.load(yaml)
    detector = cfg.ccd_detector
    pipeline = cfg.pipeline
    processor = Processor(detector, pipeline)
    calibration = cfg.calibration
    mf = ModelFitting(processor, calibration.parameters)

    configure(mf, cfg)

    lbd_expected = [1.0, -3.0, -3.0, -2.0, -2.0, 0.0, 10.0]
    ubd_expected = [10.0, 0.3010299956639812, 0.3010299956639812, 1.0, 1.0, 1.0, 200.0]
    assert mf.lbd == lbd_expected
    assert mf.ubd == ubd_expected
    ll, uu = mf.get_bounds()
    assert ll == lbd_expected
    assert uu == ubd_expected


@pytest.mark.skipif(not WITH_PYGMO, reason="Package 'pygmo' is not installed.")
@pytest.mark.parametrize(
    "simulated_data, target_data, expected_fitness",
    [
        (231, 231, 0.0),
        (231, 145, 86.0),
        (2.31, 1.45, 0.8600000000000001),
        (2.0, 1, 1.0),
        (np.array([1, 9, 45, 548, 2, 2]), np.array([1, 9, 45, 548, 2, 2]), 0.0),
        (np.array([1, 9, 45, 548, 2, 2]), np.array([1, 3, 56, 21, 235, 11]), 786.0),
        (
            np.array([[1362.0, 1378.0], [1308.0, 1309.0]]),
            np.array([[1362.0, 1378.0], [1308.0, 1309.0]]),
            0.0,
        ),
        (
            np.array([[1362.0, 1378.0], [1308.0, 1309.0]]),
            np.array([[1462.0, 1368.0], [1508.0, 1399.0]]),
            400.0,
        ),
    ],
)
def test_calculate_fitness(simulated_data, target_data, expected_fitness):
    """Test"""
    cfg = io.load("tests/data/calibrate.yaml")
    detector = cfg.ccd_detector
    pipeline = cfg.pipeline
    processor = Processor(detector, pipeline)
    calibration = cfg.calibration
    mf = ModelFitting(processor, calibration.parameters)
    configure(mf, cfg)
    fitness = mf.calculate_fitness(simulated_data, target_data)
    assert fitness == expected_fitness
    print("fitness: ", fitness)


@pytest.mark.skipif(not WITH_PYGMO, reason="Package 'pygmo' is not installed.")
@pytest.mark.parametrize(
    "yaml, factor, expected_fitness",
    [
        (Path("tests/data/calibrate_weighting.yaml"), 1, 0.0),
        (Path("tests/data/calibrate_weighting.yaml"), 2, 310815803081.51117),
        (Path("tests/data/calibrate_weighting.yaml"), 3, 621631606163.0223),
    ],
)
def test_weighting(yaml, factor, expected_fitness):
    """Test"""
    cfg = io.load(yaml)
    detector = cfg.ccd_detector
    pipeline = cfg.pipeline
    processor = Processor(detector, pipeline)
    calibration = cfg.calibration
    mf = ModelFitting(processor, calibration.parameters)
    configure(mf, cfg)
    fitness = mf.calculate_fitness(
        mf.all_target_data[0] * factor, mf.all_target_data[0]
    )
    assert fitness == expected_fitness
    print("fitness: ", fitness)


def custom_fitness_func(simulated, target, weighting=None):
    """Custom fitness func for testing"""
    return np.sum(target * 2 - simulated / 2 + 1.0)


@pytest.mark.skipif(not WITH_PYGMO, reason="Package 'pygmo' is not installed.")
@pytest.mark.parametrize(
    "yaml, simulated, target, weighting",
    [
        ("tests/data/calibrate_custom_fitness.yaml", 1.0, 2.0, 4.5),
        (
            "tests/data/calibrate_custom_fitness.yaml",
            np.array([1.0, 2.0, 3.0]),
            np.array([2.0, 5.0, 6.0]),
            26.0,
        ),
        ("tests/data/calibrate_least_squares.yaml", 2.0, 4.0, 4.0),
    ],
)
def test_custom_fitness(yaml, simulated, target, weighting):
    """Test"""
    cfg = io.load(yaml)
    assert isinstance(cfg, Configuration)

    detector = cfg.ccd_detector
    assert isinstance(detector, CCD)

    pipeline = cfg.pipeline
    assert isinstance(pipeline, DetectionPipeline)

    processor = Processor(detector, pipeline)
    assert isinstance(processor, Processor)

    calibration = cfg.calibration
    assert isinstance(calibration, Calibration)

    assert cfg.calibration is not None

    mf = ModelFitting(processor, variables=calibration.parameters)
    configure(mf=mf, sim=cfg)

    fitness = mf.calculate_fitness(simulated, target)
    assert fitness == weighting
    print("fitness: ", fitness)


@pytest.mark.skipif(not WITH_PYGMO, reason="Package 'pygmo' is not installed.")
@pytest.mark.parametrize(
    "yaml, parameter, expected_fitness",
    [
        (
            "tests/data/calibrate_models.yaml",
            np.array(
                [
                    1.0,
                    0.5,
                    1.5,
                    -2.0,
                    -3.0,
                    4.5,
                    -4.0,
                    1.0,
                    0.5,
                    -3.5,
                    2.0,
                    -3.0,
                    -4.0,
                    0.5,
                    1.0,
                    100.0,
                ]
            ),
            88663.74498686477,
        )
    ],
)
def test_fitness(yaml, parameter, expected_fitness):
    """Test"""
    cfg = io.load(yaml)

    detector = cfg.ccd_detector
    pipeline = cfg.pipeline

    processor = Processor(detector, pipeline)
    calibration = cfg.calibration

    mf = ModelFitting(processor, calibration.parameters)

    configure(mf, cfg)

    overall_fitness = mf.fitness(parameter)
    assert overall_fitness[0] == expected_fitness

    print("fitness: ", overall_fitness[0])


@pytest.mark.skipif(not WITH_PYGMO, reason="Package 'pygmo' is not installed.")
@pytest.mark.parametrize(
    "yaml, parameter, expected_array",
    [
        (
            "tests/data/calibrate_models.yaml",
            np.arange(16.0),
            np.array(
                [
                    0.0,
                    10.0,
                    100.0,
                    1000.0,
                    10000.0,
                    5.0,
                    6.0,
                    7.0,
                    8.0,
                    1.0e09,
                    1.0e10,
                    1.0e11,
                    1.0e12,
                    13.0,
                    1.0e14,
                    15.0,
                ]
            ),
        )
    ],
)
def test_split_and_update(yaml, parameter, expected_array):
    """Test"""
    cfg = io.load(yaml)
    detector = cfg.ccd_detector
    pipeline = cfg.pipeline
    processor = Processor(detector, pipeline)
    calibration = cfg.calibration
    mf = ModelFitting(processor, calibration.parameters)
    configure(mf, cfg)
    array = mf.convert_to_parameters(parameter)
    np.testing.assert_array_equal(array, expected_array)


@pytest.mark.skipif(not WITH_PYGMO, reason="Package 'pygmo' is not installed.")
@pytest.mark.parametrize(
    "yaml, param_array",
    [
        (
            "tests/data/calibrate_models.yaml",
            np.array(
                [
                    1.0,
                    10.0,
                    100.0,
                    1000.0,
                    10000.0,
                    5.0,
                    6.0,
                    7.0,
                    8.0,
                    1.0e09,
                    1.0e10,
                    1.0e11,
                    1.0e12,
                    13.0,
                    1.0e14,
                    150,
                ]
            ),
        )
    ],
)
def test_detector_and_model_update(yaml: str, param_array: np.ndarray):
    """Test"""
    cfg = io.load(yaml)

    detector = cfg.ccd_detector
    pipeline = cfg.pipeline

    processor = Processor(detector, pipeline)
    calibration = cfg.calibration
    mf = ModelFitting(processor=processor, variables=calibration.parameters)
    configure(mf=mf, sim=cfg)
    mf.processor = mf.update_processor(param_array, processor)

    assert mf.processor.pipeline.charge_transfer is not None
    assert mf.processor.pipeline.charge_measurement is not None

    attributes = [
        mf.processor.detector.characteristics.amp,
        mf.processor.pipeline.charge_transfer.models[0].arguments["tr_p"],
        mf.processor.pipeline.charge_transfer.models[0].arguments["nt_p"],
        mf.processor.pipeline.charge_transfer.models[0].arguments["sigma_p"],
        mf.processor.pipeline.charge_transfer.models[0].arguments["beta_p"],
        mf.processor.pipeline.charge_measurement.models[1].arguments["std_deviation"],
        mf.processor.detector.environment.temperature,
    ]
    a = 0
    for idx, attr in enumerate(attributes):
        if isinstance(attr, np.ndarray):
            b = len(attr)
            np.testing.assert_array_equal(attr, param_array[a : a + b])
        else:
            b = 1
            assert attr == param_array[a]
        a += b
