"""
Polymeric properties
====================
.. moduleauthor:: Benjamin Ye <GitHub: @bbye98>

This module contains classes to determine the structural
and dynamical properties of polymeric systems.
"""

from typing import Union
import warnings

import MDAnalysis as mda
from MDAnalysis.lib.log import ProgressBar
import numba
import numpy as np
from scipy import optimize, special

from . import structure
from .base import Hash, DynamicAnalysisBase, NumbaAnalysisBase
from .. import FOUND_OPENMM, Q_, ureg
from ..algorithm import accelerated, correlation
from ..algorithm.molecule import center_of_mass, radius_of_gyration
from ..algorithm.topology import unwrap, unwrap_edge
from ..algorithm.unit import strip_unit
from ..algorithm.utility import get_closest_factors
from ..fit.exponential import stretched_exp

if FOUND_OPENMM:
    from openmm import unit


class _PolymerAnalysisBase(DynamicAnalysisBase):
    r"""
    An analysis base object for polymer systems.

    Parameters
    ----------
    groups : `MDAnalysis.AtomGroup` or array-like
        Groups of polymers to be analyzed.

        .. note::

           All polymers in each group must have the same chain length.

    groupings : `str` or array-like, default: :code:`"atoms"`
        Determines whether the centers of mass are used in lieu of
        individual atom positions. If `groupings` is a `str`, the same
        value is used for all `groups`.

        .. note::

           In a standard trajectory file, segments (or chains) contain
           residues (or molecules), and residues contain atoms. This
           heirarchy must be adhered to for this analysis module to
           function correctly. If your trajectory file does not contain
           the correct residue or segment information, provide the
           number of chains and chain lengths in `n_chains` and
           `n_monomers`, respectively.

        .. container::

           **Valid values**:

           * :code:`"atoms"`: Atom positions (generally or for
             coarse-grained simulations).
           * :code:`"residues"`: Residues' centers of mass (for
             atomistic simulations).

    n_chains : `int` or array-like, optional
        Number of chains :math:`M` in each polymer group. Must be
        provided if the trajectory does not adhere to the standard
        container heirarchy. If an `int` is provided, the same value is
        used for all groups.

        **Shape**: :math:`(N_\mathrm{groups},)`.

    n_monomers : `int` or array-like, optional
        Number of monomers :math:`N` in each chain in each polymer
        group. Must be provided if the trajectory does not adhere to the
        standard container heirarchy. If an `int` is provided, the same
        value is used for all groups.

        **Shape**: :math:`(N_\mathrm{groups},)`.

    unwrap : `bool`, keyword-only, default: :code:`False`
        Determines whether atom positions are unwrapped.

    parallel : `bool`, keyword-only, default: :code:`False`
        Determines whether the analysis is performed in parallel.

    verbose : `bool`, keyword-only, default: :code:`True`
        Determines whether detailed progress is shown.

    **kwargs
        Additional keyword arguments to pass to
        :class:`MDAnalysis.analysis.base.AnalysisBase`.

    Attributes
    ----------
    universe : `MDAnalysis.Universe`
        :class:`MDAnalysis.core.universe.Universe` object containing all
        information describing the system.
    """

    def __init__(
        self,
        groups: Union[mda.AtomGroup, tuple[mda.AtomGroup]],
        groupings: Union[str, tuple[str]] = "atoms",
        n_chains: Union[int, tuple[int]] = None,
        n_monomers: Union[int, tuple[int]] = None,
        *,
        unwrap: bool = False,
        parallel: bool = False,
        verbose: bool = True,
        **kwargs,
    ) -> None:

        self._groups = [groups] if isinstance(groups, mda.AtomGroup) else groups
        self.universe = self._groups[0].universe
        super().__init__(self.universe.trajectory, parallel, verbose, **kwargs)

        self._dimensions = self.universe.dimensions
        if self._dimensions is not None:
            self._dimensions = self._dimensions[:3].copy()

        GROUPINGS = {"atoms", "residues"}
        self._n_groups = len(self._groups)
        if isinstance(groupings, str):
            if groupings not in GROUPINGS:
                emsg = (
                    f"Invalid grouping '{groupings}'. Valid values: "
                    "'" + "', '".join(GROUPINGS) + "'."
                )
                raise ValueError(emsg)
            self._groupings = self._n_groups * [groupings]
        else:
            if self._n_groups != len(groupings):
                emsg = (
                    "The shape of 'groupings' is incompatible with " "that of 'groups'."
                )
                raise ValueError(emsg)
            for gr in groupings:
                if gr not in GROUPINGS:
                    emsg = (
                        f"Invalid grouping '{gr}'. Valid "
                        "values: '" + "', '".join(GROUPINGS) + "'."
                    )
                    raise ValueError(emsg)
            self._groupings = groupings

        if n_chains is None or n_monomers is None:
            self._internal = True
            self._n_chains = np.empty(self._n_groups, dtype=int)
            self._n_monomers = np.empty_like(self._n_chains)
            for i, ag in enumerate(self._groups):
                self._n_chains[i] = ag.segments.n_segments
                self._n_monomers[i] = ag.n_atoms // self._n_chains[i]
        else:
            self._internal = False
            if isinstance(n_chains, (int, np.integer)):
                self._n_chains = n_chains * np.ones(self._n_groups, dtype=int)
            elif self._n_groups == len(n_chains):
                self._n_chains = n_chains
            else:
                emsg = (
                    "The shape of 'n_chains' is incompatible with " "that of 'groups'."
                )
                raise ValueError(emsg)
            if isinstance(n_monomers, (int, np.integer)):
                self._n_monomers = n_monomers * np.ones(n_monomers, dtype=int)
            elif self._n_groups == len(n_monomers):
                self._n_monomers = n_monomers
            else:
                emsg = (
                    "The shape of 'n_monomers' is incompatible "
                    "with that of 'groups'."
                )
                raise ValueError(emsg)

        self._unwrap = unwrap
        self._verbose = verbose


class Gyradius(_PolymerAnalysisBase):
    """
    Serial and parallel implementations to calculate the radius of
    gyration :math:`R_\\mathrm{g}` of a polymer.

    The radius of gyration is used to describe the dimensions of a
    polymer chain, and is defined as

    .. math::

        R_\\mathrm{g}=\\sqrt{
        \\frac{\\sum_i^N m_i\\|\\mathbf{r}_i
        -\\mathbf{R}_\\mathrm{com}\\|^2}{\\sum_i^N m_i}}

    where :math:`m_i` and :math:`\\mathbf{r}_i` are the mass and
    position, respectively, of particle :math:`i`, and
    :math:`\\mathbf{R}_\\mathrm{com}` is the center of mass.

    Parameters
    ----------
    groups : `MDAnalysis.AtomGroup` or array-like
        Groups of polymers to be analyzed.

        .. note::

           All polymers in each group must have the same chain length.

    groupings : `str` or array-like, default: :code:`"atoms"`
        Determines whether the centers of mass are used in lieu of
        individual atom positions. If `groupings` is a `str`, the same
        value is used for all `groups`.

        .. note::

           In a standard trajectory file, segments (or chains) contain
           residues (or molecules), and residues contain atoms. This
           heirarchy must be adhered to for this analysis module to
           function correctly. If your trajectory file does not contain
           the correct residue or segment information, provide the
           number of chains and chain lengths in `n_chains` and
           `n_monomers`, respectively.

        .. container::

           **Valid values**:

           * :code:`"atoms"`: Atom positions (generally or for
             coarse-grained simulations).
           * :code:`"residues"`: Residues' centers of mass (for
             atomistic simulations).

    n_chains : `int` or array-like, optional
        Number of chains :math:`M` in each polymer group. Must be
        provided if the trajectory does not adhere to the standard
        container heirarchy. If an `int` is provided, the same value is
        used for all groups.

        **Shape**: :math:`(N_\\mathrm{groups},)`.

    n_monomers : `int` or array-like, optional
        Number of monomers :math:`N` in each chain in each polymer
        group. Must be provided if the trajectory does not adhere to the
        standard container heirarchy. If an `int` is provided, the same
        value is used for all groups.

        **Shape**: :math:`(N_\\mathrm{groups},)`.

    components : `bool`, keyword-only, default: :code:`False`
        Specifies whether the components of the radii of gyration are
        calculated and returned instead.

    unwrap : `bool`, keyword-only, default: :code:`False`
        Determines whether atom positions are unwrapped.

    parallel : `bool`, keyword-only, default: :code:`False`
        Determines whether the analysis is performed in parallel.

    verbose : `bool`, keyword-only, default: :code:`True`
        Determines whether detailed progress is shown.

    **kwargs
        Additional keyword arguments to pass to
        :class:`MDAnalysis.analysis.base.AnalysisBase`.

    Attributes
    ----------
    universe : `MDAnalysis.Universe`
        :class:`MDAnalysis.core.universe.Universe` object containing all
        information describing the system.

    results.units : `dict`
        Reference units for the results. For example, to get the
        reference units for `results.gyradii`, call
        :code:`results.units["gyradii"]`.

    results.gyradii : `numpy.ndarray`
        Radii of gyration.

        **Shape**: :math:`(N_\\mathrm{groups},\\,N_\\mathrm{frames})` or
        :math:`(N_\\mathrm{groups},\\,N_\\mathrm{frames},\\,3)`.

        **Reference unit**: :math:`\\mathrm{Å}`.

    Example
    --------
    First, this analysis class must be imported:

    >>> from mdcraft.analysis.polymer import Gyradius

    Then, after loading a simulation trajectory:

    >>> universe = mda.Universe("simulation.nc", "topology.cif")

    We must then select the atom-groups to be analyzed:

    >>> ag = universe.select_atoms("resname POLY")

    The gyradius class can then be built and run, remembering to specify the number of chains and monomers per chain:

    >>> gyradius = Gyradius(ag, n_chains=10, n_monomers=100)
    >>> gyradius.run()

    The results can be obtained under the `results` attribute:

    >>> gyradius.results.gyradii

    """

    def __init__(
        self,
        groups: Union[mda.AtomGroup, tuple[mda.AtomGroup]],
        groupings: Union[str, tuple[str]] = "atoms",
        n_chains: Union[int, tuple[int]] = None,
        n_monomers: Union[int, tuple[int]] = None,
        *,
        components: bool = False,
        unwrap: bool = False,
        parallel: bool = False,
        verbose: bool = True,
        **kwargs,
    ):

        super().__init__(
            groups,
            groupings,
            n_chains,
            n_monomers,
            unwrap=unwrap,
            parallel=parallel,
            verbose=verbose,
            **kwargs,
        )

        if self.universe.dimensions is None and unwrap:
            emsg = (
                "System dimensions were not found, but are "
                "required when 'unwrap=True'."
            )
            raise ValueError(emsg)

        self._Ns_p = self._n_chains * self._n_monomers
        self._N_p = self._Ns_p.sum()
        self._slices = []
        _ = 0
        for N_p in self._Ns_p:
            self._slices.append(slice(_, _ + N_p))
            _ += N_p

        self._components = components

    def _prepare(self) -> None:

        if self._unwrap:

            # Navigate to first frame in analysis
            ts = self._sliced_trajectory[0]

            # Get system dimensions
            dimensions = ts.dimensions[:3]

            # Preallocate arrays to determine the number of periodic
            # boundary crossings for each entity
            self._positions_old = np.empty((self._N_p, 3))
            for ag, gr, s, M, N in zip(
                self._groups,
                self._groupings,
                self._slices,
                self._n_chains,
                self._n_monomers,
            ):
                if self._internal and gr == "residues":
                    ag.unwrap()
                    self._positions_old[s] = center_of_mass(ag, gr)
                else:
                    positions = unwrap_edge(
                        positions=ag.positions,
                        bonds=np.array(
                            [
                                (i * N + j, i * N + j + 1)
                                for i in range(M)
                                for j in range(N - 1)
                            ]
                        ),
                        dimensions=dimensions,
                        masses=ag.masses,
                    )
                    self._positions_old[s] = (
                        positions
                        if gr == "atoms"
                        else center_of_mass(
                            positions=positions.reshape(M, N, -1, 3),
                            masses=ag.masses.reshape(M, N, -1),
                        )
                    )
            self._images = np.zeros((self._N_p, 3), dtype=int)

            # Store unwrapped entity positions in a shared memory array
            # for parallel analysis
            if self._parallel:
                self._positions = np.empty((self.n_frames, self._N_p, 3))
                for i, _ in enumerate(self._sliced_trajectory):
                    for ag, gr, s, M, N in zip(
                        self._groups,
                        self._groupings,
                        self._slices,
                        self._n_chains,
                        self._n_monomers,
                    ):
                        if self._internal and gr == "residues":
                            self._positions[i, s] = center_of_mass(ag, gr)
                        else:
                            self._positions[i, s] = (
                                ag.positions
                                if gr == "atoms"
                                else center_of_mass(
                                    positions=ag.positions.reshape(M, N, -1, 3),
                                    masses=ag.masses.reshape(M, N, -1),
                                )
                            )

                    # Globally unwrap entity positions for correct
                    # centers of mass and radii of gyration
                    unwrap(
                        self._positions[i],
                        self._positions_old,
                        dimensions,
                        thresholds=dimensions / 2,
                        images=self._images,
                    )

        # Preallocate arrays to hold radii of gyration
        if not self._parallel:
            shape = [self._n_groups, self.n_frames]
            if self._components:
                shape.append(3)
            self.results.gyradii = np.empty(shape)

        # Store reference units
        self.results.units = Hash({"gyradii": ureg.angstrom})

    def _single_frame(self) -> None:

        for i, (ag, gr, M, N, s) in enumerate(
            zip(
                self._groups,
                self._groupings,
                self._n_chains,
                self._n_monomers,
                self._slices,
            )
        ):

            # Store atom or center-of-mass positions in the current frame
            if self._internal and gr == "residues":
                positions = center_of_mass(ag, gr)
            else:
                positions = (
                    ag.positions
                    if gr == "atoms"
                    else center_of_mass(
                        positions=ag.positions.reshape(M, N, -1, 3),
                        masses=ag.masses.reshape(M, N, -1),
                    )
                )

            # Globally unwrap entity positions for correct centers of
            # mass and radii of gyration
            if self._unwrap:
                dimensions = self._ts.dimensions[:3]
                unwrap(
                    positions,
                    self._positions_old[s],
                    dimensions,
                    thresholds=dimensions / 2,
                    images=self._images[s],
                )

            # Compute the radii of gyration
            self.results.gyradii[i, self._frame_index] = radius_of_gyration(
                grouping="segments",
                positions=positions.reshape((M, N, 3)),
                masses=ag.masses.reshape((M, N)),
                components=self._components,
            ).mean(axis=0)

    def _single_frame_parallel(self, index: int) -> tuple[int, np.ndarray[float]]:

        # Set current trajectory frame
        self._sliced_trajectory[index]

        # Preallocate array to store radii of gyration
        shape = [self._n_groups]
        if self._components:
            shape.append(3)
        gyradii = np.empty(shape)

        for i, (ag, gr, M, N, s) in enumerate(
            zip(
                self._groups,
                self._groupings,
                self._n_chains,
                self._n_monomers,
                self._slices,
            )
        ):

            # Calculate or get entity positions
            if self._unwrap:
                positions = self._positions[index, s]
            elif self._internal and gr == "residues":
                positions = center_of_mass(ag, gr)
            else:
                positions = (
                    ag.positions
                    if gr == "atoms"
                    else center_of_mass(
                        positions=ag.positions.reshape(M, N, -1, 3),
                        masses=ag.masses.reshape(M, N, -1),
                    )
                )

            # Compute the radii of gyration
            gyradii[i] = radius_of_gyration(
                grouping="segments",
                positions=positions.reshape((M, N, 3)),
                masses=ag.masses.reshape((M, N)),
                components=self._components,
            ).mean(axis=0)

        return index, gyradii

    def _conclude(self) -> None:

        # Consolidate parallel results and clean up memory by deleting
        # arrays that will not be reused
        if self._parallel:
            self._results = sorted(self._results)
            self.results.gyradii = np.stack([r[1] for r in self._results], axis=1)

            del self._results
            if self._unwrap:
                del self._positions
        if self._unwrap:
            del self._positions_old, self._images


def correlation_fft(
    x: np.ndarray[Union[float, complex]],
    y: np.ndarray[Union[float, complex]] = None,
    /,
    axis: int = None,
    *,
    average: bool = False,
    double: bool = False,
    vector: bool = False,
) -> np.ndarray[Union[float, complex]]:
    r"""
    Evaluates the autocorrelation functions (ACF)
    :math:`\mathrm{R_\mathbf{XX}}(\tau)` or cross-correlation functions
    (CCF) :math:`\mathrm{R_\mathbf{XY}}(\tau)` of time series
    :math:`\mathbf{X}(t)` and :math:`\mathbf{Y}(t)` using fast Fourier
    transforms (FFT).

    .. seealso::

       This function is an alias for
       :func:`mdcraft.algorithm.correlation.correlation_fft`.
    """

    return correlation.correlation_fft(
        x, y, axis, average=average, double=double, vector=vector
    )


def correlation_shift(
    x: np.ndarray[Union[float, complex]],
    y: np.ndarray[Union[float, complex]] = None,
    /,
    axis: int = None,
    *,
    average: bool = False,
    double: bool = False,
    vector: bool = False,
) -> np.ndarray[Union[float, complex]]:
    r"""
    Evaluates the autocorrelation functions (ACF)
    :math:`\mathrm{R_\mathbf{XX}}(\tau)` or cross-correlation functions
    (CCF) :math:`\mathrm{R_\mathbf{XY}}(\tau)` of time series
    :math:`\mathbf{X}(t)` and :math:`\mathbf{Y}(t)` directly by using
    sliding windows.

    .. seealso::

       This function is an alias for
       :func:`mdcraft.algorithm.correlation.correlation_shift`.
    """

    return correlation.correlation_shift(
        x, y, axis, average=average, double=double, vector=vector
    )


def calculate_relaxation_time(
    times: np.ndarray[float], acf: np.ndarray[float]
) -> float:
    r"""
    Calculates the orientational relaxation time :math:`\tau_\mathrm{r}`
    of polymers using end-to-end vector autocorrelation function (ACF)
    time series :math:`C_\mathrm{ee}(t)`.

    A stretched exponential function with :math:`\tau` and :math:`\beta`
    as coefficients,

    .. math::

       C_\mathrm{ee}(t)=\exp{\left[-(t/\tau)^\beta\right]}

    is fitted to the ACF time series, and the relaxation time is
    estimated using the first moment of :math:`C_\mathrm{ee}`,

    .. math::

       \tau_\mathrm{r}=\int_0^\infty C_\mathrm{ee}(t)\,dt
       =\frac{\tau}{\beta}\Gamma\left(\frac{1}{\beta}\right)

    where :math:`\Gamma` is the gamma function.

    Parameters
    ----------
    times : `numpy.ndarray`
        Changes in time :math:`t-t_0`.

        **Shape**: :math:`(N_\mathrm{frames},)`.

        **Reference unit**: :math:`\mathrm{ps}`.

    acf : `numpy.ndarray`
        End-to-end vector ACF :math:`C_\mathrm{ee}(t)`.

        **Shape**:
        :math:`(N_\mathrm{groups},\,N_\mathrm{bins},\,N_\mathrm{frames})`.

    Returns
    -------
    relaxation_time : `float`
        Average orientational relaxation time :math:`\tau_\mathrm{r}`.

        **Reference unit**: :math:`\mathrm{ps}`.
    """

    tau_r, beta = optimize.curve_fit(
        stretched_exp, times / times[1], acf, bounds=(0, np.inf)
    )[0]
    return tau_r * times[1] * special.gamma((beta + 1) / beta)


class EndToEndVector(_PolymerAnalysisBase):
    """
    A serial implementation to calculate the end-to-end vector
    autocorrelation function (ACF) :math:`C_\\mathrm{ee}(t)` and
    estimate the orientational relaxation time :math:`\\tau_\\mathrm{r}`
    of a polymer.

    The end-to-end vector ACF is defined as

    .. math::

       C_\\mathrm{ee}(t)=\\frac{\\langle\\mathbf{R}_\\mathrm{ee}(t)
       \\cdot\\mathbf{R}_\\mathrm{ee}(0)\\rangle}
       {\\langle\\mathbf{R}_\\mathrm{ee}^2\\rangle}

    where :math:`\\mathbf{R}_\\mathrm{ee}=\\mathbf{r}_N-\\mathbf{r}_1`
    is the end-to-end vector.

    The orientational relaxation time can then be estimated by fitting
    a stretched exponential function

    .. math::

       C_\\mathrm{ee}=\\exp{\\left[-(t/\\tau)^\\beta\\right]}

    to the end-to-end vector ACF and evaluating

    .. math::

       \\tau_\\mathrm{r}=\\int_0^\\infty C_\\mathrm{ee}\\,dt
       =\\frac{\\tau}{\\beta}\\Gamma\\left(\\frac{1}{\\beta}\\right)

    Parameters
    ----------
    groups : `MDAnalysis.AtomGroup` or array-like
        Groups of polymers to be analyzed.

        .. note::

           All polymers in each group must have the same chain length.

    groupings : `str` or array-like, default: :code:`"atoms"`
        Determines whether the centers of mass are used in lieu of
        individual atom positions. If `groupings` is a `str`, the same
        value is used for all `groups`.

        .. note::

           In a standard trajectory file, segments (or chains) contain
           residues (or molecules), and residues contain atoms. This
           heirarchy must be adhered to for this analysis module to
           function correctly. If your trajectory file does not contain
           the correct residue or segment information, provide the
           number of chains and chain lengths in `n_chains` and
           `n_monomers`, respectively.

        .. container::

           **Valid values**:

           * :code:`"atoms"`: Atom positions (generally or for
             coarse-grained simulations).
           * :code:`"residues"`: Residues' centers of mass (for
             atomistic simulations).

    n_chains : `int` or array-like, optional
        Number of chains :math:`M` in each polymer group. Must be
        provided if the trajectory does not adhere to the standard
        container heirarchy. If an `int` is provided, the same value is
        used for all groups.

        **Shape**: :math:`(N_\\mathrm{groups},)`.

    n_monomers : `int` or array-like, optional
        Number of monomers :math:`N` in each chain in each polymer
        group. Must be provided if the trajectory does not adhere to the
        standard container heirarchy. If an `int` is provided, the same
        value is used for all groups.

        **Shape**: :math:`(N_\\mathrm{groups},)`.

    n_blocks : `int`, keyword-only, default: :code:`1`
        Number of blocks to split the trajectory into.

    dt : `float`, `openmm.unit.Quantity`, or `pint.Quantity`, \
    keyword-only, optional
        Time between frames :math:`\\Delta t`. While this is normally
        determined from the trajectory, the trajectory may not have the
        correct information if the data is in reduced units. For
        example, if the reduced timestep is :math:`0.01` and trajectory
        data was outputted every :math:`10,000` timesteps, then
        :math:`\\Delta t=100`.

        **Reference unit**: :math:`\\mathrm{ps}`.

    fft : `bool`, keyword-only, default: :code:`True`
        Determines whether fast Fourier transforms (FFT) are used to
        evaluate the ACFs.

    unwrap : `bool`, keyword-only, default: :code:`False`
        Determines whether atom positions are unwrapped.

    verbose : `bool`, keyword-only, default: :code:`True`
        Determines whether detailed progress is shown.

    **kwargs
        Additional keyword arguments to pass to
        :class:`MDAnalysis.analysis.base.AnalysisBase`.

    Attributes
    ----------
    universe : `MDAnalysis.Universe`
        :class:`MDAnalysis.core.universe.Universe` object containing all
        information describing the system.

    results.units : `dict`
        Reference units for the results. For example, to get the
        reference units for `results.times`, call
        :code:`results.units["times"]`.

    results.times : `numpy.ndarray`
        Changes in time :math:`t-t_0`.

        **Shape**: :math:`(N_\\mathrm{frames},)`.

        **Reference unit**: :math:`\\mathrm{ps}`.

    results.acf : `numpy.ndarray`
        End-to-end vector ACFs :math:`C_\\mathrm{ee}(t)`.

        **Shape**:
        :math:`(N_\\mathrm{groups},\\,N_\\mathrm{blocks},\\,N_\\mathrm{frames})`.

    results.relaxation_times : `numpy.ndarray`
        Average orientational relaxation times :math:`\\tau_\\mathrm{r}`.

        **Shape**: :math:`(N_\\mathrm{groups},\\,N_\\mathrm{blocks})`.

        **Reference units**: :math:`\\mathrm{ps}`.
    
    Example
    --------
    First, this analysis class must be imported:

    >>> from mdcraft.analysis.polymer import EndToEndVector

    Then, after loading a simulation trajectory:

    >>> universe = mda.Universe("simulation.nc", "topology.cif")

    We must then select the atom-groups to be analyzed:

    >>> ag = universe.select_atoms("resname POLY")

    The `EndToEndVector` class can then be built and run, remembering to specify the number of chains and monomers per chain:
    
    >>> e2e = EndToEndVector(ag, n_chains=10, n_monomers=100)
    >>> e2e.run()

    The results can be obtained under the `results` attribute:
    
    >>> e2e.results.acf
    
    One can also obtain the relaxation times by calling the following:
    
    >>> e2e.calculate_relaxation_time()
    >>> e2e.results.relaxation_times
    
    These can be saved by doing the following:
    
    >>> e2e.save("e2e")
    
    """

    def __init__(
        self,
        groups: Union[mda.AtomGroup, tuple[mda.AtomGroup]],
        groupings: Union[str, tuple[str]] = "atoms",
        n_chains: Union[int, tuple[int]] = None,
        n_monomers: Union[int, tuple[int]] = None,
        *,
        n_blocks: int = 1,
        dt: Union[float, "unit.Quantity", Q_] = None,
        fft: bool = True,
        unwrap: bool = False,
        verbose: bool = True,
        **kwargs,
    ) -> None:

        # Disable parallel support for this analysis class
        if "parallel" in kwargs:
            wmsg = (
                "The 'EndToEndVector' analysis class does not "
                "support multithreading."
            )
            warnings.warn(wmsg)
            del kwargs["parallel"]

        super().__init__(
            groups,
            groupings,
            n_chains,
            n_monomers,
            unwrap=unwrap,
            verbose=verbose,
            **kwargs,
        )

        self._M = self._n_chains.sum()
        self._slices = []
        _ = 0
        for M in self._n_chains:
            self._slices.append(slice(_, _ + M))
            _ += M

        self._n_blocks = n_blocks
        self._dt = strip_unit(dt, "ps")[0] or self._trajectory.dt
        self._fft = fft

    def _prepare(self) -> None:

        # Ensure frames are evenly spaced and proceed forward in time
        if hasattr(self._sliced_trajectory, "frames"):
            dfs = np.diff(self._sliced_trajectory.frames)
            if (df := dfs[0]) <= 0 or not np.allclose(dfs, df):
                emsg = (
                    "The selected frames must be evenly spaced and "
                    "proceed forward in time."
                )
                raise ValueError(emsg)
        elif (df := self.step) <= 0:
            raise ValueError("The analysis must proceed forward in time.")

        # Determine number of frames used when the trajectory is split
        # into blocks
        self._n_frames_block = self.n_frames // self._n_blocks
        self._n_frames = self._n_blocks * self._n_frames_block
        if (extra_frames := self.n_frames - self._n_frames) > 0:
            wmsg = (
                f"The trajectory is not divisible into {self._n_blocks:,} "
                f"blocks, so the last {extra_frames:,} frame(s) will be "
                "discarded. To maximize performance, set appropriate "
                "starting and ending frames in run() so that the number "
                "of frames to be analyzed is divisible by the number of "
                "blocks."
            )
            warnings.warn(wmsg)

        # Preallocate arrays to store end-to-end vectors
        self._e2e = np.empty((self.n_frames, self._M, 3))

        if self._unwrap:

            # Navigate to first frame in analysis
            self._sliced_trajectory[0]

            # Preallocate arrays to determine the number of periodic
            # boundary crossings for the first and last monomer in each
            # chain
            self._positions_end_old = np.empty((self._M, 2, 3))
            for ag, gr, s, M, N in zip(
                self._groups,
                self._groupings,
                self._slices,
                self._n_chains,
                self._n_monomers,
            ):
                if self._internal and gr == "residues":
                    ag.unwrap()
                    self._positions_end_old[s] = np.stack(
                        [
                            center_of_mass(s.residues[[0, -1]].atoms, "residues")
                            for s in ag.segments
                        ]
                    )
                else:
                    positions = unwrap_edge(
                        positions=ag.positions,
                        bonds=np.array(
                            [
                                (i * N + j, i * N + j + 1)
                                for i in range(M)
                                for j in range(N - 1)
                            ]
                        ),
                        dimensions=self._dimensions,
                        masses=ag.masses,
                    ).reshape(M, N, -1, 3)[:, (0, -1)]
                    self._positions_end_old[s] = (
                        positions[:, :, 0]
                        if gr == "atoms"
                        else center_of_mass(
                            positions=positions,
                            masses=ag.masses.reshape(M, N, -1)[:, (0, -1)],
                        )
                    )
            self._images = np.zeros((self._M, 2, 3), dtype=int)
            self._thresholds = self._dimensions / 2

        # Preallocate array to store end-to-end vector ACFs
        self.results.acf = np.empty(
            (self._n_groups, self._n_blocks, self._n_frames_block)
        )

        # Store time changes
        self.results.times = df * self._dt * np.arange(self._n_frames_block)

        # Store reference units
        self.results.units = {"times": ureg.picosecond}

    def _single_frame(self) -> None:

        for ag, gr, s, M, N in zip(
            self._groups,
            self._groupings,
            self._slices,
            self._n_chains,
            self._n_monomers,
        ):

            # Store ending monomer or center-of-mass positions in the
            # current frame
            if self._internal and gr == "residues":
                positions_end = np.stack(
                    center_of_mass(s.residues[[0, -1]].atoms, "residues")
                    for s in ag.segments
                )
            else:
                positions_end = ag.positions.reshape(M, N, -1, 3)[:, (0, -1)]
                positions_end = (
                    positions_end[:, :, 0]
                    if gr == "atoms"
                    else center_of_mass(
                        positions=positions_end,
                        masses=ag.masses.reshape(M, N, -1)[:, (0, -1)],
                    )
                )

            # Globally unwrap entity positions for correct end-to-end
            # vectors
            if self._unwrap:
                unwrap(
                    positions_end,
                    self._positions_end_old[s],
                    self._dimensions,
                    thresholds=self._thresholds,
                    images=self._images[s],
                )

            # Compute the end-to-end vectors
            self._e2e[self._frame_index, s] = np.diff(positions_end, axis=1)[:, 0]

    def _conclude(self) -> None:

        # Clean up memory by deleting arrays that will not be reused
        if self._unwrap:
            del self._positions_end_old, self._images, self._thresholds

        # Compute end-to-end vector ACFs
        _acf = correlation_fft if self._fft else correlation_shift
        for i, (s, M) in ProgressBar(enumerate(zip(self._slices, self._n_chains))):
            self.results.acf[i] = _acf(
                (
                    self._e2e[:, s]
                    / np.linalg.norm(self._e2e[:, s], axis=-1, keepdims=True)
                ).reshape(self._n_blocks, -1, M, 3),
                average=True,
                vector=True,
            )

    def calculate_relaxation_times(self) -> None:
        """
        Calculates the orientational relaxation times.
        """

        if not hasattr(self.results, "acf"):
            emsg = (
                "Call EndToEndVector.run() before "
                "EndToEndVector.calculate_relaxation_times()."
            )
            raise RuntimeError(emsg)

        self.results.relaxation_times = np.empty((self._n_groups, self._n_blocks))
        self.results.units["relaxation_times"] = ureg.picosecond

        for i, block in enumerate(self.results.acf):
            for j, acf in enumerate(block):
                valid = np.where(acf >= 0)[0]
                self.results.relaxation_times[i, j] = calculate_relaxation_time(
                    self.results.times[valid], acf[valid]
                )


class SingleChainStructureFactor(NumbaAnalysisBase, _PolymerAnalysisBase):
    """
    Serial and parallel implementations to calculate the single-chain
    structure factor :math:`S_\\mathrm{sc}(q)` of a homopolymer.

    It is defined as

    .. math::

       S_{\\mathrm{sc}}(q)=\\frac{1}{MN}\\left\\langle
       \\sum_{m=1}^M\\sum_{j=1}^{N}\\sum_{k=1}^{N}\\exp{
       [i\\mathbf{q}\\cdot(\\mathbf{r}_j-\\mathbf{r}_k)]}\\right\\rangle

    where :math:`M` is the number of chains, :math:`N` is the chain
    length, :math:`\\mathbf{q}` and :math:`q` are the scattering
    wavevector and its magnitude, respectively, and
    :math:`\\mathbf{r}_i` is the position of the :math:`i`-th monomer.

    .. container::

       The single-chain structure factor reveals information about the
       characteristic length scales of the polymer:

       * In the Guinier regime (:math:`qR_g\\ll1`),
         :math:`S_{\\mathrm{sc}}(q)^{-1}\\approx N(1-(qR_g)^2/3)`
         can be used to determine the radius of gyration :math:`R_g`.
       * In the Porod regime (:math:`qR_g\\gg1`),
         :math:`S_{\\mathrm{sc}}(q)=1` since the only contribution is
         the self-scattering of the monomers.
       * In the intermediate regime, the slope :math:`s` of the log-log
         plot of :math:`S_{\\mathrm{sc}}(q)` is related to the scaling
         exponent :math:`\\nu` via :math:`\\nu=-1/s`.

    Parameters
    ----------
    groups : `MDAnalysis.AtomGroup` or array-like
        Groups of polymers to be analyzed.

        .. note::

           All polymers in each group must have the same chain length.

    groupings : `str` or array-like, default: :code:`"atoms"`
        Determines whether the centers of mass are used in lieu of
        individual atom positions. If `groupings` is a `str`, the same
        value is used for all `groups`.

        .. note::

           In a standard trajectory file, segments (or chains) contain
           residues (or molecules), and residues contain atoms. This
           heirarchy must be adhered to for this analysis module to
           function correctly. If your trajectory file does not contain
           the correct residue or segment information, provide the
           number of chains and chain lengths in `n_chains` and
           `n_monomers`, respectively.

        .. container::

           **Valid values**:

           * :code:`"atoms"`: Atom positions (generally or for
             coarse-grained simulations).
           * :code:`"residues"`: Residues' centers of mass (for
             atomistic simulations).

    n_chains : `int` or array-like, optional
        Number of chains :math:`M` in each polymer group. Must be
        provided if the trajectory does not adhere to the standard
        container heirarchy. If an `int` is provided, the same value is
        used for all groups.

        **Shape**: :math:`(N_\\mathrm{groups},)`.

    n_monomers : `int` or array-like, optional
        Number of monomers :math:`N` in each chain in each polymer
        group. Must be provided if the trajectory does not adhere to the
        standard container heirarchy. If an `int` is provided, the same
        value is used for all groups.

        **Shape**: :math:`(N_\\mathrm{groups},)`.

    form : `str`, keyword-only, default: :code:`"exp"`
        Expression used to evaluate the single-chain structure factors.

        .. container::

           **Valid values**:

           * :code:`"exp"`: Exponential form. Slightly faster due to
             fewer mathematical operations.
           * :code:`"trig"`: Trigonometric form. Slightly slower but
             doesn't have overflow issues.

    dimensions : array-like, `openmm.unit.Quantity`, or \
    `pint.Quantity`, keyword-only, optional
        System dimensions :math:`(L_x,\\,L_y,\\,L_z)`. Only used to
        determine the scattering wavevectors when `wavevectors` is not
        specified.

        **Shape**: :math:`(3,)`.

        **Reference unit**: :math:`\\mathrm{Å}`.

    n_points : `int`, keyword-only, default: :code:`32`
        Number of points :math:`n_\\mathrm{points}` in the scattering
        wavevector grid to generate

        .. math::

           \\mathbf{q}=2\\pi\\left(\\frac{a}{L_x},\\,\\frac{b}{L_y},\\,
           \\frac{c}{L_z}\\right)

        where :math:`a`, :math:`b`, and :math:`c` are integers from
        :math:`0` up to :math:`n_\\mathrm{points}-1`.

        Additional wavevectors can be introduced via `n_surfaces` and
        `n_surface_points` for more accurate structure factors at small
        wavenumbers. Alternatively, the desired wavevectors can be
        specified directly in `wavevectors`.

    n_surfaces : `int`, keyword-only, optional
        Number of spherical surfaces in the first octant that intersect
        with the grid wavevectors along the three coordinate axes for
        which to introduce extra wavevectors for more accurate structure
        factor values. Only available if the system is perfectly cubic.

    n_surface_points : `int`, keyword-only, default: :code:`8`
        Number of extra wavevectors to introduce per spherical surface.
        Has no effect if `n_surfaces` is not specified.

    q_max : `float`, `openmm.unit.Quantity`, or `pint.Quantity`, \
    keyword-only, optional
        Maximum wavenumber :math:`q_\\mathrm{max}`.

        **Reference unit**: :math:`\\mathrm{Å}^{-1}`.

    wavevectors : array-like, `openmm.unit.Quantity`, or `pint.Quantity`, \
    keyword-only, optional
        Scattering wavevectors for which to compute structure factors.
        Has precedence over `n_points`, `n_surfaces`, and
        `n_surface_points` if specified.

        **Shape**: :math:`(N_q,\\,3)`.

        **Reference unit**: :math:`\\mathrm{Å}^{-1}`.

    sort : `bool`, keyword-only, default: :code:`True`
        Determines whether the results are sorted by the wavenumbers.

    unique : `bool`, keyword-only, default: :code:`True`
        Determines whether structure factors for the same wavenumber
        are grouped and averaged.

    parallel : `bool`, keyword-only, default: :code:`False`
        Determines whether the analysis is performed in parallel.

    verbose : `bool`, keyword-only, default: :code:`True`
        Determines whether detailed progress is shown.

    **kwargs
        Additional keyword arguments to pass to
        :class:`MDAnalysis.analysis.base.AnalysisBase`.

    Attributes
    ----------
    universe : `MDAnalysis.Universe`
        :class:`MDAnalysis.core.universe.Universe` object containing all
        information describing the system.

    results.units : `dict`
        Reference units for the results. For example, to get the
        reference units for :code:`results.wavenumbers`, call
        :code:`results.units["wavenumbers"]`.

    results.wavenumbers : `numpy.ndarray`
        Wavenumbers :math:`q`.

        **Shape**: :math:`(N_q,)`.

        **Reference unit**: :math:`\\mathrm{Å}^{-1}`.

    results.scsf : `numpy.ndarray`
        Single-chain structure factors :math:`S_\\mathrm{sc}(q)`.

        **Shape**: :math:`(N_q,)`.
        
    Example
    --------
    First, this analysis class must be imported:

    >>> from mdcraft.analysis.polymer import SingleChainStructureFactor

    Then, after loading a simulation trajectory:

    >>> universe = mda.Universe("simulation.nc", "topology.cif")

    We must then select the atom-groups to be analyzed:

    >>> ag = universe.select_atoms("resname POLY")

    The `SingleChainStructureFactor` class can then be built and run, remembering to specify the number of chains and monomers per chain:
    
    >>> scsf = SingleChainStructureFactor(ag, n_chains=10, n_monomers=100)
    >>> scsf.run()

    The results can be obtained under the `results` attribute:
    
    >>> scsf.results.scsf
    
    """

    def __init__(
        self,
        groups: mda.AtomGroup,
        groupings: str = "atoms",
        n_chains: int = None,
        n_monomers: int = None,
        *,
        form: str = "exp",
        dimensions: Union[np.ndarray[float], "unit.Quantity", Q_] = None,
        n_points: int = 32,
        n_surfaces: int = None,
        n_surface_points: int = 8,
        q_max: Union[float, "unit.Quantity", Q_] = None,
        wavevectors: np.ndarray[float] = None,
        sort: bool = True,
        unique: bool = True,
        unwrap: bool = False,
        parallel: bool = False,
        verbose: bool = True,
        **kwargs,
    ) -> None:

        # Specify 'parallel=False' to use the NumbaAnalysisBase.run method
        # instead of the ParallelAnalysisBase.run method
        _PolymerAnalysisBase.__init__(
            self,
            groups,
            groupings,
            n_chains,
            n_monomers,
            unwrap=unwrap,
            parallel=False,
            verbose=verbose,
            **kwargs,
        )

        if wavevectors is not None:
            self._wavevectors = wavevectors
        else:
            dimensions = strip_unit(dimensions or self.universe.dimensions[:3], "Å")[0]
            if dimensions is None:
                emsg = (
                    "System dimensions were not found, but are "
                    "required when 'wavevectors' is not specified."
                )
                raise ValueError(emsg)
            elif len(dimensions) != 3:
                raise ValueError("'dimensions' must have length 3.")
            dimensions = np.asarray(dimensions)

            if np.allclose(dimensions, dimensions[0]):
                grid = 2 * np.pi * np.arange(n_points) / dimensions[0]
                self._wavevectors = np.stack(np.meshgrid(grid, grid, grid), -1).reshape(
                    -1, 3
                )
                if n_surfaces:
                    n_theta, n_phi = get_closest_factors(
                        n_surface_points, 2, reverse=True
                    )
                    theta = np.linspace(
                        np.pi / (2 * n_theta + 4),
                        np.pi / 2 - np.pi / (2 * n_theta + 4),
                        n_theta,
                    )
                    phi = np.linspace(
                        np.pi / (2 * n_phi + 4),
                        np.pi / 2 - np.pi / (2 * n_phi + 4),
                        n_phi,
                    )
                    self._wavevectors = np.vstack(
                        (
                            self._wavevectors,
                            np.einsum(
                                "o,tpd->otpd",
                                grid[1 : n_surfaces + 1],
                                np.stack(
                                    (
                                        np.sin(theta) * np.cos(phi)[:, None],
                                        np.sin(theta) * np.sin(phi)[:, None],
                                        np.tile(np.cos(theta)[None, :], (n_phi, 1)),
                                    ),
                                    axis=-1,
                                ),
                            ).reshape((n_surfaces * n_surface_points, 3)),
                        )
                    )
            else:
                self._wavevectors = np.stack(
                    np.meshgrid(
                        *[2 * np.pi * np.arange(n_points) / L for L in dimensions]
                    ),
                    axis=-1,
                ).reshape(-1, 3)
        self._wavenumbers = np.linalg.norm(self._wavevectors, axis=1)

        if q_max is not None:
            q_max = strip_unit(q_max, "Å^-1")[0]
            keep = self._wavenumbers <= q_max
            self._wavevectors = self._wavevectors[keep]
            self._wavenumbers = self._wavenumbers[keep]

        self._njit = lambda s: numba.njit(s, fastmath=True, parallel=parallel)
        self._delta_fourier_transform_sum = self._njit(
            structure.delta_fourier_transform_sum
        )
        self._ssf_trigonometric = self._njit(structure.ssf_trigonometric)
        self._psf_trigonometric = self._njit(structure.psf_trigonometric)
        self._inner = (
            accelerated.numba_inner_parallel if parallel else accelerated.numba_inner
        )

        self._form = form
        self._sort = sort
        self._unique = unique
        self._verbose = verbose

    def _prepare(self) -> None:

        # Preallocate array to store single-chain structure factors
        if not self._parallel:
            self.results.scsf = np.zeros((self._n_groups, len(self._wavenumbers)))

        # Determine the unique wavenumbers
        self.results.wavenumbers = np.unique(self._wavenumbers.round(11))

        # Store reference units
        self.results.units = Hash({"results.wavenumbers": ureg.angstrom**-1})

    def _single_frame(self) -> None:

        for ig, (ag, gr, M, N) in enumerate(
            zip(self._groups, self._groupings, self._n_chains, self._n_monomers)
        ):

            # Get entity positions in the current frame
            if self._internal and gr == "residues":
                positions = center_of_mass(ag, gr)
            else:
                positions = (
                    ag.positions
                    if gr == "atoms"
                    else center_of_mass(
                        positions=ag.positions.reshape(M, N, -1, 3),
                        masses=ag.masses.reshape(M, N, -1),
                    )
                )

            # Calculate single-chain structure factor contributions
            if self._form == "exp":
                for chain in positions.reshape((M, N, 3)):
                    rhos = self._delta_fourier_transform_sum(self._wavevectors, chain)
                    self.results.scsf[ig] += (rhos * rhos.conj()).real
            elif self._form == "trig":
                for chain in positions.reshape((M, N, 3)):
                    self.results.scsf[ig] += self._ssf_trigonometric(
                        self._inner(self._wavevectors, chain)
                    )

    def _conclude(self) -> None:

        # Normalize the single-chain structure factor by the number of
        # monomers and timesteps
        self.results.scsf /= self._n_chains * self._n_monomers * self.n_frames

        # Combine values sharing the same wavenumber, if desired
        if self._unique:
            self.results.scsf = np.hstack(
                [
                    self.results.scsf[:, np.isclose(q, self._wavenumbers)].mean(
                        axis=1, keepdims=True
                    )
                    for q in self.results.wavenumbers
                ]
            )

        # Sort the results by wavenumber, if desired
        if self._sort:
            order = np.argsort(self.results.wavenumbers)
            self.results.wavenumbers = self.results.wavenumbers[order]
            self.results.scsf = self.results.scsf[:, order]
