import logging
from typing import Dict
from typing import List

from PyMca5.PyMcaPhysics.xrf import Elements
from PyMca5.PyMcaPhysics.xrf import KShell
from PyMca5.PyMcaPhysics.xrf import LShell
from PyMca5.PyMcaPhysics.xrf import MShell

from ..model import PyMcaXrfConfiguration

logger = logging.getLogger(__name__)


def split_mlines(configuration: PyMcaXrfConfiguration, mlines: Dict[str, List[str]]):
    """
    :param configuration:
    :param dict mlines: for example `{"Pb":["M4", "M5]}}`
    """
    split = False
    for element, new_lines in mlines.items():
        if element in configuration.peaks:
            old_lines = configuration.peaks[element]
            if "M" in old_lines:
                lines = [group for group in old_lines if group != "M"]
                configuration.peaks[element] = lines + new_lines
                logger.warning(
                    "PyMca configuration: modified peak groups of %s to %s",
                    element,
                    configuration.peaks[element],
                )
                split = True

    if split:
        _patch_mlines()
    else:
        _unpatch_mlines()


def _patch_mlines():
    global _ORIGINAL_NAMES, _ORIGINAL_TRANSITIONS, _ORIGINAL_RATES
    if Elements.ElementXrays == _PATCHED_NAMES:
        return
    if _ORIGINAL_NAMES is None:
        _ORIGINAL_NAMES = Elements.ElementXray
        _ORIGINAL_TRANSITIONS = Elements.ElementShellTransitions
        _ORIGINAL_RATES = Elements.ElementShellRates

    Elements.ElementXrays = _PATCHED_NAMES
    Elements.ElementShellTransitions = _PATCHED_TRANSITIONS
    Elements.ElementShellRates = _PATCHED_RATES


def _unpatch_mlines():
    if Elements.ElementXrays == _ORIGINAL_NAMES or _ORIGINAL_NAMES is None:
        return
    Elements.ElementXrays = _ORIGINAL_NAMES
    Elements.ElementShellTransitions = _ORIGINAL_TRANSITIONS
    Elements.ElementShellRates = _ORIGINAL_RATES


_ORIGINAL_NAMES = None
_ORIGINAL_TRANSITIONS = None
_ORIGINAL_RATES = None


_PATCHED_NAMES = [
    "K xrays",
    "Ka xrays",
    "Kb xrays",
    "L xrays",
    "L1 xrays",
    "L2 xrays",
    "L3 xrays",
    "M xrays",
    "M1 xrays",
    "M2 xrays",
    "M3 xrays",
    "M4 xrays",
    "M5 xrays",
]
_PATCHED_TRANSITIONS = [
    KShell.ElementKShellTransitions,
    KShell.ElementKAlphaTransitions,
    KShell.ElementKBetaTransitions,
    LShell.ElementLShellTransitions,
    LShell.ElementL1ShellTransitions,
    LShell.ElementL2ShellTransitions,
    LShell.ElementL3ShellTransitions,
    [s + "*" for s in MShell.ElementMShellTransitions],
    MShell.ElementM1ShellTransitions,
    MShell.ElementM2ShellTransitions,
    MShell.ElementM3ShellTransitions,
    MShell.ElementM4ShellTransitions,
    MShell.ElementM5ShellTransitions,
]
_PATCHED_RATES = [
    KShell.ElementKShellRates,
    KShell.ElementKAlphaRates,
    KShell.ElementKBetaRates,
    LShell.ElementLShellRates,
    LShell.ElementL1ShellRates,
    LShell.ElementL2ShellRates,
    LShell.ElementL3ShellRates,
    MShell.ElementMShellRates,
    MShell.ElementM1ShellRates,
    MShell.ElementM2ShellRates,
    MShell.ElementM3ShellRates,
    MShell.ElementM4ShellRates,
    MShell.ElementM5ShellRates,
]
