import re

import numpy
from PyMca5.PyMcaPhysics.xrf import ClassMcaTheory
from PyMca5.PyMcaPhysics.xrf.XRFBatchFitOutput import OutputBuffer

from ..types import NotAvailable
from ..types import XRFBatchFitResult
from .calibration import channel_info
from .calibration import energy_info


def parse_native_fit_output(
    xrf_spectra: numpy.ndarray,
    mca_theory: ClassMcaTheory,
    outbuffer: OutputBuffer,
    quantification: bool,
    diagnostics: bool,
    fast_fitting: bool,
) -> XRFBatchFitResult:
    """Convert native PyMca batch fitting result to `XRFBatchFitResult`."""
    channels, raw_channels = channel_info(mca_theory)
    if not fast_fitting:
        start_channel = numpy.round(channels[0]).astype(int)
        stop_channel = numpy.round(channels[-1]).astype(int) + 1
        channels = raw_channels[start_channel:stop_channel]

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

    parameters = outbuffer["parameters"]
    uncertainties = outbuffer["uncertainties"]
    parameter_names = outbuffer.labels("parameters")

    if quantification:
        massfractions_merged = outbuffer["massfractions"]

        mass_fraction_names = []
        layer_indices = {}
        for idx, name in enumerate(outbuffer.labels("massfractions")):
            if isinstance(name, str):
                mass_fraction_names.append(name)
                key = 0
            else:
                _, layer = name
                key = int(layer.lower().split("layer")[1])
            layer_indices.setdefault(key, []).append(idx)

        if len(layer_indices) == 1:
            mass_fractions = massfractions_merged
            mass_fractions_per_layer = []
        else:
            mass_fractions = massfractions_merged[layer_indices[0]]
            mass_fractions_per_layer = [
                massfractions_merged[layer_indices[key]]
                for key in sorted(layer_indices)[1:]
            ]
    else:
        mass_fractions = None
        mass_fraction_names = None
        mass_fractions_per_layer = None

    quantification_parameters = NotAvailable()
    if diagnostics:
        spectra = xrf_spectra
        model = outbuffer["model"]
        residuals = outbuffer["residuals"]
    else:
        spectra = None
        model = None
        residuals = None

    if fast_fitting:
        derivatives = outbuffer["derivatives"]
    else:
        derivatives = None

        parameters = _flatten_3d_last(parameters)
        uncertainties = _flatten_3d_last(uncertainties)
        if quantification:
            mass_fractions = _flatten_3d_last(mass_fractions)
            mass_fractions_per_layer = [
                _flatten_3d_last(arr) for arr in mass_fractions_per_layer
            ]
        if diagnostics:
            model = _flatten_3d_first(model)
            residuals = _flatten_3d_first(residuals)

    return XRFBatchFitResult(
        parameters=parameters,
        uncertainties=uncertainties,
        derivatives=derivatives,
        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=quantification_parameters,
    )


def parse_native_legacy_fit_output(
    xrf_spectra: numpy.ndarray,
    mca_theory: ClassMcaTheory,
    outputdict: dict,
    quantification: bool,
    diagnostics: bool,
    fast_fitting: bool,
) -> XRFBatchFitResult:
    """Convert native PyMca legacy batch fitting result to `XRFBatchFitResult`."""
    channels, raw_channels = channel_info(mca_theory)
    energies = energy_info(mca_theory, channels)
    raw_energies = energy_info(mca_theory, raw_channels)

    parameters = _flatten_3d_last(outputdict["parameters"])
    uncertainties = _flatten_3d_last(outputdict["uncertainties"])
    names = outputdict["names"]
    if quantification:
        noffset = len(outputdict["parameters"])
        parameter_names = names[:noffset]

        massfractions_merged = _flatten_3d_last(outputdict["concentrations"])

        mass_fraction_names = []
        layer_indices = {}
        for idx, name in enumerate(names[noffset:]):
            if "layer" not in name.lower():
                key = 0
                name = re.search(r"\((.*?)\)", name).group(1)
                mass_fraction_names.append(name)
            else:
                key = int(name.lower().split("-layer")[1])
            layer_indices.setdefault(key, []).append(idx)

        if len(layer_indices) == 1:
            mass_fractions = massfractions_merged
            mass_fractions_per_layer = []
        else:
            mass_fractions = massfractions_merged[layer_indices[0]]
            mass_fractions_per_layer = [
                massfractions_merged[layer_indices[key]]
                for key in sorted(layer_indices)[1:]
            ]

    else:
        parameter_names = names
        mass_fractions = None
        mass_fraction_names = None
        mass_fractions_per_layer = None

    if diagnostics:
        spectra = xrf_spectra
    else:
        spectra = None

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


def _flatten_3d_last(arr: numpy.ndarray) -> numpy.ndarray:
    """Reshape array with shape (A, B, C) into shape (A, B*C)"""
    if arr.ndim != 3:
        raise ValueError(f"Numpy is {arr.ndim}D instead of 3D")
    return arr.reshape(arr.shape[0], -1)


def _flatten_3d_first(arr: numpy.ndarray) -> numpy.ndarray:
    """Reshape array with shape (A, B, C) into shape (A*B, C)"""
    if arr.ndim != 3:
        raise ValueError(f"Numpy is {arr.ndim}D instead of 3D")
    return arr.reshape(-1, arr.shape[-1])
