import sys
from typing import Dict
from typing import Sequence
from typing import Union

if sys.version_info < (3, 12):
    import importlib_resources as resources
else:
    from importlib import resources

import numpy

from ..xrffit import pymca_config


def test_round_trip():
    files = resources.files("PyMca5.PyMcaData")
    with resources.as_file(files) as pkg_path:
        for cfg_file in pkg_path.glob("*.cfg"):
            config_dict_org = pymca_config.pymca_configdict_from_file(cfg_file)
            model = pymca_config.pymca_configdict_to_model(config_dict_org)
            config_dict = pymca_config.pymca_configdict_from_model(model)
            assert config_dict == config_dict_org, cfg_file


def _generate_pymca_config(fit_config: Dict[str, Sequence[Union[float, int]]]):
    files = resources.files("PyMca5.PyMcaData")
    with resources.as_file(files) as pkg_path:
        config_dict = pymca_config.pymca_configdict_from_file(pkg_path / "Steel.cfg")
        config_dict["fit"].update(fit_config)
        return pymca_config.pymca_configdict_to_model(config_dict)


def test_replace_energy():
    configuration = _generate_pymca_config(
        {
            "energy": [2, 5, 10, 20, 50],
            "energyflag": [1, 0, 1, 0, 0],
            "energyscatter": [1, 0, 0, 1, 0],
            "energyweight": [1.0, 0.25, 0.2, 0.1, 0.1],
        }
    )
    energy = 2.5
    energy_multiplier = 10

    pymca_config.set_beam_energy(configuration, energy, energy_multiplier)

    numpy.testing.assert_allclose(
        configuration.fit.energy, numpy.array([2.5, 12.5, 25])
    )
    numpy.testing.assert_equal(configuration.fit.energyflag, numpy.array([1, 1, 1]))
    numpy.testing.assert_equal(configuration.fit.energyscatter, numpy.array([1, 0, 0]))
    numpy.testing.assert_allclose(
        configuration.fit.energyweight, numpy.array([1.0, 0.2, 1e-10])
    )


def test_replace_energy_all_defined():
    configuration = _generate_pymca_config(
        {
            "energy": [2],
            "energyflag": [1],
            "energyscatter": [1],
            "energyweight": [1.0],
        }
    )
    energy = 4
    energy_multiplier = 10

    pymca_config.set_beam_energy(configuration, energy, energy_multiplier)

    numpy.testing.assert_allclose(configuration.fit.energy, numpy.array([4, 40]))
    numpy.testing.assert_equal(configuration.fit.energyflag, numpy.array([1, 1]))
    numpy.testing.assert_equal(configuration.fit.energyscatter, numpy.array([1, 0]))
    numpy.testing.assert_allclose(
        configuration.fit.energyweight, numpy.array([1.0, 1e-10])
    )


def test_replace_energy_no_multiplier():
    configuration = _generate_pymca_config(
        {
            "energy": [2, 5, 10, 20, 50],
            "energyflag": [1, 0, 1, 0, 0],
            "energyscatter": [1, 0, 0, 1, 0],
            "energyweight": [1.0, 0.25, 0.2, 0.1, 0.1],
        }
    )
    energy = 2.5
    energy_multiplier = None

    pymca_config.set_beam_energy(configuration, energy, energy_multiplier)

    numpy.testing.assert_allclose(configuration.fit.energy, numpy.array([2.5, 12.5]))
    numpy.testing.assert_equal(configuration.fit.energyflag, numpy.array([1, 1]))
    numpy.testing.assert_equal(configuration.fit.energyscatter, numpy.array([1, 0]))
    numpy.testing.assert_allclose(
        configuration.fit.energyweight, numpy.array([1.0, 0.2])
    )


def test_replace_energy_not_defined():
    configuration = _generate_pymca_config(
        {
            "energy": [2, 5, 10, 20, 50],
            "energyflag": [0, 0, 0, 0, 0],
            "energyscatter": [1, 0, 0, 1, 0],
            "energyweight": [1.0, 0.25, 0.2, 0.1, 0.1],
        }
    )
    energy = 6
    energy_multiplier = 3

    pymca_config.set_beam_energy(configuration, energy, energy_multiplier)

    numpy.testing.assert_allclose(
        configuration.fit.energy, numpy.array([energy, energy_multiplier * energy])
    )
    numpy.testing.assert_equal(configuration.fit.energyflag, numpy.array([1, 1]))
    numpy.testing.assert_equal(configuration.fit.energyscatter, numpy.array([0, 0]))
    numpy.testing.assert_allclose(
        configuration.fit.energyweight, numpy.array([1.0, 1e-10])
    )


def test_replace_energy_not_defined_no_multiplier():
    configuration = _generate_pymca_config(
        {
            "energy": [2, 5, 10, 20, 50],
            "energyflag": [0, 0, 0, 0, 0],
            "energyscatter": [1, 0, 0, 1, 0],
            "energyweight": [1.0, 0.25, 0.2, 0.1, 0.1],
        }
    )
    energy = 7.1
    energy_multiplier = None

    pymca_config.set_beam_energy(configuration, energy, energy_multiplier)

    numpy.testing.assert_allclose(configuration.fit.energy, numpy.array([energy]))
    numpy.testing.assert_equal(configuration.fit.energyflag, numpy.array([1]))
    numpy.testing.assert_equal(configuration.fit.energyscatter, numpy.array([0]))
    numpy.testing.assert_allclose(configuration.fit.energyweight, numpy.array([1.0]))


def test_replace_multiplier_only():
    configuration = _generate_pymca_config(
        {
            "energy": [4],
            "energyflag": [1],
            "energyscatter": [1],
            "energyweight": [1.0],
        }
    )
    energy_multiplier = 10

    pymca_config.set_beam_energy(configuration, None, energy_multiplier)

    numpy.testing.assert_allclose(configuration.fit.energy, numpy.array([4, 40]))
    numpy.testing.assert_equal(configuration.fit.energyflag, numpy.array([1, 1]))
    numpy.testing.assert_equal(configuration.fit.energyscatter, numpy.array([1, 0]))
    numpy.testing.assert_allclose(
        configuration.fit.energyweight, numpy.array([1.0, 1e-10])
    )


def test_replace_multiplier_only_no_energy():
    configuration = _generate_pymca_config(
        {
            "energy": [],
            "energyflag": [],
            "energyscatter": [],
            "energyweight": [],
        }
    )
    energy = None
    energy_multiplier = 10

    pymca_config.set_beam_energy(configuration, energy, energy_multiplier)

    numpy.testing.assert_allclose(configuration.fit.energy, numpy.array([]))
    numpy.testing.assert_equal(configuration.fit.energyflag, numpy.array([]))
    numpy.testing.assert_equal(configuration.fit.energyscatter, numpy.array([]))
    numpy.testing.assert_allclose(configuration.fit.energyweight, numpy.array([]))
