from contextlib import contextmanager
from typing import Dict
from typing import List
from typing import NamedTuple
from typing import Optional
from typing import Sequence
from typing import Tuple
from typing import Union

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

from ...xrffit.pymca_config import PyMcaXrfConfiguration
from ...xrffit.pymca_config import pymca_configdict_from_model
from ...xrffit.pymca_config import pymca_configdict_to_model


class EmissionLineGroup(NamedTuple):
    element: str
    name: str
    counts: Union[float, Sequence]

    @property
    def counts_asarray(self):
        return get_counts_array(self.counts)


class ScatterLineGroup(NamedTuple):
    name: str
    counts: Union[float, Sequence]
    prefix: str = "Scatter"

    @property
    def counts_asarray(self):
        return get_counts_array(self.counts)


def get_counts_array(counts: Union[float, Sequence]):
    if isinstance(counts, Sequence):
        return numpy.asarray(counts)
    return numpy.array((counts))


def xrf_spectra(
    linegroups: Sequence[EmissionLineGroup],
    scattergroups: Sequence[ScatterLineGroup],
    nchannels: int = 1024,
    energy: float = 7.5,
    angle_in: float = 62.0,
    angle_out: float = 49.0,
    flux: float = 1e10,
    elapsed_time: float = 1.0,
    fit_background: bool = False,
    quantification: Optional[str] = None,
    num_layers: int = 1,
) -> Tuple[numpy.ndarray, PyMcaXrfConfiguration]:
    """The returned spectra have the shape `counts.shape + (nchannels,)`"""
    peaks = {linegroup.element: linegroup.name for linegroup in linegroups}
    counts: Dict[str, numpy.ndarray] = {
        f"{group.element} {group.name}": group.counts_asarray for group in linegroups
    }
    counts.update(
        {f"Scatter {group.name}": group.counts_asarray for group in scattergroups}
    )

    # MCA parameters
    zero = 0  # keV
    max_energy = energy + 1.5  # keV
    gain = (max_energy - zero) / nchannels

    # Initialize mca_theory
    mca_theory = McaTheory()

    if quantification == "ref":
        ref_idx = numpy.argmax(
            [linegroup.counts_asarray.mean() for linegroup in linegroups]
        )
        quantification = linegroups[ref_idx].element

    elements = [linegroup.element for linegroup in linegroups]
    configuration = _get_configuration(
        peaks,
        xmin=1,
        xmax=nchannels - 3,
        zero=zero,
        gain=gain,
        energy=energy,
        angle_in=angle_in,
        angle_out=angle_out,
        flux=flux,
        elapsed_time=elapsed_time,
        fit_background=fit_background,
        quantification=quantification,
        elements=elements,
        num_layers=num_layers,
    )
    configuration = pymca_configdict_to_model(
        mca_theory.configure(pymca_configdict_from_model(configuration))
    )

    # Calculate emission lines
    x = numpy.arange(nchannels)
    mca_theory.setData(x=x, y=numpy.zeros(nchannels))
    with _enable_linear_derivatives(mca_theory):
        mca_theory.estimate()

        # Analytical XRF spectrum
        nglobal = mca_theory.NGLOBAL
        if counts:
            shape = next(iter(counts.values())).shape
        else:
            shape = tuple()
        y = numpy.zeros(shape + (nchannels,))
        for i, name in enumerate(mca_theory.PARAMETERS[nglobal:], nglobal):
            normalized_group = mca_theory.linearMcaTheoryDerivative(
                numpy.asarray(mca_theory.parameters), i, x
            )
            assert isinstance(normalized_group, numpy.ndarray)
            if shape:
                y += counts[name][..., None] * normalized_group[None, ...]
            else:
                y += counts[name] * normalized_group
        if fit_background:
            if shape:
                ymin = numpy.clip(numpy.min(y, axis=1), -numpy.inf, 0)
                y += -ymin[:, numpy.newaxis] + 10
            else:
                y += -min(min(y), 0) + 10

    # Configuration for fitting
    mca_theory.config["fit"]["use_limit"] = 1
    mca_theory.config["fit"]["linearfitflag"] = 0

    return y, pymca_configdict_to_model(mca_theory.config)


@contextmanager
def _enable_linear_derivatives(mca_theory: McaTheory):
    linearfitflag = mca_theory.config["fit"]["linearfitflag"]
    mca_theory.config["fit"]["linearfitflag"] = 1
    mca_theory.enableOptimizedLinearFit()
    try:
        yield
    finally:
        mca_theory.disableOptimizedLinearFit()
        mca_theory.config["fit"]["linearfitflag"] = linearfitflag


def _get_configuration(
    peaks: dict,
    xmin: int = 0,
    xmax: int = 0,
    zero: float = 0,
    gain: float = 0.005,
    energy: float = 0,
    angle_in: float = 62.0,
    angle_out: float = 49.0,
    flux: float = 1e10,
    elapsed_time: float = 1.0,
    fit_background: bool = False,
    quantification: Optional[str] = None,
    elements: Optional[List[str]] = None,
    num_layers: int = 1,
) -> PyMcaXrfConfiguration:
    material_config = get_material_config(
        elements=elements,
        num_layers=num_layers,
        angle_in=angle_in,
        angle_out=angle_out,
        quantification=quantification,
    )
    return PyMcaXrfConfiguration(
        peaks=peaks,
        fit=get_fit_config(
            xmin=xmin, xmax=xmax, energy=energy, fit_background=fit_background
        ),
        concentrations=get_quantification_config(
            flux=flux, elapsed_time=elapsed_time, quantification=quantification
        ),
        detector=get_detector(zero=zero, gain=gain),
        peakshape=get_peak_shape(),
        tube=get_tube(),
        **material_config,
    )


def get_quantification_config(
    flux: float = 1e10, elapsed_time: float = 1.0, quantification: Optional[str] = None
) -> dict:
    if quantification:
        if quantification == "fundamental":
            usematrix = 0
            reference = "Auto"
        elif quantification == "auto":
            usematrix = 1
            reference = "Auto"
        else:
            # This is a fix element
            usematrix = 1
            reference = quantification
    else:
        usematrix = 0
        reference = "Auto"
    return {
        "usematrix": usematrix,
        "useattenuators": 1,
        "usemultilayersecondary": 0,
        "usexrfmc": 0,
        "mmolarflag": 0,
        "flux": flux,
        "time": elapsed_time,
        "area": 0.774671,
        "distance": 4.15114,
        "reference": reference,
        "useautotime": 0,
    }


def get_peak_shape():
    return {
        "st_arearatio": 0.0697778,
        "deltast_arearatio": 0.03,
        "fixedst_arearatio": 1,
        "st_sloperatio": 0.0105551,
        "deltast_sloperatio": 0.49,
        "fixedst_sloperatio": 1,
        "lt_arearatio": 0.1,
        "deltalt_arearatio": 0.15,
        "fixedlt_arearatio": 1,
        "lt_sloperatio": 0.1,
        "deltalt_sloperatio": 0.08,
        "fixedlt_sloperatio": 1,
        "step_heightratio": 0.1,
        "deltastep_heightratio": 5e-05,
        "fixedstep_heightratio": 1,
        "eta_factor": 0.02,
        "deltaeta_factor": 0.02,
        "fixedeta_factor": 1,
    }


def get_tube() -> dict:
    return {
        "transmission": 0,
        "voltage": 30.0,
        "anode": "Ag",
        "anodethickness": 0.0002,
        "anodedensity": 10.5,
        "window": "Be",
        "windowthickness": 0.0125,
        "windowdensity": 1.848,
        "filter1": "He",
        "filter1thickness": 0.0,
        "filter1density": 0.000118,
        "alphax": 90.0,
        "alphae": 90.0,
        "deltaplotting": 0.1,
    }


def get_fit_config(
    xmin: int = 0, xmax: int = 0, energy: float = 0, fit_background: bool = False
) -> dict:
    return {
        "deltaonepeak": 0.01,
        "strategy": "SingleLayerStrategy",
        "strategyflag": 0,
        "fitfunction": 0,
        "continuum": 1 if fit_background else 0,
        "fitweight": 1,
        "stripalgorithm": 1,  # SNIP
        "linpolorder": 5,
        "exppolorder": 6,
        "stripconstant": 1.0,
        "snipwidth": 65,
        "stripiterations": 400,
        "stripwidth": 5,
        "stripfilterwidth": 7,
        "stripanchorsflag": 0,
        "maxiter": 50,
        "deltachi": 0.001,
        "xmin": xmin,
        "xmax": xmax,
        "linearfitflag": 1,  # set to 0 after calculating the theoretical spectrum
        "use_limit": 0,  # set to 1 after calculating the theoretical spectrum
        "stripflag": 0 if fit_background else 1,
        "escapeflag": 0,
        "sumflag": 0,
        "scatterflag": 1,
        "hypermetflag": 1,
        "stripanchorslist": [0, 0, 0, 0],
        "energy": [energy],
        "energyweight": [1.0],
        "energyflag": [1],
        "energyscatter": [1],
    }


def get_material_config(
    elements: Optional[List[str]] = None,
    num_layers: int = 1,
    angle_in: float = 62.0,
    angle_out: float = 49.0,
    quantification: Optional[str] = None,
) -> dict:
    if elements is None:
        elements = []
    fixed_internal_reference = quantification not in (None, "fundamental", "auto")
    # WARNING: when quantification=="auto" at least one element in the fit must be in the matrix
    if fixed_internal_reference:
        if quantification not in elements:
            raise ValueError(
                f"'quantification={quantification}' is not in the list of elements {elements}"
            )
        # Put the internal reference in "compound1"
        elements.insert(0, elements.pop(elements.index(quantification)))

    # Split list of elements in layers
    num_elements_per_layer = max(len(elements) // num_layers, 1)
    element_layers = [
        elements[i : i + num_elements_per_layer]
        for i in range(0, len(elements), num_elements_per_layer)
    ]
    num_layers = len(element_layers)

    # Get materials database and define attenuator/filter/detector materials
    materials = get_materials(element_layers)
    attenuators = get_attenuators()

    # Define the matrix
    if num_layers > 1:
        matrix_material = "MULTILAYER"
    else:
        matrix_material = "compound1"
    attenuators["Matrix"] = [
        1,  # enable
        matrix_material,
        2.111,  # density (g/cm^3)
        0.1,  # thickness (cm)
        angle_in,  # incident angle w.r.t sample surface  (deg)
        angle_out,  # outgoing angle w.r.t sample surface (deg)
        0,  # enable scattering angle
        angle_in + angle_out,  # outgoing angle w.r.t primary beam
    ]

    # Define the MULTILAYER composition
    multilayer = {
        f"Layer{i}": [1, f"compound{i}", 2.111, 0.1] for i in range(1, num_layers + 1)
    }

    # Define the iterative fit strategy
    single_layer_strategy = {
        "layer": "Auto",
        "iterations": 3,
        "completer": "-",
        "peaks": "",
        "materials": ["-", "-", "-", "-", "-", "-", "-", "-", "-", "-"],
        "flags": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    }

    return {
        "attenuators": attenuators,
        "materials": materials,
        "multilayer": multilayer,
        "SingleLayerStrategy": single_layer_strategy,
    }


def get_detector(zero: float = 0, gain: float = 0.005) -> dict:
    return {
        "detele": "Si",
        "nthreshold": 4,
        "zero": zero,
        "deltazero": 0.1,
        "fixedzero": 0,  # only free non-linear parameter
        "gain": gain,
        "deltagain": 0.002,
        "fixedgain": 1,
        "noise": -0.0955892,
        "deltanoise": 0.05,
        "fixednoise": 1,
        "fano": 0.0484433,
        "deltafano": 0.114,
        "fixedfano": 1,
        "sum": 0.0,
        "deltasum": 2e-06,
        "fixedsum": 1,
        "ignoreinputcalibration": 0,
    }


def get_attenuators() -> dict:
    return {
        "kapton": [1, "Mylar", 1.4, 0.0004, 1.0],
        "atmosphere": [0, "-", 0.0, 0.0, 1.0],
        "deadlayer": [0, "Si1", 2.33, 0.002, 1.0],
        "absorber": [0, "-", 1.0, 0.1, 1.0],
        "window": [1, "Be", 1.848, 0.0025, 1.0],
        "contact": [0, "moxtekAP3.3", 5.45e-05, 1.0, 1.0],
        "Filter 6": [0, "-", 0.0, 0.0, 1.0],
        "Filter 7": [0, "-", 1.0, 0.1, 1.0],
        "BeamFilter0": [0, "-", 0.0, 0.0, 1.0],
        "BeamFilter1": [0, "-", 0.0, 0.0, 1.0],
        "Detector": [1, "Si1", 2.33, 0.5, 1.0],
    }


def get_materials(element_layers: List[List[str]]) -> dict:
    materials = {**_POLYMERS, **_OTHER_COMPOUNDS}
    for i, elements in enumerate(element_layers, 1):
        num_elements = len(elements)
        materials[f"compound{i}"] = {
            "Comment": f"Compound with elements {' '.join(elements)}",
            "Thickness": 0.0025,
            "Density": 1.42,
            "CompoundFraction": [1.0 / num_elements] * num_elements,
            "CompoundList": [f"{el}1" for el in elements],
        }
    return materials


_POLYMERS = {
    "Mylar": {
        "Comment": "Mylar (Polyethylene Terephthalate) density=1.40 g/cm3",
        "Density": 1.4,
        "CompoundFraction": [0.041959, 0.625017, 0.333025],
        "CompoundList": ["H1", "C1", "O1"],
    },
    "Kapton": {
        "Comment": "Kapton 100 HN 25 micron density=1.42 g/cm3",
        "Thickness": 0.0025,
        "Density": 1.42,
        "CompoundFraction": [0.628772, 0.066659, 0.304569],
        "CompoundList": ["C1", "N1", "O1"],
    },
    "Teflon": {
        "Comment": "Teflon density=2.2 g/cm3",
        "Density": 2.2,
        "CompoundFraction": [0.240183, 0.759817],
        "CompoundList": ["C1", "F1"],
    },
    "Viton": {
        "Comment": "Viton Fluoroelastomer density=1.8 g/cm3",
        "Density": 1.8,
        "CompoundFraction": [0.009417, 0.280555, 0.710028],
        "CompoundList": ["H1", "C1", "F1"],
    },
    "moxtekAP3.3": {
        "Comment": "New Material",
        "Thickness": 1.0,
        "Density": 5.45e-05,
        "CompoundFraction": [0.06, 0.62, 0.055, 0.155, 0.106],
        "CompoundList": ["B", "C", "N", "O", "Al"],
    },
}


_OTHER_COMPOUNDS = {
    "Air": {
        "Comment": "Dry Air (Near sea level) density=0.001204790 g/cm3",
        "Thickness": 1.0,
        "Density": 0.0012048,
        "CompoundFraction": [0.000124, 0.75527, 0.23178, 0.012827, 3.2e-06],
        "CompoundList": ["C1", "N1", "O1", "Ar1", "Kr1"],
    },
    "Water": {
        "Comment": "Water density=1.0 g/cm3",
        "CompoundFraction": 1.0,
        "Thickness": 1.0,
        "Density": 1.0,
        "CompoundList": "H2O1",
    },
    "SingleLayerStrategyMaterial": {
        "Comment": "Last Single Layer Strategy iteration",
        "Thickness": 0.0025,
        "Density": 1.42,
        "CompoundFraction": [],
        "CompoundList": [],
    },
}
