# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
#
# lipyphilic --- lipyphilic.readthedocs.io
#
# Released under the GNU Public Licence, v2 or any higher version
#

r"""
Lateral diffusion --- :mod:`lipyphilic.analysis.lateral_diffusion`
==================================================================

This module contains methods for calculating the lateral diffusion coefficient
of lipids in a bilayer.

The class :class:`lipyphilic.analysis.lateral_diffusion.MSD` calculates the two-dimensional
mean squared displacent (MSD) of lipids in a bilayer. The `Fast Correlation Algorithm
<https://www.sciencedirect.com/science/article/pii/001046559500048K>`__, implemented
in `tidynamics <http://lab.pdebuyl.be/tidynamics/>`__ is used to calculate the MSD of
each lipid, with optional removal of the center of mass motion of the bilayer.

:class:`lipyphilic.analysis.lateral_diffusion.MSD` also contains a method for calculating the
lateral diffusion coefficient, :math:`D_{xy}`, via the Einstein relation:

.. math::

    D_{xy} = \frac{1}{4} \lim_{t\to\infty} \frac{d}{dt} \displaystyle \Bigg\langle \frac{1}{N} \sum_{i=1}^{N} \left  | r_i(t_0 + \Delta t) - r_i(t_0) \right |^2 \displaystyle \Bigg\rangle_{t_0}

where :math:`N` is the number of lipids, :math:`r_i(t0)` is the center of mass in :math:`xy`
of lipids :math:`i` at a time origin `t_0`, :math:`r_i(t0 + \Delta t)` is the same lipid's
center of mass at a lagtime :math:`\Delta t`, and the angular brackets denote an average
over all time origins, :math:`t_0`.

Typically, the MSD is averaged over all molecules. However, :class:`lipyphilic.analysis.lateral_diffusion.MSD`
will return the MSD for each individual lipid. This makes it simple to later calculate the diffusion
coefficient using a subset of the lipids, such as a specific lipid species or lipids near a protein.

Input
-----

Required:
  - *universe* : an MDAnalysis Universe object
  - *lipid_sel* : atom selection for calculating the MSD

Optional:
  - *com_removal_sel* : atom selection for center of mass removal from the MSD
  - *dt* : time period betwen consecutive frames in the MSD analysis

Output
------

  - *msd* : the mean squared displacement of each lipid at each lagtime, :math:`\Delta t`, in :math:`nm^2`:
  - *lagtimes* : a NumPy array of lagtimes (in :math:`ns`)

The data are stored in the :attr:`MSD.msd` and :attr:`MSD.lagtimes` attributes.


Warning
-------

Before using `lipyphilic.analysis.lateral_diffusion.MSD` you *must* ensure that the coordinates have
been unwrapped using, for example, :class:`MDAnalysis.transformations.NoJump`.


Example usage of :class:`MSD`
-----------------------------

To calculate the MSD of each lipid in a bilayer we must first load a trajectory using
MDAnalysis::

  import MDAnalysis as mda
  from lipyphilic.analysis.lateral_diffusion import MSD

  u = mda.Universe(tpr, trajectory)

If we have used the MARTINI forcefield to study a phospholipid/cholesterol mixture,
we can calculate the MSD of each lipid as follows::

  msd = MSD(
    universe=u,
    lipid_sel="resname DPPC DOPC CHOL"
  )

We then select which frames of the trajectory to analyse (`None` will use every
frame) and choose to display a progress bar (`verbose=True`)::

  msd.run(
    start=None,
    stop=None,
    step=None,
    verbose=True
  )

The results are then available in the :attr:`msd.MSD` attribute as a
:class:`numpy.ndarray`. Each row corresponds to an individual lipid and each column
to a different lagtime`.

Center of mass removal
----------------------

During your simulation, it is likely that you removed the center of mass motion of your
bilayer in the :math:`z` direction. However, it is not possible to remove the :math:`x`
and :math:`y` center of mass motions until you have unwrapped your lipid positions.

You may select which lipids to use for the center of mass motion removal using the
:attr:`com_removal_sel` keyword::

  msd = MSD(
    universe=u,
    lipid_sel="resname DPPC",
    com_removla_sel="resname DPPC DOPC CHOL"
  )

In this case, the MSD of DPPC will be calculated with the center of mass motion of the
bilayer will be subtracted from it.

Plotting the MSD of each species
--------------------------------

If you have calculated the MSD of DPPC, DOPC and cholesterol as in the first example, you
can plot the MSD of each species as follows::

  for species in ["DPPC", "DOPC", "CHOL"]:

    plt.loglog(
      msd.lagtimes,
      np.mean(msd.msd[msd.membrane.residues.resnames == species], axis=0),
      label=species
    )

  plt.legend()

The linear part of the log-log plot can be used for fitting a line and calculating the
diffusion coefficient.

Calculating the lateral diffusion coefficient
---------------------------------------------

After calculating the MSD and identifying the linear portion of the plot, the
:func: `diffusion_coefficient` method of :class:`lipyphilic.analysis.lateral_diffusion.MSD`
can be used to calculate :math:`D_{xy}`. We need to pass the time at which to start and
stop the linear fit::

  d, sem = msd.diffusion_coefficient(
    start_fit=400,
    end_fit=600
  )

This will calculate a diffusion coefficient for each individual lipid and return the mean
and standard error of the distribution of coefficients.

To calculate the diffusion coefficient of a subset of lipids we can use the `:attr:lipid_sel`
keyword::

  d, sem = msd.diffusion_coefficient(
    start_fit=400,
    end_fit=600,
    lipid_sel="resname CHOL"
  )

which will calculate the lateral diffusion coefficient for cholesterol, using a fit to the MSD
curve from lagtime :math:`\Delta t = 400` to lagtime :math:`\Delta t = 600`.

.. autoclass:: MSD
    :members:

"""

from attrs import define
from MDAnalysis import Universe
from MDAnalysis.analysis.base import AnalysisBase
import numpy as np
import scipy.stats
import tidynamics

__all__ = [
    "MSD",
]


@define(auto_attribs=True, auto_detect=True, eq=False)
class MSD(AnalysisBase):
    """
    Calculate the mean-squared lateral displacement (MSD) of lipids in a bilayer.

    The MSD is calculated in units of :math:`nm^2/ns`.

    Parameters
    ----------
    universe : Universe
        MDAnalysis Universe object
    lipid_sel : str
        Selection string for calculating the mean-squared displacement. IF multiple atoms
        per lipid are selected, the center-of-mass of these atoms will be used for
        calculating the MSD.
    com_removal_sel : str, optional
        The MSD of the center of mass of atoms in this selection will be subtracted from
        all individual lipid MSDs. The default is `None`, in which case no center of mass
        motion removal is performed.
    dt : float, optional
        The time, in picoseconds, between consecutive frames in `universe.trajectory`.
        The default is `None`, in which case `dt` is taken to be `universe.trajectory.dt`.
    """

    universe: Universe
    lipid_sel: str
    com_removal_sel: str | None = None
    dt: float | None = None

    def __attrs_post_init__(self):
        super().__init__(self.universe.trajectory)

        self.membrane = self.universe.select_atoms(self.lipid_sel, updating=False)
        self.com_removal = self.universe.select_atoms(self.com_removal_sel)
        self.dt = self.dt if self.dt is not None else self.universe.trajectory.dt

        self.results.msd = None
        self.results.lagtimes = None

    @property
    def msd(self):
        return self.results.msd

    @property
    def lagtimes(self):
        return self.results.lagtimes

    def _prepare(self):
        self.lipid_com_pos = np.full(
            (self.membrane.n_residues, self.n_frames, 2),
            fill_value=np.NaN,
            dtype=np.float64,
        )

        self.results.msd = np.full(
            (self.membrane.n_residues, self.n_frames),
            fill_value=np.NaN,
        )

    def _single_frame(self):
        self.lipid_com_pos[:, self._frame_index] = self.membrane.center_of_mass(compound="residues")[:, :2]

        # Remove COM motion if necessary
        if self.com_removal.n_atoms > 0:
            self.lipid_com_pos[:, self._frame_index] -= self.com_removal.center_of_mass()[:2]

    def _conclude(self):
        self.results.lagtimes = self.frames * self.dt

        # lagtimes must start from zero, and should be in ns
        self.results.lagtimes -= self.results.lagtimes[0]
        self.results.lagtimes /= 1000

        for lipid_index in range(self.membrane.n_residues):
            self.results.msd[lipid_index] = tidynamics.msd(self.lipid_com_pos[lipid_index])

        # MSD must start at 0
        self.results.msd[:, 0] = 0.0

        # Convert A^2 to nm^2
        self.results.msd /= 100

    def diffusion_coefficient(
        self,
        start_fit: float | None = None,
        stop_fit: float | None = None,
        lipid_sel: str | None = None,
    ) -> tuple[float, float]:
        """Calculate the lateral diffusion coefficient via the Einstein relation.

        A diffusion is calculated for each lipid through a linear fit to its MSD curve.
        The mean and standard error of the diffusion coefficient is returned.

        Parameters
        ----------

        start_fit : float, optional
            The time at which to start the linear fit to the MSD curve. The default is
            `None`, in which case the fit will exclude the first 20% of the MSD data.
        stop_fit : float, optional
            The time at which to stop the linear fit to the MSD curve. The default is
            `None`, in which case the fit will exclude the final 20% of the MSD data.
        lipid_sel : str, optional
            Selection string for lipids to include in calculating the diffusion
            coefficient.

        Returns
        -------

        d : float
            The mean lateral diffusion coefficient, in
            :math:`cm^2/s`., averaged over all lipids in `lipid_sel`.
        sem : float
            The standard error of the diffusion coefficients.
        """
        if start_fit is None:
            start_fit_index = self.results.lagtimes.size * 20 // 100
        else:
            start_fit_index = np.searchsorted(self.results.lagtimes, start_fit)

        if stop_fit is None:
            stop_fit_index = self.results.lagtimes.size * 80 // 100
        else:
            stop_fit_index = np.searchsorted(self.results.lagtimes, stop_fit)

        if lipid_sel is None:
            mask = np.full(self.membrane.n_residues, fill_value=True, dtype=bool)
        else:
            keep_lipids = self.universe.select_atoms(lipid_sel)
            mask = np.in1d(self.membrane.residues.resindices, keep_lipids.residues.resindices)

        all_coeffs = np.full(sum(mask), fill_value=np.NaN)
        for index, msd in enumerate(self.results.msd[mask, start_fit_index:stop_fit_index]):
            linear_fit = scipy.stats.linregress(self.results.lagtimes[start_fit_index:stop_fit_index], msd)
            slope = linear_fit.slope
            d = slope * 1 / 4 * 1e-5  # 1e-5 converts nm^2/ns to cm^2/s
            all_coeffs[index] = d

        return np.mean(all_coeffs), scipy.stats.sem(all_coeffs)
