from typing import Optional

import numpy
from PyMca5.PyMcaPhysics.xrf.LegacyMcaAdvancedFitBatch import (
    McaAdvancedFitBatch as LegacyMcaAdvancedFitBatch,
)
from PyMca5.PyMcaPhysics.xrf.McaAdvancedFitBatch import McaAdvancedFitBatch

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:
    """Slow fitting refers to fitting XRF spectra one-by-one.
    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

    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")

    batch = McaAdvancedFitBatch(
        None,
        filelist=xrf_spectra,
        concentrations=int(quantification),
        fitfiles=0,
        fitconcfile=0,
        nosave=True,
        saveResiduals=True,
        saveFit=True,
        saveData=True,
        diagnostics=True,
        saveFOM=True,
    )

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

    _ = batch.mcafit.configure(pymca_configdict_from_model(configuration))

    with numpy.errstate(over="ignore"):  # Stores results in float32
        batch.processList()

    return parse_native_fit_output(
        xrf_spectra,
        batch.mcafit,
        batch.outbuffer,
        quantification,
        diagnostics,
        fast_fitting=False,
    )


class _LegacyMcaAdvancedFitBatch(LegacyMcaAdvancedFitBatch):
    def __log(self, text):
        # Base class does `print`
        pass


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:
    """Slow fitting refers to fitting XRF spectra one-by-one.
    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

    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")

    batch = _LegacyMcaAdvancedFitBatch(
        None,
        filelist=xrf_spectra,
        concentrations=int(quantification),
        fitfiles=0,
        nosave=True,
    )

    class Dummy:
        def getConcentrationsAsAscii(self, _):
            # Prevent saving of _concentrations.txt file
            return []

    batch._toolConversion = Dummy()

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

    _ = batch.mcafit.configure(pymca_configdict_from_model(configuration))

    with numpy.errstate(over="ignore"):  # Stores results in float32
        batch.processList()

    images = batch._McaAdvancedFitBatch__images
    sigmas = batch._McaAdvancedFitBatch__sigmas

    parameter_names = list(sigmas)
    parameters = numpy.array([images[name] for name in parameter_names])
    uncertainties = numpy.array([sigmas[name] for name in parameter_names])
    outputdict = {"parameters": parameters, "uncertainties": uncertainties}

    names = list(parameter_names)
    if quantification:
        concentrations = []
        for name in parameter_names:
            key = f"{name} mass fraction"
            if key in images:
                concentrations.append(images[key])
                names.append(f"C({name})")

        layeri = 0
        has_layeri = True
        while has_layeri:
            layeri += 1
            has_layeri = False
            for name in parameter_names:
                key = f"{name} Layer{layeri}"
                if key in images:
                    concentrations.append(images[key])
                    names.append(f"C({name})-Layer{layeri}")
                    has_layeri = True

        outputdict["concentrations"] = numpy.array(concentrations)

    outputdict["names"] = names

    return parse_native_legacy_fit_output(
        xrf_spectra,
        batch.mcafit,
        outputdict,
        quantification,
        diagnostics,
        fast_fitting=False,
    )
