"""
.. _lhc-queries:

**Querying Utilities**

The functions below are settings query utilities for the ``LHC``.
"""
from logging import getLogger
from typing import Dict, Sequence, Union

import tfs

from cpymad.madx import Madx

from cpymadtools import twiss
from cpymadtools.constants import (
    LHC_KCD_KNOBS,
    LHC_KCO_KNOBS,
    LHC_KCOSX_KNOBS,
    LHC_KCOX_KNOBS,
    LHC_KCS_KNOBS,
    LHC_KCSSX_KNOBS,
    LHC_KCSX_KNOBS,
    LHC_KCTX_KNOBS,
    LHC_KO_KNOBS,
    LHC_KQS_KNOBS,
    LHC_KQSX_KNOBS,
    LHC_KQTF_KNOBS,
    LHC_KSF_KNOBS,
    LHC_KSS_KNOBS,
)
from cpymadtools.lhc._setup import lhc_orbit_variables
from cpymadtools.utils import _get_k_strings

logger = getLogger(__name__)


def get_magnets_powering(
    madx: Madx, patterns: Sequence[str] = [r"^mb\.", r"^mq\.", r"^ms\."], brho: Union[str, float] = None, **kwargs
) -> tfs.TfsDataFrame:
    r"""
    .. versionadded:: 0.17.0

    Gets the twiss table with additional defined columns for the given *patterns*.

    .. note::
        Here are below certain useful patterns for the ``LHC`` and their meaning:

        * ``^mb\.`` :math:`\rightarrow` main bends.
        * ``^mq\.`` :math:`\rightarrow` main quadrupoles.
        * ``^ms\.`` :math:`\rightarrow` main sextupoles.
        * ``^mb[rswx]`` :math:`\rightarrow` separation dipoles.
        * ``^mq[mwxy]`` :math:`\rightarrow` insertion quads.
        * ``^mqt.1[23]`` :math:`\rightarrow` short tuning quads (12 & 13).
        * ``^mqtl`` :math:`\rightarrow` long  tuning quads.
        * ``^mcbx`` :math:`\rightarrow` crossing scheme magnets.
        * ``^mcb[cy]`` :math:`\rightarrow` crossing scheme magnets.

        To make no selection, one can give ``patterns=[""]`` and this will give back
        the results for *all* elements. One can also give a specific magnet's exact
        name to include it in the results.

    .. note::
        The ``TWISS`` flag will be fully cleared after running this function.

    Args:
        madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object.
        patterns (Sequence[str]): a list of regex patterns to define which elements
            should be selected and included in the returned table. Defaults to selecting
            the main bends, quads and sextupoles. See the note admonition above for
            useful patterns to select specific ``LHC`` magnet families.
        brho (Union[str, float]): optional, an explicit definition for the magnetic
            rigidity in :math:`Tm^{-1}`. If not given, it will be assumed that
            a ``brho`` quantity is defined in the ``MAD-X`` globals.
        **kwargs: any keyword argument will be passed to `~.twiss.get_pattern_twiss` and
            later on to the ``TWISS`` command executed in ``MAD-X``.

    Returns:
        A `~tfs.TfsDataFrame` of the ``TWISS`` table, with the relevant newly defined columns
        and including the elements matching the regex *patterns* that were provided.

    Example:
        .. code-block:: python

            >>> sextupoles_powering = get_magnets_powering(madx, patterns=[r"^ms\."])
    """
    logger.debug("Computing magnets field and powering limits proportions")
    NEW_COLNAMES = ["name", "keyword", "ampere", "imax", "percent", "kn", "kmax", "integrated_field", "L"]
    NEW_COLNAMES = list(set(NEW_COLNAMES + kwargs.pop("columns", [])))  # in case user gives explicit columns
    _list_field_currents(madx, brho=brho)
    return twiss.get_pattern_twiss(madx, columns=NEW_COLNAMES, patterns=patterns, **kwargs)


def query_arc_correctors_powering(madx: Madx) -> Dict[str, float]:
    """
    .. versionadded:: 0.15.0

    Queries for the arc corrector strengths and returns their values as a percentage of
    their max powering. This is a port of one of the **corr_value.madx** file's macros

    Args:
        madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object with an
            active (HL)LHC sequence.

    Returns:
        A `dict` with the percentage for each corrector.

    Example:
        .. code-block:: python

            >>> arc_knobs = query_arc_correctors_powering(madx)
    """
    logger.debug("Querying triplets correctors powering")
    result: Dict[str, float] = {}

    logger.debug("Querying arc tune trim quadrupole correctors (MQTs) powering")
    k_mqt_max = 120 / madx.globals.brho  # 120 T/m
    result.update({knob: 100 * _knob_value(madx, knob) / k_mqt_max for knob in LHC_KQTF_KNOBS})

    logger.debug("Querying arc short straight sections skew quadrupole correctors (MQSs) powering")
    k_mqs_max = 120 / madx.globals.brho  # 120 T/m
    result.update({knob: 100 * _knob_value(madx, knob) / k_mqs_max for knob in LHC_KQS_KNOBS})

    logger.debug("Querying arc sextupole correctors (MSs) powering")
    k_ms_max = 1.280 * 2 / 0.017 ** 2 / madx.globals.brho  # 1.28 T @ 17 mm
    result.update({knob: 100 * _knob_value(madx, knob) / k_ms_max for knob in LHC_KSF_KNOBS})

    logger.debug("Querying arc skew sextupole correctors (MSSs) powering")
    k_mss_max = 1.280 * 2 / 0.017 ** 2 / madx.globals.brho  # 1.28 T @ 17 mm
    result.update({knob: 100 * _knob_value(madx, knob) / k_mss_max for knob in LHC_KSS_KNOBS})

    logger.debug("Querying arc spool piece (skew) sextupole correctors (MCSs) powering")
    k_mcs_max = 0.471 * 2 / 0.017 ** 2 / madx.globals.brho  # 0.471 T @ 17 mm
    result.update({knob: 100 * _knob_value(madx, knob) / k_mcs_max for knob in LHC_KCS_KNOBS})

    logger.debug("Querying arc spool piece (skew) octupole correctors (MCOs) powering")
    k_mco_max = 0.040 * 6 / 0.017 ** 3 / madx.globals.brho  # 0.04 T @ 17 mm
    result.update({knob: 100 * _knob_value(madx, knob) / k_mco_max for knob in LHC_KCO_KNOBS})

    logger.debug("Querying arc spool piece (skew) decapole correctors (MCDs) powering")
    k_mcd_max = 0.100 * 24 / 0.017 ** 4 / madx.globals.brho  # 0.1 T @ 17 mm
    result.update({knob: 100 * _knob_value(madx, knob) / k_mcd_max for knob in LHC_KCD_KNOBS})

    logger.debug("Querying arc short straight sections octupole correctors (MOs) powering")
    k_mo_max = 0.29 * 6 / 0.017 ** 3 / madx.globals.brho  # 0.29 T @ 17 mm
    result.update({knob: 100 * _knob_value(madx, knob) / k_mo_max for knob in LHC_KO_KNOBS})
    return result


def query_triplet_correctors_powering(madx: Madx) -> Dict[str, float]:
    """
    .. versionadded:: 0.15.0

    Queries for the triplet corrector strengths and returns their values as a percentage of
    their max powering. This is a port of one of the **corr_value.madx** file's macros.

    Args:
        madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object with an
            active (HL)LHC sequence.

    Returns:
        A `dict` with the percentage for each corrector.

    Example:
        .. code-block:: python

            >>> triplet_knobs = query_triplet_correctors_powering(madx)
    """
    logger.debug("Querying triplets correctors powering")
    result: Dict[str, float] = {}

    logger.debug("Querying triplet skew quadrupole correctors (MQSXs) powering")
    k_mqsx_max = 1.360 / 0.017 / madx.globals.brho  # 1.36 T @ 17mm
    result.update({knob: 100 * _knob_value(madx, knob) / k_mqsx_max for knob in LHC_KQSX_KNOBS})

    logger.debug("Querying triplet sextupole correctors (MCSXs) powering")
    k_mcsx_max = 0.028 * 2 / 0.017 ** 2 / madx.globals.brho  # 0.028 T @ 17 mm
    result.update({knob: 100 * _knob_value(madx, knob) / k_mcsx_max for knob in LHC_KCSX_KNOBS})

    logger.debug("Querying triplet skew sextupole correctors (MCSSXs) powering")
    k_mcssx_max = 0.11 * 2 / 0.017 ** 2 / madx.globals.brho  # 0.11 T @ 17 mm
    result.update({knob: 100 * _knob_value(madx, knob) / k_mcssx_max for knob in LHC_KCSSX_KNOBS})

    logger.debug("Querying triplet octupole correctors (MCOXs) powering")
    k_mcox_max = 0.045 * 6 / 0.017 ** 3 / madx.globals.brho  # 0.045 T @ 17 mm
    result.update({knob: 100 * _knob_value(madx, knob) / k_mcox_max for knob in LHC_KCOX_KNOBS})

    logger.debug("Querying triplet skew octupole correctors (MCOSXs) powering")
    k_mcosx_max = 0.048 * 6 / 0.017 ** 3 / madx.globals.brho  # 0.048 T @ 17 mm
    result.update({knob: 100 * _knob_value(madx, knob) / k_mcosx_max for knob in LHC_KCOSX_KNOBS})

    logger.debug("Querying triplet decapole correctors (MCTXs) powering")
    k_mctx_max = 0.01 * 120 / 0.017 ** 5 / madx.globals.brho  # 0.010 T @ 17 mm
    result.update({knob: 100 * _knob_value(madx, knob) / k_mctx_max for knob in LHC_KCTX_KNOBS})
    return result


def get_current_orbit_setup(madx: Madx) -> Dict[str, float]:
    """
    .. versionadded:: 0.8.0

    Get the current values for the (HL)LHC orbit variables. Initial implementation credits go to
    :user:`Joschua Dilly <joschd>`.

    Args:
        madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object.

    Returns:
        A `dict` of all orbit variables set, and their values as set in the ``MAD-X`` globals.

    Example:
        .. code-block:: python

            >>> orbit_setup = get_current_orbit_setup(madx)
    """
    logger.debug("Extracting orbit variables from global table")
    variables, specials = lhc_orbit_variables()
    return {orbit_variable: madx.globals[orbit_variable] for orbit_variable in variables + list(specials.keys())}


# ----- Helpers ----- #


def _list_field_currents(madx: Madx, brho: Union[str, float] = None) -> None:
    """
    Creates additional columns for the ``TWISS`` table with the magnets' total fields
    and currents, to help later on determine which proportion of their maximum powering
    the current setting is using. This is an implementation of the old utility script
    located at **/afs/cern.ch/eng/lhc/optics/V6.503/toolkit/list_fields_currents.madx**.

    .. important::
        Certain quantities are assumed to be defined in the ``MAD-X`` globals, such as
        ``brho``, or available in the magnets definition, such as ``calib``. For this
        reason, this script most likely only works for the ``(HL)LHC`` sequences where
        those are defined.

    Args:
        madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object.
        brho (Union[str, float]): optional, an explicit definition for the magnetic
            rigidity in :math:`Tm^{-1}`. If not given, it will be assumed that
            a ``brho`` quantity is defined in the ``MAD-X`` globals and this one will
            be used.
    """
    logger.debug("Creating additional TWISS table columns for magnets' fields and currents")

    if brho is not None:
        logger.debug(f"Setting 'brho' to explicitely defined '{brho}'")
        madx.globals["brho"] = brho

    # Define strength := table(twiss, k0l) + ... + table(twiss, k5sl) +  table(twiss, hkick)  +  table(twiss, vkick);
    madx.globals["strength"] = (
        " + ".join(f"table(twiss, {a.lower()})" for a in _get_k_strings(stop=6))
        + " +  table(twiss, hkick)  +  table(twiss, vkick)"
    )

    # All here are given as strings to make it deferred expressions in MAD-X
    madx.globals["epsilon"] = 1e-20  # to avoid divisions by zero
    madx.globals["length"] = "table(twiss, l) + table(twiss, lrad) + epsilon"
    madx.globals["kmaxx"] = "table(twiss, kmax) + epsilon"
    madx.globals["calibration"] = "table(twiss, calib) + epsilon"
    madx.globals["kn"] = "abs(strength) / length"
    # madx.globals["rho"] = "kn / (kn + epsilon) / (kn + epsilon)"

    madx.globals["field"] = "kn * brho"
    madx.globals["percent"] = "field * 100 / (kmaxx + epsilon)"
    madx.globals["ampere"] = "field / calibration"
    madx.globals["imax"] = "kmaxx / calibration"
    madx.globals["integrated_field"] = "field * length"


def _knob_value(madx: Madx, knob: str) -> float:
    """
    Queryies the current value of a given *knob* name in the ``MAD-X`` process, and defaults
    to 0 (as ``MAD-X`` does) in case that knob has not been defined in the current process.

    Args:
        madx (cpymad.madx.Madx): an instanciated `~cpymad.madx.Madx` object.
        knob (str): the name the knob.

    Returns:
        The knob value if it was defined, otherwise 0.

    Example:
        .. code-block:: python

            >>> _knob_value(madx, knob="underfined_for_sure")
            0
    """
    try:
        return madx.globals[knob]
    except KeyError:  # cpymad gives a 'Variable not defined: var_name'
        return 0
