from __future__ import annotations

import copy
import functools
import logging
import pathlib
import typing
from typing import (
    Any,
    Dict,
    Generator,
    List,
    NamedTuple,
    Optional,
    Sequence,
    Tuple,
    Type,
    TypeVar,
    Union,
)

import lark
import matplotlib
import matplotlib.axes
import matplotlib.pyplot as plt
import pydantic
import pydantic.alias_generators
from pmd_beamphysics.units import nice_array
from typing_extensions import override

from ...errors import (
    NotFlatError,
    RecursiveLineError,
)
from ..types import (
    AnyPath,
    BaseModel,
    BeamlineElement,
    ReprTableData,
    ValueType,
)
from . import _lattice as auto_lattice
from . import parsers
from ..interfaces.bmad import genesis4_elements_and_line_from_tao


from typing import Literal


AnyBeamlineElement = Union[
    auto_lattice.AutogeneratedBeamlineElement,
    "Line",
]


T_BeamlineElement = TypeVar("T_BeamlineElement", bound=BeamlineElement)

logger = logging.getLogger(__name__)
LineItem = Union[str, "DuplicatedLineItem", "PositionedLineItem", AnyBeamlineElement]


class ElementListProxy(List[T_BeamlineElement]):
    """
    A list proxy class for input elements.

    Getting or setting an attribute on an instance of this class will get or
    set that attribute on each element of the list.

    May be used as a normal list with indexing and standard methods such as `.append()`.
    """

    # NOTE: not all of these will be available - it depends on the element type.
    # this is primarily to satisfy the type checker and give some tab completion options.
    # Generated by way of generate_attr_list_ below.
    L: List[float] | float
    aw: List[float] | float
    ax: List[float] | float
    ay: List[float] | float
    cx: List[float] | float
    cy: List[float] | float
    delay: List[float] | float
    dumpbeam: List[int] | int
    dumpfield: List[int] | int
    elements: List[List] | List
    gradx: List[float] | float
    grady: List[float] | float
    helical: List[bool] | bool
    k1: List[float] | float
    kx: List[float] | float
    ky: List[float] | float
    label: List[str] | str
    lambdau: List[float] | float
    lb: List[float] | float
    ld: List[float] | float
    nwig: List[int] | int
    phi: List[float] | float
    sort_: List[int] | int
    stop: List[int] | int
    x_offset: List[float] | float
    y_offset: List[float] | float

    @staticmethod
    def _generate_attr_list_():
        eles = typing.get_args(AnyBeamlineElement)[:-1] + (Line,)  # forwardref
        # eles = typing.get_args(AnyBeamlineElement)
        all_fields = {}
        for ele in eles:
            for name, fld in type(ele).model_fields.items():
                all_fields.setdefault(name, []).append(fld)

        for name, flds in sorted(all_fields.items()):
            annotation = list(set(fld.annotation for fld in flds))
            if len(annotation) == 1:
                (cls,) = annotation
                print(f"{name}: list[{cls.__name__}] | {cls.__name__}")

    def __getattr__(self, attr: str):
        if attr == "sort_":
            attr = "sort"
        return [getattr(element, attr) for element in self]

    def __setattr__(self, attr: str, value):
        # Allow setting of "internal" attributes on this proxy object.
        if attr.startswith("_"):
            object.__setattr__(self, attr, value)
        else:
            if attr == "sort_":
                attr = "sort"

            if not len(self):
                return

            if not isinstance(value, Sequence):
                value = [value] * len(self)

            for element, val in zip(self, value):
                setattr(element, attr, val)


class ZElement(NamedTuple):
    """A tuple of a Z position and its corresponding beamline element."""

    zend: float
    element: AnyBeamlineElement


class _InspectElement(NamedTuple):
    """A tuple of a Z position and its corresponding beamline element."""

    z: Optional[float]
    context: List[LineItem]
    element: AnyBeamlineElement


class ElementPlotSettings(NamedTuple):
    color: str
    width: float
    alpha: float = 1.0


class _BarPlotBar(NamedTuple):
    x: float
    width: float
    height: float
    color: str
    edgecolor: str
    element: Union[auto_lattice.Undulator, auto_lattice.Quadrupole]


class DuplicatedLineItem(BaseModel):
    """
    A Genesis 4 lattice Line item which is at a certain position.

    This corresponds to Genesis 4 `"COUNT*LABEL"`.

    Attributes
    ----------
    label : str
        The name/label of the line item.
    count : int
        The number of times to repeat the line item.
    """

    label: str
    count: int

    @classmethod
    def from_string(cls, value: str) -> DuplicatedLineItem:
        count, label = value.split("*", 1)
        return cls(
            label=label.strip(),
            count=int(count),
        )

    def to_genesis(self) -> str:
        return str(self)

    @override
    def __str__(self) -> str:
        return f"{self.count}*{self.label}"


class PositionedLineItem(BaseModel):
    """
    A Genesis 4 lattice Line item which is at a certain position.

    This corresponds to Genesis 4 `"LABEL@POSITION"`.

    Attributes
    ----------
    label : str
        The name/label of the line item.
    position : float
        The position of the element.
    """

    label: str
    position: float

    @classmethod
    def from_string(cls, value: str) -> PositionedLineItem:
        label, position = value.split("@", 1)
        return cls(
            label=label.strip(),
            position=float(position.strip()),
        )

    def to_genesis(self) -> str:
        return str(self)

    @override
    def __str__(self) -> str:
        return f"{self.label}@{self.position}"


def _fix_line_item(line_item: LineItem) -> LineItem:
    """Make the appropriate dataclass for a serialized line item, if necessary."""
    if isinstance(line_item, (DuplicatedLineItem, PositionedLineItem)):
        return line_item
    if isinstance(line_item, BeamlineElement):
        return line_item
    if "@" in line_item:
        return PositionedLineItem.from_string(line_item)
    if "*" in line_item:
        return DuplicatedLineItem.from_string(line_item)
    return line_item


class Line(BeamlineElement):
    """
    A Genesis 4 beamline Line.

    Attributes
    ----------
    elements : list of LineItem
        Elements contained in the line.
    label : str, optional
        An optional label to attach to the line.
    """

    type: Literal["line"] = "line"
    elements: List[LineItem] = pydantic.Field(default_factory=list, kw_only=False)
    label: str = pydantic.Field(default="")

    @override
    def model_post_init(self, __context: Any) -> None:
        self.elements = [_fix_line_item(item) for item in self.elements]

    @property
    def flattened(self) -> List[Tuple[Optional[float], str]]:
        """Flattened elements by name, including Z location if available."""
        res = []
        for item in self.names:
            item = _fix_line_item(item)
            if isinstance(item, PositionedLineItem):
                res.append((item.position, item.label))
            elif isinstance(item, DuplicatedLineItem):
                for _ in range(item.count):
                    res.append((None, item.label))
            else:
                res.append((None, str(item)))
        return res

    @property
    def names(self) -> List[str]:
        return [
            elem.label if isinstance(elem, BeamlineElement) else str(elem)
            for elem in self.elements
        ]

    @override
    def to_genesis(self) -> str:
        elements = ", ".join(self.names)
        return "".join(
            (
                self.label,
                ": LINE = {",
                elements,
                "};",
            )
        )

    @classmethod
    def from_labels(
        cls,
        elements: Dict[str, AnyBeamlineElement],
        *element_labels: str,
        label: str = "",
    ) -> Line:
        """
        Create a `Line` from labeled beamline elements.

        Parameters
        ----------
        elements : Dict[str, BeamlineElement]
            The element dictionary.
        *element_labels : str
            Labels of elements.  May be either a single label per argument or a
            whitespace-delimited string of labels.
        label : str
            Label name for the line.

        Returns
        -------
        Line
        """
        try:
            return cls(
                elements=[
                    elements[label]
                    for labels in element_labels
                    for label in labels.strip().split()
                ],
                label=label,
            )
        except KeyError as ex:
            raise ValueError(
                f"Label {ex} is not present in the beamline element dictionary. "
                f"The following are valid: {tuple(elements.keys())}"
            )


def lattice_elements_from_list(
    elements: Sequence[AnyBeamlineElement],
) -> Dict[str, AnyBeamlineElement]:
    """
    Make a label-to-lattice-element dictionary from a sequence of elements.

    Parameters
    ----------
    elements : Sequence[AnyBeamlineElement]

    Returns
    -------
    Dict[str, AnyBeamlineElement]
    """
    res = {}
    for idx, element in enumerate(elements):
        label = element.label
        if not label:
            raise ValueError(f"Element #{idx} does not have a label: ({element})")
        if label in res:
            raise ValueError(
                f"Element #{idx} has a duplicate label to another element. "
                f"Previous: {res[label]} "
                f"element[{idx}] = {element}"
            )
        res[label] = element
    return res


def _nice_bar_plot(
    ax: matplotlib.axes.Axes,
    bars: Sequence[_BarPlotBar],
    units: str,
    label: str,
    alpha: float = 1.0,
    show_labels: bool = True,
):
    if not bars:
        return None
    values = [bar.height for bar in bars]
    y, _, prefix = nice_array(values)
    ylabel = f"{label} ({prefix}{units})"
    ax.bar(
        x=[bar.x + bar.width / 2 for bar in bars],
        height=y,
        width=[bar.width for bar in bars],
        color=[bar.color for bar in bars],
        edgecolor=[bar.edgecolor for bar in bars],
        alpha=alpha,
    )
    ax.set_ylabel(ylabel)

    if not show_labels:
        return
    for bar in bars:
        mid = bar.x + (bar.width / 2.0)
        annotation = ax.annotate(
            bar.element.label,
            xy=(mid, 0.0),
        )
        annotation.set_rotation(90)
        annotation.set_fontsize("x-small")


class Lattice(BaseModel):
    """
    A Genesis 4 beamline Lattice configuration.

    Attributes
    ----------
    elements : list of BeamlineElement or Line
        Elements contained in the lattice.
    filename : pathlib.Path or None
        The filename from which this lattice was loaded.
    """

    elements: Dict[str, AnyBeamlineElement] = pydantic.Field(
        default_factory=dict,
    )
    filename: Optional[pathlib.Path] = None

    def __init__(
        self,
        elements: Union[
            Dict[str, AnyBeamlineElement], List[AnyBeamlineElement], None
        ] = None,
        *,
        filename: Optional[pathlib.Path] = None,
    ) -> None:
        if elements is None:
            elements = {}
        if not isinstance(elements, dict):
            elements = lattice_elements_from_list(elements)
        super().__init__(elements=elements, filename=filename)

    def taper_custom(
        self,
        beamline: str,
        *,
        aw0: float,
        zstart: float,
        zend: float,
    ) -> List[ZElement]:
        """
        Get a list of undulators to perform a custom taper on.

        Only elements that fit within the range [zstart, zend] will be
        included.

        Parameters
        ----------
        beamline : str
            The beamline containing the undulators.
        aw0 : float
            The starting undulator strength.
        zstart : float
            The starting Z position, relative to the first element on the line.
        zend : float
            The ending Z position, relative to the first element on the line.

        Returns
        -------
        list of ZElement
            List of ZElement instances, which contain a normalized Z location
            (zstart=0.0, zend=1.0) and an Undulator element.
        """
        if not self.is_flat(beamline):
            raise NotFlatError(
                f"The beamline {beamline!r} is not flat: it contains reused "
                f"elements.  Before performing a taper operation, first flatten "
                f"the beamline with `.flatten_line()`"
            )

        if zend < zstart:
            zstart, zend = zend, zstart

        undulators = self.by_z_location(beamline, limit_to=auto_lattice.Undulator)
        norm_z = []
        for zelem in undulators:
            und = zelem.element
            assert isinstance(und, auto_lattice.Undulator)
            if zstart <= zelem.zend <= zend:
                norm_z.append(
                    ZElement(
                        zend=(zelem.zend - zstart) / (zend - zstart),
                        element=und,
                    )
                )
            else:
                und.aw = aw0

        return norm_z

    def taper_array(
        self,
        beamline: str,
        aws: Sequence[float],
    ) -> None:
        """
        Apply a list of undulator strengths to set the taper.

        Parameters
        ----------
        beamline : str
            The beamline name.
        aws : list of floats (or equivalent, np.ndarray)
            The undulator strengths, one for each undulator in the beamline.
            Undulators are sorted based on Z position.
        """
        if not self.is_flat(beamline):
            raise NotFlatError(
                f"The beamline {beamline!r} is not flat: it contains reused "
                f"elements.  Before performing a taper operation, first flatten "
                f"the beamline with `.flatten_line()`"
            )

        undulators = self.by_z_location(beamline, limit_to=auto_lattice.Undulator)
        aws = list(aws)
        if len(aws) != len(undulators):
            raise ValueError(
                f"The number of undulators ({len(undulators)}) is not equal to "
                f"the number of aw parameters provided ({len(aws)})"
            )
        for aw, (_z, und) in zip(aws, undulators):
            und = typing.cast(auto_lattice.Undulator, und)
            und.aw = aw

    def taper_linear(
        self,
        beamline: str,
        *,
        aw0: float,
        zstart: float,
        zend: float,
        taper_start: float,
        taper_end: float,
    ) -> List[float]:
        """
        Perform a linear taper on undulators in the given beamline.

        Only elements that fit within the range [zstart, zend] will be
        included.

        Parameters
        ----------
        beamline : str
            The beamline containing the undulators.
        aw0 : float
            The starting undulator strength.
        zstart : float
            The starting Z position, relative to the first element on the line.
        zend : float
            The ending Z position, relative to the first element on the line.
        taper_start : float
            The starting taper multiplier (applied to `aw0`) at `zstart`.
        taper_end : float
            The final taper multiplier (applied to `aw0`) at `zend`.

        Returns
        -------
        list of float
            List of undulator strengths.
        """
        # taper_custom will give normalized Z values for [zstart, zend] -> [0, 1]
        undulators = self.taper_custom(beamline, aw0=aw0, zstart=zstart, zend=zend)
        slope = aw0 * (taper_end - taper_start) / (1.0 - 0.0)  # / (zend - zstart)
        for norm_z, und in undulators:
            assert isinstance(und, auto_lattice.Undulator)
            und.aw = aw0 + slope * norm_z

        return [zelem.element.aw for zelem in undulators]

    def _check_beamline_name(self, name: Optional[str] = None) -> str:
        self.fix_labels()
        options = ", ".join(line.label for line in self.lines)
        if name is not None:
            el = self.elements.get(name, None)
            if el is None:
                raise ValueError(
                    f"{name!r} is not a valid beamline.  Options are: {options}"
                )
            if not isinstance(el, Line):
                raise ValueError(
                    f"{name!r} is not a valid beamline.  It's a {type(el).__name__}"
                )

            return name

        if len(self.lines) == 1:
            return self.lines[0].label

        raise ValueError(
            f"You must specify a beamline to plot.  Options are: {options}"
        )

    def plot(
        self,
        beamline: Optional[str] = None,
        *,
        ax: Optional[matplotlib.axes.Axes] = None,
        show_labels: bool = True,
        show: bool = True,
        normalize_aw: bool = False,
        figsize: Optional[Tuple[float, float]] = None,
        xlim: Optional[Tuple[float, float]] = None,
        ylim: Optional[Tuple[float, float]] = None,
    ):
        """
        Plot the layout of the given beamline.

        Parameters
        ----------
        beamline : str
            The name of the beamline.
        ax : matplotlib.axes.Axes or None
            Plot on the given axis, or create a new one.
        show_labels : bool
            Show labels for each undulator.
        show : bool
            Show the plot.
        normalize_aw : bool, default=False
            Normalize undulator strengths with respect to the first undulator
            ($aw0$): $aw / aw0 - 1$
        figsize : Tuple[float, float], optional
            The figure size to create.  Used if ``ax`` is not provided.
        xlim : Tuple[float, float], optional
            X axis limits for the plot.
        ylim : Tuple[float, float], optional
            Y axis limits for the plot.
        """
        beamline = self._check_beamline_name(beamline)
        elements = self.by_z_location(beamline)
        undulators = [
            (elem.zend, elem.element)
            for elem in elements
            if isinstance(elem.element, auto_lattice.Undulator)
        ]
        quads = [
            (elem.zend, elem.element)
            for elem in elements
            if isinstance(elem.element, auto_lattice.Quadrupole)
        ]
        if ax is None:
            _, ax = plt.subplots(figsize=figsize)
        assert ax is not None

        aw0 = undulators[0][1].aw if undulators else 1.0
        if aw0 == 0.0:
            # Avoid dividing by zero when the strengths aren't yet set
            normalize_aw = False

        bars = [
            _BarPlotBar(
                x=zend - und.L,
                width=und.L,
                height=und.aw / aw0 - 1 if normalize_aw else und.aw,
                color="red",
                edgecolor="red",
                element=und,
            )
            for zend, und in undulators
        ]
        _nice_bar_plot(
            ax,
            bars,
            units="1",
            alpha=0.5,
            label="$aw/aw0 - 1$" if normalize_aw else "$aw$",
            show_labels=show_labels,
        )

        bars = [
            _BarPlotBar(
                x=zend - quad.L,
                width=quad.L,
                height=quad.k1,
                color="blue",
                edgecolor="blue",
                element=quad,
            )
            for zend, quad in quads
        ]
        _nice_bar_plot(
            ax.twinx(),
            bars,
            units="$1/m^2$",
            alpha=0.5,
            label="Quad $k$",
            show_labels=show_labels,
        )

        ax.set_xlabel("$z$ (m)")
        ax.set_title(f"{beamline} Layout")
        if xlim is not None:
            ax.set_xlim(xlim)
        if ylim is not None:
            ax.set_xlim(ylim)
        if show:
            plt.show()

    def is_flat(self, beamline: str) -> bool:
        """Checks if the given `beamline` does not reuse elements."""
        names = self.inspect_line_by_name(beamline)
        return len(set(names)) == len(names)

    def _inspect(self, beamline: str) -> List[_InspectElement]:
        """
        This recursively inspects any LINEs inside the beamline.

        Returns
        -------
        list of _InspectElement
            List of names.
        """
        self.fix_labels()
        stack = []

        def _inspect(name: str) -> Generator[_InspectElement, None, None]:
            if name in stack:
                raise RecursiveLineError(
                    f"Recursion of beamline elements in line {beamline} detected: {name}"
                )
            stack.append(name)
            element = self.elements[name]
            if isinstance(element, Line):
                yield _InspectElement(
                    z=None,
                    context=[name],
                    element=element,
                )
                for z, el in element.flattened:
                    for inner in _inspect(el):
                        yield _InspectElement(
                            z=inner.z if inner.z is not None else z,
                            context=[name] + inner.context,
                            element=inner.element,
                        )
            else:
                yield _InspectElement(z=None, context=[], element=element)
            stack.pop(-1)

        return list(_inspect(beamline))

    def inspect_line_by_name(self, beamline: str) -> List[str]:
        """
        Returns a list of all element names in the given beamline.

        This recursively inspects any LINEs inside the beamline.

        Returns
        -------
        list of str
            List of names.
        """
        self.fix_labels()
        return [
            insp.element.label
            for insp in self._inspect(beamline)
            if not isinstance(insp.element, Line)
        ]

    def inspect_line(self, beamline: str) -> List[AnyBeamlineElement]:
        """
        Returns a list of all elements in the given beamline.

        This recursively inspects any LINEs inside the beamline.

        Returns
        -------
        list of BeamlineElement
            List of beamline element instances.
        """
        return [
            insp.element
            for insp in self._inspect(beamline)
            if not isinstance(insp.element, Line)
        ]

    def flatten_line(
        self,
        beamline: str,
        *,
        count: int = 1,
        start: int = 0,
        prefix: str = "",
        suffix: str = "",
        delimiter: str = "_",
        line_format: str = "{name}{index}",
        last_format: str = "",
        in_place: bool = True,
        new_beamline: str = "",
    ):
        """
        Flatten a beamline such that any nested beamlines are made into distinct elements.

        For example, take a beamline ``LN`` which contains 2 nested beamlines
        "LA" (containing UND) and "LB" (containing QUAD).

        Flattening this with the default line format of ``{name}{index}`` and
        the default label format of the same ``{name}{index}`` would result
        with replacing ``LN`` as follows.

        Using start=0, count=1:

        * ``LN0_LA0_UND0`` and ``LN0_LB0_QUAD0``

        Using start=1, count=2:

        * ``LN1_LA1_UND1``, ``LN1_LB1_QUAD1``, ``LN2_LA1_UND1``, ``LN2_LB1_QUAD1``

        Parameters
        ----------
        beamline : str
            The beamline name to flatten.
        count : int, default=1
            The number of times to duplicate the flattened beamline.
        start : int, default=0
            The starting index for relabeling elements.
        prefix : str, default=""
            A prefix to add to the renamed elements.
        suffix : str, default=""
            A suffix to add to the renamed elements.
        delimiter : str, default="_"
            The delimiter to use between the name and index in the new element names.
        line_format : str, default="{name}{index}"
            The string format to use for renaming beamline elements.
        last_format : str, default=line_format
            An optional format string to use for the last element in each segment.
            Defaults to using the 'line_format'.
        in_place : bool, default=True
            Perform the flattening in-place. Remove unused elements and add
            the new elements directly to the lattice.
        new_beamline : str
            For in-place mode, use this as the name for the new beamline.

        Returns
        -------
        [segment(start), segment(start+1), ... segment(start+count-1)]
            Each segment corresponds to a single index of the flattened line,
            from ``start`` to ``start + count - 1``. Each segment is a list of
            created elements based on that index.
        """
        self.fix_labels()
        last_format = last_format or line_format
        line_stack = []

        def _inspect(
            name: str,
            seen: Dict[str, int],
        ) -> Generator[Tuple[str, AnyBeamlineElement], None, None]:
            if name in line_stack:
                raise RecursiveLineError(
                    f"Recursion of beamline elements in line {beamline} detected: {name}"
                )

            line_stack.append(name)

            element = self.elements[name]
            seen_count = seen.get(name, start - 1) + 1
            seen[name] = seen_count

            format = line_format if isinstance(element, Line) else last_format

            line_prefix = format.format(
                name=name,
                index=seen_count,
                beamline=beamline,
                delimiter=delimiter,
            )

            if isinstance(element, Line):
                seen = {}
                for _, el in element.flattened:
                    for inner, sub_ele in _inspect(el, seen=seen):
                        yield f"{line_prefix}{delimiter}{inner}", sub_ele
            else:
                yield line_prefix, element
            line_stack.pop(-1)

        to_duplicate = self.inspect_line_by_name(beamline)
        new_elements = {}
        sections = []
        seen: Dict[str, int] = {}
        for _ in range(start, start + count):
            section = []
            for new_name, elem in list(_inspect(beamline, seen=seen)):
                new_name = f"{prefix}{new_name}{suffix}"
                if new_name in new_elements:
                    raise ValueError(
                        f"Corrected name for element conflicts with another new element: {new_name}"
                    )

                new_element = copy.deepcopy(elem)
                new_element.label = new_name
                new_elements[new_name] = new_element
                section.append(new_element)
            sections.append(section)

        if in_place:
            for label in to_duplicate:
                if label not in new_elements:
                    self.elements.pop(label, None)
            if new_beamline != beamline and beamline not in new_elements:
                self.elements.pop(beamline, None)

            self.elements[new_beamline or beamline] = Line(elements=list(new_elements))
            self.elements.update(new_elements)

        return sections

    def by_z_location(
        self,
        beamline: str,
        limit_to: Union[
            Type[BeamlineElement],
            Tuple[Type[BeamlineElement], ...],
            None,
        ] = None,
    ) -> List[ZElement]:
        """
        Get all (flattened) beamline elements by Z location.

        Returns
        -------
        list of (zend, element)
            Each list item is a ZElement, a namedtuple which has `.zend` and
            `.element` that is also usable as a normal tuple.
        """
        z = 0.0
        by_z = []
        for insp in self._inspect(beamline):
            if isinstance(insp.element, Line):
                # Line is already expanded
                continue

            if insp.z is not None:
                # A located lattice element has a user-specified, explicit
                # location and is exempt from our element-by-element counting:
                by_z.append(ZElement(zend=insp.z, element=insp.element))
                continue

            if isinstance(
                insp.element,
                (
                    auto_lattice.Drift,
                    auto_lattice.Quadrupole,
                    auto_lattice.Corrector,
                    auto_lattice.Chicane,
                    auto_lattice.PhaseShifter,
                    auto_lattice.Undulator,
                ),
            ):
                z += insp.element.L
            elif isinstance(insp.element, auto_lattice.Marker):
                ...
            else:
                raise NotImplementedError(type(insp.element))

            if limit_to is None or isinstance(insp.element, limit_to):
                by_z.append(ZElement(zend=z, element=insp.element))

        def sort_z(item: ZElement):
            return item.zend

        return sorted(by_z, key=sort_z)

    @override
    def _repr_table_data_(self) -> ReprTableData:
        return {
            "obj": self.elements,
            "annotations": None,
            "descriptions": None,
        }

    @override
    def __str__(self) -> str:
        return self.to_string("repr")

    @override
    def _repr_html_(self) -> str:
        repr_ = self.to_string("repr")
        return f"<pre>{repr_}</pre>"

    def to_genesis(self) -> str:
        """Generate a Genesis 4-compatible lattice input string."""
        self.fix_labels()
        return "\n".join(element.to_genesis() for element in self.elements.values())

    def fix_labels(self) -> None:
        """Fix labels of elements based on their dictionary key in `.elements`."""
        for label, element in self.elements.items():
            if element.label != label:
                if element.label:
                    logger.warning(
                        "Renaming beamline element in lattice from %s to %s",
                        element.label,
                        label,
                    )
                element.label = label

    @property
    def by_element(
        self,
    ) -> Dict[Type[T_BeamlineElement], ElementListProxy[T_BeamlineElement]]:
        """Get beamline elements organized by their class."""
        by_element = {}
        for element in self.elements.values():
            by_element.setdefault(type(element), ElementListProxy[type(element)]())
            by_element[type(element)].append(element)
        return by_element

    @property
    def undulators(self) -> ElementListProxy[auto_lattice.Undulator]:
        """List of all Undulator instances."""
        return self.by_element.get(auto_lattice.Undulator, ElementListProxy())

    @property
    def drifts(self) -> ElementListProxy[auto_lattice.Drift]:
        """List of all Drift instances."""
        return self.by_element.get(auto_lattice.Drift, ElementListProxy())

    @property
    def quadrupoles(self) -> ElementListProxy[auto_lattice.Quadrupole]:
        """List of all Quadrupole instances."""
        return self.by_element.get(auto_lattice.Quadrupole, ElementListProxy())

    @property
    def correctors(self) -> ElementListProxy[auto_lattice.Corrector]:
        """List of all Corrector instances."""
        return self.by_element.get(auto_lattice.Corrector, ElementListProxy())

    @property
    def chicanes(self) -> ElementListProxy[auto_lattice.Chicane]:
        """List of all Chicane instances."""
        return self.by_element.get(auto_lattice.Chicane, ElementListProxy())

    @property
    def phase_shifters(self) -> ElementListProxy[auto_lattice.PhaseShifter]:
        """List of all PhaseShifter instances."""
        return self.by_element.get(auto_lattice.PhaseShifter, ElementListProxy())

    @property
    def markers(self) -> ElementListProxy[auto_lattice.Marker]:
        """List of all Marker instances."""
        return self.by_element.get(auto_lattice.Marker, ElementListProxy())

    @property
    def lines(self) -> ElementListProxy[Line]:
        """List of all Line instances."""
        return self.by_element.get(Line, ElementListProxy())

    @override
    def model_dump(self, **kwargs: Any) -> Dict[str, Any]:
        """Serialize this lattice to a list of dictionaries."""
        self.fix_labels()
        return super().model_dump(**kwargs)

    @classmethod
    def from_contents(
        cls, contents: str, filename: Optional[AnyPath] = None
    ) -> Lattice:
        """
        Load a lattice from its file contents.

        Parameters
        ----------
        contents : str
            The contents of the lattice file.
        filename : AnyPath or None, optional
            The filename of the lattice, if known.

        Returns
        -------
        Lattice
        """
        parser = parsers.new_lattice_parser()
        filename = filename or "unknown"
        try:
            tree = parser.parse(contents)
        except Exception:
            if "\n" not in contents:
                raise ValueError(
                    f"Unable to parse the provided input in Genesis4 lattice "
                    f"format. It looks like this might have been a filename: "
                    f"{contents!r}"
                )
            raise
        return _LatticeTransformer(filename).transform(tree)

    @classmethod
    def from_file(cls, filename: AnyPath) -> Lattice:
        """
        Load a lattice file from disk.

        Parameters
        ----------
        filename : AnyPath
            The filename to load.

        Returns
        -------
        Lattice
        """
        with open(filename) as fp:
            contents = fp.read()
        return cls.from_contents(contents, filename=filename)

    @classmethod
    def from_elements(
        cls, elements: Sequence[AnyBeamlineElement], filename: Optional[AnyPath] = None
    ) -> Lattice:
        """
        Create a Lattice from a list of labeled elements.

        Parameters
        ----------
        elements : list of BeamlineElement
            Elements of the beamline.  These must have a label set.
        filename : AnyPath or None, optional
            The filename of the lattice, if applicable.

        Returns
        -------
        Lattice
        """
        elements = list(elements)
        for el in elements:
            if not el.label:
                raise ValueError(f"Beamline element is missing a label: {el}")
        elements_by_name = {el.label: el for el in elements}
        if len(elements_by_name) != len(elements):
            raise ValueError("One or more duplicate labels exist in the input")
        return cls(
            elements_by_name,
            filename=pathlib.Path(filename) if filename else None,
        )

    @classmethod
    def from_tao(
        cls,
        tao,
        ele_start="beginning",
        ele_end="end",
        universe: int = 1,
        branch: int = 0,
        line_label: str | None = None,
    ):
        """
        Creates a Genesis4 Lattice from a running PyTao Tao instance.

        The lattice will contain one Line element.

        Parameters
        ----------
        tao : Tao
            The Tao instance containing the element definitions.
        ele_start : str, optional
            Element to start. Defaults to "beginning".
        ele_end : str, optional
            Element to end. Defaults to "end".
        universe: int, optional = 1
            Tao universe to use
        branch: int, optional = 0
            Tao branch to use
        line_label: str, default = None
            Name of the line element.
            If not given, the Tao branch name will be used.

        Returns
        -------
        Lattice

        """
        elements, line_labels = genesis4_elements_and_line_from_tao(
            tao, ele_start=ele_start, ele_end=ele_end, universe=universe, branch=branch
        )

        if line_label is None:
            line_label = tao.branch1(universe, branch)["name"]
        elements[line_label] = Line(elements=line_labels)
        return cls(elements=elements)

    def to_file(
        self,
        filename: AnyPath,
    ) -> None:
        """
        Write the lattice input file, in Genesis format, to ``filename``.

        Parameters
        ----------
        filename : str or pathlib.Path
        """
        main_config = self.to_genesis()
        with open(filename, "wt") as fp:
            print(main_config, file=fp)

        logger.debug(
            "Wrote lattice to %s:\n%s",
            main_config,
            filename,
        )


def get_short_type(type_: str) -> str:
    # Per the Genesis documents, only the first 4 characters are compared.
    # For example, this maps the uppercase 'DRIF' to the Drift dataclass.
    return type_[:4].upper()


@functools.lru_cache(maxsize=50)
def element_class_from_genesis(
    type_: str, sequence_type: Optional[str]
) -> Type[BeamlineElement]:
    def get_key(cls: Type[BeamlineElement]) -> Tuple[str, Optional[str]]:
        short_type = get_short_type(cls.model_fields["type"].default)
        # Sequences share the same 'type' and must be differentiated by their
        # sequence type, which we have to special case:
        sequence_type = (
            cls.model_fields["sequence_type"].default.upper()
            if "sequence_type" in cls.model_fields
            else None
        )
        return (short_type, sequence_type)

    to_class = {get_key(cls): cls for cls in BeamlineElement.__subclasses__()}
    key = (get_short_type(type_), sequence_type.upper() if sequence_type else None)
    try:
        return to_class[key]
    except KeyError:
        raise ValueError(
            f"Unknown element type '{type_}'"
            + (f" with sequence_type '{sequence_type}'" if sequence_type else "")
        )


class _LatticeTransformer(lark.visitors.Transformer_InPlaceRecursive):
    """
    Grammar transformer which takes lark objects and makes a :class:`Lattice`.

    Attributes
    ----------
    _filename : str
        Filename source of the input.
    """

    _filename: Optional[pathlib.Path]

    def __init__(self, filename: AnyPath) -> None:
        super().__init__()
        self._filename = pathlib.Path(filename)

    @lark.v_args(inline=True)
    def line(
        self,
        label: lark.Token,
        _: lark.Token,  # line
        element_list: List[LineItem],
    ) -> Tuple[str, Line]:
        return str(label), Line(elements=element_list)

    @lark.v_args(inline=True)
    def parameter_set(
        self,
        parameter: lark.Token,
        value: lark.Token,
    ) -> Tuple[str, lark.Token]:
        return str(parameter), value

    def parameter_list(
        self, sets: List[Tuple[str, ValueType]]
    ) -> List[Tuple[str, ValueType]]:
        return list(sets)

    @lark.v_args(inline=True)
    def beamline_element(
        self,
        label: lark.Token,
        type_: lark.Token,
        parameter_list: Optional[List[Tuple[str, lark.Token]]],
    ) -> Tuple[str, BeamlineElement]:
        param_dict = dict(parameter_list or [])
        cls = element_class_from_genesis(
            type_=str(type_).strip(),
            sequence_type=param_dict.pop("type", None),
        )
        parameters, unknown = parsers.fix_parameters(cls, param_dict)
        if unknown:
            raise ValueError(
                f"Beamline element {label} received unexpected parameter(s): {unknown}"
            )

        parameters["type"] = cls.model_fields["type"].default.lower()
        element = cls.model_validate(parameters)
        return str(label), element

    @lark.v_args(inline=True)
    def duplicate_item(
        self,
        count: lark.Token,
        label: lark.Token,
    ) -> DuplicatedLineItem:
        return DuplicatedLineItem(
            label=str(label),
            count=int(count),
        )

    @lark.v_args(inline=True)
    def positioned_item(
        self,
        label: lark.Token,
        position: lark.Token,
    ) -> PositionedLineItem:
        return PositionedLineItem(
            label=str(label),
            position=float(position),
        )

    @lark.v_args(inline=True)
    def line_item(
        self, item: Union[lark.Token, DuplicatedLineItem, PositionedLineItem]
    ) -> Union[str, DuplicatedLineItem, PositionedLineItem]:
        if isinstance(item, lark.Token):
            return str(item)
        return item

    def element_list(
        self, items: List[Union[lark.Token, DuplicatedLineItem, PositionedLineItem]]
    ) -> List[Union[lark.Token, DuplicatedLineItem, PositionedLineItem]]:
        return items

    def lattice(self, elements: List[Tuple[str, AnyBeamlineElement]]) -> Lattice:
        return Lattice(
            elements=dict(elements),
            filename=self._filename,
        )
