from typing import Callable
from typing import List
from typing import Tuple
from typing import Union

import numpy
from PyMca5.PyMcaPhysics.xrf.ClassMcaTheory import ClassMcaTheory
from PyMca5.PyMcaPhysics.xrf.ConcentrationsTool import ConcentrationsTool

from ...pymca_config import PyMcaXrfConfiguration
from ..types import QuantificationParameters


def calculate_massfractions(
    mca_theory: ClassMcaTheory,
    peak_group_areas: numpy.ndarray,
    peak_group_names: List[str],
    configuration: PyMcaXrfConfiguration,
    norm_factors: Union[float, numpy.ndarray],
) -> Tuple[numpy.ndarray, List[numpy.ndarray], List[str], QuantificationParameters]:
    """Calculate mass fractions of the fit of many spectra based on the fit of a reference spectrum and
    fundamental parameters or an internal reference.

    Without internal reference (fundamental parameters):

        mass_fraction(group) = norm_factors * peak_area(group)  / ref_fit_peak_area(group)  * ref_fit_mass_fraction(group)

    With internal reference:

        mass_fraction(group) = peak_area(group) / peak_area(int) * ref_fit_peak_area(int) / ref_fit_peak_area(group) * ref_fit_mass_fraction(group)

    :param mca_theory:
    :param peak_group_areas: shape ``(num_parameters, num_profiles)``
    :param peak_group_names:
    :param configuration:
    :param norm_factors: scalar or shape ``(num_profiles, )`` (only for fundamental parameters)
    :returns: mass fractions with shape ``(num_parameters, num_profiles)``,
              mass fractions per layer with shape ``(num_layers, num_parameters, num_profiles)``,
              mass fraction names,
              internal reference element
    """

    ctool = ConcentrationsTool()
    ctool_config = ctool.configure()
    ctool_config.update(configuration.concentrations.model_dump(mode="python"))

    # Prepare fit results
    usematrix = configuration.concentrations.usematrix
    reference = configuration.concentrations.reference
    select_fitreference = usematrix and reference.upper() == "AUTO"
    configuration = configuration.model_dump(mode="python")  # copy
    if select_fitreference:
        # Select internal reference
        _ = mca_theory.startfit(digest=0, linear=True)
        fitresult = {"result": mca_theory.imagingDigestResult()}
        fitresult["result"]["config"] = configuration
        _, add_info = ctool.processFitResult(
            config=ctool_config,
            fitresult=fitresult,
            elementsfrommatrix=False,
            fluorates=mca_theory._fluoRates,
            addinfo=True,
        )
        for group in fitresult["result"]["groups"]:
            if fitresult["result"][group]["fitarea"] <= 0.0:
                fitresult["result"][group]["fitarea"] = 1.0e-6
        configuration["concentrations"]["reference"] = add_info["ReferenceElement"]
    else:
        fitresult = {"result": {"config": configuration, "groups": peak_group_names}}
        for name in peak_group_names:
            # We just need the ratio between peak area and concentration
            fitresult["result"][name] = {"fitarea": 1.0, "sigmaarea": 1.0}

    # Calculate mass fractions of reference fit
    wresult, add_info = ctool.processFitResult(
        config=ctool_config,
        fitresult=fitresult,
        elementsfrommatrix=False,
        fluorates=mca_theory._fluoRates,
        addinfo=True,
    )
    quantification_parameters = QuantificationParameters(**add_info)

    # Choose calculation mode
    if quantification_parameters.ReferenceElement in ["", None, "None"]:
        # Fundamental parameters

        def _calc_mass_fraction(
            name, peak_area, ref_fit_peak_area, ref_fit_mass_fraction, layer=None
        ):
            return _calc_mass_fraction_fp(
                peak_area, ref_fit_peak_area, ref_fit_mass_fraction, norm_factors
            )

        mass_fractions, mass_fractions_per_layer, mass_fraction_names = (
            _calc_mass_fractions(
                fitresult["result"]["groups"],
                peak_group_areas,
                wresult,
                fitresult,
                _calc_mass_fraction,
            )
        )
    else:
        # Internal reference
        reference_group = (
            quantification_parameters.ReferenceElement
            + " "
            + quantification_parameters.ReferenceTransitions.split()[0]
        )
        idx_int = fitresult["result"]["groups"].index(reference_group)
        name_int = fitresult["result"]["groups"][idx_int]
        peak_area_int = peak_group_areas[idx_int]
        ref_fit_peak_area_int = fitresult["result"][name_int]["fitarea"]

        def _calc_mass_fraction(
            name, peak_area, ref_fit_peak_area, ref_fit_mass_fraction, layer=None
        ):
            if name == name_int:
                if layer is None:
                    return [wresult["mass fraction"][name]] * len(peak_area)
                return [wresult[layer]["mass fraction"][name]] * len(peak_area)

            return _calc_mass_fraction_int(
                peak_area,
                ref_fit_peak_area,
                ref_fit_mass_fraction,
                peak_area_int,
                ref_fit_peak_area_int,
            )

        mass_fractions, mass_fractions_per_layer, mass_fraction_names = (
            _calc_mass_fractions(
                fitresult["result"]["groups"],
                peak_group_areas,
                wresult,
                fitresult,
                _calc_mass_fraction,
            )
        )

    # Convert lists to numpy arrays
    mass_fractions_per_layer = [numpy.asarray(lst) for lst in mass_fractions_per_layer]
    return (
        numpy.asarray(mass_fractions),
        mass_fractions_per_layer,
        mass_fraction_names,
        quantification_parameters,
    )


def _calc_mass_fraction_fp(
    peak_area: numpy.ndarray,
    ref_fit_peak_area: float,
    ref_fit_mass_fraction: float,
    norm_factors: Union[float, numpy.ndarray],
) -> numpy.ndarray:
    """Mass fraction calculation for fundamental parameters."""
    factor = ref_fit_mass_fraction / ref_fit_peak_area
    return factor * (norm_factors * peak_area)


def _calc_mass_fraction_int(
    peak_area: numpy.ndarray,
    ref_fit_peak_area: float,
    ref_fit_mass_fraction: float,
    peak_area_int: float,
    ref_fit_peak_area_int: float,
) -> numpy.ndarray:
    """Mass fraction calculation for internal reference."""
    factor_num = ref_fit_mass_fraction * ref_fit_peak_area_int
    factor_denom = peak_area_int * ref_fit_peak_area
    return (factor_num / factor_denom) * peak_area


def _calc_mass_fractions(
    groups: List[str],
    peak_group_areas: numpy.ndarray,
    wresult: dict,
    fitresult: dict,
    calc_mass_fraction: Callable,
) -> Tuple[List[numpy.ndarray], List[List[numpy.ndarray]], List[str]]:
    """Iterate over groups and calculate mass fractions + per-layer fractions."""
    mass_fraction_names = []
    mass_fractions = []
    nlayers = len(wresult["layerlist"])
    mass_fractions_per_layer = [[] for _ in range(nlayers)]

    for name, peak_area in zip(groups, peak_group_areas):
        if name.lower().startswith("scatter"):
            continue

        mass_fraction_names.append(name)
        ref_fit_peak_area = fitresult["result"][name]["fitarea"]

        # Overall mass fraction
        ref_fit_mass_fraction = wresult["mass fraction"][name]
        mass_fractions.append(
            calc_mass_fraction(
                name, peak_area, ref_fit_peak_area, ref_fit_mass_fraction
            )
        )

        # Mass fraction per layer
        for layer, layer_mass_fractions in zip(
            wresult["layerlist"], mass_fractions_per_layer
        ):
            ref_fit_mass_fraction = wresult[layer]["mass fraction"][name]
            layer_mass_fractions.append(
                calc_mass_fraction(
                    name,
                    peak_area,
                    ref_fit_peak_area,
                    ref_fit_mass_fraction,
                    layer=layer,
                )
            )

    return mass_fractions, mass_fractions_per_layer, mass_fraction_names
