from typing import Optional

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

from ..pymca_config import PyMcaXrfConfiguration
from ..pymca_config import pymca_configdict_from_model
from ..pymca_config import pymca_configdict_to_model
from ._utils.calibration import channel_info
from ._utils.calibration import energy_info
from ._utils.limits import expand_data
from .types import NotAvailable
from .types import XRFBatchFitResult


def fit_xrf_spectra(
    xrf_spectra: numpy.ndarray,
    configuration: PyMcaXrfConfiguration,
    quantification: Optional[bool] = None,
    individual_weights: Optional[bool] = None,
    live_times: Optional[numpy.ndarray] = None,
    positive_peak_areas: Optional[bool] = None,
    diagnostics: Optional[bool] = None,
) -> XRFBatchFitResult:
    """Slow fitting refers to fitting XRF spectra one-by-one.
    This function does not use the high-level API from PyMca.
    It uses the low-level API to

    - calculate elemental spectra
    - calculate mass fractions from fit results
    - strip spectral background
    - and perform linear least-squares fit.

    :param xrf_spectra: shape `(num_spectra, num_channels)`
    :param configuration: PyMca configuration.
    :param quantification: Calculate mass fractions from peak area's.
    :param individual_weights: When fitting with weights, use the weight of each
                                spectrum (slow) instead of the average weight.
    :param live_times: 1D array with one value for each XRF spectrum.
                       Only used for mass fractions with fundamental parameters.
    :param positive_peak_areas:
    :param diagnostics: fit model and residuals.
    """

    if quantification is None:
        quantification = True
    if individual_weights is None:
        individual_weights = True
    if positive_peak_areas is None:
        positive_peak_areas = True
    if diagnostics is None:
        diagnostics = False

    if not positive_peak_areas:
        raise NotImplementedError("'positive_peak_areas' must be True")
    if live_times is not None:
        raise NotImplementedError("'live_times' is not supported")
    if not individual_weights:
        raise NotImplementedError("'individual_weights' must be True")

    mca_theory = ClassMcaTheory()
    configuration = pymca_configdict_to_model(
        mca_theory.configure(pymca_configdict_from_model(configuration))
    )

    parameters = []
    uncertainties = []
    model = []
    residuals = []
    parameter_names = mca_theory.PARAMETERS
    if quantification:
        mass_fractions = []
        mass_fraction_names = None
        mass_fractions_per_layer = None

        ctool = ConcentrationsTool()
        ctool_config = ctool.configure()
        ctool_config.update(configuration.concentrations.model_dump(mode="python"))
        for spectrum in xrf_spectra:
            mca_theory.setData(y=spectrum)
            mca_theory.estimate()
            fitresult, result = mca_theory.startfit(digest=1)

            fittedpar, chisq, sigmapar, niter, lastdeltachi = fitresult
            parameters.append(fittedpar)
            uncertainties.append(sigmapar)
            if diagnostics:
                model.append(result["yfit"])
                residuals.append(result["yfit"] - result["ydata"])

            digest = {}
            digest["fitresult"] = fitresult
            digest["result"] = result
            digest["result"]["config"] = mca_theory.configure()  # deepcopy
            wresult = ctool.processFitResult(
                config=ctool_config,
                fitresult=digest,
                elementsfrommatrix=False,
                fluorates=mca_theory._fluoRates,
            )

            if mass_fraction_names is None:
                nlayers = len(wresult["layerlist"])
                mass_fraction_names = wresult["groups"]
                mass_fractions_per_layer = [[] for _ in range(nlayers)]

            mass_fractions.append(
                [wresult["mass fraction"][name] for name in mass_fraction_names]
            )
            for layer, lst_dest in zip(wresult["layerlist"], mass_fractions_per_layer):
                lst_dest.append(
                    [
                        wresult[layer]["mass fraction"][name]
                        for name in mass_fraction_names
                    ]
                )

        mass_fractions = numpy.asarray(mass_fractions).T
        mass_fractions_per_layer = [
            numpy.asarray(mf_layer).T for mf_layer in mass_fractions_per_layer
        ]
    else:
        mass_fractions = None
        mass_fractions_per_layer = None
        mass_fraction_names = None

        for spectrum in xrf_spectra:
            mca_theory.setData(y=spectrum)
            mca_theory.estimate()
            fitresult, result = mca_theory.startfit(digest=1)

            fittedpar, chisq, sigmapar, niter, lastdeltachi = fitresult
            parameters.append(fittedpar)
            uncertainties.append(sigmapar)

            if diagnostics:
                model.append(result["yfit"])
                residuals.append(result["yfit"] - result["ydata"])

    parameters = numpy.stack(parameters, axis=1)
    uncertainties = numpy.stack(uncertainties, axis=1)

    channels, raw_channels = channel_info(mca_theory)
    num_fit_observations = len(channels)
    num_raw_observations = len(raw_channels)

    if diagnostics:
        spectra = xrf_spectra
        model = numpy.array(model)
        residuals = numpy.array(residuals)
        if num_fit_observations != num_raw_observations:
            # Expand with NaN's so the model, residuals and base profiles match to
            # orginal number of observations.
            model = expand_data(channels, model, num_raw_observations, axis=1)
            residuals = expand_data(channels, residuals, num_raw_observations, axis=1)
    else:
        spectra = None
        model = None
        residuals = None

    energies = energy_info(mca_theory, channels)
    raw_energies = energy_info(mca_theory, raw_channels)

    return XRFBatchFitResult(
        parameters=parameters,
        uncertainties=uncertainties,
        derivatives=None,
        parameter_names=parameter_names,
        channels=channels,
        energies=energies,
        raw_channels=raw_channels,
        raw_energies=raw_energies,
        spectra=spectra,
        model=model,
        residuals=residuals,
        mass_fractions=mass_fractions,
        mass_fractions_per_layer=mass_fractions_per_layer,
        mass_fraction_names=mass_fraction_names,
        quantification_parameters=NotAvailable(),
    )
