import logging
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
from typing import Union

import numpy

from ..model import PyMcaXrfConfiguration
from .matrix import ensure_matrix

logger = logging.getLogger(__name__)


def set_beam_energy(
    configuration: PyMcaXrfConfiguration,
    energy: Optional[float],
    energy_multiplier: Optional[float],
):
    parameters = _get_beam_energy_parameters(configuration, energy, energy_multiplier)
    if parameters:
        _apply_beam_energy_parameters(configuration, parameters)


def _get_beam_energy_parameters(
    configuration: PyMcaXrfConfiguration,
    energy: Optional[float],
    energy_multiplier: Optional[float],
) -> Dict[str, List[Union[float, int]]]:
    if not _is_positive_finite_number(energy) and energy_multiplier is None:
        return {}

    fit = configuration.fit
    old_cfg = {
        "energy": _as_list(fit.energy),
        "energyweight": _as_list(fit.energyweight),
        "energyflag": _as_list(fit.energyflag),
        "energyscatter": _as_list(fit.energyscatter),
    }

    valid_items = [
        (i, e)
        for i, (e, energyflag) in enumerate(
            zip(old_cfg["energy"], old_cfg["energyflag"])
        )
        if energyflag == 1 and _is_positive_finite_number(e)
    ]
    indices = [i for i, _ in sorted(valid_items, key=lambda tpl: tpl[1])]

    modified = False

    if indices:

        def extract_defined_energies(name: str, dtype):
            energies = [old_cfg[name][i] for i in indices]
            return numpy.array(energies, dtype=dtype).tolist()

        parameters: Dict[str, list] = {
            "energy": extract_defined_energies("energy", float),
            "energyflag": extract_defined_energies("energyflag", int),
            "energyweight": extract_defined_energies("energyweight", float),
            "energyscatter": extract_defined_energies("energyscatter", int),
        }

        if _is_positive_finite_number(energy):
            first_energy = parameters["energy"][0]
            if not numpy.isclose(first_energy, energy):
                # Change the first energy and renormalize the others
                parameters["energy"] = [
                    e * energy / first_energy for e in parameters["energy"]
                ]
                modified = True

                logger.warning(
                    "PyMca configuration: modify primary beam energies to %s keV",
                    parameters["energy"],
                )
    else:

        if not _is_positive_finite_number(energy):
            return {}

        parameters: Dict[str, list] = {
            "energy": [energy],
            "energyflag": [1],
            "energyweight": [1.0],
            "energyscatter": [0],
        }
        modified = True

        logger.warning("PyMca configuration: set primary beam energy to %s keV", energy)

    # Add faraway line at the end (with small weight to not be fitted)
    # to be sure to include all lines in the energy range.
    first_energy = parameters["energy"][0]
    last_energy = parameters["energy"][-1]
    if _is_positive_finite_number(energy_multiplier):
        extra_energy = first_energy * energy_multiplier
        if last_energy < extra_energy:
            parameters["energy"].append(extra_energy)
            parameters["energyflag"].append(1)
            parameters["energyweight"].append(1e-10)
            parameters["energyscatter"].append(0)
            modified = True

            logger.warning(
                "PyMca configuration: append primary beam energy at %s keV (negligible contribution)",
                extra_energy,
            )

    if modified:
        return parameters
    return {}


def _apply_beam_energy_parameters(
    configuration: PyMcaXrfConfiguration,
    parameters: Dict[str, List[Union[float, int]]],
):
    fit = configuration.fit
    for key, value in parameters.items():
        setattr(fit, key, value)

    # Dummy matrix (apparently needed for multi-energy)
    if len(parameters["energy"]) > 1:
        ensure_matrix(configuration)


def _as_list(value: Any) -> List[Any]:
    if isinstance(value, list):
        return value
    return [value]


def _is_positive_finite_number(value: Any) -> bool:
    if not isinstance(value, (int, float)):
        return False
    return value > 0 and numpy.isfinite(value)
