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 energy is None or not numpy.isfinite(energy):
        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),
    }

    indices = [
        i
        for i, (energy, energyflag) in enumerate(
            zip(old_cfg["energy"], old_cfg["energyflag"])
        )
        if energyflag == 1 and isinstance(energy, (int, float)) and energy > 0
    ]

    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),
        }
    else:
        parameters: Dict[str, list] = {
            "energy": [energy],
            "energyflag": [1],
            "energyweight": [1.0],
            "energyscatter": [0],
        }

    # Renormalize energies
    e0 = parameters["energy"][0]
    parameters["energy"] = [e * energy / e0 for e in parameters["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
    if energy_multiplier and parameters["energy"][-1] < energy * energy_multiplier:
        parameters["energy"].append(energy * energy_multiplier)
        parameters["energyflag"].append(1)
        parameters["energyweight"].append(1e-10)
        parameters["energyscatter"].append(0)
        logger.warning(
            "PyMca configuration: set beam energy to %s keV (+ negligible contribution at %s keV)",
            energy,
            energy * energy_multiplier,
        )
    else:
        logger.warning("PyMca configuration: set beam energy to %s keV", energy)

    return parameters


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]
