r"""Definition of :py:class:`.Observable` objects used in matching and
response matrices.

Observables are ways to monitor various lattice parameters and use them in
matching and correction procedures.

Class hierarchy
---------------

:py:class:`Observable`\ (evalfun, name, target, weight, bounds, ...)
    :py:class:`RingObservable`\ (fun, ...)

    :py:func:`GlobalOpticsObservable`\ (attrname, plane, name, ...)

    :py:class:`EmittanceObservable`\ (attrname, plane, name, ...)

    :py:class:`ElementObservable`\ (fun, ...)
        :py:class:`OrbitObservable`\ (refpts, axis, name, ...)

        :py:class:`TrajectoryObservable`\ (refpts, axis, name, ...)

        :py:class:`MatrixObservable`\ (refpts, axis, name, ...)

        :py:class:`LocalOpticsObservable`\ (refpts, attrname, axis, name, ...)

        :py:class:`LatticeObservable`\ (refpts, attrname, index, name, ...)

        :py:class:`GeometryObservable`\ (refpts, attrname, name, ...)

        :py:class:`.RDTObservable`\ (refpts, RDTname, name, ...)

:py:class:`.Observable`\ s are usually not evaluated directly, but through a
container which performs the required optics computation and feeds each
:py:class:`.Observable` with its specific data. After evaluation, each
:py:class:`.Observable` provides the following properties:

- :py:attr:`~Observable.value`
- :py:attr:`~Observable.weight`
- :py:attr:`~Observable.weighted_value`: :pycode:`value / weight`
- :py:attr:`~Observable.deviation`:  :pycode:`value - target`
- :py:attr:`~Observable.residual`:  :pycode:`((value - target)/weight)**2`
"""

from __future__ import annotations

__all__ = [
    "ElementObservable",
    "EmittanceObservable",
    "GeometryObservable",
    "GlobalOpticsObservable",
    "LatticeObservable",
    "LocalOpticsObservable",
    "MatrixObservable",
    "Need",
    "Observable",
    "OrbitObservable",
    "RingObservable",
    "TrajectoryObservable",
]

from collections.abc import Callable, Set as AbstractSet, Mapping
from functools import partial
from enum import Enum
from itertools import repeat
from typing import ClassVar, Any

import numpy as np
import numpy.typing as npt

from ..lattice import AtError, AxisDef, axis_, plane_
from ..lattice import Lattice, Refpts, End

RefIndex = int | tuple[int, ...] | slice


# Observables must be pickleable. For this, the evaluation function must be a
# module-level function. No inner, nested function is allowed. So nested
# functions are replaced be module-level callable class instances:


class _Convolve:
    def __init__(self, modfun, fun, *args, **kwargs):
        self.modfun = modfun
        self.fun = fun
        self.args = args
        self.kwargs = kwargs

    def __call__(self, *a, **__):
        return self.modfun(self.fun(*a), *self.args, **self.kwargs)


class _ArrayAccess:
    """Access a selected item in an array."""

    def __init__(self, index):
        self.index = _all_rows(index)

    def __call__(self, data, **__):
        index = self.index
        return data if index is None else data[self.index]


def _record_access(param, index, data, **__):
    """Access a selected item in a record array."""
    val = getattr(data, param)
    return val if index is None else val[index]


def _fun_access(fun, index, data, **kwargs):
    """Access a selected item in the output of a user-defined function."""
    val = fun(data, **kwargs)
    return val if index is None else val[index]


def _muf_access(_, index, data, **kwargs):
    mu = _record_access("mu", index, data, **kwargs)
    return np.remainder(mu, 2.0 * np.pi)


def _mu2pi_access(_, index, data, **kwargs):
    mu = _record_access("mu", index, data, **kwargs)
    return mu / 2.0 / np.pi


def _mu2pif_access(_, index, data, **kwargs):
    mu = _record_access("mu", index, data, **kwargs)
    return np.remainder(mu / 2.0 / np.pi, 1.0)


_opdata = {
    "muf": _muf_access,
    "mu2pi": _mu2pi_access,
    "mu2pif": _mu2pif_access,
}

_arrayproc = {
    None: None,
    "real": np.real,
    "imag": np.imag,
    "abs": np.absolute,
    "angle": np.angle,
    "log": np.log,
    "exp": np.exp,
    "sqrt": np.sqrt,
}

_statproc = {
    None: None,
    "mean": np.mean,
    "std": np.std,
    "var": np.var,
    "min": np.min,
    "max": np.max,
}


def _all_rows(index: RefIndex | None):
    """Prepends "all rows" (":") to an index tuple."""
    if index is None:
        return None
    if isinstance(index, tuple):
        return slice(None), *index
    else:
        return slice(None), index


def _get_fun(fname, fdict) -> Callable | None:
    """Get a processing from its name."""
    if callable(fname):
        return fname
    else:
        return fdict[fname]


def _subscript(plane):
    idx = axis_(plane, key="index")
    if isinstance(idx, tuple):
        return "".join(str(i) for i in idx)
    else:
        return ""


def _mod_name(name, fun, *args):
    if name:
        if fun:
            name = f"{fun.__name__}({name})"
        if args:
            name = _mod_name(name, *args)
    return name


def _tune(data, **__):
    return data.mu[-1] / 2.0 / np.pi


class _Ring:
    """Get an attribute of a lattice element."""

    def __init__(self, attrname, index, refpts):
        self.get_val = partial(_record_access, attrname, index)
        self.refpts = refpts

    def __call__(self, lattice, **__):
        vals = [self.get_val(el) for el in lattice.select(self.refpts)]
        return np.array(vals)


class Need(Enum):
    """Defines the computation requirements for an :py:class:`Observable`."""

    #:  Provides the *ring* data to the evaluation function
    RING = 1
    #:  Specify :py:func:`.find_orbit` computation and provide its *orbit* output
    #:  to the evaluation function
    ORBIT = 2
    #:  Specify :py:func:`.find_m44` or :py:func:`.find_m44` computation and
    #:  provide its *m44* or *m66* output to the evaluation function
    MATRIX = 3
    #:  Specify :py:func:`.get_optics` computation and provide its *ringdata* output
    #:  to the evaluation function
    GLOBALOPTICS = 4
    #:  Specify :py:func:`.get_optics` computation and provide its *elemdata* output
    #:  to the evaluation function
    LOCALOPTICS = 5
    #:  Specify :py:func:`.lattice_pass` computation and provide its *r_out*
    #:  to the evaluation function
    TRAJECTORY = 6
    #:  Specify :py:func:`.envelope_parameters` computation and provide its
    #:  *params* output to the evaluation function
    EMITTANCE = 7
    #:  Associated with LOCALOPTICS, require local optics computation at all
    #:  points: slower but avoids jumps in phase advance
    ALL_POINTS = 8
    #:  Associated with LOCALOPTICS, require the *get_chrom* keyword
    CHROMATICITY = 9
    #:  Associated with LOCALOPTICS, require the *get_w* keyword
    W_FUNCTIONS = 10
    #:  Specify :py:meth:`~.Lattice.get_geometry` computation and provide its output
    #:  to the evaluation function
    GEOMETRY = 11
    #:  Specify :py:func:`RDT <.get_rdts>` computation and provide its output to
    #:  the evaluation function
    RDT = 12
    #:  Associated with RDT: request 2nd order calculation
    RDT_2ND_ORDER = 13
    #: Specify :py:func:`.ohmi_envelope` computation and provide its *beamdata* output
    #: to the evaluation function
    GLOBALEMIT = 14
    #: Specify :py:func:`.ohmi_envelope` computation and provide its *emit* output
    #: to the evaluation function
    LOCALEMIT = 15
    #: For observables needing a 4D lattice
    NEED_4D = 16
    #: For observables needing a 6D lattice
    NEED_6D = 17


class Observable:
    """Base class for Observables. Can be used for user-defined observables."""

    # Class attributes
    _default: ClassVar[tuple] = ("{param}[{plane}]", "{param}", lambda x: x)
    _pinfo: ClassVar[dict] = {}

    # Instance attributes
    name: str  #: Observable name. A default name is automatically generated.
    #: Label used in plot legends.
    #:
    #: It may contain LaTeX math code. Example: ``"$\beta_x$"`` will appear as
    #: :math:`\beta_x`.
    #:
    #: By default the label is identical to the name of the Observable.
    #: Labels starting with a ``_`` will not appear in the legends.
    label: str
    #: Line formatting used when plotting the Observable. See
    #: :py:meth:`~matplotlib.axes.Axes.plot` for a description of line formatting.
    #: *plot_fmt* may be a :py:class:`str` for simple formatting (ex.: ``"o-"``) or a
    #: :py:class:`dict` for detailed formatting (ex.: :pycode:`{"linewidth": 3.0}`).
    plot_fmt: str | Mapping
    fun: Callable  #: Evaluation function.
    needs: AbstractSet[Need]  #: Set of requirements.
    target: npt.ArrayLike | None  #: Target value.
    eval_args: tuple[Any, ...]  #: Evaluation arguments
    eval_kw: dict[str, Any]  #: Evaluation keywords

    def __init__(
        self,
        fun: Callable,
        *eval_args,
        name: str | None = None,
        target: npt.ArrayLike | None = None,
        weight: npt.ArrayLike = 1.0,
        bounds=(0.0, 0.0),
        needs: AbstractSet[Need] | None = None,
        postfun: Callable | str | None = None,
        plot_fmt: str | Mapping | None = None,
        label: str | None = None,
        **eval_kw,
    ):
        r"""Args:
            name:           Observable name. If :py:obj:`None`, an explicit
              name will be generated
            fun:            :ref:`evaluation function <base_eval>`
            *eval_args:          Arguments provided to the evaluation function
            target:         Target value for a constraint. If :py:obj:`None`
              (default), the residual will always be zero.
            weight:         Weight factor: the residual is
              :pycode:`((value-target)/weight)**2`
            bounds:         Tuple of lower and upper bounds. The parameter
              is constrained in the interval
              [*target*\ +\ *low_bound* *target*\ +\ *up_bound*]
            postfun:        Post-processing function. It can be any numpy ufunc or a
              function name in {"real", "imag", "abs", "angle", "log", "exp", "sqrt"}.
            needs:          Set of requirements. This selects the data provided
              to the evaluation function. *needs* items are members of the
              :py:class:`Need` enumeration.
            plot_fmt:       Line formatting used when plotting the Observable. See
              :py:meth:`~matplotlib.axes.Axes.plot` for a description of line
              formatting. *plot_fmt* may be a :py:class:`str` for simple formatting
              (ex.: ``"o-"``) or a :py:class:`dict` for detailed formatting (ex.:
              :pycode:`{"linewidth": 3.0}`).


        Keyword Args:
            \*\*eval_kw:    Keyword arguments provided to the evaluation function

        The *target*, *weight* and *bounds* inputs must be broadcastable to the
        shape of *value*.

        .. _base_eval:
        .. rubric:: User-defined evaluation function

        The general form is:

        :pycode:`value = fun(*data, *eval_args, **eval_kw)`

        - *data* depends on the *needs* argument, and by default is empty. If several
          needed values are specified, their order is: *ring*, *orbit*, *m44/m66*,
          *ringdata*, *elemdata*, *r_out*, *params*, *geomdata*,
        - *eval_args* are the evaluation positional arguments provided to the observable
           constructor,
        - *eval_kw* are the evaluation keywords provided to the observable constructor,
          to the constructor of the enclosing :py:class:`.ObservableList` and to the
          :py:meth:`~.ObservableList.evaluate` method.
        - *value* is the value of the observable.

        For user-defined evaluation functions using linear optics data or
        emittance data, it is recommended to use
        :py:class:`LocalOpticsObservable`, :py:obj:`GlobalOpticsObservable`
        or :py:class:`EmittanceObservable` which provide the corresponding
        *data* argument.
        """
        name = name or fun.__name__
        postfun = _get_fun(postfun, _arrayproc)
        if postfun:
            fun = _Convolve(postfun, fun)
        name = _mod_name(name, postfun)
        label = _mod_name(label, postfun)
        self.fun = fun
        self.needs = needs or set()
        self.name = name
        self.target = target
        self.w: npt.NDArray[float] = np.asarray(weight, dtype=float)
        self.lbound, self.ubound = bounds
        self.initial: npt.NDArray[float] | None = None
        self._value: npt.NDArray[float] | Exception | None = None
        self._shape: tuple[int, ...] | None = None
        if plot_fmt is not None:
            self.plot_fmt = plot_fmt
        self.label = label or name
        self._axis_label = eval_kw.pop("axis_label", None)
        self.eval_args = eval_args
        self.eval_kw = eval_kw

    @classmethod
    def _ax_lab(cls, param, plane) -> str | None:
        if callable(param):
            return None
        else:
            _, fmt, code = cls._pinfo.get(param, cls._default)
            return fmt.format(plane=code(plane))

    @classmethod
    def _pl_lab(cls, param, plane) -> str | None:
        if callable(param):
            return None
        else:
            fmt, _, code = cls._pinfo.get(param, cls._default)
            return fmt.format(plane=code(plane))

    @property
    def axis_label(self):
        """Label used for the y-axes of plots (read only)."""
        return self._axis_label

    def __str__(self):
        """Return the string representation of the Observable."""
        return "\n".join((self._header(), self._all_lines()))

    @staticmethod
    def _header():
        """Header line."""
        fstring = "\n    {:<12} {:>16}  {:>16}  {:>16}  {:>16}  {:>16} "
        return fstring.format(
            "location", "Initial", "Actual", "Low bound", "High bound", "deviation"
        )

    @staticmethod
    def _line(loc, *items):
        def pitem(v):
            if v is None:
                return f" {'-':>14}   "
            elif isinstance(v, Exception):
                return f" {type(v).__name__:>16} "
            elif isinstance(v, np.ndarray) and v.ndim > 0:
                if v.size == 0:
                    return f" {'[]':16} "
                else:
                    return f" [{v.flat[0]:< 10.4} ...] "
            else:
                return f" {v: 16.6} "

        its = [pitem(v) for v in items]
        return f"    {loc:<12}" + "".join(its)

    def _all_lines(self):
        vnow = self._value
        if vnow is None or isinstance(vnow, Exception):
            deviation = None
            vmin = None
            vmax = None
        else:
            deviation = self.deviation
            if self.target is None:
                vmin = None
                vmax = None
            else:
                target = np.broadcast_to(self.target, vnow.shape)  # type: ignore
                vmin = target + self.lbound
                vmax = target + self.ubound
        values = self._line("", self.initial, vnow, vmin, vmax, deviation)
        return "\n".join((self.name, values))

    def _setup(self, ring: Lattice):
        """Setup function called when the observable is added to a list."""

    def evaluate(
        self, *data, initial: bool = False, **evalkw
    ) -> npt.NDArray[float] | Exception:
        """Compute and store the value of the observable.

        The direct evaluation of a single :py:class:`Observable` is normally
        not used. This method is called by the :py:class:`.ObservableList`
        container which provides the *data* arguments.

        Args:
            *data:      Raw data, provided by :py:class:`.ObservableList` and
              sent to the evaluation function
            initial:    It :py:obj:`None`, store the result as the initial
              value

        Returns:
            value:      The value of the observable or the error in evaluation
        """
        for d in data:
            if isinstance(d, Exception):
                message = f"Evaluation of {self.name} failed: {d.args[0]}"
                err = type(d)(message).with_traceback(d.__traceback__)
                self._value = err
                return err

        val = np.asarray(self.fun(*data, *self.eval_args, **(self.eval_kw | evalkw)))
        if initial:
            self.initial = val
        self._shape = val.shape
        self._value = val
        return val

    def check(self) -> bool:
        """Check if evaluation is done.

        Returns:
            ok: :py:obj:`True` if evaluation is done, :py:obj:`False` otherwise.

        Raises:
            AtError:    if the value is doubtful: evaluation failed, empty value…
        """
        return self.value is not None

    @staticmethod
    def check_value(value: npt.NDArray[float] | Exception) -> npt.NDArray[float]:
        if isinstance(value, Exception):
            raise type(value)(value.args[0]) from value
        return value

    @property
    def value(self) -> npt.NDArray[float]:
        """Value of the observable."""
        return self.check_value(self._value)

    @property
    def weight(self) -> npt.NDArray[float]:
        """Observable weight."""
        return np.broadcast_to(self.w, self._value.shape)  # type: ignore

    @weight.setter
    def weight(self, w: npt.ArrayLike):
        self.w = np.asarray(w, dtype=float)

    @property
    def weighted_value(self) -> npt.NDArray[float]:
        """Weighted value of the Observable, computed as
        :pycode:`weighted_value = value/weight`.
        """
        return self.value / self.w

    @property
    def deviation(self) -> npt.NDArray[float]:
        """Deviation from target value, computed as
        :pycode:`deviation = value-target`.

        If *target* is :py:obj:`None`, the deviation is zero for any value.
        """
        vnow = self.value
        if vnow is None:
            deviation = None
        elif self.target is None:
            deviation = np.broadcast_to(0.0, vnow.shape)
        else:
            vsh = vnow.shape
            diff = np.atleast_1d(vnow - np.broadcast_to(self.target, vsh))
            lb = diff - self.lbound
            ub = diff - self.ubound
            lb[lb >= 0] = 0
            ub[ub <= 0] = 0
            deviation = (lb + ub).reshape(vsh)
        return deviation

    @property
    def weighted_deviation(self) -> npt.NDArray[float]:
        """:pycode:`weighted_deviation = (value-target)/weight`.

        If *target* is :py:obj:`None`, the weighted deviation is zero for any value.
        """
        return self.deviation / self.w

    @property
    def residual(self) -> npt.NDArray[float]:
        """residual, computed as :pycode:`residual = ((value-target)/weight)**2`.

        If *target* is :py:obj:`None`, the residual is zero for any value.
        """
        # absolute necessary for complex data
        return np.absolute(self.weighted_deviation) ** 2

    @staticmethod
    def _set_name(name, param, index):
        """Compute a default observable names."""
        if name is None:
            if (
                index is Ellipsis
                or index is None
                or (isinstance(index, str) and index in {":", "..."})
            ):
                subscript = ""
            elif isinstance(index, tuple):
                ids = ", ".join(str(k) for k in index)
                subscript = f"[{ids}]"
            else:
                subscript = f"[{index}]"
            if callable(param):
                try:
                    base = param.__name__
                except AttributeError:
                    base = "<function>"
            else:
                base = param
            name = base + subscript
        return name


class RingObservable(Observable):
    """Observe any user-defined property of a ring."""

    def __init__(
        self,
        fun: Callable,
        name: str | None = None,
        needs: Need | None = None,
        **eval_kw,
    ):
        r"""Args:
            fun:            :ref:`user-defined evaluation function <ring_eval>`
            name:           Observable name. If :py:obj:`None`, an explicit
              name will be generated.
            needs:          :py:obj:`.Need.NEED_4D` or :py:obj:`.Need.NEED_6D`. Ensure
              that the  lattice provided to the evaluation function has the desired
              property. If :py:obj:`None`, the lattice is the one given to
              :py:meth:`~.ObservableList.evaluate` unmodified.

        Keyword Args:
            target:         Target value for a constraint. If :py:obj:`None`
              (default), the residual will always be zero.
            weight:         Weight factor: the residual is
              :pycode:`((value-target)/weight)**2`
            bounds:         Tuple of lower and upper bounds. The parameter
              is constrained in the interval
              [*target*\ +\ *low_bound* *target*\ +\ *up_bound*]
            postfun:        Post-processing function. It can be any numpy ufunc or a
              function name in {"real", "imag", "abs", "angle", "log", "exp", "sqrt"}.

        The *target*, *weight* and *bounds* inputs must be broadcastable to the
        shape of *value*.

        .. rubric:: Evaluation keywords

        These values must be provided to the :py:meth:`~.ObservableList.evaluate`
        method. Default values may be given at instantiation.

        * **ring** –    Lattice description,
        * **dp** –      Momentum deviation. Defaults to :py:obj:`None`,
        * **dct** –     Path lengthening. Defaults to :py:obj:`None`,
        * **df** –      Deviation from the nominal RF frequency.
          Defaults to :py:obj:`None`,

        .. _ring_eval:
        .. rubric:: User-defined evaluation function

        It is called as:

        :pycode:`value = fun(ring, **eval_kw)`

        - *ring* is the lattice description,
        - *eval_kw* are the evaluation keywords provided to the observable constructor,
          to the constructor of the enclosing :py:class:`.ObservableList` and to the
          :py:meth:`~.ObservableList.evaluate` method.
        - *value* is the value of the Observable.

        Examples:
            >>> def circumference(ring, **__):
            ...     return ring.get_s_pos(len(ring))[0]
            >>> obs = RingObservable(circumference)

            Defines an Observable for the ring circumference.

            >>> def momentum_compaction(ring, dp=0.0, **__):
            ...     return ring.get_mcf(dp=dp)
            >>> obs = RingObservable(momentum_compaction, needs=Need.NEED_4D)

            Defines an Observable for the momentum compaction factornad make sure that
            the lattice provided to :py:meth:`~.Lattice.get_mcf` is 4D.
        """
        nds = {Need.RING}
        if needs is not None:
            nds.add(needs)
        name = self._set_name(name, fun, None)
        super().__init__(fun, name=name, needs=nds, **eval_kw)


class ElementObservable(Observable):
    """Base class for Observables linked to a position in the lattice."""

    def __init__(
        self,
        fun: Callable,
        refpts: Refpts,
        name: str | None = None,
        statfun: Callable | str | None = None,
        postfun: Callable | str | None = None,
        summary: bool = False,
        label: str | None = None,
        **eval_kw,
    ):
        r"""Args:
            fun:            :ref:`evaluation function <base_eval>`
            refpts:         Observation points.
              See ":ref:`Selecting elements in a lattice <refpts>`"
            name:           Observable name. If :py:obj:`None`, an explicit
              name will be generated
            postfun:        Post-processing function. It can be any numpy ufunc or a
              function name in {"real", "imag", "abs", "angle", "log", "exp", "sqrt"}.
            statfun:        Statistics post-processing function. it can be a numpy
              function or a function name in {"mean", "std", "var", "min", "max"}.
              Example: :pycode:`statfun=numpy.mean`.

        Keyword Args:
            target:         Target value for a constraint. If :py:obj:`None`
              (default), the residual will always be zero.
            weight:         Weight factor: the residual is
              :pycode:`((value-target)/weight)**2`
            bounds:         Tuple of lower and upper bounds. The parameter
              is constrained in the interval
              [*target*\ +\ *low_bound* *target*\ +\ *up_bound*]

        .. rubric:: Shape of the value

        A :pycode:`nrefs` dimension, with :pycode:`nrefs` the number
        of reference points, is prepended to the shape of the output of *fun*,
        **even if nrefs == 1**. The *target*, *weight* and *bounds* inputs must be
        broadcastable to the shape of *value*.
        """
        name = name or fun.__name__
        postfun = _get_fun(postfun, _arrayproc)
        if postfun:
            fun = _Convolve(postfun, fun)
        statfun = _get_fun(statfun, _statproc)
        if statfun:
            fun = _Convolve(statfun, fun, axis=0)
            summary = True
        name = _mod_name(name, postfun, statfun)
        label = _mod_name(label, postfun, statfun)
        super().__init__(fun, name=name, label=label, **eval_kw)
        self.summary = summary
        self.refpts = refpts
        self._boolrefs = None
        self._excluded = None
        self._locations = [""]

    def check(self) -> bool:
        ok = super().check()
        shp = self._shape
        if ok and shp and shp[0] <= 0:
            msg = f"Observable {self.name!r}: No location selected in the lattice."
            raise AtError(msg)
        return ok

    def _all_lines(self):
        if self.summary:
            return super()._all_lines()
        else:
            vnow = self._value
            if vnow is None or isinstance(vnow, Exception):
                vnow = repeat(vnow)
                deviation = repeat(None)
                vmin = repeat(None)
                vmax = repeat(None)
            else:
                deviation = self.deviation
                if self.target is None:
                    vmin = repeat(None)
                    vmax = repeat(None)
                else:
                    target = np.broadcast_to(self.target, vnow.shape)  # type: ignore
                    vmin = target + self.lbound
                    vmax = target + self.ubound
            vini = self.initial
            if vini is None:
                vini = repeat(None)
            viter = zip(
                self._locations, vini, vnow, vmin, vmax, deviation, strict=False
            )
            values = "\n".join(self._line(*vv) for vv in viter)
            return "\n".join((self.name, values))

    def _setup(self, ring: Lattice):
        boolrefs = ring.get_bool_index(self.refpts)
        excluded = ring.get_bool_index(self._excluded)
        boolrefs &= ~excluded
        self._boolrefs = boolrefs
        locs = [el.FamName for el in ring.select(self._boolrefs[:-1])]
        if boolrefs[-1]:
            locs.append("End")
        self._locations = locs


class GeometryObservable(ElementObservable):
    """Observe the geometrical parameters of the reference trajectory.

    Process the *geomdata* output of :py:func:`.get_geometry`.
    """

    # Class attributes
    _pinfo: ClassVar[dict] = {"x": "x [m]", "y": "y [m]", "angle": "angle"}

    def __init__(
        self,
        refpts: Refpts,
        param: str,
        name: str | None = None,
        label: str | None = None,
        **eval_kw,
    ):
        # noinspection PyUnresolvedReferences
        r"""Args:
            refpts:         Observation points.
              See ":ref:`Selecting elements in a lattice <refpts>`"
            param:          Geometry parameter name: one in {'x', 'y', 'angle'}
            name:           Observable name. If :py:obj:`None`, an explicit
              name will be generated.

        Keyword Args:
            statfun:        Post-processing function called on the value of the
              observable. Example: :pycode:`statfun=numpy.mean`
            target:         Target value for a constraint. If :py:obj:`None`
              (default), the residual will always be zero.
            weight:         Weight factor: the residual is
              :pycode:`((value-target)/weight)**2`
            bounds:         Tuple of lower and upper bounds. The parameter
              is constrained in the interval
              [*target*\ +\ *low_bound* *target*\ +\ *up_bound*]

        The *target*, *weight* and *bounds* inputs must be broadcastable to the
        shape of *value*.

        .. rubric:: Evaluation keywords

        These values must be provided to the :py:meth:`~.ObservableList.evaluate`
        method. Default values may be given at instantiation.

        * **ring** –    Lattice description,

        Example:
            >>> obs = GeometryObservable(at.Monitor, param="x")

            Observe x coordinate of monitors
        """

        if param not in self._pinfo:
            msg = f"Expected {param!r} to be one of {self._pinfo.keys()!r}"
            raise ValueError(msg)
        name = self._set_name(name, "geometry", param)
        fun = partial(_record_access, param, None)
        needs = {Need.GEOMETRY}
        super().__init__(
            fun,
            refpts,
            needs=needs,
            name=name,
            label=param,
            axis_label=self._pinfo[param],
            **eval_kw,
        )
        if label:
            self.label = label


class OrbitObservable(ElementObservable):
    """Observe the closed orbit coordinates at selected locations.

    Process the *orbit* output of :py:func:`.find_orbit`.
    """

    # Class attributes
    _plist: ClassVar = ["x [m]", "$p_x$", "y [m]", "$p_y$", "angle"]

    def __init__(
        self,
        refpts: Refpts,
        axis: AxisDef = None,
        name: str | None = None,
        label: str | None = None,
        **eval_kw,
    ):
        # noinspection PyUnresolvedReferences
        r"""Args:
            refpts:         Observation points.
              See ":ref:`Selecting elements in a lattice <refpts>`"
            axis:           Index in the orbit vector, If :py:obj:`None`,
              the whole vector is specified
            name:           Observable name. If :py:obj:`None`, an explicit
              name will be generated.

        Keyword Args:
            postfun:        Post-processing function. It can be any numpy ufunc or a
              function name in {"real", "imag", "abs", "angle", "log", "exp", "sqrt"}.
            statfun:        Statistics post-processing function. it can be a numpy
              function or a function name in {"mean", "std", "var", "min", "max"}.
              Example: :pycode:`statfun=numpy.mean`.
            target:         Target value for a constraint. If :py:obj:`None`
              (default), the residual will always be zero.
            weight:         Weight factor: the residual is
              :pycode:`((value-target)/weight)**2`
            bounds:         Tuple of lower and upper bounds. The parameter
              is constrained in the interval
              [*target*\ +\ *low_bound* *target*\ +\ *up_bound*]

        .. rubric:: Evaluation keywords

        These values must be provided to the :py:meth:`~.ObservableList.evaluate`
        method. Default values may be given at instantiation.

        * **ring** –    Lattice description,
        * **dp** –      Momentum deviation. Defaults to :py:obj:`None`,
        * **dct** –     Path lengthening. Defaults to :py:obj:`None`,
        * **df** –      Deviation from the nominal RF frequency.
          Defaults to :py:obj:`None`,

        .. rubric:: Shape of the value

        If *axis* is :py:obj:`None` (whole orbit vector), then *value* has shape
        :pycode:`(nrefs, 6)` with :pycode:`nrefs` number of reference points: a
        :pycode:`nrefs` dimension is prepended to the shape of the orbit vector,
        **even if nrefs == 1**. A single coordinate has shape :pycode:`(nrefs,)`.
        The *target*, *weight* and *bounds* inputs must be broadcastable to the shape
        of *value*.

        Example:
            >>> obs = OrbitObservable(at.Monitor, axis=0)

            Observe the horizontal closed orbit at monitor locations
        """
        descr = axis_(axis)
        name = self._set_name(name, "orbit", descr["code"])
        fun = _ArrayAccess(descr["index"])
        needs = {Need.ORBIT}
        super().__init__(
            fun,
            refpts,
            needs=needs,
            name=name,
            label=descr["label"],
            axis_label="".join((descr["label"], descr["unit"])),
            **eval_kw,
        )
        if label:
            self.label = label


class MatrixObservable(ElementObservable):
    """Observe coefficients of the transfer matrix.

    Process the result of calling :py:func:`.find_m44` or :py:func:`.find_m44`
    depending upon :py:meth:`~.Lattice.is_6d`.
    """

    def __init__(
        self,
        refpts: Refpts,
        axis: AxisDef = Ellipsis,
        name: str | None = None,
        label: str | None = None,
        **eval_kw,
    ):
        # noinspection PyUnresolvedReferences
        r"""Args:
            refpts:         Observation points.
              See ":ref:`Selecting elements in a lattice <refpts>`"
            axis:           Index in the transfer matrix, If :py:obj:`Ellipsis`,
              the whole matrix is specified
            name:           Observable name. If :py:obj:`None`, an explicit
              name will be generated.

        Keyword Args:
            postfun:        Post-processing function. It can be any numpy ufunc or a
              function name in {"real", "imag", "abs", "angle", "log", "exp", "sqrt"}.
            statfun:        Statistics post-processing function. it can be a numpy
              function or a function name in {"mean", "std", "var", "min", "max"}.
              Example: :pycode:`statfun=numpy.mean`.
            target:         Target value for a constraint. If :py:obj:`None`
              (default), the residual will always be zero.
            weight:         Weight factor: the residual is
              :pycode:`((value-target)/weight)**2`
            bounds:         Tuple of lower and upper bounds. The parameter
              is constrained in the interval
              [*target*\ +\ *low_bound* *target*\ +\ *up_bound*]

        The *target*, *weight* and *bounds* inputs must be broadcastable to the
        shape of *value*.

        .. rubric:: Evaluation keywords

        These values must be provided to the :py:meth:`~.ObservableList.evaluate`
        method. Default values may be given at instantiation.

        * **ring** –    Lattice description,
        * **dp** –      Momentum deviation. Defaults to :py:obj:`None`,
        * **dct** –     Path lengthening. Defaults to :py:obj:`None`,
        * **df** –      Deviation from the nominal RF frequency.
          Defaults to :py:obj:`None`,
        * **orbit** –   Initial orbit. Avoids looking for the closed orbit if it is
          already known,

        Example:
            >>> obs = MatrixObservable(at.Monitor, axis=("x", "px"))

            Observe the transfer matrix from origin to monitor locations and
            extract T[0,1]
        """
        name = self._set_name(name, "matrix", axis_(axis, key="code"))
        fun = _ArrayAccess(axis_(axis, key="index"))
        needs = {Need.MATRIX}
        super().__init__(
            fun,
            refpts,
            needs=needs,
            name=name,
            label=f"$T_{{{_subscript(axis)}}}$",
            axis_label="T [m]",
            **eval_kw,
        )
        if label:
            self.label = label


class _GlobalOpticsObservable(Observable):
    # Class attributes
    _pinfo: ClassVar[dict] = {
        "tune": (r"$\nu_{{{plane}}}$", "Tune", partial(plane_, key="label")),
        "chromaticity": (
            r"$\xi_{{{plane}}}$",
            "Chromaticity",
            partial(plane_, key="label"),
        ),
        "damping_time": (
            r"$\tau_{{{plane}}}$",
            "Damping time [s]",
            partial(plane_, key="label"),
        ),
    }

    def __init__(
        self,
        param: str | Callable,
        plane: AxisDef = None,
        name: str | None = None,
        label: str | None = None,
        **eval_kw,
    ):
        # noinspection PyUnresolvedReferences
        r"""Args:
            param:          Optics parameter name (see :py:func:`.get_optics`)
              or user-defined evaluation function called as:
              :pycode:`value = fun(ringdata, **eval_kw)` and returning the
              value of the Observable
            plane:          Index in the parameter array, If :py:obj:`None`,
              the whole array is specified
            name:           Observable name. If :py:obj:`None`, an explicit
              name will be generated.

        Keyword Args:
            target:         Target value for a constraint. If :py:obj:`None`
              (default), the residual will always be zero.
            weight:         Weight factor: the residual is
              :pycode:`((value-target)/weight)**2`
            bounds:         Tuple of lower and upper bounds. The parameter
              is constrained in the interval
              [*target*\ +\ *low_bound* *target*\ +\ *up_bound*]
            postfun:        Post-processing function. It can be any numpy ufunc or a
              function name in {"real", "imag", "abs", "angle", "log", "exp", "sqrt"}.

        The *target*, *weight* and *bounds* inputs must be broadcastable to the
        shape of *value*.

        .. rubric:: Evaluation keywords

        These values must be provided to the :py:meth:`~.ObservableList.evaluate`
        method. Default values may be given at instantiation.

        * **ring** –    Lattice description,
        * **dp** –      Momentum deviation. Defaults to :py:obj:`None`,
        * **dct** –     Path lengthening. Defaults to :py:obj:`None`,
        * **df** –      Deviation from the nominal RF frequency.
          Defaults to :py:obj:`None`,
        * **orbit** –   Initial orbit. Avoids looking for the closed orbit if it is
          already known,
        """

        needs = {Need.GLOBALOPTICS}
        name = self._set_name(name, param, plane_(plane, key="code"))
        index = plane_(plane, key="index")
        if callable(param):
            fun = partial(_fun_access, param, index)
            needs.add(Need.CHROMATICITY)
        else:
            fun = partial(_record_access, param, index)
            if param == "chromaticity":
                needs.add(Need.CHROMATICITY)

        super().__init__(
            fun,
            needs=needs,
            name=name,
            label=self._pl_lab(param, plane),
            axis_label=self._ax_lab(param, plane),
            **eval_kw,
        )
        if label:
            self.label = label


class LocalOpticsObservable(ElementObservable):
    """Observe a local optics parameter at selected locations.

    Process the local output of :py:func:`.get_optics`.
    """

    # Class attributes
    _pinfo: ClassVar[dict] = {
        "alpha": (r"$\alpha_{{{plane}}}$", "alpha", partial(plane_, key="label")),
        "beta": (r"$\beta_{{{plane}}}$", "beta [m]", partial(plane_, key="label")),
        "gamma": (r"$\gamma_{{{plane}}}$", "gamma", partial(plane_, key="label")),
        "mu": (
            r"$\mu_{{{plane}}}$",
            "phase advance [rad]",
            partial(plane_, key="label"),
        ),
        "muf": (
            r"$\mu_{{{plane}}}$",
            "phase advance [rad]",
            partial(plane_, key="label"),
        ),
        "mu2pi": (
            r"$\mu_{{{plane}}}/2\pi$",
            "phase advance",
            partial(plane_, key="label"),
        ),
        "mu2pif": (
            r"$\mu_{plane}/2 \pi$",
            "phase advance",
            partial(plane_, key="label"),
        ),
        "closed_orbit": (
            r"${{{plane}}}_{{co}}$",
            "closed orbit",
            partial(axis_, key="code"),
        ),
        "dispersion": (
            r"$\eta_{{{plane}}}$",
            r"dispersion [m]",
            partial(axis_, key="code"),
        ),
        "s_pos": ("s", "s [m]", partial(plane_, key="label")),
        "M": (r"$M_{{{plane}}}$", "M", _subscript),
        "A": (r"$A_{{{plane}}}$", "A", _subscript),
        "B": (r"$B_{{{plane}}}$", "B", _subscript),
        "C": (r"$C_{{{plane}}}$", "C", _subscript),
        "R": (r"$R_{{{plane}}}$", "R", _subscript),
        "W": (r"$W_{{{plane}}}$", "W", partial(plane_, key="label")),
        "Wp": (r"$Wp_{{{plane}}}$", "Wp", partial(plane_, key="label")),
        "dalpha": (
            r"$\partial \alpha_{{{plane}}}/ \partial \delta$",
            r"$\partial \alpha / \partial \delta$",
            partial(plane_, key="label"),
        ),
        "dbeta": (
            r"$\partial \beta_{{{plane}}}/ \partial \delta$",
            r"$\partial \beta/ \partial \delta$ [m]",
            partial(plane_, key="label"),
        ),
        "dmu": (
            r"$\partial \mu_{{{plane}}}/ \partial \delta$",
            r"$\partial \mu/ \partial \delta$ [rad]",
            partial(plane_, key="label"),
        ),
        "ddispersion": (
            r"$\partial \eta_{{{plane}}}/ \partial \delta$",
            r"$\partial \eta/ \partial \delta$ [m]",
            partial(axis_, key="code"),
        ),
        "dR": (
            r"$\partial R_{{{plane}}}/ \partial \delta$",
            r"$\partial R/ \partial \delta$",
            _subscript,
        ),
    }

    _default = ("{param}[{plane}]", "{param}", lambda x: x)

    def __init__(
        self,
        refpts: Refpts,
        param: str | Callable,
        *,
        plane: AxisDef = Ellipsis,
        name: str | None = None,
        all_points: bool = False,
        summary: bool = False,
        label: str | None = None,
        **eval_kw,
    ):
        # noinspection PyUnresolvedReferences
        r"""Args:
            refpts:         Observation points.
              See ":ref:`Selecting elements in a lattice <refpts>`"
            param:          :ref:`Optics parameter name <localoptics_param>`
              or :ref:`user-defined evaluation function <localoptics_eval>`
            plane:          Index in the parameter array, If :py:obj:`Ellipsis`,
              the whole array is specified
            name:           Observable name. If :py:obj:`None`, an explicit
              name will be generated
            all_points:     Compute the local optics at all elements. This avoids
              discontinuities in phase advances. This is automatically set for the
              'mu' parameter, but may need to be specified for user-defined evaluation
              functions using the phase advance.
            summary:        Set to :py:obj:`True` if the user-defined
             evaluation function returns a single item (see below) instead of one item
             per refpoint.

        Keyword Args:
            postfun:        Post-processing function. It can be any numpy ufunc or a
              function name in {"real", "imag", "abs", "angle", "log", "exp", "sqrt"}.
            statfun:        Statistics post-processing function. it can be a numpy
              function or a function name in {"mean", "std", "var", "min", "max"}.
              Example: :pycode:`statfun=numpy.mean`.
            target:         Target value for a constraint. If :py:obj:`None`
              (default), the residual will always be zero.
            weight:         Weight factor: the residual is
              :pycode:`((value-target)/weight)**2`
            bounds:         Tuple of lower and upper bounds. The parameter
              is constrained in the interval
              [*target*\ +\ *low_bound* *target*\ +\ *up_bound*]

        .. rubric:: Evaluation keywords

        These values must be provided to the :py:meth:`~.ObservableList.evaluate`
        method. Default values may be given at instantiation.

        * **ring** –    Lattice description,
        * **dp** –      Momentum deviation. Defaults to :py:obj:`None`,
        * **dct** –     Path lengthening. Defaults to :py:obj:`None`,
        * **df** –      Deviation from the nominal RF frequency.
          Defaults to :py:obj:`None`,
        * **orbit** –   Initial orbit. Avoids looking for the closed orbit if it is
          already known,
        * **twiss_in** – Initial conditions for transfer line optics.
          See :py:func:`.get_optics`,
        * **method** –  Method for linear optics. Default: :py:obj:`~.linear.linopt6`.

        .. rubric:: Shape of the value

        If the requested attribute has shape :pycode:`shp`, then
        *value* has shape :pycode:`(nrefs,) + shp` with :pycode:`nrefs` number of
        reference points: a  :pycode:`nrefs` dimension is prepended to the shape of
        the attribute, **even if nrefs == 1**. The *target*, *weight* and *bounds*
        inputs must be broadcastable to the shape of *value*. For instance, a *target*
        with shape :pycode:`shp` will automatically broadcast and apply  to all
        reference points.

        .. _localoptics_param:
        .. rubric:: Optics parameter name

        In addition to :py:func:`.get_optics` parameter names, LocalOpticsObservable
        adds 3 parameters: *muf*, *mu2pi* and *mu2pif*:

        ================    ===================================================
        **s_pos**           longitudinal position [m]
        **M**               (6, 6) transfer matrix M from the beginning of ring
                            to the entrance of the element
        **closed_orbit**    (6,) closed orbit vector
        **dispersion**      (4,) dispersion vector
        **A**               (6, 6) A-matrix
        **R**               (3, 6, 6) R-matrices
        **beta**            :math:`\left[ \beta_x,\beta_y \right]` vector
        **alpha**           :math:`\left[ \alpha_x,\alpha_y \right]` vector
        **mu**              :math:`\left[ \mu_x,\mu_y \right]`, betatron phase
        **mu2pi**           :math:`\left[ \mu_x,\mu_y \right]/2\pi`, reduced betatron
                            phase
        **muf**             :math:`\left[ \mu_x,\mu_y \right]`, betatron phase
                            (modulo :math:`2\pi`)
        **mu2pif**          :math:`\mathrm{frac}(\left[ \mu_x,\mu_y \right]/2\pi)`,
                            fractional part of the reduced betatron phase
        **W**               :math:`\left[ W_x,W_y \right]` only if *get_w*
                            is :py:obj:`True`: chromatic amplitude function
        **Wp**              :math:`\left[ Wp_x,Wp_y \right]` only if *get_w*
                            is :py:obj:`True`: chromatic phase function
        **dalpha**          (2,) alpha derivative vector
                            (:math:`\Delta \alpha/ \delta_p`)
        **dbeta**           (2,) beta derivative vector
                            (:math:`\Delta \beta/ \delta_p`)
        **dmu**             (2,) mu derivative vector
                            (:math:`\Delta \mu/ \delta_p`)
        **ddispersion**     (4,) dispersion derivative vector
                            (:math:`\Delta D/ \delta_p`)
        **dR**              (3, 6, 6) R derivative vector
                            (:math:`\Delta R/ \delta_p`)
        ================    ===================================================

        .. _localoptics_eval:
        .. rubric:: User-defined evaluation function

        The observable value is computed as:

        :pycode:`value = fun(elemdata, **eval_kw)[plane]`

        - *elemdata* is the output of :py:func:`.get_optics`, evaluated at the *refpts*
          of the observable,
        - *eval_kw* are the evaluation keywords provided to the observable constructor,
          to the constructor of the enclosing :py:class:`.ObservableList` and to the
          :py:meth:`~.ObservableList.evaluate` method,
        - *value* is the value of the Observable and must have one line per
          refpoint. Alternatively, it may be a single line, but then the
          *summary* keyword must be set to :py:obj:`True`,
        - the *plane* keyword then selects the desired values in the function output.

        Examples:
            Observe the beta in both planes at all :py:class:`.Monitor`
            locations:

            >>> obs = LocalOpticsObservable(at.Monitor, "beta")

            Observe the maximum vertical beta in Quadrupoles:

            >>> obs = LocalOpticsObservable(
            ...     at.Quadrupole, "beta", plane="y", statfun=np.max
            ... )

            The user-defined evaluation function computes the phase-advance
            between the 1st and last given reference points, here the elements
            33 and 101 of the lattice

            >>> def phase_advance(elemdata, **__):
            ...     mu = elemdata.mu
            ...     return mu[-1] - mu[0]
            >>>
            >>> obs = LocalOpticsObservable(
            ...     [33, 101],
            ...     phase_advance,
            ...     plane="y",
            ...     all_points=True,
            ...     summary=True,
            ... )

        A user-defined function may accept any evaluation keyword. Here is a function
        which computes the beam envelope for given emittances and energy spread. We
        define *emit* and *sigma_e* as evaluation keywords:

        >>> def beam_size(elemdata, emit=None, sigma_e=None, **__):
        ...     return np.sqrt(
        ...         elemdata.beta*emit + (elemdata.dispersion[:, [0, 2]] * sigma_e)**2
        ...     )

        We instantiate a :py:class:`LocalOpticsObservable` using this function and
        with default values for emittances and energy spread:

        >>> obs = LocalOpticsObservable(
        ...     0, beam_size, emit=[130.0e-12, 10.0e-12], sigma_e=0.9e-3
        ... )
        >>> allobs = ObservableList([obs])

        We can evaluate *obs* with the default emittance values:

        >>> allobs.evaluate(ring=ring)
        array([2.9990243e-05, 5.14264471e-06])

        We can then evaluate *obs* with arbitrary emittances:

        >>> allobs.evaluate(ring=ring, emit=[140.0e-12, 20.0e-12])
        array([3.11193609e-05, 7.2727979e-06])

        """

        if param in {"M", "closed_orbit", "dispersion", "A", "R"}:
            ax_ = axis_
        else:
            ax_ = plane_

        needs = {Need.LOCALOPTICS}
        name = self._set_name(name, param, ax_(plane, key="code"))
        index = _all_rows(ax_(plane, key="index"))
        if callable(param):
            if summary:
                fun = partial(_fun_access, param, ax_(plane, key="index"))
            else:
                fun = partial(_fun_access, param, index)
        else:
            fun = partial(_opdata.get(param, _record_access), param, index)
            if param in {"mu", "mu2pi"} or all_points:
                needs.add(Need.ALL_POINTS)
            if param in {"W", "Wp", "dalpha", "dbeta", "dmu", "ddispersion", "dR"}:
                needs.add(Need.W_FUNCTIONS)

        super().__init__(
            fun,
            refpts,
            needs=needs,
            name=name,
            summary=summary,
            label=self._pl_lab(param, plane),
            axis_label=self._ax_lab(param, plane),
            **eval_kw,
        )
        if label:
            self.label = label


class LatticeObservable(ElementObservable):
    """Observe an attribute of selected lattice elements."""

    def __init__(
        self,
        refpts: Refpts,
        attrname: str,
        index: int | None = None,
        name: str | None = None,
        label: str | None = None,
        **eval_kw,
    ):
        # noinspection PyUnresolvedReferences
        r"""Args:
            refpts:         Elements to be observed
              See ":ref:`Selecting elements in a lattice <refpts>`"
            attrname:       Attribute name
            index:          Index in the attribute array. If :py:obj:`None`,
              the whole array is specified
            name:           Observable name. If :py:obj:`None`, an explicit
              name will be generated.

        Keyword Args:
            postfun:        Post-processing function. It can be any numpy ufunc or a
              function name in {"real", "imag", "abs", "angle", "log", "exp", "sqrt"}.
            statfun:        Statistics post-processing function. it can be a numpy
              function or a function name in {"mean", "std", "var", "min", "max"}.
              Example: :pycode:`statfun=numpy.mean`.

        .. rubric:: Evaluation keywords

        These values must be provided to the :py:meth:`~.ObservableList.evaluate`
        method. Default values may be given at instantiation.

        * **ring** –    Lattice description,

        Example:
            >>> obs = LatticeObservable(
            ...     at.Sextupole, "KickAngle", index=0, statfun=np.sum
            ... )

            Observe the sum of horizontal kicks in Sextupoles
        """
        fun = _Ring(attrname, index, refpts)
        needs = {Need.RING}
        name = self._set_name(name, attrname, index)
        super().__init__(fun, refpts, needs=needs, name=name, **eval_kw)
        if label:
            self.label = label


class TrajectoryObservable(ElementObservable):
    """Observe trajectory coordinates at selected locations.

    Process the *r_out* output if :py:meth:`.Lattice.track`
    """

    # Class attributes
    _pinfo: ClassVar[dict] = {
        "x": (r"$x_{{{plane}}}$", "position [m]", lambda x: x),
        "px": (r"$p_{{x{plane}}}$", "transverse momentum", lambda x: x),
        "y": (r"$y_{{{plane}}}$", "position [m]", lambda x: x),
        "py": (r"$p_{{y{plane}}}$", "transverse momentum", lambda x: x),
        "dp": (r"$\delta_{{{plane}}}$", "off-momentum", lambda x: x),
        "ct": (r"$\beta c \tau_{{{plane}}}$", "path lengthening [m]", lambda x: x),
    }

    def __init__(
        self,
        refpts: Refpts,
        axis: AxisDef = Ellipsis,
        npart: int = 0,
        name: str | None = None,
        label: str | None = None,
        **eval_kw,
    ):
        r"""Args:
            refpts:         Observation points.
              See ":ref:`Selecting elements in a lattice <refpts>`"
            axis:           Index in the orbit array, If :py:obj:`Ellipsis`,
              the whole array is specified
            npart:          Particle number,
            name:           Observable name. If :py:obj:`None`, an explicit
              name will be generated.

        Keyword Args:
            postfun:        Post-processing function. It can be any numpy ufunc or a
               function name in {"real", "imag", "abs", "angle", "log", "exp", "sqrt"}.
            statfun:        Statistics post-processing function. it can be a numpy
              function or a function name in {"mean", "std", "var", "min", "max"}.
              Example: :pycode:`statfun=numpy.mean`.
            target:         Target value for a constraint. If :py:obj:`None`
              (default), the residual will always be zero.
            weight:         Weight factor: the residual is
              :pycode:`((value-target)/weight)**2`
            bounds:         Tuple of lower and upper bounds. The parameter
              is constrained in the interval
              [*target*\ +\ *low_bound* *target*\ +\ *up_bound*]

        The *target*, *weight* and *bounds* inputs must be broadcastable to the
        shape of *value*.

        .. rubric:: Evaluation keywords

        These values must be provided to the :py:meth:`~.ObservableList.evaluate`
        method. Default values may be given at instantiation.

        * **ring** –    Lattice description,
        * **r_in** –    Initial coordinates of one or several tracked particles.
        """
        descr = axis_(axis)
        name = self._set_name(name, "trajectory", descr["code"])
        fun = _ArrayAccess((npart, descr["index"]))
        needs = {Need.TRAJECTORY}
        super().__init__(
            fun,
            refpts,
            needs=needs,
            name=name,
            label=self._pl_lab(descr["code"], npart),
            axis_label=self._ax_lab(descr["code"], npart),
            **eval_kw,
        )
        if label:
            self.label = label


class EmittanceObservable(Observable):
    """Observe emittance-related parameters.

    Output of :py:func:`.radiation_parameters` or :py:func:`.envelope_parameters`.

    For 4D lattices, the data is extracted form the output of
    :py:func:`.radiation_parameters`: emittances are computed  with the radiation
    integrals.

    For 6D lattices, the data is extracted form the output of
    :py:func:`.envelope_parameters`: emittances are computed by
    :py:func:`.ohmi_envelope`
    """

    # Class attributes
    _pinfo: ClassVar[dict] = {
        "emittances": (
            r"$\epsilon_{{{plane}}}$",
            "Emittance [m]",
            partial(plane_, key="label"),
        ),
        "J": (
            r"$\mathrm{{J}}_{{{plane}}}$",
            "Damping partition number",
            partial(plane_, key="label"),
        ),
        "Tau": (
            r"$\tau_{{{plane}}}$",
            "Damping time [s]",
            partial(plane_, key="label"),
        ),
        "sigma_e": (r"$\sigma_e$", "Energy spread", partial(plane_, key="label")),
        "sigma_l": (r"$\sigma_l$", "Bunch length [m]", partial(plane_, key="label")),
        "phi_s": (r"$\phi_s$", "Synchronous phase [rad]", partial(plane_, key="label")),
        "f_s": (
            r"$f_s$",
            "Synchrotron frequency [Hz]",
            partial(plane_, key="label"),
        ),
        "tunes6": (r"$\nu_{{{plane}}}$", "Tune", partial(plane_, key="label")),
        "U0": (r"$U_0$", "Energy loss / turn [eV]", partial(plane_, key="label")),
        "voltage": ("V", "V [V]", partial(plane_, key="label")),
    }

    def __init__(
        self,
        param: str | Callable,
        plane: AxisDef = None,
        name: str | None = None,
        label: str | None = None,
        **eval_kw,
    ):
        r"""Args:
            param:          Parameter name (see :py:func:`.envelope_parameters`) or
              :ref:`user-defined evaluation function <emittance_eval>`
            plane:          One out of {0, 'x', 'h', 'H'} for horizontal plane,
             one out of {1, 'y', 'v', 'V'} for vertical plane or one out of
             {2, 'z', 'l', 'L'} for longitudinal plane
            name:           Observable name. If :py:obj:`None`, an explicit
              name will be generated.

        Keyword Args:
            postfun:        Post-processing function. It can be any numpy ufunc or a
              function name in {"real", "imag", "abs", "angle", "log", "exp", "sqrt"}.
            statfun:        Statistics post-processing function. it can be a numpy
              function or a function name in {"mean", "std", "var", "min", "max"}.
              Example: :pycode:`statfun=numpy.mean`.
            target:         Target value for a constraint. If :py:obj:`None`
              (default), the residual will always be zero.
            weight:         Weight factor: the residual is
              :pycode:`((value-target)/weight)**2`
            bounds:         Tuple of lower and upper bounds. The parameter
              is constrained in the interval
              [*target*\ +\ *low_bound* *target*\ +\ *up_bound*]

        .. rubric:: Evaluation keywords

        These values must be provided to the :py:meth:`~.ObservableList.evaluate`
        method. Default values may be given at instantiation.

        * **ring** –    Lattice description,

        .. _emittance_eval:
        .. rubric:: User-defined evaluation function

        It is called as:

        :pycode:`value = fun(paramdata, **eval_kw)`

        - *paramdata* if the :py:class:`.RingParameters` object returned by
          :py:func:`.envelope_parameters`.
        - *eval_kw* are the evaluation keywords provided to the observable constructor,
          to the constructor of the enclosing :py:class:`.ObservableList` and to the
          :py:meth:`~.ObservableList.evaluate` method,
        - *value* is the value of the Observable.

        Example:
            >>> EmittanceObservable("emittances", plane="h")

            Observe the horizontal emittance
        """

        name = self._set_name(name, param, plane_(plane, key="code"))
        if callable(param):
            fun = param
        else:
            fun = partial(_record_access, param, plane_(plane, key="index"))
        needs = {Need.EMITTANCE}
        super().__init__(
            fun,
            needs=needs,
            name=name,
            label=self._pl_lab(param, plane),
            axis_label=self._ax_lab(param, plane),
            **eval_kw,
        )
        if label:
            self.label = label


# noinspection PyPep8Naming
def GlobalOpticsObservable(
    param: str,
    *,
    plane: AxisDef = Ellipsis,
    name: str | None = None,
    use_integer: bool = False,
    **eval_kw,
):
    # noinspection PyUnresolvedReferences
    r"""Observe a global optics parameter.

    Process the *ringdata* output of :py:func:`.get_optics`.

    Args:
        param:          Optics parameter name (see :py:func:`.get_optics`)
          or :ref:`user-defined evaluation function <globaloptics_eval>`
        plane:          Index in the parameter array, If :py:obj:`Ellipsis`,
          the whole array is specified
        name:           Observable name. If :py:obj:`None`, an explicit
          name will be generated
        use_integer:    For *'tune'* parameter: derive the tune from the
          phase advance to avoid skipping integers. Slower than looking only
          at the fractional part

    Keyword Args:
        target:         Target value for a constraint. If :py:obj:`None`
          (default), the residual will always be zero.
        weight:         Weight factor: the residual is
          :pycode:`((value-target)/weight)**2`
        bounds:         Tuple of lower and upper bounds. The parameter
          is constrained in the interval
          [*target*\ +\ *low_bound* *target*\ +\ *up_bound*]
        postfun:        Post-processing function. It can be any numpy ufunc or a
              function name in {"real", "imag", "abs", "angle", "log", "exp", "sqrt"}.

    The *target*, *weight* and *bounds* inputs must be broadcastable to the
    shape of *value*.

    .. rubric:: Evaluation keywords

    These values must be provided to the :py:meth:`~.ObservableList.evaluate`
    method. Default values may be given at instantiation.

    * **ring** –    Lattice description,
    * **dp** –      Momentum deviation. Defaults to :py:obj:`None`,
    * **dct** –     Path lengthening. Defaults to :py:obj:`None`,
    * **df** –      Deviation from the nominal RF frequency.
      Defaults to :py:obj:`None`,
    * **orbit** –   Initial orbit. Avoids looking for the closed orbit if it is
      already known,

    .. _globaloptics_eval:
    .. rubric:: User-defined evaluation function

    It is called as:

    :pycode:`value = fun(ring, ringdata, **eval_kw)`

    - *ringdata* is the output of :py:func:`.get_optics`,
    - *eval_kw* are the evaluation keywords provided to the observable constructor,
      to the constructor of the enclosing :py:class:`.ObservableList` and to the
      :py:meth:`~.ObservableList.evaluate` method.
    - *value* is the value of the Observable.

    Examples:
        >>> obs = GlobalOpticsObservable("tune", use_integer=True)

        Observe the tune in both planes, including the integer part (slower)

        >>> obs = GlobalOpticsObservable("chromaticity", plane="v")

        Observe the vertical chromaticity
    """
    if param == "tune" and use_integer:
        # noinspection PyProtectedMember
        name = ElementObservable._set_name(name, param, plane_(plane, key="code"))
        return LocalOpticsObservable(
            End,
            _tune,
            plane=plane,
            name=name,
            summary=True,
            all_points=True,
            **eval_kw,
        )
    else:
        return _GlobalOpticsObservable(param, plane=plane, name=name, **eval_kw)
