from __future__ import annotations

import logging
import os
import pathlib
import shlex
import shutil
from contextlib import contextmanager
from inspect import isclass
from typing import (
    Any,
    Dict,
    Generator,
    List,
    Optional,
    Sequence,
    Tuple,
    Type,
    TypeVar,
    Union,
)

import h5py
import lark
import numpy as np
import pydantic
import pydantic.alias_generators
from lume import tools as lume_tools
from pmd_beamphysics import ParticleGroup
from pmd_beamphysics.units import c_light
from typing_extensions import override

from ... import tools
from ...errors import (
    MultipleNamelistsError,
    NamelistAccessError,
    NoSetupNamelistError,
    NoSuchNamelistError,
)
from .. import archive as _archive
from ..field import FieldFile
from ..interfaces.bmad import genesis4_namelists_from_tao
from ..particles import Genesis4ParticleData
from ..types import (
    AnyPath,
    BaseModel,
    BeamlineElement,
    NameList,
    NDArray,
    PydanticParticleGroup,
    ReprTableData,
    ValueType,
)
from .lattice import Lattice
from . import _main as auto_main
from . import parsers, util
from ._main import (
    ImportBeam,
    ImportDistribution,
    ImportField,
    ImportTransformation,
    ProfileFile,
    ProfileFileMulti,
    SequenceFilelist,
    Setup,
)

from typing import Literal


AnyNameList = Union[
    auto_main.AutogeneratedNameList,
    "ProfileArray",
]


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

logger = logging.getLogger(__name__)


class ProfileArray(NameList):
    r"""
    ProfileArray is a lume-genesis convenience class for generating
    ``profile_file`` namelists.

    Attributes
    ----------
    x_label : str
        Name of the profile, which is used to refer to it in later calls of namelists
    xdata : list of float or np.ndarray
        The `s`-position for the look-up table.
    ydata : list of float or np.ndarray
        The function values of the look-up table.
    isTime : bool, default=False
        If true the `s`-position is a time variable and therefore multiplied with the
        speed of light `c` to get the position in meters.
    reverse : bool, default=False
        if true the order in the look-up table is reverse. This is sometimes needed
        because time and spatial coordinates differ sometimes by a minus sign.
    filename : str, optional
        By default, this is a randomly-generated filename that lume-genesis
        manages for you.  If desirable, you may set a fixed filename relative
        to the main input file.  Path delimiters (such as ``/``) are not
        allowed.
    x_label : str, optional
        The X label (key) to use for the saved data.
    y_label : str, optional
        The Y label (key) to use for the saved data.
    """

    type: Literal["ProfileArray"] = "ProfileArray"
    label: str
    xdata: NDArray
    ydata: NDArray
    isTime: bool = False
    reverse: bool = False
    autoassign: bool = False
    filename: str = ""
    x_label: str = "x"
    y_label: str = "y"

    def write(self, base_path: AnyPath) -> pathlib.Path:
        if not self.filename:
            self.filename = util.get_temporary_filename(
                prefix=type(self).__name__, extension=".h5"
            )

        if "/" in self.filename:
            raise ValueError(
                f"Filename is not allowed to contain the path separator "
                f"forward slash (/).  Genesis 4 interprets these as part of "
                f"the HDF group. "
                f"Got: {self.filename!r}"
            )
        path = pathlib.Path(base_path) / self.filename
        with h5py.File(path, "w") as fp:
            # for key, value in self.hdf_data.items():
            #     fp[key] = value
            fp.update(self.get_hdf_data())
        logger.info("Saved %s to %s", self.label, path)
        return path

    def get_hdf_data(self) -> Dict[str, np.ndarray]:
        """Get all HDF5 data to be written for Genesis 4."""
        return {
            self.x_label: self.xdata,
            self.y_label: self.ydata,
        }

    def to_profile_file(self) -> ProfileFile:
        """Convert this ProfileArray into a 'profile_file' namelist for Genesis 4."""
        return ProfileFile(
            label=self.label,
            xdata=f"{self.filename}/{self.x_label}",
            ydata=f"{self.filename}/{self.y_label}",
            isTime=self.isTime,
            reverse=self.reverse,
            autoassign=self.autoassign,
        )

    @override
    def to_genesis(self) -> str:
        return self.to_profile_file().to_genesis()


def get_particles_slen(particles: ParticleGroup) -> float:
    """
    Length of the time window in meters.

    Note that for parallel jobs this might be adjusted towards larger
    values.

    This can be used to update the Time namelist with information
    from this initial particle distribution.
    """
    return max(
        c_light * np.ptp(particles.t),
        np.ptp(particles.z),
    )


class MainInput(BaseModel):
    """
    A Genesis 4 main input configuration file.

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

    namelists: List[AnyNameList]
    filename: Optional[pathlib.Path]

    def __init__(
        self,
        namelists: Optional[List[AnyNameList]] = None,
        *,
        filename: Optional[pathlib.Path] = None,
    ) -> None:
        if namelists is None:
            namelists = []
        super().__init__(namelists=namelists, filename=filename)

    def to_genesis(self) -> str:
        return "\n\n".join(namelist.to_genesis() for namelist in self.namelists)

    @override
    def _repr_table_data_(self) -> ReprTableData:
        return {
            "obj": {f"{idx}": obj for idx, obj in enumerate(self.namelists)},
            "annotations": None,
            "descriptions": None,
        }

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

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

    @property
    def by_namelist(self) -> Dict[Type[T_NameList], List[T_NameList]]:
        """Get namelists organized by their class."""
        by_namelist = {}
        for namelist in self.namelists:
            by_namelist.setdefault(type(namelist), [])
            by_namelist[type(namelist)].append(namelist)
        return by_namelist

    @property
    def setup(self) -> Setup:
        """Get the required setup namelist."""
        setups = self.by_namelist.get(Setup, [])
        if len(setups) == 0:
            raise NoSuchNamelistError("Setup is not defined in the input.")
        if len(setups) > 1:
            raise MultipleNamelistsError(
                "Multiple setup namelists were defined in the input."
            )
        return setups[0]

    def to_dicts(
        self,
        exclude_defaults: bool = True,
        by_alias: bool = True,
        **kwargs: Any,
    ) -> List[Dict]:
        """Serialize this main input to a list of dictionaries."""
        return [
            {
                "type": type(namelist).model_fields["type"].default,
                **namelist.model_dump(
                    exclude_defaults=exclude_defaults, by_alias=by_alias, **kwargs
                ),
            }
            for namelist in self.namelists
        ]

    @classmethod
    def from_dicts(
        cls, contents: Sequence[Dict], filename: Optional[pathlib.Path] = None
    ) -> MainInput:
        """
        Load main input from a list of serialized dictionaries.

        Parameters
        ----------
        contents : sequence of dict
            The serialized contents of the main input file.

        Returns
        -------
        MainInput
        """

        if isinstance(contents, dict):
            contents = [contents]

        adapter = pydantic.TypeAdapter(AnyNameList)

        def make_class(idx: int, dct: dict) -> AnyNameList:
            type_ = dct.get("type", None)
            if type_ is None:
                raise ValueError(f"'type' missing from dictionary #{idx}")

            try:
                return adapter.validate_python(dct)
            except pydantic.ValidationError as ex:
                raise ValueError(
                    f"Dictionary #{idx} with type {type_!r} did not pass pydantic validation. "
                    f"Are all parameters for the given class valid?"
                ) from ex

        return cls(
            namelists=[make_class(idx, dct) for idx, dct in enumerate(contents)],
            filename=filename,
        )

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

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

        Returns
        -------
        MainInput
        """
        parser = parsers.new_main_input_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 main input "
                    f"format. It looks like this might have been a filename: "
                    f"{contents!r}"
                )
            raise
        return _MainInputTransformer(filename).transform(tree)

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

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

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

    def to_file(
        self,
        filename: AnyPath,
    ) -> None:
        """
        Write the main 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 main config to %s:\n%s",
            main_config,
            filename,
        )

    @classmethod
    def from_tao(
        cls, tao, ele_start: str = "beginning", branch: int = 0, universe: int = 1
    ):
        """
        Creates Genesis4 MainInput from a Tao instance, using specified parameters to
        extract relevant beamline, beam, and field configurations.

        Parameters
        ----------
        tao : pytao.Tao
            A running Tao instance, providing access to element attributes, Twiss parameters,
            and orbit data.
        ele_start : str, optional
            Element to start. Defaults to "beginning".
        branch : int, optional
            The branch index within the specified Tao universe. Defaults to 0.
        universe : int, optional
            The universe index within the Tao object. Defaults to 1.

        Returns
        -------
        MainInput

        Notes
        -----
        - The generated `beamline` name is based on the Tao universe and branch configuration,
          with `gamma0` calculated from the total energy.

        Examples
        --------
        >>> tao = pytao.Tao(...)
        >>> input = MainInput.from_tao(tao, ele_start='beginning', branch=0, universe=1)

        """
        namelists = genesis4_namelists_from_tao(
            tao, ele_start=ele_start, branch=branch, universe=universe
        )
        return cls(namelists=namelists)

    def _check_for_mistakes(self) -> None:
        """Check and fix simple mistakes to be friendly to the end-user."""
        try:
            setup = self.setup
        except NoSuchNamelistError:
            raise NoSetupNamelistError(
                "No setup namelist: the configuration is invalid"
            ) from None

        if not setup.rootname:
            logger.debug("Setting setup.rootname as the user left it blank")
            setup.rootname = "output"

    def write_files(
        self,
        workdir: AnyPath,
        main_filename: str = "genesis4.in",
        source_path: AnyPath = pathlib.Path(),
        rename: bool = True,
    ) -> List[pathlib.Path]:
        """
        Write the main input file, arrays and other necessary files to run
        Genesis to ``workdir``.

        This may modify ``setup`` or other namelists as required to generate
        valid Genesis4 input.

        Parameters
        ----------
        workdir : pathlib.Path or str
            The work directory where Genesis is to be run.
        main_filename : str, default = "genesis4.in"
            Filename to use for the main input file.
        source_path : pathlib.Path or str, optional
            The source directory, where any referenced HDF5 files will be
            found.  Defaults to the current directory.
        rename : bool, optional
            Adjust temporary filenames of HDF5 files, replacing random
            characters with sensible namelist-prefixed names.

        Returns
        -------
        List[pathlib.Path]
            The additional filenames that were written.
        """
        self._check_for_mistakes()

        workdir = pathlib.Path(workdir).resolve()
        source_path = pathlib.Path(source_path).resolve()
        if not workdir.exists():
            raise ValueError(f"The configured work directory {workdir} does not exist")
        if not workdir.is_dir():
            raise ValueError(
                f"The configured work directory {workdir} is not a directory"
            )

        # We expect that external files are provided by the user in the
        # same directory as the input files.  Verify they exist first.
        for namelist, filename in self._get_files_to_symlink(workdir):
            source_file = source_path / filename
            if not source_file.exists():
                raise FileNotFoundError(f"{source_file} (referenced in {namelist})")

        paths = []
        # Write out arrays to temporary HDF5 files so the user doesn't have to:
        for cls in (ProfileArray,):
            for idx, namelist in enumerate(self.by_namelist.get(cls, [])):
                if rename or not namelist.filename:
                    namelist.filename = f"{cls.__name__}_{idx}.h5"
                paths.append(namelist.write(workdir))

        # Symlink user-referenced pre-existing HDF5 files:
        for _, filename in self._get_files_to_symlink(workdir):
            source_file = source_path / filename
            symlink = workdir / filename
            _symlink_or_copy(symlink=symlink, file=source_file)

            paths.append(symlink)

        # Write the main input file last, as temporary filenames may have
        # changed above.
        self.to_file(workdir / main_filename)
        paths.append(workdir / main_filename)
        return paths

    def _get_files_to_symlink(
        self, workdir: pathlib.Path
    ) -> Generator[Tuple[AnyNameList, str], None, None]:
        for namelist in self.namelists:
            for attr in ("xdata", "ydata", "file"):
                file = getattr(namelist, attr, None)
                if file is None or isinstance(file, np.ndarray):
                    continue
                file = _split_file(namelist, attr, file)
                if not file:
                    continue

                work_file = workdir / file
                if work_file.exists() and not work_file.is_symlink():
                    # The file was written externally, likely by the
                    # initial_particles handler
                    continue

                yield namelist, file

    def remove(
        self, item_or_class: Union[AnyNameList, Type[NameList], Sequence[AnyNameList]]
    ) -> List[AnyNameList]:
        """
        Remove a item from the namelist by instance or class.

        Parameters
        ----------
        item_or_class : NameList class, instance, or list

        Examples
        --------

        Remove a single instance of "Track" namelists.

        >>> G.input.main.remove(G.input.main.track)

        Remove a single instance of "Track" namelists.

        >>> G.input.main.remove(G.input.main.track)

        Remove all instances of "Track" namelists using the class:

        >>> G.input.main.remove(Track)
        """
        to_remove: List[AnyNameList]
        if isinstance(item_or_class, NameList):
            to_remove = [item_or_class]
        elif isclass(item_or_class):
            to_remove = [
                obj for obj in list(self.namelists) if isinstance(obj, item_or_class)
            ]
        else:
            to_remove = list(item_or_class)

        for obj in to_remove:
            self.namelists.remove(obj)
        return to_remove

    def insert_initial_field(
        self,
        field: FieldFile,
        harmonic: int = 1,
        time: bool = True,
    ) -> int:
        """Insert a FieldFile instance as as an initial field with importfield."""
        to_insert = ImportField(
            file="initial_field.h5",
            harmonic=harmonic,
            time=time,
        )

        logger.debug("Inserting initial field file: %s", field)
        try:
            previous = self.import_field
        except NamelistAccessError:
            if self.fields:
                # Before Field, if it exists
                insert_pos = self.namelists.index(self.fields[0])
            elif self.times:
                # Otherwise, it needs to come after Time
                insert_pos = self.namelists.index(self.times[-1]) + 1
            else:
                logger.warning(
                    "Unable to determine where to insert the importfield; "
                    "placing it after 'setup'"
                )
                insert_pos = self.namelists.index(self.setup) + 1
        else:
            insert_pos = self.namelists.index(previous)
            self.remove(previous)

        self.namelists.insert(insert_pos, to_insert)
        return insert_pos

    def insert_initial_particles(
        self,
        particles: Union[ParticleGroup, Genesis4ParticleData],
        update_slen: bool,
    ) -> int:
        """Insert a ParticleGroup instance as as an initial particle distribution."""

        if isinstance(particles, ParticleGroup):
            to_insert = ImportDistribution(
                file="initial_particles.h5",
                charge=particles.charge,
            )
            to_remove = self.import_beams + self.beams
        else:
            to_insert = ImportBeam(file="initial_particles.h5")
            to_remove = self.import_distributions + self.beams

        for remove in to_remove:
            logger.warning(f"Removing existing {type(remove).__name__}")
            self.remove(remove)

        try:
            if isinstance(to_insert, ImportDistribution):
                previous = self.import_distribution
            else:
                previous = self.import_beam
        except NamelistAccessError:
            try:
                # Insert it after the time namelist:
                insert_pos = self.namelists.index(self.time) + 1
            except NamelistAccessError:
                # Last attempt: put it after setup
                insert_pos = self.namelists.index(self.setup) + 1
        else:
            insert_pos = self.namelists.index(previous)
            self.remove(previous)

        if update_slen and isinstance(particles, ParticleGroup):
            particles_slen = get_particles_slen(particles)
            for time_ in self.times:
                was = time_.slen
                if was < particles_slen:
                    time_.slen = particles_slen
                    logger.warning(
                        "Updating time namelist slen: %f (was %f) to span all particles",
                        time_.slen,
                        was,
                    )

        self.namelists.insert(insert_pos, to_insert)
        return insert_pos

    @property
    def alter_setups(self) -> List[auto_main.AlterSetup]:
        """List of all AlterSetup instances."""
        return self.by_namelist.get(auto_main.AlterSetup, [])

    @property
    def lattices(self) -> List[auto_main.LatticeNamelist]:
        """List of all Lattice instances."""
        return self.by_namelist.get(auto_main.LatticeNamelist, [])

    @property
    def times(self) -> List[auto_main.Time]:
        """List of all Time instances."""
        return self.by_namelist.get(auto_main.Time, [])

    @property
    def profile_consts(self) -> List[auto_main.ProfileConst]:
        """List of all ProfileConst instances."""
        return self.by_namelist.get(auto_main.ProfileConst, [])

    @property
    def profile_gausses(self) -> List[auto_main.ProfileGauss]:
        """List of all ProfileGauss instances."""
        return self.by_namelist.get(auto_main.ProfileGauss, [])

    @property
    def profile_steps(self) -> List[auto_main.ProfileStep]:
        """List of all ProfileStep instances."""
        return self.by_namelist.get(auto_main.ProfileStep, [])

    @property
    def profile_polynoms(self) -> List[auto_main.ProfilePolynom]:
        """List of all ProfilePolynom instances."""
        return self.by_namelist.get(auto_main.ProfilePolynom, [])

    @property
    def profile_files(self) -> List[auto_main.ProfileFile]:
        """List of all ProfileFile instances."""
        return self.by_namelist.get(auto_main.ProfileFile, [])

    @property
    def sequence_consts(self) -> List[auto_main.SequenceConst]:
        """List of all SequenceConst instances."""
        return self.by_namelist.get(auto_main.SequenceConst, [])

    @property
    def sequence_polynoms(self) -> List[auto_main.SequencePolynom]:
        """List of all SequencePolynom instances."""
        return self.by_namelist.get(auto_main.SequencePolynom, [])

    @property
    def sequence_powers(self) -> List[auto_main.SequencePower]:
        """List of all SequencePower instances."""
        return self.by_namelist.get(auto_main.SequencePower, [])

    @property
    def sequence_randoms(self) -> List[auto_main.SequenceRandom]:
        """List of all SequenceRandom instances."""
        return self.by_namelist.get(auto_main.SequenceRandom, [])

    @property
    def beams(self) -> List[auto_main.Beam]:
        """List of all Beam instances."""
        return self.by_namelist.get(auto_main.Beam, [])

    @property
    def fields(self) -> List[auto_main.Field]:
        """List of all Field instances."""
        return self.by_namelist.get(auto_main.Field, [])

    @property
    def import_distributions(self) -> List[auto_main.ImportDistribution]:
        """List of all ImportDistribution instances."""
        return self.by_namelist.get(auto_main.ImportDistribution, [])

    @property
    def import_beams(self) -> List[auto_main.ImportBeam]:
        """List of all ImportBeam instances."""
        return self.by_namelist.get(auto_main.ImportBeam, [])

    @property
    def import_fields(self) -> List[auto_main.ImportField]:
        """List of all ImportField instances."""
        return self.by_namelist.get(auto_main.ImportField, [])

    @property
    def import_transformations(self) -> List[auto_main.ImportTransformation]:
        """List of all ImportTransformation instances."""
        return self.by_namelist.get(auto_main.ImportTransformation, [])

    @property
    def efields(self) -> List[auto_main.Efield]:
        """List of all Efield instances."""
        return self.by_namelist.get(auto_main.Efield, [])

    @property
    def sponrads(self) -> List[auto_main.Sponrad]:
        """List of all Sponrad instances."""
        return self.by_namelist.get(auto_main.Sponrad, [])

    @property
    def wakes(self) -> List[auto_main.Wake]:
        """List of all Wake instances."""
        return self.by_namelist.get(auto_main.Wake, [])

    @property
    def writes(self) -> List[auto_main.Write]:
        """List of all Write instances."""
        return self.by_namelist.get(auto_main.Write, [])

    @property
    def tracks(self) -> List[auto_main.Track]:
        """List of all Track instances."""
        return self.by_namelist.get(auto_main.Track, [])

    @property
    def alter_fields(self) -> List[auto_main.AlterField]:
        """List of all AlterField instances."""
        return self.by_namelist.get(auto_main.AlterField, [])

    @property
    def profile_file_multis(self) -> List[auto_main.ProfileFileMulti]:
        """List of all ProfileFileMulti instances."""
        return self.by_namelist.get(auto_main.ProfileFileMulti, [])

    @property
    def sequence_lists(self) -> List[auto_main.SequenceList]:
        """List of all SequenceList instances."""
        return self.by_namelist.get(auto_main.SequenceList, [])

    @property
    def sequence_filelists(self) -> List[auto_main.SequenceFilelist]:
        """List of all SequenceFilelist instances."""
        return self.by_namelist.get(auto_main.SequenceFilelist, [])

    @property
    def profile_arrays(self) -> List[ProfileArray]:
        """List of all ProfileArray instances."""
        return self.by_namelist.get(ProfileArray, [])

    def _get_only_one(self, cls: Type[T_NameList]) -> T_NameList:
        items = self.by_namelist.get(cls, [])
        if len(items) == 0:
            raise NoSuchNamelistError(f"{cls.__name__} is not defined in the input.")
        if len(items) > 1:
            plural_fix = {
                "profile_gauss": "profile_gausses",
                "lattice_namelist": "lattices",
            }
            plural = pydantic.alias_generators.to_snake(cls.__name__)
            plural = plural_fix.get(plural, f"{plural}s")
            raise MultipleNamelistsError(
                f"Multiple {cls.__name__} namelists were defined in the input. "
                f"Please use .{plural}"
            )
        return items[0]

    @property
    def alter_setup(self) -> auto_main.AlterSetup:
        """Get a single AlterSetup instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.AlterSetup)

    @property
    def lattice(self) -> auto_main.LatticeNamelist:
        """Get a single Lattice instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.LatticeNamelist)

    @property
    def time(self) -> auto_main.Time:
        """Get a single Time instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.Time)

    @property
    def profile_const(self) -> auto_main.ProfileConst:
        """Get a single ProfileConst instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.ProfileConst)

    @property
    def profile_gauss(self) -> auto_main.ProfileGauss:
        """Get a single ProfileGauss instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.ProfileGauss)

    @property
    def profile_step(self) -> auto_main.ProfileStep:
        """Get a single ProfileStep instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.ProfileStep)

    @property
    def profile_polynom(self) -> auto_main.ProfilePolynom:
        """Get a single ProfilePolynom instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.ProfilePolynom)

    @property
    def profile_file(self) -> auto_main.ProfileFile:
        """Get a single ProfileFile instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.ProfileFile)

    @property
    def sequence_const(self) -> auto_main.SequenceConst:
        """Get a single SequenceConst instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.SequenceConst)

    @property
    def sequence_polynom(self) -> auto_main.SequencePolynom:
        """Get a single SequencePolynom instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.SequencePolynom)

    @property
    def sequence_power(self) -> auto_main.SequencePower:
        """Get a single SequencePower instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.SequencePower)

    @property
    def sequence_random(self) -> auto_main.SequenceRandom:
        """Get a single SequenceRandom instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.SequenceRandom)

    @property
    def beam(self) -> auto_main.Beam:
        """Get a single Beam instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.Beam)

    @property
    def field(self) -> auto_main.Field:
        """Get a single Field instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.Field)

    @property
    def import_distribution(self) -> auto_main.ImportDistribution:
        """Get a single ImportDistribution instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.ImportDistribution)

    @property
    def import_beam(self) -> auto_main.ImportBeam:
        """Get a single ImportBeam instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.ImportBeam)

    @property
    def import_field(self) -> auto_main.ImportField:
        """Get a single ImportField instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.ImportField)

    @property
    def import_transformation(self) -> auto_main.ImportTransformation:
        """Get a single ImportTransformation instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.ImportTransformation)

    @property
    def efield(self) -> auto_main.Efield:
        """Get a single Efield instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.Efield)

    @property
    def sponrad(self) -> auto_main.Sponrad:
        """Get a single Sponrad instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.Sponrad)

    @property
    def wake(self) -> auto_main.Wake:
        """Get a single Wake instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.Wake)

    @property
    def write(self) -> auto_main.Write:
        """Get a single Write instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.Write)

    @property
    def track(self) -> auto_main.Track:
        """Get a single Track instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.Track)

    @property
    def alter_field(self) -> auto_main.AlterField:
        """Get a single AlterField instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.AlterField)

    @property
    def profile_file_multi(self) -> auto_main.ProfileFileMulti:
        """Get a single ProfileFileMulti instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.ProfileFileMulti)

    @property
    def sequence_list(self) -> auto_main.SequenceList:
        """Get a single SequenceList instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.SequenceList)

    @property
    def sequence_filelist(self) -> auto_main.SequenceFilelist:
        """Get a single SequenceFilelist instance. Errors if zero or 2+ exist."""
        return self._get_only_one(auto_main.SequenceFilelist)

    @property
    def profile_array(self) -> ProfileArray:
        """Get a single ProfileArray instance. Errors if zero or 2+ exist."""
        return self._get_only_one(ProfileArray)


def _symlink_or_copy(symlink: pathlib.Path, file: pathlib.Path):
    if os.name == "nt":
        # With Windows 10, users need Administator Privileges or run on
        # Developer mode in order to be able to create symlinks.
        # Ref: https://docs.python.org/3/library/os.html#os.symlink
        return shutil.copy2(file, symlink)  # type: ignore[unreachable]

    if symlink != file and not symlink.is_symlink():
        return symlink.symlink_to(file)


def _split_file(namelist: AnyNameList, attr: str, file: str) -> Optional[str]:
    """
    Get a filename from the namelist.

    Namelists such as ProfileFile store arrays in HDF5 files and reference
    them by way of a string in the form: ``"FILENAME/KEY"``.

    This limits Genesis4's ability to load files in absolute paths, as the
    delimiter ``/`` is used to separate the filename from the key.
    """
    if not isinstance(
        namelist,
        (
            ProfileFile,
            ProfileFileMulti,
            ImportDistribution,
            ImportBeam,
            ImportField,
            ImportTransformation,
            SequenceFilelist,
            # ProfileArray,
        ),
    ):
        return None
    if attr in ("xdata", "ydata"):
        # We can't handle paths here, although perhaps lume-genesis
        # could fix this for users... hmm.
        return file.split("/")[0]
    if attr in ("file",):
        return file
    raise NotImplementedError(attr)


class Genesis4Input(BaseModel):
    """
    All Genesis 4-related command input.

    Attributes
    ----------
    main : MainInput
        The main input. This contains all the namelists, starting with
        `&setup`.
    lattice : Lattice
        The lattice definition.
    beamline : str, optional
        Optional override for `beamline` in the setup namelist.
    lattice_name : str, optional
        Optional override for `lattice` in the setup namelist.
    seed : str, optional
        Optional override for `seed` in the setup namelist.
    output_path : Optional[AnyPath] = None
        Optional override for `rootname` in the setup namelist.
    source_path : Optional[AnyPath] = None
        When using Genesis 4-compatible input as strings, this is the directory
        where we expect to find other input HDF5 files.
    input_filename : str = "genesis4.in"
        The filename to use when writing the main input file. As a user, you
        should not need to worry about this as Genesis4Input will handle it for
        you.
    initial_particles : ParticleGroup, optional
        OpenPMD ParticleGroup to use as the initial Genesis 4 particle
        distribution.  If set, this acts as a shortcut for automatically adding
        an ``importdistribution`` namelist and writing the corresponding HDF5
        file for Genesis.
    initial_field : FieldFile, optional
        Initial Genesis 4 format field file to use in the simulation. Use
        `FieldFile.from_file` to load the file first.
    """

    main: MainInput
    lattice: Lattice
    beamline: Optional[str] = None
    lattice_name: Optional[str] = None
    seed: Optional[str] = None
    source_path: pathlib.Path = pathlib.Path(".")
    output_path: Optional[AnyPath] = None
    input_filename: str = "genesis4.in"
    initial_particles: Optional[Union[PydanticParticleGroup, Genesis4ParticleData]] = (
        None
    )
    initial_field: Optional[FieldFile] = None
    # initial_wavefront: Optional[PmdWavefront] = None

    def __init__(self, /, **data: Any) -> None:
        self.__pydantic_validator__.validate_python(
            data,
            self_instance=self,
            context={"initializing": True},
        )

    @pydantic.model_validator(mode="after")
    def _validate_initial_particles(self, info: pydantic.ValidationInfo):
        initializing = info.context and info.context.get("initializing")
        if self.initial_particles is None:
            if initializing:
                # No initial particles were specified.  Do not modify the namelists.
                return self

            if ImportDistribution in self.main.by_namelist:
                logger.warning(
                    "initial_particles cleared; removing ImportDistribution namelist."
                )
                self.main.remove(ImportDistribution)
            if ImportBeam in self.main.by_namelist:
                logger.warning(
                    "initial_particles cleared; removing ImportBeam namelist."
                )
                self.main.remove(ImportBeam)
        else:
            self.main.insert_initial_particles(self.initial_particles, update_slen=True)
        return self

    @pydantic.model_validator(mode="after")
    def _validate_initial_field(self, info: pydantic.ValidationInfo):
        initializing = info.context and info.context.get("initializing")
        if self.initial_field is None:
            if initializing:
                # No initial field was specified.  Do not modify the namelists.
                return self
            if ImportField in self.main.by_namelist:
                logger.warning("initial_field cleared; removing ImportField namelist.")
                self.main.remove(ImportField)
        else:
            self.main.insert_initial_field(self.initial_field)
        return self

    def to_genesis(self) -> str:
        return "\n".join(
            (
                "# Main input",
                self.main.to_genesis(),
                "",
                "# Lattice",
                self.lattice.to_genesis(),
            )
        )

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

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

    @property
    def arguments(self) -> List[str]:
        """
        Get all of the command-line arguments for running Genesis 4.

        Returns
        -------
        list of str
            Individual arguments to pass to Genesis 4.
        """
        optional_args = []
        if self.beamline:
            optional_args.extend(["-b", self.beamline])
        if self.lattice_name:
            optional_args.extend(["-l", self.lattice_name])
        if self.seed is not None:
            optional_args.extend(["-s", self.seed])
        if self.output_path is not None:
            optional_args.extend(["-o", str(self.output_path)])
        return [
            *optional_args,
            "-l",
            str(self.lattice_filename),
            str(self.input_filename),
        ]

    @property
    def lattice_filename(self) -> str:
        """
        The filename of the lattice, determined from the Setup namelist.

        Defaults to "genesis.lat" if not previously set.
        """
        setup = self.main.setup
        if not setup.lattice:
            setup.lattice = "genesis.lat"
        return setup.lattice

    def write(
        self,
        workdir: AnyPath = pathlib.Path("."),
        rename: bool = True,
    ) -> List[pathlib.Path]:
        """
        Write all input files for executing Genesis 4.

        Parameters
        ----------
        workdir : AnyPath
            The working directory for Genesis 4. This will be where all
            input and output files are written.
        rename : bool, optional
            Adjust temporary filenames of HDF5 files, replacing random
            characters with sensible namelist-prefixed names.

        Returns
        -------
        List[pathlib.Path]
            Paths written.
        """
        path = pathlib.Path(workdir)
        path.mkdir(parents=True, mode=0o755, exist_ok=True)

        extra_paths = []
        # Write the particle and field files prior to MainInput, as it will
        # check to ensure that they exist
        if self.initial_particles is not None:
            extra_paths.append(self.write_initial_particles(path))

        if self.initial_field is not None:
            extra_paths.append(self.write_initial_field(path))

        main_paths = self.main.write_files(
            path,
            main_filename=self.input_filename,
            source_path=self.source_path,
            rename=rename,
        )

        extra_paths.extend(main_paths)

        lattice_path = path / self.lattice_filename
        self.lattice.to_file(filename=lattice_path)
        return [path / self.input_filename, lattice_path, *extra_paths]

    def write_initial_field(self, workdir: AnyPath) -> pathlib.Path:
        """Write the initial field file."""
        if self.initial_field is None:
            raise ValueError("initial_field was not set")

        workdir = pathlib.Path(workdir)
        if not workdir.is_dir():
            raise ValueError("Provided path is not a directory: {}")

        import_field = self.main.import_field
        if not import_field.file:
            raise ValueError("import_field filename (.file) unset")

        filename = workdir / import_field.file
        self.initial_field.write_genesis4(filename)
        return filename

    def write_initial_particles(self, workdir: AnyPath) -> pathlib.Path:
        """Write the initial particles file."""
        if self.initial_particles is None:
            raise ValueError("initial_particles was not set")

        workdir = pathlib.Path(workdir)
        if not workdir.is_dir():
            raise ValueError("Provided path is not a directory: {}")

        if isinstance(self.initial_particles, Genesis4ParticleData):
            beam = self.main.import_beam
            particle_file = pathlib.Path(workdir) / beam.file
        elif isinstance(self.initial_particles, ParticleGroup):
            dist = self.main.import_distribution
            particle_file = pathlib.Path(workdir) / dist.file
            if dist.charge != self.initial_particles.charge:
                logger.warning(
                    f"Updating distribution charge: was={dist.charge} "
                    f"now={self.initial_particles.charge}"
                )
                dist.charge = self.initial_particles.charge

            particles_slen = get_particles_slen(self.initial_particles)
            for time_ in self.main.times:
                was = time_.slen
                if was < particles_slen:
                    time_.slen = particles_slen
                    logger.warning(
                        f"Updating time namelist slen: was={was} now {time_.slen} to span all particles",
                    )
        else:
            raise ValueError(
                f"Initial particles type unexpected: {type(self.initial_particles).__name__}"
            )

        self.initial_particles.write_genesis4_distribution(
            str(particle_file),
            verbose=False,
        )
        return particle_file

    def write_run_script(
        self,
        path: pathlib.Path,
        command_prefix: str = "genesis4",
    ) -> None:
        with open(path, mode="wt") as fp:
            print(shlex.join(shlex.split(command_prefix) + self.arguments), file=fp)
        lume_tools.make_executable(str(path))

    @contextmanager
    def write_context(
        self, workdir: AnyPath
    ) -> Generator[List[pathlib.Path], None, None]:
        """
        Write all input files for executing Genesis 4.

        This is a context manager. When the context manager exits, all input
        files will be cleaned up (i.e., deleted from disk).

        Parameters
        ----------
        workdir : AnyPath
            The working directory for Genesis 4. This will be where all
            input and output files are written.

        Yields
        ------
        List[pathlib.Path]
            Paths written.
        """
        paths = self.write(workdir)
        yield paths
        for path in paths:
            path.unlink(missing_ok=True)

    @classmethod
    def from_main_input(
        cls,
        main: MainInput,
        lattice: Union[Lattice, str, pathlib.Path] = "",
        *,
        source_path: Optional[pathlib.Path] = None,
    ) -> Genesis4Input:
        """
        Work with a lume-genesis MainInput and potentially an external
        lattice file or instance.

        If the input refers to files that already exist on disk, ensures
        that `source_path` is set correctly.
        """
        from .lattice import Lattice

        if source_path is None:
            source_path = pathlib.Path(".")

        if isinstance(lattice, Lattice):
            return cls(
                main=main,
                lattice=lattice,
                source_path=source_path,
            )

        if lattice:
            lattice_path, lattice = tools.read_if_path(lattice)
            if lattice_path is not None:
                try:
                    setup = main.setup
                except NoSuchNamelistError:
                    # No setup yet; skip for now
                    pass
                else:
                    main.setup.lattice = lattice_path.name
            lattice = Lattice.from_contents(lattice, filename=lattice_path)
        else:
            logger.debug("Lattice not specified; determining from main input")
            setup = None
            try:
                setup = main.setup
                assert setup is not None
                # Use the lattice filename from main input's setup:
                with open(source_path / setup.lattice) as fp:
                    lattice = fp.read()
            except Exception:
                if setup is not None:
                    logger.exception(
                        (
                            "Lattice not specified and unable to determine it from the "
                            "main input's setup. Setup.lattice=%s Lattice file should "
                            "be located here: %s"
                        ),
                        setup.lattice,
                        source_path / setup.lattice,
                    )
                raise
            lattice = Lattice.from_contents(
                lattice, filename=source_path / setup.lattice
            )

        return cls(
            main=main,
            lattice=lattice,
            source_path=source_path,
        )

    @classmethod
    def from_strings(
        cls,
        main: str,
        lattice: str = "",
        *,
        source_path: pathlib.Path = pathlib.Path("."),
        input_filename: Optional[pathlib.Path] = None,
        lattice_filename: Optional[pathlib.Path] = None,
    ) -> Genesis4Input:
        """
        Work directly with Genesis 4-compatible inputs.

        If unspecified, `lattice` will be determined from the `main` input
        settings.

        If the input refers to files that already exist on disk, ensures
        that `source_path` is set correctly.
        """
        from .lattice import Lattice

        main_config = MainInput.from_contents(main, filename=input_filename)

        if not lattice:
            return cls.from_main_input(main_config, source_path=source_path)

        return cls(
            main=main_config,
            lattice=Lattice.from_contents(lattice, filename=lattice_filename),
            source_path=source_path,
        )

    def archive(self, h5: h5py.Group) -> None:
        """
        Dump input data into the given HDF5 group.

        Parameters
        ----------
        h5 : h5py.Group
            The HDF5 file in which to write the information.
        """
        _archive.store_in_hdf5_file(h5, self)

    @classmethod
    def from_archive(cls, h5: h5py.Group) -> Genesis4Input:
        """
        Loads input from archived h5 file.

        Parameters
        ----------
        h5 : str or h5py.File
            The filename or handle on h5py.File from which to load data.
        """
        loaded = _archive.restore_from_hdf5_file(h5)
        if not isinstance(loaded, Genesis4Input):
            raise ValueError(
                f"Loaded {loaded.__class__.__name__} instead of a "
                f"Genesis4Input instance.  Was the HDF group correct?"
            )
        return loaded


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

    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)

        # This maps, e.g., "setup" to the Setup dataclass
        self.type_name_to_class = {
            cls.model_fields["type"].default: cls for cls in NameList.__subclasses__()
        }

    @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 namelist(
        self,
        name: lark.Token,
        parameter_list: List[Tuple[str, lark.Token]],
        _: lark.Token,  # end
    ) -> NameList:
        cls = self.type_name_to_class[name]
        parameters, unknown = parsers.fix_parameters(cls, dict(parameter_list))
        if unknown:
            raise ValueError(
                f"Namelist {name} received unexpected parameter(s): {unknown}"
            )
        return cls(**parameters)

    @lark.v_args(inline=True)
    def main_input(self, *namelists: AnyNameList) -> MainInput:
        return MainInput(namelists=list(namelists), filename=self._filename)
