from typing import Optional

import numpy
from PyMca5.PyMcaPhysics.xrf.FastXRFLinearFit import FastXRFLinearFit
from PyMca5.PyMcaPhysics.xrf.LegacyFastXRFLinearFit import (
    FastXRFLinearFit as LegacyFastXRFLinearFit,
)

from ..pymca_config import PyMcaXrfConfiguration
from ..pymca_config import pymca_configdict_from_model
from ._utils.parse_native import parse_native_fit_output
from ._utils.parse_native import parse_native_legacy_fit_output
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:
    """Fast fitting refers to linear fitting of several spectra by solving a single
    system of equations. This function uses the high-level API from PyMca.

    :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

    fastFit = FastXRFLinearFit()

    if configuration.fit.fitweight:
        weight = 1 + bool(individual_weights)
    else:
        weight = 0

    # WARNING: FastXRFLinearFit ignores `use_limit=False`
    #          and always uses xmin/xmax from the configuration.
    if configuration.fit.use_limit:
        xmin = None
        xmax = None
    else:
        xmin = 0
        xmax = xrf_spectra.shape[1]

    with numpy.errstate(over="ignore"):  # Stores results in float32
        outbuffer = fastFit.fitMultipleSpectra(
            y=xrf_spectra.copy(),  # background subtraction in-place
            refit=int(positive_peak_areas),
            concentrations=int(quantification),
            livetime=live_times,
            weight=weight,
            configuration=pymca_configdict_from_model(configuration),
            saveResiduals=True,
            saveFit=True,
            saveData=True,
            diagnostics=True,
            saveFOM=True,
            xmin=xmin,
            xmax=xmax,
        )

    return parse_native_fit_output(
        xrf_spectra,
        fastFit._mcaTheory,
        outbuffer,
        quantification,
        diagnostics,
        fast_fitting=True,
    )


def fit_xrf_spectra_legacy(
    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:
    """Fast fitting refers to linear fitting of several spectra by solving a single
    system of equations. This function uses the legacy high-level API from PyMca.

    :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

    fastFit = LegacyFastXRFLinearFit()

    if configuration.fit.fitweight:
        weight = 1 + bool(individual_weights)
    else:
        weight = 0

    # WARNING: FastXRFLinearFit ignores `use_limit=False`
    #          and always uses xmin/xmax from the configuration.
    if configuration.fit.use_limit:
        xmin = None
        xmax = None
    else:
        xmin = 0
        xmax = xrf_spectra.shape[1]

    y = xrf_spectra.copy()[numpy.newaxis, ...]  # background subtraction in-place

    with numpy.errstate(over="ignore"):  # Stores results in float32
        outputdict = fastFit.fitMultipleSpectra(
            y=y,
            refit=int(positive_peak_areas),
            concentrations=int(quantification),
            livetime=live_times,
            weight=weight,
            configuration=pymca_configdict_from_model(configuration),
            xmin=xmin,
            xmax=xmax,
        )

    return parse_native_legacy_fit_output(
        xrf_spectra,
        fastFit._mcaTheory,
        outputdict,
        quantification,
        diagnostics,
        fast_fitting=True,
    )
