"""cryoJAX compatibility with [RELION](https://relion.readthedocs.io/en/release-5.0/)."""

import abc
import pathlib
import re
import typing
import warnings
from collections.abc import Callable
from copy import deepcopy
from typing import Any, Literal, TypedDict, cast
from typing_extensions import Self, override

import equinox as eqx
import jax
import jax.numpy as jnp
import mrcfile
import numpy as np
import pandas as pd
from cryojax.io import write_image_stack_to_mrc
from cryojax.jax_util import NDArrayLike
from cryojax.ndimage import FourierConstant, FourierGaussian
from cryojax.simulator import (
    AstigmaticCTF,
    BasicImageConfig,
    ContrastTransferTheory,
    EulerAnglePose,
)
from jaxtyping import Array, Float, Int

from .._io import read_starfile, write_starfile
from .base_dataset import AbstractParticleDataset, AbstractParticleParameterFile


# RELION column entries
RELION_CTF_OPTICS_ENTRIES = [
    ("rlnSphericalAberration", "Float64"),
    ("rlnAmplitudeContrast", "Float64"),
]
RELION_INSTRUMENT_OPTICS_ENTRIES = [
    ("rlnImageSize", "Int64"),
    ("rlnVoltage", "Float64"),
    ("rlnImagePixelSize", "Float64"),
]
RELION_CTF_PARTICLE_ENTRIES = [
    ("rlnDefocusU", "Float64"),
    ("rlnDefocusV", "Float64"),
    ("rlnDefocusAngle", "Float64"),
    ("rlnPhaseShift", "Float64"),
]
RELION_POSE_PARTICLE_ENTRIES = [
    ("rlnOriginXAngst", "Float64"),
    ("rlnOriginYAngst", "Float64"),
    ("rlnAngleRot", "Float64"),
    ("rlnAngleTilt", "Float64"),
    ("rlnAnglePsi", "Float64"),
]
# Default entries for writing
RELION_DEFAULT_OPTICS_ENTRIES = [
    *RELION_INSTRUMENT_OPTICS_ENTRIES,
    *RELION_CTF_OPTICS_ENTRIES,
    ("rlnOpticsGroup", "Int64"),
]
RELION_DEFAULT_PARTICLE_ENTRIES = [
    *RELION_CTF_PARTICLE_ENTRIES,
    *RELION_POSE_PARTICLE_ENTRIES,
    ("rlnOpticsGroup", "Int64"),
]
# Required entries for loading
RELION_REQUIRED_OPTICS_ENTRIES = RELION_DEFAULT_OPTICS_ENTRIES
RELION_REQUIRED_PARTICLE_ENTRIES = [
    *RELION_CTF_PARTICLE_ENTRIES,
    ("rlnOpticsGroup", "Int64"),
]
RELION_SUPPORTED_PARTICLE_ENTRIES = [
    *RELION_REQUIRED_PARTICLE_ENTRIES,
    *RELION_POSE_PARTICLE_ENTRIES,
    ("rlnCtfBfactor", "Float64"),
    ("rlnCtfScalefactor", "Float64"),
]


if hasattr(typing, "GENERATING_DOCUMENTATION"):
    _ParticleParameterInfo = dict[str, Any]  # pyright: ignore[reportAssignmentType]
    _ParticleStackInfo = dict[str, Any]  # pyright: ignore[reportAssignmentType]
    _ParticleParameterLike = dict[str, Any]  # pyright: ignore[reportAssignmentType]
    _ParticleStackLike = dict[str, Any]  # pyright: ignore[reportAssignmentType]
    _Options = dict[str, Any]  # pyright: ignore[reportAssignmentType]
    _StarfileData = dict[str, Any]  # pyright: ignore[reportAssignmentType]
    _MrcfileSettings = dict[str, Any]  # pyright: ignore[reportAssignmentType]

else:

    class _ParticleParameterInfo(TypedDict):
        """Parameters for a particle stack from RELION."""

        image_config: BasicImageConfig
        pose: EulerAnglePose
        transfer_theory: ContrastTransferTheory

        metadata: pd.DataFrame | None

    class _ParticleStackInfo(TypedDict):
        """Particle stack info from RELION."""

        parameters: _ParticleParameterInfo | None
        images: Float[np.ndarray, "... y_dim x_dim"]

    _ParticleParameterLike = dict[str, Any] | _ParticleParameterInfo
    _ParticleStackLike = dict[str, Any] | _ParticleStackInfo

    class _Options(TypedDict):
        loads_metadata: bool
        loads_envelope: bool
        broadcasts_image_config: bool
        updates_optics_group: bool
        rotation_convention: Literal["object", "frame"]
        pad_options: dict[str, Any]

    class _StarfileData(TypedDict):
        optics: pd.DataFrame
        particles: pd.DataFrame

    class _MrcfileSettings(TypedDict):
        prefix: str
        output_folder: str | pathlib.Path
        n_characters: int
        delimiter: str
        overwrite: bool
        compression: str | None


class AbstractRelionParticleParameterFile(
    AbstractParticleParameterFile[_ParticleParameterInfo, _ParticleParameterLike]
):
    """Abstract class for a RELION particle parameter file.

    This class is mostly useful for wrapping
    the [`cryospax.RelionParticleParameterFile`][] into
    slightly different indexing behavior downstream.
    """

    @property
    @override
    def path_to_output(self) -> pathlib.Path:
        return self.path_to_starfile

    @path_to_output.setter
    @override
    def path_to_output(self, value: str | pathlib.Path):
        self.path_to_starfile = value

    @property
    @abc.abstractmethod
    def path_to_starfile(self) -> pathlib.Path:
        raise NotImplementedError

    @path_to_starfile.setter
    @abc.abstractmethod
    def path_to_starfile(self, value: str | pathlib.Path):
        raise NotImplementedError

    @property
    @abc.abstractmethod
    def starfile_data(self) -> _StarfileData:
        raise NotImplementedError

    @starfile_data.setter
    @abc.abstractmethod
    def starfile_data(self, value: dict[str, pd.DataFrame]):
        raise NotImplementedError

    @property
    @abc.abstractmethod
    def loads_metadata(self) -> bool:
        raise NotImplementedError

    @loads_metadata.setter
    @abc.abstractmethod
    def loads_metadata(self, value: bool):
        raise NotImplementedError

    @property
    @abc.abstractmethod
    def loads_envelope(self) -> bool:
        raise NotImplementedError

    @loads_envelope.setter
    @abc.abstractmethod
    def loads_envelope(self, value: bool):
        raise NotImplementedError

    @property
    @abc.abstractmethod
    def broadcasts_image_config(self) -> bool:
        raise NotImplementedError

    @broadcasts_image_config.setter
    @abc.abstractmethod
    def broadcasts_image_config(self, value: bool):
        raise NotImplementedError

    @property
    @abc.abstractmethod
    def updates_optics_group(self) -> bool:
        raise NotImplementedError

    @updates_optics_group.setter
    @abc.abstractmethod
    def updates_optics_group(self, value: bool):
        raise NotImplementedError

    @property
    @abc.abstractmethod
    def rotation_convention(self) -> Literal["object", "frame"]:
        raise NotImplementedError

    @rotation_convention.setter
    @abc.abstractmethod
    def rotation_convention(self, value: Literal["object", "frame"]):
        raise NotImplementedError

    def copy(self) -> Self:
        return deepcopy(self)


class RelionParticleParameterFile(AbstractRelionParticleParameterFile):
    """A dataset that wraps a RELION particle stack in
    [STAR](https://relion.readthedocs.io/en/latest/Reference/Conventions.html)
    format.
    """

    def __init__(
        self,
        path_to_starfile: str | pathlib.Path,
        mode: Literal["r", "w"] = "r",
        *,
        exists_ok: bool = False,
        selection_filter: dict[str, Callable] | None = None,
        options: dict[str, Any] = {},
    ):
        """**Arguments:**

        - `path_to_starfile`:
            The path to the RELION STAR file. If the path does not exist
            and `mode = 'w'`, an empty dataset will be created.
        - `path_to_relion_project`: The path to the RELION project directory.
        - `mode`:
            - If `mode = 'w'`, the dataset is prepared to write new
            *parameters*. This is done by storing an empty dataset in
            `RelionParticleParameterFile.starfile_data`. If a STAR file
            already exists at `path_to_starfile`, set `exists_ok = True`.
            - If `mode = 'r'`, the STAR file at `path_to_starfile` is read
            into `RelionParticleParameterFile.starfile_data`.
        - `exists_ok`:
            If the `path_to_starfile` already exists, if `True` and `mode = 'w'`
            nonetheless stores an empty `RelionParticleParameterFile.starfile_data`.
        - `selection_filter`:
            A dictionary used to include only particular dataset elements.
            The keys of this dictionary should be any data entry in the STAR
            file, while the values should be a function that takes in a
            column and returns a boolean mask for the column. For example,
            filter by class using
            `selection_filter["rlnClassNumber"] = lambda x: x == 0`.
        - `options`:
            A dictionary of options for modifying the behavior of loading/writing.
            - 'loads_metadata':
                If `True`, the resulting dict loads
                the raw metadata from the STAR file that is not otherwise included
                into a `pandas.DataFrame`.
                If this is set to `True`, note that dictionaries cannot pass through
                JIT boundaries without removing the metadata.
                By default, `False`.
            - 'broadcasts_image_config':
                If `True`, image config parameters are broadcasted with leading dimension
                as the number of particles.
                By default, `True`.
            - 'loads_envelope':
                If `True`, read in the parameters of the CTF envelope function, i.e.
                "rlnCtfScalefactor" and "rlnCtfBfactor".
                By default, `False`.
            - 'updates_optics_group':
                If `True`, when re-writing STAR file entries via
                `dataset[idx] = parameters` syntax, creates a new optics group entry.
                By default, `False`.
            - 'rotation_convention':
                If `'object'`, the loader loads/writes poses in the convention that
                the rotation is of the *object*. If `'frame'`, the rotation is of
                the frame (i.e. the rotation is inverted).
                The pose passed to [`cryojax.simulator.make_image_model`](https://michael-0brien.github.io/cryojax/api/simulator/entry-point/#cryojax.simulator.make_image_model)
                is always of the object, but advanced considerations may require
                setting `rotation_convention = 'frame'` (or manually calling
                `pose.to_inverse_rotation()`) to correctly match RELION and
                cryoJAX conventions.
                By default, `'object'`.
            - 'pad_options':
                Padding options for image simulation, passed to the `BasicImageConfig`.
                See `BasicImageConfig` for documentation.
                By default, `{}`.
        """  # noqa: E501
        # Private attributes
        self._options = _dict_to_options(options)
        self._mode = _validate_mode(mode)
        # The STAR file data
        self._path_to_starfile = pathlib.Path(path_to_starfile)
        self._starfile_data = _load_starfile_data(
            self._path_to_starfile, selection_filter, mode, exists_ok
        )

    @override
    def __getitem__(
        self, index: int | slice | Int[np.ndarray, ""] | Int[np.ndarray, " _"]
    ) -> _ParticleParameterInfo:
        """Load STAR file entries with `value = parameter_file[...]` syntax,
        where `value` is a dictionary with keys

        - 'pose':
            The [`cryojax.simulator.EulerAnglePose`](https://michael-0brien.github.io/cryojax/api/simulator/pose/#cryojax.simulator.EulerAnglePose)
        - 'image_config':
            The [`cryojax.simulator.BasicImageConfig`](https://michael-0brien.github.io/cryojax/api/simulator/config/#cryojax.simulator.BasicImageConfig)
        - 'transfer_theory':
            The [`cryojax.simulator.ContrastTransferTheory`](https://michael-0brien.github.io/cryojax/api/simulator/transfer_theory/#cryojax.simulator.ContrastTransferTheory)
        - 'metadata':
            If `loads_metadata = True`, a `pandas.DataFrame` of entries
            *not* used when loading the `pose`, `image_config`, and
            `transfer_theory`. Otherwise, `None`.
        """  # noqa: E501
        # Validate index
        n_rows = self.starfile_data["particles"].shape[0]
        _validate_dataset_index(type(self), index, n_rows)
        # ... read particle data at the requested indices
        particle_data = self.starfile_data["particles"]
        particle_data_at_index = particle_data.iloc[index]
        # ... read optics group given the particle data
        optics_data = self.starfile_data["optics"]
        optics_group = _get_optics_group_from_particle_data(
            particle_data_at_index, optics_data
        )
        # Load the image stack and STAR file parameters
        image_config, transfer_theory, pose = _make_pytrees_from_starfile(
            particle_data_at_index,
            optics_group,
            self.broadcasts_image_config,
            self.loads_envelope,
            self._options["pad_options"],
            self.rotation_convention,
        )
        if self.loads_metadata:
            # ... convert to dataframe for serialization
            if isinstance(particle_data_at_index, pd.Series):
                particle_data_at_index = particle_data_at_index.to_frame().T
            # ... no overlapping keys with loaded pytrees
            redundant_entry_labels, _ = list(zip(*RELION_SUPPORTED_PARTICLE_ENTRIES))
            columns = particle_data_at_index.columns
            remove_columns = [
                column for column in columns if column in redundant_entry_labels
            ]
            metadata = particle_data_at_index.drop(remove_columns, axis="columns")
        else:
            metadata = None

        return _ParticleParameterInfo(
            image_config=image_config,
            pose=pose,
            transfer_theory=transfer_theory,
            metadata=metadata,
        )

    @override
    def __len__(self) -> int:
        """The number of particles in the STAR file."""
        return len(self.starfile_data["particles"])

    @override
    def __setitem__(
        self,
        index: int | slice | Int[np.ndarray, ""] | Int[np.ndarray, " _"],
        value: _ParticleParameterLike,
    ):
        """Set STAR file entries with `parameter_file[...] = value` syntax,
        where `value` is a dictionary with keys

        - 'pose': The [`cryojax.simulator.EulerAnglePose`]()
        - 'image_config': The [`cryojax.simulator.BasicImageConfig`]()
        - 'transfer_theory': The [`cryojax.simulator.ContrastTransferTheory`]()
        - 'metadata':
            A `pandas.DataFrame` used to write custom entries to the STAR
            file. `len(metadata)` should be equal to the number of particles.

        All keys are optional.
        """
        # Make sure index is valid
        n_rows = self.starfile_data["particles"].shape[0]
        _validate_dataset_index(type(self), index, n_rows)
        # ... also, the parameters too
        _validate_parameters(value, force_keys=False)
        # Invert the pose if desired
        if "pose" in value:
            if self.rotation_convention == "frame":
                value["pose"] = _invert_rotation(value["pose"])
        # Grab the current and new optics and particle data
        if self.updates_optics_group:
            optics_group_index = _make_optics_group_index(self.starfile_data["optics"])
            particle_data_for_update = _parameters_to_particle_data(
                value, optics_group_index
            )
            optics_data_to_append = _parameters_to_optics_data(value, optics_group_index)
            optics_data = pd.concat(
                [self.starfile_data["optics"], optics_data_to_append], ignore_index=True
            )
        else:
            particle_data_for_update = _parameters_to_particle_data(value)
            optics_data = self.starfile_data["optics"]
        particle_data = self.starfile_data["particles"]
        # Set new empty columns in the particle data, if the update data includes this
        new_columns = list(
            set(particle_data_for_update.columns) - set(particle_data.columns)
        )
        for column in new_columns:
            dtype = pd.api.types.pandas_dtype(particle_data_for_update[column].dtype)
            particle_data[column] = pd.Series(dtype=dtype)
        # Finally, set the updated data
        if isinstance(index, (int, np.ndarray)):
            index = np.atleast_1d(index)
        particle_data.loc[
            particle_data.index[index], particle_data_for_update.columns
        ] = particle_data_for_update.values
        self._starfile_data = _StarfileData(optics=optics_data, particles=particle_data)

    @override
    def append(self, value: _ParticleParameterLike):
        """Add an entry or entries to the STAR file with
        `parameter_file.append(value)` syntax, where `value` is
        a dictionary with keys

        - 'pose': The [`cryojax.simulator.EulerAnglePose`](). This key
           is required.
        - 'image_config': The [`cryojax.simulator.BasicImageConfig`](). This key
           is required.
        - 'transfer_theory': The [`cryojax.simulator.ContrastTransferTheory`]().
           This key is required.
        - 'metadata':
            A `pandas.DataFrame` used to write custom entries to the STAR
            file. `len(metadata)` should be equal to the number of particles.
            This key is optional.
        """
        # Make sure parameters are valid
        _validate_parameters(value, force_keys=True)
        # Invert the pose if desired
        if "pose" in value:
            if self.rotation_convention == "frame":
                value["pose"] = _invert_rotation(value["pose"])
        # Make new optics group
        optics_group_index = _make_optics_group_index(self.starfile_data["optics"])
        optics_data, optics_data_to_append = (
            self.starfile_data["optics"],
            _parameters_to_optics_data(value, optics_group_index),
        )
        # Make new particle entries
        particle_data, particle_data_to_append = (
            self.starfile_data["particles"],
            _parameters_to_particle_data(value, optics_group_index),
        )
        # Concatenate and set new entries
        optics_data = (
            pd.concat([optics_data, optics_data_to_append], ignore_index=True)
            if len(optics_data) > 0
            else optics_data_to_append
        )
        particle_data = (
            pd.concat([particle_data, particle_data_to_append], ignore_index=True)
            if len(particle_data) > 0
            else particle_data_to_append
        )
        self._starfile_data = _StarfileData(optics=optics_data, particles=particle_data)

    @override
    def save(
        self,
        *,
        overwrite: bool = False,
        **kwargs: Any,
    ):
        """Save the STAR file at the current `parameter_file.path_to_starfile`."""
        path_to_starfile = self.path_to_starfile
        path_exists = path_to_starfile.exists()
        if path_exists and not overwrite:
            raise FileExistsError(
                f"Tried saving STAR file, but file {str(path_to_starfile)} "
                "already exists. To overwrite existing STAR file, set "
                f"`{type(self).__name__}.overwrite = True`."
            )
        else:
            if not path_to_starfile.parent.exists():
                path_to_starfile.parent.mkdir(parents=True)
            write_starfile(self.starfile_data, path_to_starfile, **kwargs)

    @property
    @override
    def path_to_starfile(self) -> pathlib.Path:
        """The STAR file path"""
        return self._path_to_starfile

    @path_to_starfile.setter
    @override
    def path_to_starfile(self, value: str | pathlib.Path):
        self._path_to_starfile = pathlib.Path(value)

    @property
    def mode(self) -> Literal["r", "w"]:
        """Whether or not the `parameter_file` was
        instantiated in reading ('r') or writing ('w') mode.
        """
        return self._mode  # type: ignore

    @property
    @override
    def starfile_data(self) -> _StarfileData:
        """The `pandas.DataFrame` of STAR file entries."""
        return self._starfile_data

    @starfile_data.setter
    @override
    def starfile_data(self, value: dict[str, pd.DataFrame]):
        if "particles" in value and "optics" in value:
            particle_data, optics_data = value["particles"], value["optics"]
            if isinstance(particle_data, pd.DataFrame) and isinstance(
                optics_data, pd.DataFrame
            ):
                self._starfile_data = _StarfileData(
                    optics=optics_data, particles=particle_data
                )
            else:
                raise ValueError(
                    "STAR file data must be a dictionary "
                    "of pandas DataFrames, with keys equal to "
                    "'particles' and 'optics'. Found that the "
                    f"particle data was type `{type(particle_data).__name__}`"
                    "and the optics data was type "
                    f"`{type(optics_data).__name__}`."
                )
        else:
            raise ValueError(
                "STAR file data must be a dictionary "
                "of pandas DataFrames, with keys equal to "
                "'particles' and 'optics'. Tried setting "
                f"with a dictionary with keys {list(value.keys())}."
            )

    @property
    @override
    def loads_metadata(self) -> bool:
        """Whether or not to load the following information:

        ```python
        parameter_info = parameter_file[index]
        assert isinstance(parameter["metadata"], pandas.DataFrame)
        ```

        The `metadata` is a `pandas.DataFrame` that includes the
        information of
        """
        return self._options["loads_metadata"]

    @loads_metadata.setter
    @override
    def loads_metadata(self, value: bool):
        self._options["loads_metadata"] = value

    @property
    @override
    def loads_envelope(self) -> bool:
        """Whether or not to load the following information:

        ```python
        parameter_info = parameter_file[index]
        assert parameter["transfer_theory"].envelope is not None  # True
        ```
        """
        return self._options["loads_envelope"]

    @loads_envelope.setter
    @override
    def loads_envelope(self, value: bool):
        self._options["loads_envelope"] = value

    @property
    @override
    def broadcasts_image_config(self) -> bool:
        """Whether or not parameters in the loaded `image_config`
        have a batch dimension.

        ```python
        # If `True`:
        parameter_info = parameter_file[1:10]
        assert parameter["image_config"].pixel_size.shape == (10,)  # True
        ...
        # If `False`
        parameter_info = parameter_file[1:10]
        assert parameter["image_config"].pixel_size.shape == ()  # True
        ```
        """
        return self._options["broadcasts_image_config"]

    @broadcasts_image_config.setter
    @override
    def broadcasts_image_config(self, value: bool):
        self._options["broadcasts_image_config"] = value

    @property
    @override
    def updates_optics_group(self) -> bool:
        """If `True`, write a new optics group entry with
        every call to `parameter_file[...] = value`.
        """
        return self._options["updates_optics_group"]

    @updates_optics_group.setter
    @override
    def updates_optics_group(self, value: bool):
        self._options["updates_optics_group"] = value

    @property
    def rotation_convention(self) -> Literal["object", "frame"]:
        """Whether or not parameters are loaded with convention
        that the rotation is of the *object* (i.e. the protein)
        or the *frame* (i.e. a voxel map).

        !!! info "Convert between conventions"

            If you load in one convention, it is straightforward
            to convert to the other.

            ```python
            # Load in 'object' convention
            parameter_file.rotation_convention = "object"
            parameter_info = parameter_file[0]
            pose_object = parameter_info["pose"]
            # Convert to 'frame' convention
            pose_frame = pose_object.to_inverse_rotation()
            ```
        """
        return self._options["rotation_convention"]

    @rotation_convention.setter
    def rotation_convention(self, value: Literal["object", "frame"]):
        self._options["rotation_convention"] = _validate_rotation_convention(value)


class RelionParticleDataset(
    AbstractParticleDataset[_ParticleStackInfo, _ParticleStackLike]
):
    """A dataset that wraps a RELION particle stack in
    [STAR](https://relion.readthedocs.io/en/latest/Reference/Conventions.html) format.
    """

    def __init__(
        self,
        parameter_file: AbstractRelionParticleParameterFile,
        path_to_relion_project: str | pathlib.Path,
        mode: Literal["r", "w"] = "r",
        *,
        mrcfile_settings: dict[str, Any] = {},
        loads_parameters: bool = True,
    ):
        """**Arguments:**

        - `path_to_relion_project`:
            In RELION STAR files, only a relative path is added to the
            'rlnImageName' column. This is relative to the path to the
            "project", which is given by this parameter.
        - `parameter_file`:
            The `RelionParticleParameterFile`.
        - `mode`:
            - If `mode = 'w'`, the dataset is prepared to write new
            *images*. This is done by removing 'rlnImageName' from
            `parameter_file.starfile_data`, if it exists at all.
            does not have a column 'rlnImageName' and image files
            are not yet written.
            - If `mode = 'r'`, images are read from the 'rlnImageName'
            stored in the `parameter_file.starfile_data`.
        - `mrcfile_settings`:
            A dictionary with the following keys:
            - 'prefix':
                A `str` which acts as the prefix to the filenames. If this
                is equal to `"f"`, then the filename for image stack 0 will
                be called "f-00000.mrcs", for `delimiter = '-'` and
                `n_characters = 5`.
                are of format "filenam"
            - 'output_folder':
                A `str` or `pathlib.Path` type where to write MRC files,
                relative to the `path_to_relion_project`.
            - 'n_characters':
                An `int` for the number of characters to write the filename
                number string. If this is equal to `5`, then the filename
                for image stack 0 will be called "f-00000.mrcs", for
                `delimiter = '-'` and `prefix = 'f'`.
             - 'delimiter':
                A `str` for the delimiter between the filename prefix
                and number string. If this is equal to `'-'`, then the
                filename for image stack 0 will be called "f-00000.mrcs",
                for `n_characters = 5` and `prefix = 'f'`.
            - 'overwrite':
                If `True`, overwrite existing MRC file path if it exists.
        - `loads_parameters`:
            If `True`, load parameters and images. Otherwise, load only images.
        """
        # Set properties. First, core properties of the dataset, starting
        # with the `mode``
        self._mode = _validate_mode(mode)
        # ... then, the `parameter_file`. If `mode = 'w'` but
        # the images already exist, we should make a copy in case
        # those images are being used elsewhere
        particle_data, optics_data = (
            parameter_file.starfile_data["particles"],
            parameter_file.starfile_data["optics"],
        )
        self._parameter_file = parameter_file
        # ... properties common to reading and writing images
        self._path_to_relion_project = pathlib.Path(path_to_relion_project)
        # ... properties for reading images
        self._loads_parameters = loads_parameters
        # ... properties for writing images
        self._mrcfile_settings = _dict_to_mrcfile_settings(mrcfile_settings)
        # Now, initialize for `mode = 'r'` vs `mode = 'w'`
        images_exist = "rlnImageName" in particle_data.columns
        project_exists = self.path_to_relion_project.exists()
        if mode == "w":
            # Write empty "rlnImageName" column (defaults to NaN values)
            particle_data["rlnImageName"] = pd.Series(dtype=str)
            self.parameter_file.starfile_data = dict(
                optics=optics_data, particles=particle_data
            )
            # Make the project directory, if it does not yet exist
            if not project_exists:
                self.path_to_relion_project.mkdir(parents=True, exist_ok=False)
        else:
            if not images_exist:
                raise OSError(
                    "Could not find column 'rlnImageName' in the STAR file. "
                    "When using `mode = 'r'`, the STAR file must have this "
                    "column. To write images in a STAR file, "
                    "set `mode = 'w'`."
                )
            if not project_exists:
                raise FileNotFoundError(
                    "`RelionParticleDataset` opened in "
                    "'mode = `r`', but the RELION project directory "
                    "`path_to_relion_project` does not exist. "
                    "To write images in a STAR file in a new RELION project, "
                    "set `mode = 'w'`."
                )

    @override
    def __getitem__(
        self, index: int | slice | Int[np.ndarray, ""] | Int[np.ndarray, " N"]
    ) -> _ParticleStackInfo:
        if self.loads_parameters:
            # Load images and parameters. First, read parameters
            # and metadata from the STAR file
            loads_metadata = self.parameter_file.loads_metadata
            self.parameter_file.loads_metadata = True
            # ... read parameters
            parameters = self.parameter_file[index]
            # ... validate the metadata
            particle_data_at_index = cast(pd.DataFrame, parameters["metadata"])
            _validate_rln_image_name_exists(particle_data_at_index, index)
            # ... reset boolean to original value
            self.parameter_file.loads_metadata = loads_metadata
            if not loads_metadata:
                parameters["metadata"] = None
            # ... grab shape
            shape = parameters["image_config"].shape
            # ... load stack of images
            images = _load_image_stack_from_mrc(
                shape, particle_data_at_index, self.path_to_relion_project
            )
            # ... make sure images and parameters have same leading dim
            if parameters["pose"].offset_x_in_angstroms.ndim == 0:
                images = np.squeeze(images)

            return _ParticleStackInfo(parameters=parameters, images=images)
        else:
            # Otherwise, do not read parameters to more efficiently read images. First,
            # validate the dataset index.
            n_rows = self.parameter_file.starfile_data["particles"].shape[0]
            _validate_dataset_index(type(self), index, n_rows)
            # ... read particle data at the requested indices
            particle_data = self.parameter_file.starfile_data["particles"]
            particle_data_at_index = particle_data.iloc[index]
            if isinstance(particle_data_at_index, pd.Series):
                particle_data_at_index = particle_data_at_index.to_frame().T
            _validate_rln_image_name_exists(particle_data_at_index, index)
            # ... grab shape by reading the optics group
            optics_data = self.parameter_file.starfile_data["optics"]
            optics_group = _get_optics_group_from_particle_data(
                particle_data_at_index, optics_data
            )
            image_size = int(optics_group["rlnImageSize"])
            shape = (image_size, image_size)
            # ... load stack of images
            images = _load_image_stack_from_mrc(
                shape, particle_data_at_index, self.path_to_relion_project
            )
            # ... make sure image leading dim matches with index query
            if isinstance(index, int) or (
                isinstance(index, np.ndarray) and index.size == 0
            ):
                images = np.squeeze(images)

            return _ParticleStackInfo(parameters=None, images=images)

    @override
    def __len__(self) -> int:
        """The number of particles in the STAR file."""
        return len(self.parameter_file)

    @override
    def __setitem__(
        self, index: int | slice | Int[np.ndarray, ""], value: _ParticleStackLike
    ):
        """Set dataset entries with `parameter_file[...] = value` syntax,
        where `value` is a dictionary with keys

        - 'images':
            An image or image stack to write to an MRC file. This
            key is required.
        - 'parameters':
            See [`cryospax.RelionParticleParameterFile`][] for more
            information. This key is optional.
        """
        if isinstance(index, Int[np.ndarray, "_"]):  # type: ignore
            raise ValueError(
                "When setting `dataset[index] = ...`, "
                "it is not supported to pass `index` as a 1D numpy-array."
            )
        if not isinstance(value, dict):
            raise TypeError(
                "When setting `dataset[index] = foo`, "
                "foo must be a dictionary with key "
                "'images' and optionally key 'parameters'."
            )
        images, parameters = _unpack_particle_stack_dict(value)
        if parameters is not None:
            self.parameter_file[index] = parameters
        n_particles = len(self.parameter_file)
        index_array = np.atleast_1d(_index_to_array(index, n_particles))
        self.write_images(index_array, images, parameters=parameters)

    @override
    def append(self, value: _ParticleStackLike):
        """Add an entry or entries to the dataset with
        `dataset.append(value)` syntax, where `value`
        is a dictionary with keys

        - 'images':
            An image or image stack to write to the MRC file.
        - 'parameters':
            See [`cryospax.RelionParticleParameterFile`][] for more
            information.

        Both keys are required.
        """
        if not isinstance(value, dict):
            raise TypeError(
                "When appending `dataset.append(foo)`, "
                "`foo` must be a dictionary with keys "
                "'images' and 'parameters'."
            )
        images, parameters = _unpack_particle_stack_dict(value)
        if parameters is None:
            raise ValueError(
                "When appending dictionary `foo` as `dataset.append(foo)`, "
                "`foo` must have key 'parameters'."
            )
        start = len(self.parameter_file.starfile_data["particles"])
        # Append parameters. This automatically sets the 'rlnImageName'
        # column to NaNs
        self.parameter_file.append(parameters)
        # Write images
        stop = len(self.parameter_file.starfile_data["particles"])
        index_array = np.arange(start, stop, dtype=int)
        self.write_images(index_array, images, parameters=parameters)

    @override
    def write_images(
        self,
        index_array: Int[np.ndarray, " _"],
        images: Float[NDArrayLike, "... _ _"],
        parameters: _ParticleParameterLike | None = None,
    ):
        """Write images to the dataset given `images` and `parameters`.

        This function is wrapped by [`cryospax.RelionParticleDataset.append`][]
        and [`cryospax.RelionParticleDataset.__setitem__`][] and in most cases
        these APIs are preferred.
        """
        # Get relevant metadata
        particle_data = self.parameter_file.starfile_data["particles"]
        optics_data = self.parameter_file.starfile_data["optics"]
        if parameters is None:
            optics_group = _get_optics_group_from_particle_data(
                particle_data.iloc[index_array], optics_data
            )
            pixel_size, dim = (
                float(optics_group["rlnImagePixelSize"]),
                int(optics_group["rlnImageSize"]),
            )
        else:
            pixel_size, dim = (
                float(np.atleast_1d(parameters["image_config"].pixel_size)[0]),
                parameters["image_config"].shape[0],
            )
        if not (images.ndim in [2, 3] and images.shape[-2:] == (dim, dim)):
            raise ValueError(
                "Image(s) must be of "
                "shape `(n_images, dim, dim)` or `(dim, dim)`. "
                f"Tried writing image(s) of "
                f"shape {images.shape}."
            )
        # Prepare to write images
        if images.ndim == 2:
            images = images[None, ...]
        n_images, _ = images.shape[0], images.shape[1]
        n_particles = len(self.parameter_file)

        # Convert index into 1D ascending numpy array
        n_indices = index_array.size
        if n_images != n_indices:
            raise ValueError(
                "Tried to set dataset elements with an inconsistent number "
                f"of images. Found that the number of images was {n_images}, "
                f"while the number of dataset indices was {n_indices}."
            )
        # Get absolute path to the filename, as well as the 'rlnImageName'
        # column
        path_to_filename, rln_image_names = _make_image_filename(
            index_array,
            particle_data,
            n_particles,
            self.mrcfile_settings,
            self.path_to_relion_project,
        )
        # Set the STAR file column
        particle_data.loc[particle_data.index[index_array], "rlnImageName"] = (
            rln_image_names
        )
        self.parameter_file.starfile_data = dict(
            particles=particle_data, optics=optics_data
        )
        # ... and write the images to disk
        try:
            write_image_stack_to_mrc(
                images.astype(jnp.float32),
                pixel_size,
                path_to_filename,
                overwrite=self.mrcfile_settings["overwrite"],
                compression=self.mrcfile_settings["compression"],
            )
        except Exception as err:
            raise OSError(
                "Error occurred when writing image stack to MRC "
                "file. Most likely, the filename the writer "
                f"chose ({str(path_to_filename)}) already "
                "exists. Try changing the "
                "`RelionParticleDataset.mrcfile_settings`. "
                f"The error message was:\n{err}"
            )

    @property
    @override
    def parameter_file(self) -> AbstractRelionParticleParameterFile:
        return self._parameter_file

    @property
    @override
    def mode(self) -> Literal["r", "w"]:
        return self._mode  # type: ignore

    @property
    def path_to_relion_project(self) -> pathlib.Path:
        return self._path_to_relion_project

    @property
    def mrcfile_settings(self) -> _MrcfileSettings:
        """Settings for writing MRC files with. See
        [`cryojax.RelionParticleDataset.__init__`][]
        for more information.
        """
        return self._mrcfile_settings

    @mrcfile_settings.setter
    def mrcfile_settings(self, value: dict[str, Any]):
        self._mrcfile_settings = _dict_to_mrcfile_settings(value)

    @property
    def loads_parameters(self) -> bool:
        """If `False`, load images and *not* parameters. This gives
        better performance when it is not necessary to load parameters.

        ```python
        dataset.loads_parameters = False
        particle_info = dataset[0]
        assert particle_info["parameters"] is None  # True
        ```
        """
        return self._loads_parameters

    @loads_parameters.setter
    def loads_parameters(self, value: bool):
        self._loads_parameters = value


def _load_starfile_data(
    path_to_starfile: pathlib.Path,
    selection_filter: dict[str, Callable] | None,
    mode: Literal["r", "w"],
    exists_ok: bool,
) -> _StarfileData:
    if mode == "r":
        if path_to_starfile.exists():
            starfile_data = read_starfile(path_to_starfile)
            _validate_starfile_data(starfile_data)
            if selection_filter is not None:
                starfile_data = _select_particles(starfile_data, selection_filter)
        else:
            raise FileNotFoundError(
                f"Set `mode = '{mode}'`, but STAR file {str(path_to_starfile)} does not "
                "exist. To write a new STAR file, set `mode = 'w'`."
            )
    else:
        if path_to_starfile.exists() and not exists_ok:
            raise FileExistsError(
                f"Set `mode = 'w'`, but STAR file {str(path_to_starfile)} already "
                "exists. To read an existing STAR file, set `mode = 'r'` or "
                "to erase an existing STAR file, set `mode = 'w'` and "
                "`exists_ok=True`."
            )
        else:
            if selection_filter is None:
                starfile_data = dict(
                    optics=pd.DataFrame(
                        data={
                            column: pd.Series(dtype=dtype)
                            for column, dtype in RELION_DEFAULT_OPTICS_ENTRIES
                        }
                    ),
                    particles=pd.DataFrame(
                        data={
                            column: pd.Series(dtype=dtype)
                            for column, dtype in RELION_DEFAULT_PARTICLE_ENTRIES
                        }
                    ),
                )
            else:
                raise ValueError(
                    "Initialized a `RelionParticleParameterFile` in `mode = 'w'` "
                    "but also passed a `selection_filter`. Selection is only used "
                    "in `mode = 'r'`."
                )

    return _StarfileData(
        optics=starfile_data["optics"], particles=starfile_data["particles"]
    )


def _validate_mode(mode: str) -> Literal["r", "w"]:
    if mode not in ["r", "w"]:
        raise ValueError(
            f"Passed unsupported `mode = {mode}`. Supported modes are 'r' and 'w'."
        )
    return mode  # type: ignore


def _validate_rotation_convention(mode: str) -> Literal["frame", "object"]:
    if mode not in ["frame", "object"]:
        raise ValueError(
            f"Passed unsupported `rotation_convention = {mode}`. "
            "Supported modes are 'object' and 'frame'."
        )
    return mode  # type: ignore


def _select_particles(
    starfile_data: dict[str, pd.DataFrame], selection_filter: dict[str, Callable]
) -> dict[str, pd.DataFrame]:
    particle_data = starfile_data["particles"]
    boolean_mask = pd.Series(True, index=particle_data.index)
    for key in selection_filter:
        if key in particle_data.columns:
            fn = selection_filter[key]
            column = particle_data[key]
            base_error_message = (
                f"Error filtering key '{key}' in the `selection_filter`. "
                f"To filter the STAR file entries, `selection_filter['{key}']`"
                "must be a function that takes in an array and returns a "
                "boolean mask."
            )
            if isinstance(selection_filter[key], Callable):
                try:
                    mask_at_column = fn(column)
                except Exception as err:
                    raise ValueError(
                        f"{base_error_message} "
                        "When calling the function, caught an error:\n"
                        f"{err}"
                    )
                if not pd.api.types.is_bool_dtype(mask_at_column):
                    raise ValueError(
                        f"{base_error_message} "
                        "Found that the function did not return "
                        "a boolean dtype."
                    )
            else:
                raise ValueError(base_error_message)
            # Update mask
            boolean_mask = mask_at_column & boolean_mask
        else:
            raise ValueError(
                f"Included key '{key}' in the `selection_filter`, "
                "but this entry could not be found in the STAR file. "
                "The `selection_filter` must be a dictionary whose "
                "keys are strings in the STAR file and whose values "
                "are functions that take in columns and return boolean "
                "masks."
            )
    # Select particles using mask
    starfile_data["particles"] = particle_data[boolean_mask]

    return starfile_data


#
# STAR file reading
#
def _make_pytrees_from_starfile(
    particle_data,
    optics_data,
    broadcasts_image_config,
    loads_envelope,
    pad_options,
    rotation_convention,
) -> tuple[BasicImageConfig, ContrastTransferTheory, EulerAnglePose]:
    float_dtype = jax.dtypes.canonicalize_dtype(float)
    # Load CTF parameters. First from particle data
    defocus_in_angstroms = (
        np.asarray(particle_data["rlnDefocusU"], dtype=float_dtype)
        + np.asarray(particle_data["rlnDefocusV"], dtype=float_dtype)
    ) / 2
    astigmatism_in_angstroms = np.asarray(
        particle_data["rlnDefocusU"], dtype=float_dtype
    ) - np.asarray(particle_data["rlnDefocusV"], dtype=float_dtype)
    astigmatism_angle = np.asarray(particle_data["rlnDefocusAngle"], dtype=float_dtype)
    phase_shift = np.asarray(particle_data["rlnPhaseShift"], dtype=float_dtype)
    # Then from optics data
    batch_shape = (
        () if defocus_in_angstroms.ndim == 0 else (defocus_in_angstroms.shape[0],)
    )
    spherical_aberration_in_mm = np.full(
        batch_shape, np.asarray(optics_data["rlnSphericalAberration"], dtype=float_dtype)
    )
    amplitude_contrast_ratio = np.full(
        batch_shape, np.asarray(optics_data["rlnAmplitudeContrast"], dtype=float_dtype)
    )
    ctf_params = (
        defocus_in_angstroms,
        astigmatism_in_angstroms,
        astigmatism_angle,
        spherical_aberration_in_mm,
        amplitude_contrast_ratio,
        phase_shift,
    )
    # Envelope parameters
    if loads_envelope:
        b_factor, scale_factor = (
            (
                np.asarray(particle_data["rlnCtfBfactor"], dtype=float_dtype)
                if "rlnCtfBfactor" in particle_data.keys()
                else None
            ),
            (
                np.asarray(particle_data["rlnCtfScalefactor"], dtype=float_dtype)
                if "rlnCtfScalefactor" in particle_data.keys()
                else None
            ),
        )
    else:
        b_factor, scale_factor = None, None
    # Image config parameters
    pixel_size = np.asarray(optics_data["rlnImagePixelSize"], dtype=float_dtype)
    voltage_in_kilovolts = np.asarray(optics_data["rlnVoltage"], dtype=float_dtype)
    if broadcasts_image_config and len(batch_shape) > 0:
        pixel_size = np.full(batch_shape, pixel_size)
        voltage_in_kilovolts = np.full(batch_shape, voltage_in_kilovolts)
    # Pose parameters. Values for the pose are optional,
    # so look to see if each key is present
    particle_keys = particle_data.keys()
    # Read the pose. first, xy offsets
    rln_origin_x_angst = (
        particle_data["rlnOriginXAngst"] if "rlnOriginXAngst" in particle_keys else 0.0
    )
    rln_origin_y_angst = (
        particle_data["rlnOriginYAngst"] if "rlnOriginYAngst" in particle_keys else 0.0
    )
    # ... rot angle
    rln_angle_rot = (
        particle_data["rlnAngleRot"] if "rlnAngleRot" in particle_keys else 0.0
    )
    # ... tilt angle
    if "rlnAngleTilt" in particle_keys:
        rln_angle_tilt = particle_data["rlnAngleTilt"]
    elif "rlnAngleTiltPrior" in particle_keys:  # support for helices
        rln_angle_tilt = particle_data["rlnAngleTiltPrior"]
    else:
        rln_angle_tilt = 0.0
    # ... psi angle
    if "rlnAnglePsi" in particle_keys:
        # Relion uses -999.0 as a placeholder for an un-estimated in-plane
        # rotation
        if isinstance(particle_data["rlnAnglePsi"], pd.Series):
            # ... check if all values are equal to -999.0. If so, just
            # replace the whole pandas.Series with 0.0
            if (
                particle_data["rlnAnglePsi"].nunique() == 1
                and particle_data["rlnAnglePsi"].iloc[0] == -999.0
            ):
                rln_angle_psi = 0.0
            # ... otherwise, replace -999.0 values with 0.0
            else:
                rln_angle_psi = particle_data["rlnAnglePsi"].where(
                    lambda x: x != -999.0, 0.0
                )
        else:
            # ... if the column is just equal to a float, then
            # directly check if it is equal to -999.0
            rln_angle_psi = (
                0.0
                if particle_data["rlnAnglePsi"] == -999.0
                else particle_data["rlnAnglePsi"]
            )
    elif "rlnAnglePsiPrior" in particle_keys:  # support for helices
        rln_angle_psi = particle_data["rlnAnglePsiPrior"]
    else:
        rln_angle_psi = 0.0
    # Now, flip the sign of the translations and transpose rotations.
    # RELION's convention thinks about the translation as "undoing" the translation
    # and rotation in the image
    maybe_make_full = lambda param: (
        np.full(batch_shape, param)
        if len(batch_shape) > 0 and param.shape == ()
        else param
    )
    pose_params = tuple(
        maybe_make_full(x)
        for x in (
            -np.asarray(rln_origin_x_angst, dtype=float_dtype),
            -np.asarray(rln_origin_y_angst, dtype=float_dtype),
            -np.asarray(rln_angle_rot, dtype=float_dtype),
            -np.asarray(rln_angle_tilt, dtype=float_dtype),
            -np.asarray(rln_angle_psi, dtype=float_dtype),
        )
    )
    # Now, create cryojax objects. Do this on the CPU
    cpu_device = jax.devices(backend="cpu")[0]
    with jax.default_device(cpu_device):
        # First, create the `BasicImageConfig`
        image_size = int(optics_data["rlnImageSize"])
        image_shape = (image_size, image_size)
        image_config = _make_config(
            image_shape, pixel_size, voltage_in_kilovolts, pad_options
        )
        # ... now the `ContrastTransferTheory`
        envelope = (
            _make_envelope_function(scale_factor, b_factor) if loads_envelope else None
        )
        transfer_theory_params = (*ctf_params, envelope)
        transfer_theory = _make_transfer_theory(*transfer_theory_params)  # type: ignore
        # ... finally the `EulerAnglePose`
        pose = _make_pose(*pose_params)
        if rotation_convention == "frame":
            pose = _invert_rotation(pose)
    # Now, convert arrays to numpy in case the user wishes to do preprocessing
    pytree_dynamic, pytree_static = eqx.partition(
        (image_config, transfer_theory, pose), eqx.is_array
    )
    pytree_dynamic = jax.tree.map(lambda x: np.asarray(x), pytree_dynamic)
    image_config, transfer_theory, pose = eqx.combine(pytree_dynamic, pytree_static)

    return image_config, transfer_theory, pose


def _make_config(
    image_shape,
    pixel_size,
    voltage_in_kilovolts,
    pad_options,
):
    return eqx.tree_at(
        lambda x: (x.pixel_size, x.voltage_in_kilovolts),
        BasicImageConfig(image_shape, 1.0, 1.0, pad_options=pad_options),
        (pixel_size, voltage_in_kilovolts),
    )


def _make_pose(offset_x, offset_y, phi, theta, psi):
    _make_fn = lambda _x, _y, _phi, _theta, _psi: EulerAnglePose(
        _x, _y, _phi, _theta, _psi
    )
    if offset_x.ndim == 1:
        _make_fn = eqx.filter_vmap(_make_fn)
    return _make_fn(offset_x, offset_y, phi, theta, psi)


def _make_envelope_function(amp, b_factor):
    if b_factor is None and amp is None:
        warnings.warn(
            "`loads_envelope` was set to True, but no envelope parameters were found. "
            "Setting envelope as None. "
            "Make sure your starfile is correctly formatted or set "
            "`loads_envelope=False`."
        )
        return None

    elif b_factor is None and amp is not None:
        return eqx.tree_at(lambda x: x.value, FourierConstant(1.0), amp)
    else:
        if amp is None:
            amp = np.asarray(1.0) if b_factor.ndim == 0 else np.ones_like(b_factor)
        return eqx.tree_at(
            lambda x: (x.amplitude, x.b_factor),
            FourierGaussian(1.0, 1.0),
            (amp, b_factor),
        )


def _make_transfer_theory(defocus, astig, angle, sph, ac, ps, env=None):
    ctf = eqx.tree_at(
        lambda x: (
            x.defocus_in_angstroms,
            x.astigmatism_in_angstroms,
            x.astigmatism_angle,
            x.spherical_aberration_in_mm,
        ),
        AstigmaticCTF(),
        (defocus, astig, angle, sph),
    )
    transfer_theory = ContrastTransferTheory(
        ctf, envelope=env, amplitude_contrast_ratio=0.1, phase_shift=0.0
    )

    return eqx.tree_at(
        lambda x: (x.amplitude_contrast_ratio, x.phase_shift), transfer_theory, (ac, ps)
    )


def _invert_rotation(pose: EulerAnglePose) -> EulerAnglePose:
    negate_angle = lambda angle: ((-angle + 180) % 360) - 180
    return eqx.tree_at(
        lambda x: (x.phi_angle, x.theta_angle, x.psi_angle),
        pose,
        (
            negate_angle(pose.psi_angle),
            negate_angle(pose.theta_angle),
            negate_angle(pose.phi_angle),
        ),
    )


def _load_image_stack_from_mrc(
    shape: tuple[int, int],
    particle_dataframe_at_index: pd.DataFrame,
    path_to_relion_project: str | pathlib.Path,
) -> Float[np.ndarray, "... y_dim x_dim"]:
    # Load particle image stack rlnImageName
    rln_image_names = particle_dataframe_at_index["rlnImageName"].reset_index(drop=True)
    if rln_image_names.convert_dtypes().dtype != "string":
        raise TypeError(
            "The 'rlnImageName' column was expected to be type string. "
            f"Instead, found type {rln_image_names.dtype}."
        )
    # Split the pandas.Series into two: one for the image index
    # and another for the filename
    split = rln_image_names.str.split("@").str
    # ... relion convention starts indexing at 1, not 0
    mrc_index, filenames = split[0], split[1]
    if filenames.isnull().any() or mrc_index.isnull().any():
        raise ValueError(
            "Could not parse filenames from the 'rlnImageName' "
            "column. Check to make sure these are the correct format, "
            "for example '00000@path/to/image-00000.mrcs'."
        )
    mrc_index = mrc_index.astype(int) - 1
    # Allocate memory for stack
    n_images = len(filenames)
    image_stack = np.empty((n_images, *shape), dtype=float)
    # Loop over filenames to fill stack
    unique_filenames = filenames.unique()
    for filename in unique_filenames:
        # Get the MRC indices
        mrc_index_at_filename = mrc_index[filename == filenames]
        particle_index_at_filename = mrc_index_at_filename.index
        path_to_filename = pathlib.Path(path_to_relion_project, filename)
        with mrcfile.mmap(path_to_filename, mode="r", permissive=True) as mrc:
            mrc_data = np.asarray(mrc.data)
            mrc_ndim = mrc_data.ndim
            mrc_shape = mrc_data.shape if mrc_ndim == 2 else mrc_data.shape[1:]

            if shape != mrc_shape:
                raise ValueError(
                    f"The shape of the MRC with filename {filename} "
                    "was found to not have the same shape loaded from "
                    "the 'rlnImageSize'. Check your MRC files and also "
                    "the STAR file optics group formatting."
                )
            image_stack[particle_index_at_filename] = (
                mrc_data if mrc_ndim == 2 else mrc_data[mrc_index_at_filename]
            )

    return image_stack


def _validate_dataset_index(cls, index, n_rows):
    index_error_msg = lambda idx: (
        f"The index at which the `{cls.__name__}` was accessed was out of bounds! "
        f"The number of rows in the dataset is {n_rows}, but you tried to "
        f"access the index {idx}."
    )
    # ... pandas has bad error messages for its indexing
    if isinstance(index, (int, np.integer)):  # type: ignore
        if index > n_rows - 1:
            raise IndexError(index_error_msg(index))
    elif isinstance(index, slice):
        if index.start is not None and index.start > n_rows - 1:
            raise IndexError(index_error_msg(index.start))
    elif isinstance(index, np.ndarray):
        if index.size == 0:
            raise IndexError(
                "Found that the index passed to the dataset "
                "was an empty numpy array. Please pass a "
                "supported index."
            )
    else:
        raise IndexError(
            f"Indexing with the type {type(index)} is not supported by "
            f"`{cls.__name__}`. Indexing by integers is supported, one-dimensional "
            "fancy indexing is supported, and numpy-array indexing is supported. "
            "For example, like `particle = particle_dataset[0]`, "
            "`particle_stack = particle_dataset[0:5]`, "
            "or `particle_stack = dataset[np.array([1, 4, 3, 2])]`."
        )


def _validate_starfile_data(starfile_data: dict[str, pd.DataFrame]):
    if "particles" not in starfile_data.keys():
        raise ValueError("Missing key 'particles' in `starfile.read` output.")
    else:
        required_particle_keys, _ = zip(*RELION_REQUIRED_PARTICLE_ENTRIES)
        if not set(required_particle_keys).issubset(
            set(starfile_data["particles"].keys())
        ):
            raise ValueError(
                "Missing required keys in starfile 'particles' group. "
                f"Required keys are {required_particle_keys}."
            )
    if "optics" not in starfile_data.keys():
        raise ValueError("Missing key 'optics' in `starfile.read` output.")
    else:
        required_optics_keys, _ = zip(*RELION_REQUIRED_OPTICS_ENTRIES)
        if not set(required_optics_keys).issubset(set(starfile_data["optics"].keys())):
            raise ValueError(
                "Missing required keys in starfile 'optics' group. "
                f"Required keys are {required_optics_keys}."
            )


def _validate_rln_image_name_exists(particle_data, index):
    if "rlnImageName" not in particle_data.columns:
        raise OSError(
            "Tried to read STAR file for "
            f"`RelionParticleDataset` index = {index}, "
            "but no entry found for 'rlnImageName'."
        )


#
# Working with pytrees for I/O
#
def _unpack_particle_stack_dict(
    value: _ParticleStackLike,
) -> tuple[Float[NDArrayLike, "... y_dim x_dim"], _ParticleParameterLike | None]:
    if "images" in value:
        images = value["images"]
    else:
        raise ValueError(
            "When passing dictionary `foo` as `dataset.append(foo)` or "
            "`dataset[index] = foo`, `foo` must have key `images`."
        )
    if "parameters" in value:
        parameters = value["parameters"]
    else:
        parameters = None

    return images, parameters


#
# STAR file writing. First, functions for writing parameters
#
def _validate_parameters(parameters: _ParticleParameterLike, force_keys: bool = False):
    if force_keys:
        if not {"image_config", "transfer_theory", "pose"}.issubset(parameters):
            raise ValueError(
                "When passing dictionary `foo` as `parameter_file.append(foo)` "
                "`foo` must have keys 'pose', 'transfer_theory', and 'image_config'."
            )
    if "image_config" in parameters:
        if not isinstance(parameters["image_config"], BasicImageConfig):
            raise TypeError(
                "Found that dict key 'image_config' was "
                "not type `cryojax.simulator.BasicImageConfig`. "
                f"Instead, it was type "
                f"{type(parameters['image_config']).__name__}."
            )
    if "transfer_theory" in parameters:
        if not isinstance(parameters["transfer_theory"], ContrastTransferTheory):
            raise TypeError(
                "Found that dict key 'transfer_theory' was "
                "not type `cryojax.simulator.ContrastTransferTheory`. "
                f"Instead, it was type "
                f"{type(parameters['transfer_theory']).__name__}."
            )
    if "pose" in parameters:
        if not isinstance(parameters["pose"], EulerAnglePose):
            raise TypeError(
                "Found that dict key 'pose' was "
                "not type `cryojax.simulator.EulerAnglePose`. "
                f"Instead, it was type "
                f"{type(parameters['pose']).__name__}."
            )


def _parameters_to_optics_data(
    parameters: _ParticleParameterLike, optics_group_index: int
) -> pd.DataFrame:
    if {"image_config", "transfer_theory"}.issubset(parameters):
        shape = parameters["image_config"].shape
        if shape[0] == shape[1]:
            dim = shape[0]
        else:
            raise ValueError(
                "When adding optics group to STAR file, found "
                "non-square shape in `image_config.shape`. Only "
                "square shapes are supported."
            )
        pixel_size = parameters["image_config"].pixel_size
        voltage_in_kilovolts = parameters["image_config"].voltage_in_kilovolts
        amplitude_contrast_ratio = parameters["transfer_theory"].amplitude_contrast_ratio
        if isinstance(parameters["transfer_theory"].ctf, AstigmaticCTF):
            spherical_aberration_in_mm = getattr(
                parameters["transfer_theory"].ctf, "spherical_aberration_in_mm"
            )
        else:
            raise ValueError(
                "When adding optics group or particle to STAR file, "
                "`transfer_theory.ctf` must be type "
                "`AstigmaticCTF`. Instead, got type "
                f"{type(parameters['transfer_theory'].ctf).__name__}."
            )
        optics_group_dict = {
            "rlnOpticsGroup": optics_group_index,
            "rlnImageSize": dim,
            "rlnImagePixelSize": pixel_size,
            "rlnVoltage": voltage_in_kilovolts,
            "rlnSphericalAberration": spherical_aberration_in_mm,
            "rlnAmplitudeContrast": amplitude_contrast_ratio,
        }
        for k, v in optics_group_dict.items():
            if isinstance(v, Array | np.ndarray):
                arr = np.atleast_1d(np.asarray(v))
                if arr.size > 1:
                    if np.unique(arr).size > 1:
                        raise ValueError(
                            "Tried to fill a RELION optics group entry with an array "
                            "that has multiple unique values. Optics group compatible "
                            "arrays such as `BasicImageConfig.pixel_size` "
                            "must be either scalars or arrays all with the same value. "
                            f"Error occurred when filling '{k}' with array {v}."
                        )
                optics_group_dict[k] = arr.ravel()[0, None]
            else:
                optics_group_dict[k] = [v]
        optics_data = pd.DataFrame.from_dict(optics_group_dict)

        return optics_data
    else:
        raise ValueError(
            "Tried to add optics group to the STAR file, but "
            "parameter dictionary did not include 'image_config' "
            "and 'transfer_theory' keys. If you are setting parameters "
            "`parameter_file[index] = dict(pose=pose)` make sure "
            "`parameter_file.updates_optics_group = False`."
        )


def _parameters_to_particle_data(
    parameters: _ParticleParameterLike,
    optics_group_index: int | None = None,
) -> pd.DataFrame:
    particles_dict = {}
    if "pose" in parameters:
        # Now, pose parameters
        pose = parameters["pose"]
        if not isinstance(pose, EulerAnglePose):
            raise ValueError(
                "When adding particle to STAR file, "
                "`pose` must be type "
                "`EulerAnglePose`. Instead, got type "
                f"{type(pose).__name__}."
            )
        if pose.offset_in_angstroms.ndim == 2:
            particles_dict["rlnOriginXAngst"] = -pose.offset_in_angstroms[:, 0]
            particles_dict["rlnOriginYAngst"] = -pose.offset_in_angstroms[:, 1]
        elif pose.offset_in_angstroms.ndim == 1:
            particles_dict["rlnOriginXAngst"] = -pose.offset_in_angstroms[0]
            particles_dict["rlnOriginYAngst"] = -pose.offset_in_angstroms[1]
        else:
            raise RuntimeError(
                "Internal `cryojax` error when loading translations to STAR file."
            )
        particles_dict["rlnAngleRot"] = -pose.phi_angle
        particles_dict["rlnAngleTilt"] = -pose.theta_angle
        particles_dict["rlnAnglePsi"] = -pose.psi_angle
        # Now, broadcast parameters to same dimension
        n_particles = pose.offset_x_in_angstroms.size
        for k, v in particles_dict.items():
            if v.shape == ():
                particles_dict[k] = np.full((n_particles,), np.asarray(v))
            elif v.size == n_particles:
                particles_dict[k] = v.ravel()
            else:
                raise ValueError(
                    "Found inconsistent number of particles "
                    "when adding particle to STAR file. "
                    "When running `parameter_file[index] = foo` "
                    "or `parameter_file.append(foo)`, make sure `foo` "
                    "has arrays that are scalars or all have the same "
                    "number of dimensions."
                )
    else:
        raise ValueError(
            "Tried to modify parameters in STAR file, but "
            "parameter dictionary did not include 'pose' "
            "key. If you are setting parameters "
            "`parameter_file[index] = foo` make sure "
            "`foo` is, for example, `foo = dict(pose=EulerAnglePose(...))`."
        )
    # Fill CTF parameters
    if "transfer_theory" in parameters:
        transfer_theory = parameters["transfer_theory"]
        if isinstance(transfer_theory.ctf, AstigmaticCTF):
            if pose.offset_z_in_angstroms is None:
                defocus_offset = 0.0
            else:
                defocus_offset = pose.offset_z_in_angstroms
            particles_dict["rlnDefocusU"] = (
                transfer_theory.ctf.defocus_in_angstroms
                + defocus_offset
                + transfer_theory.ctf.astigmatism_in_angstroms / 2
            )
            particles_dict["rlnDefocusV"] = (
                transfer_theory.ctf.defocus_in_angstroms
                + defocus_offset
                - transfer_theory.ctf.astigmatism_in_angstroms / 2
            )
            particles_dict["rlnDefocusAngle"] = transfer_theory.ctf.astigmatism_angle
        else:
            raise ValueError(
                "When adding particle to STAR file, "
                "`transfer_theory.ctf` must be type "
                "`AstigmaticCTF`. Instead, got type "
                f"{type(parameters['transfer_theory'].ctf).__name__}."
            )

        if isinstance(transfer_theory.envelope, FourierGaussian):
            particles_dict["rlnCtfBfactor"] = transfer_theory.envelope.b_factor
            particles_dict["rlnCtfScalefactor"] = transfer_theory.envelope.amplitude
        elif isinstance(transfer_theory.envelope, FourierConstant):
            particles_dict["rlnCtfScalefactor"] = transfer_theory.envelope.value
        elif transfer_theory.envelope is None:
            pass
        else:
            raise ValueError(
                "When adding particle to STAR file, "
                "`transfer_theory.envelope` must be "
                "type `cryojax.ndimage.FourierGaussian` "
                "or `cryojax.ndimage.FourierConstant`, or `None`. Got "
                f"{type(transfer_theory.envelope).__name__}."
            )
        particles_dict["rlnPhaseShift"] = transfer_theory.phase_shift
    # Now, miscellaneous parameters
    if optics_group_index is not None:
        particles_dict["rlnOpticsGroup"] = np.full(
            (n_particles,), optics_group_index, dtype=int
        )
    # Make sure parameters are numpy arrays
    particles_dict = jax.tree.map(
        lambda x: np.asarray(x) if isinstance(x, Array) else x, particles_dict
    )
    particle_data = pd.DataFrame.from_dict(particles_dict)
    # Finally, see if the particle parameters has metadata and if so,
    # add this
    if "metadata" in parameters:
        if parameters["metadata"] is not None:
            metadata = cast(pd.DataFrame, parameters["metadata"])
            if n_particles != metadata.index.size:
                raise ValueError(
                    "When adding custom metadata to STAR file "
                    "with `parameter_file[index] = foo` or `parameter_file.append(foo)`, "
                    "found the number of particles "
                    "in `foo['metadata']` was inconsistent with the "
                    "number of particles in `foo['pose']`."
                )
            # Add metadata to dataframe
            particle_data = pd.concat(
                [particle_data, metadata], axis="columns", verify_integrity=True
            )
    return particle_data


def _make_optics_group_index(optics_data: pd.DataFrame) -> int:
    optics_group_indices = np.asarray(optics_data["rlnOpticsGroup"], dtype=int)
    last_optics_group_index = (
        0 if optics_group_indices.size == 0 else int(optics_group_indices[-1])
    )
    return last_optics_group_index + 1


def _parse_optics_group_index(particle_data_at_index: pd.DataFrame | pd.Series) -> int:
    # ... read optics data
    optics_group_indices = np.unique(
        np.atleast_1d(np.asarray(particle_data_at_index["rlnOpticsGroup"]))
    )
    if optics_group_indices.size > 1:
        raise NotImplementedError(
            "Tried to read multiple particles at once that belong "
            "to different optics groups, but this is not yet "
            "implemented. In the meantime, try reading one particle "
            "at a time."
        )
    optics_group_index = optics_group_indices[0]

    return int(optics_group_index)


def _get_optics_group_from_index(
    optics_data: pd.DataFrame, optics_group_index: int
) -> pd.Series:
    return optics_data[optics_data["rlnOpticsGroup"] == optics_group_index].iloc[0]


def _get_optics_group_from_particle_data(
    particle_data_at_index: pd.DataFrame | pd.Series, optics_data: pd.DataFrame
) -> pd.Series:
    optics_group_index = _parse_optics_group_index(particle_data_at_index)
    return _get_optics_group_from_index(optics_data, optics_group_index)


#
# Now, functions for writing image files
#
def _dict_to_mrcfile_settings(d: dict[str, Any]) -> _MrcfileSettings:
    prefix = d["prefix"] if "prefix" in d else ""
    output_folder = d["output_folder"] if "output_folder" in d else ""
    delimiter = d["delimiter"] if "delimiter" in d else "_"
    n_characters = d["n_characters"] if "n_characters" in d else 6
    overwrite = d["overwrite"] if "overwrite" in d else False
    compression = d["compression"] if "compression" in d else None
    return _MrcfileSettings(
        prefix=prefix,
        output_folder=output_folder,
        delimiter=delimiter,
        n_characters=n_characters,
        overwrite=overwrite,
        compression=compression,
    )


def _dict_to_options(d: dict[str, Any]) -> _Options:
    loads_metadata = d["loads_metadata"] if "loads_metadata" in d else False
    broadcasts_image_config = (
        d["broadcasts_image_config"] if "broadcasts_image_config" in d else True
    )
    loads_envelope = d["loads_envelope"] if "loads_envelope" in d else False
    updates_optics_group = (
        d["updates_optics_group"] if "updates_optics_group" in d else False
    )
    rotation_convention = (
        _validate_rotation_convention(d["rotation_convention"])
        if "rotation_convention" in d
        else "object"
    )
    pad_options = d["pad_options"] if "pad_options" in d else {}
    return _Options(
        loads_metadata=loads_metadata,
        broadcasts_image_config=broadcasts_image_config,
        loads_envelope=loads_envelope,
        updates_optics_group=updates_optics_group,
        rotation_convention=rotation_convention,
        pad_options=pad_options,
    )


def _index_to_array(indices: slice | int | np.ndarray, size: int) -> np.ndarray:
    if isinstance(indices, slice):
        return np.asarray(range(*indices.indices(size)))
    else:
        return np.asarray(indices, dtype=int)


def _make_image_filename(
    index: Int[np.ndarray, " _"],
    particle_data: pd.DataFrame,
    n_particles: int,
    mrcfile_settings: _MrcfileSettings,
    path_to_relion_project: pathlib.Path,
) -> tuple[pathlib.Path, list[str]]:
    # Get the file number for this MRC file
    if n_particles == 0:
        file_number = 0
    else:
        last_index = index[0] - 1
        if last_index == -1:
            file_number = 0
        else:
            last_filename = particle_data["rlnImageName"].iloc[last_index].split("@")[1]
            if pd.isna(last_filename):
                raise OSError(
                    "Tried to assign a number to the MRC file while writing "
                    "images, but could not grab the previous file number at "
                    f"index {int(last_index)}. At this index, found that the "
                    "filename was NaN."
                )
            else:
                file_number = _parse_filename_for_number(last_filename) + 1
    # Unpack settings
    prefix = mrcfile_settings["prefix"]
    output_folder = mrcfile_settings["output_folder"]
    delimiter = mrcfile_settings["delimiter"]
    n_characters = mrcfile_settings["n_characters"]
    # Generate filename
    file_number_fmt = _format_number_for_filename(file_number, n_characters=n_characters)
    if prefix == "":
        relative_path_to_filename = str(
            pathlib.Path(output_folder, file_number_fmt + ".mrcs")
        )
    else:
        relative_path_to_filename = str(
            pathlib.Path(output_folder, prefix + delimiter + file_number_fmt + ".mrcs")
        )
    # Finally, generate the 'rln_image_name' column, which includes the particle index
    rln_image_names = [
        _format_number_for_filename(int(i + 1), n_characters)
        + "@"
        + relative_path_to_filename
        for i in range(index.size)
    ]
    # Finally, the path to the filename
    path_to_filename = pathlib.Path(path_to_relion_project, relative_path_to_filename)

    return path_to_filename, rln_image_names


def _parse_filename_for_number(filename: str) -> int:
    match = re.search(r"(\d+)\.[^.]+$", filename)
    try:
        file_number = int(match.group(1))  # type: ignore
    except Exception as err:
        raise OSError(
            f"Could not get the file number from file {filename} "
            "Files must be enumerated with the trailing part of the "
            "filename as the file number, like so: '/path/to/file-0000.txt'. "
            f"When extracting the file number and converting it to an integer, "
            f"found error:\n\t{err}"
        )
    return file_number


def _format_number_for_filename(file_number: int, n_characters: int = 6):
    if file_number == 0:
        return "0" * n_characters
    else:
        n_digits = int(np.log10(file_number)) + 1
        return "0" * (n_characters - n_digits) + str(file_number)
