# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
#
# MDAnalysis --- https://www.mdanalysis.org
# Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
# (see the file AUTHORS for the full list of names)
#
# Released under the Lesser GNU Public Licence, v2.1 or any higher version
#
# Please cite your use of MDAnalysis in published work:
#
# R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
# D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
# MDAnalysis: A Python package for the rapid analysis of molecular dynamics
# simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
# Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
# doi: 10.25080/majora-629e541a-00e
#
# N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
# MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
# J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
#

r"""
Updated nucleic acid analysis --- :mod:`MDAnalysis.analysis.nucleicacids`
=========================================================================

:Author: Alia Lescoulie
:Year: 2022-2023
:copyright: LGPLv2.1

The module provides classes for analyzing nucleic acids structures.
This is an updated, higher performance version of previous nucleic acid tools.
For applications see :footcite:p:`Denning2011,Denning2012`.

.. rubric:: References

.. footbibliography::

Distances
---------

.. autoclass:: NucPairDist
    :members:
    :inherited-members:

.. autoclass:: WatsonCrickDist
    :members:
    :exclude-members: select_strand_atoms
    :inherited-members:

.. autoclass:: MinorPairDist
    :members:
    :exclude-members: select_strand_atoms
    :inherited-members:

.. autoclass:: MajorPairDist
    :members:
    :exclude-members: select_strand_atoms
    :inherited-members:

.. versionadded 2.2.0

"""

from typing import List, Dict, Tuple, Union
import warnings

import numpy as np

import MDAnalysis as mda
from .distances import calc_bonds
from .base import AnalysisBase, ResultsGroup
from MDAnalysis.core.groups import Residue, ResidueGroup


# Deprecation: In 3.0.0 change type to just
# ResidueClass = ResidueGroup
ResidueClass = Union[List[Residue], ResidueGroup]
r"""A type alias for :code:`Union[List[Residue], ResidueGroup]`

Used as an alias for methods where either class is acceptable.
"""


class NucPairDist(AnalysisBase):
    r"""Atom pair distance calculation base class.

    Takes two lists of :class:`~MDAnalysis.core.groups.AtomGroup` and
    computes the distances between them over a trajectory. Used as a
    superclass for the other nucleic acid distances classes. The distance
    will be measured between atoms sharing an index in the two lists of
    :class:`~MDAnalysis.core.groups.AtomGroup`.

    Parameters
    ----------
    selection1: List[AtomGroup]
        List of :class:`~MDAnalysis.core.groups.AtomGroup` containing an atom
        of each nucleic acid being analyzed.
    selection2: List[AtomGroup]
        List of :class:`~MDAnalysis.core.groups.AtomGroup` containing an atom
        of each nucleic acid being analyzed.
    kwargs: dict
        Arguments for :class:`~MDAnalysis.analysis.base.AnalysisBase`


    Attributes
    ----------
    results.pair_distances: numpy.ndarray
        2D array of pair distances. First dimension is simulation time,
        second dimension contains the pair distances for each each entry
        pair in selection1 and selection2.

        .. versionadded:: 2.4.0

        .. note::
            `results.pair_distances` is slated for deprecation in
            version 3.0.0, use `results.distances` instead.
        .. deprecated:: 2.7.0
            `results.pair_distances` will be removed in
            version 3.0.0, use :attr:`results.distances` instead.

    results.distances: numpy.ndarray
        stored in a 2d numpy array with first index selecting the
        Residue pair, and the second index selecting the frame number
        Distances are stored in a 2d numpy array with axis 0 (first index)
        indexing the trajectory frame and axis 1 (second index) selecting the
        Residue pair.

        .. versionadded:: 2.7.0

    times: numpy.ndarray
        Simulation times for analysis.


    Raises
    ------
    ValueError
        If the selections given are not the same length
    ValueError
        An :class:`~MDAnalysis.core.groups.AtomGroup` in one of the
        strands not a valid nucleic acid
    ValueError
        If a given residue pair from the provided strands returns an empty
        :class:`~MDAnalysis.core.groups.AtomGroup` when selecting the atom
        pairs used in the distance calculations


    *Version Info*

    .. versionchanged:: 2.5.0
       The ability to access by passing selection indices to :attr:`results`
       is now removed as of MDAnalysis version 2.5.0. Please use
       :attr:`results.pair_distances` instead.
       The :attr:`results.times` was deprecated and is now removed as of
       MDAnalysis 2.5.0.
       Please use the class attribute :attr:`times` instead.

    .. versionchanged:: 2.7.0
        Added static method :attr:`select_strand_atoms` as a
        helper for selecting atom pairs for distance analysis.

    .. versionchanged:: 2.9.0
       Enabled **parallel execution** with the ``multiprocessing`` and ``dask``
       backends; use the new method :meth:`get_supported_backends` to see all
       supported backends.
    """

    _analysis_algorithm_is_parallelizable = True

    @classmethod
    def get_supported_backends(cls):
        return ("serial", "multiprocessing", "dask")

    _s1: mda.AtomGroup
    _s2: mda.AtomGroup
    _n_sel: int

    def __init__(
        self,
        selection1: List[mda.AtomGroup],
        selection2: List[mda.AtomGroup],
        **kwargs,
    ) -> None:
        super(NucPairDist, self).__init__(
            selection1[0].universe.trajectory, **kwargs
        )

        if len(selection1) != len(selection2):
            raise ValueError("Selections must be same length")

        self._n_sel: int = len(selection1)

        self._s1 = selection1[0]
        self._s2 = selection2[0]

        for i in range(1, self._n_sel):
            self._s1 += selection1[i]
            self._s2 += selection2[i]

    @staticmethod
    def select_strand_atoms(
        strand1: ResidueGroup,
        strand2: ResidueGroup,
        a1_name: str,
        a2_name: str,
        g_name: str = "G",
        a_name: str = "A",
        u_name: str = "U",
        t_name: str = "T",
        c_name: str = "C",
    ) -> Tuple[List[mda.AtomGroup], List[mda.AtomGroup]]:
        r"""
        A helper method for nucleic acid pair distance analyses.
        Used for selecting specific atoms from two strands of nucleic acids.


        Parameters
        ----------
        strand1: List[Residue]
            The first nucleic acid strand
        strand2: List[Residue]
            The second nucleic acid strand
        a1_name: str
            The selection for the purine base of the strand pair
        a2_name: str
            the selection for the pyrimidine base of the strand pair
        g_name: str (optional)
            Name of Guanine in topology, by default assigned to G
        a_name: str (optional)
            Name of Adenine in topology, by default assigned to A
        u_name: str (optional)
            Name of Uracil in topology, by default assigned to U
        t_name: str (optional)
            Name of Thymine in topology, by default assigned to T
        c_name: str (optional)
            Name of Cytosine in topology, by default assigned to C

        Returns
        -------
        Tuple[List[AtomGroup], List[AtomGroup]]
            returns a tuple containing two lists of
            :class:`~MDAnalysis.core.groups.AtomGroup`\s
            corresponding to the provided selections from each strand.

        Raises
        ------
        ValueError:
            An :class:`~MDAnalysis.core.groups.AtomGroup`
            in one of the strands not a valid nucleic acid
        ValueError:
            An :class:`~MDAnalysis.core.groups.Residue` returns an empty
            :class:`~MDAnalysis.core.groups.AtomGroup`
            with the provided selection


        .. versionadded:: 2.7.0
        """
        pyrimidines: List[str] = [c_name, t_name, u_name]
        purines: List[str] = [a_name, g_name]

        sel1: List[mda.AtomGroup] = []
        sel2: List[mda.AtomGroup] = []

        for pair in zip(strand1.residues, strand2.residues):
            if pair[0].resname[0] in pyrimidines:
                a1, a2 = a2_name, a1_name
            elif pair[0].resname[0] in purines:
                a1, a2 = a1_name, a2_name
            else:
                raise ValueError(
                    f"AtomGroup in {pair} is not a valid nucleic acid"
                )

            ag1 = pair[0].atoms.select_atoms(f"name {a1}")
            ag2 = pair[1].atoms.select_atoms(f"name {a2}")

            if not all(len(ag) > 0 for ag in [ag1, ag2]):
                err_info: Tuple[Residue, str] = (
                    (pair[0], a1) if len(ag1) == 0 else (pair[1], a2)
                )

                raise ValueError(
                    (
                        f"{err_info[0]} returns an empty AtomGroup"
                        'with selection string "name {a2}"'
                    )
                )

            sel1.append(ag1)
            sel2.append(ag2)

        return (sel1, sel2)

    def _prepare(self) -> None:
        self.results.distances: np.ndarray = np.zeros(
            [self.n_frames, self._n_sel]
        )

    def _single_frame(self) -> None:
        dist: np.ndarray = calc_bonds(self._s1.positions, self._s2.positions)

        self.results.distances[self._frame_index, :] = dist

    def _conclude(self) -> None:
        self.results["pair_distances"] = self.results["distances"]
        # TODO: remove pair_distances in 3.0.0

    def _get_aggregator(self):
        return ResultsGroup(
            lookup={
                "distances": ResultsGroup.ndarray_vstack,
            }
        )


class WatsonCrickDist(NucPairDist):
    r"""
    Watson-Crick base pair distance for selected
    residues over a trajectory.

    Takes two :class:`~MDAnalysis.core.groups.ResidueGroup`
    objects or two lists of :class:`~MDAnalysis.core.groups.Residue`
    and calculates the distance between the nitrogen atoms in the
    Watson-Crick hydrogen bond over the trajectory. Bases are matched
    either by their index in the two
    :class:`~MDAnalysis.core.groups.ResidueGroup` provided as arguments,
    or based on the indices of the provided lists of
    :class:`~MDAnalysis.core.groups.Residue` objects depending
    on which is provided.

    .. note::
        Support for :class:`~MDAnalysis.core.groups.Residue` is slated for
        deprecation and will raise a warning when used. It still works but
        :class:`~MDAnalysis.core.groups.ResidueGroup` is preferred.

    Parameters
    ----------
    strand1: ResidueClass
        First list of bases

        .. deprecated:: 2.7.0
           Using a list of :class:`~MDAnalysis.core.groups.Residue` will
           be removed in 3.0.0. Use a
           :class:`~MDAnalysis.core.groups.ResidueGroup`.

    strand2: ResidueClass
        Second list of bases

        .. deprecated:: 2.7.0
           Using a list of :class:`~MDAnalysis.core.groups.Residue` will
           be removed in 3.0.0. Use a
           :class:`~MDAnalysis.core.groups.ResidueGroup`.

    n1_name: str (optional)
        Name of Nitrogen 1 of nucleic acids, by default assigned to "N1"
    n3_name: str (optional)
        Name of Nitrogen 3 of nucleic acids, by default assigned to "N3"
    g_name: str (optional)
        Name of Guanine in topology, by default assigned to "G"
    a_name: str (optional)
        Name of Adenine in topology, by default assigned to "A"
    u_name: str (optional)
        Name of Uracil in topology, by default assigned to "U"
    t_name: str (optional)
        Name of Thymine in topology, by default assigned to "T"
    c_name: str (optional)
        Name of Cytosine in topology, by default assigned to C
    **kwargs: dict
        Key word arguments for
        :class:`~MDAnalysis.analysis.base.AnalysisBase`

    Attributes
    ----------
    results.distances: numpy.ndarray
        Distances are stored in a 2d numpy array with axis 0 (first index)
        indexing the trajectory frame and axis 1 (second index) selecting the
        Residue pair.

        .. versionadded:: 2.7.0

    results.pair_distances: numpy.ndarray
        2D array of pair distances. First dimension is
        simulation time, second dimension contains the
        pair distances for each each entry pair in
        selection1 and selection2.

        .. versionadded:: 2.4.0

        .. deprecated:: 2.7.0
            `results.pair_distances` will be removed in version 3.0.0,
            use :attr:`results.distances` instead.

    times: numpy.ndarray
        Simulation times for analysis.

    Raises
    ------
    TypeError
        If the provided list of :class:`~MDAnalysis.core.Residue` contains
        non-Residue elements

        .. deprecated:: 2.7.0
           Starting with version 3.0.0, this exception will no longer
           be raised because only
           :class:`~MDAnalysis.core.groups.ResidueGroup` will be allowed.

    ValueError
        If `strand1` and `strand2` are not the same length
    ValueError:
        An :class:`~MDAnalysis.core.groups.AtomGroup`
        in one of the strands not a valid nucleic acid
    ValueError
        If a given residue pair from the provided strands returns an empty
        :class:`~MDAnalysis.core.groups.AtomGroup` when selecting the atom
        pairs used in the distance calculations


    *Version Info*

    .. versionchanged:: 2.5.0
       Accessing results by passing strand indices to :attr:`results`
       was deprecated and is now removed as of MDAnalysis version 2.5.0.
       Please use :attr:`results.pair_distances` instead.
       The :attr:`results.times` was deprecated and is now removed as of
       MDAnalysis 2.5.0. Please use the class attribute
       :attr:`times` instead.

    .. versionchanged:: 2.7.0
        `strand1` and `strand2` now also accept a
        :class:`~MDAnalysis.core.groups.ResidueGroup` as input.
        The previous input type, ``List[Residue]`` is still supported,
        but it is **deprecated** and will be removed in release 3.0.0.
    """

    def __init__(
        self,
        strand1: ResidueClass,
        strand2: ResidueClass,
        n1_name: str = "N1",
        n3_name: str = "N3",
        g_name: str = "G",
        a_name: str = "A",
        u_name: str = "U",
        t_name: str = "T",
        c_name: str = "C",
        **kwargs,
    ) -> None:

        def verify_strand(strand: ResidueClass) -> ResidueGroup:
            # Helper method to verify the strands

            if isinstance(strand, list):  # Checking if a list is given
                # verify list is only Residues
                if not all(isinstance(resid, Residue) for resid in strand):
                    raise TypeError(f"{strand} contains non-Residue elements")

                warnings.warn(
                    DeprecationWarning(
                        (
                            f"ResidueGroup should be used for {strand} instead"
                            "of giving a Residue list"
                        )
                    )
                )
                # Convert to a ResidueGroup
                strand: ResidueGroup = ResidueGroup(strand)

            return strand

        strand1: ResidueGroup = verify_strand(strand1)
        strand2: ResidueGroup = verify_strand(strand2)

        strand_atomgroups: Tuple[List[mda.AtomGroup], List[mda.AtomGroup]] = (
            self.select_strand_atoms(
                strand1,
                strand2,
                n1_name,
                n3_name,
                g_name=g_name,
                a_name=a_name,
                t_name=t_name,
                u_name=u_name,
                c_name=c_name,
            )
        )

        super(WatsonCrickDist, self).__init__(
            strand_atomgroups[0], strand_atomgroups[1], **kwargs
        )


class MinorPairDist(NucPairDist):
    r"""Minor-Pair basepair distance for selected residues over a trajectory.

    Takes two :class:`~MDAnalysis.core.groups.ResidueGroup` objects and
    calculates the Minor-groove hydrogen bond length between the
    nitrogen and oxygen atoms over the trajectory. Bases are
    matched by their index in the two
    :class:`~MDAnalysis.core.groups.ResidueGroup` provided as arguments.

    Parameters
    ----------
    strand1: List[Residue]
        First list of bases
    strand2: List[Residue]
        Second list of bases
    o2_name: str (optional)
        Name of Oxygen 2 of nucleic acids;
        by default assigned to "O2";
    c2_name: str (optional)
        Name of Carbon 2 of nucleic acids;
        by default assigned to "C2";
    g_name: str (optional)
        Name of Guanine in topology;
        by default assigned to "G";
    a_name: str (optional)
        Name of Adenine in topology
        by default assigned to "A";
    u_name: str (optional)
        Name of Uracil in topology;
        by default assigned to "U";
    t_name: str (optional)
        Name of Thymine in topology;
        by default assigned to "T";
    c_name: str (optional)
        Name of Cytosine in topology;
        by default assigned to "C";
    **kwargs:
        keyword arguments for
        :class:`~MDAnalysis.analysis.base.AnalysisBase`

    Attributes
    ----------
    results.distances: numpy.ndarray
        stored in a 2d numpy array with first index selecting
        the Residue pair, and the second index selecting the frame number
    times: numpy.ndarray
        Simulation times for analysis.

    Raises
    ------
    ValueError
        If the selections given are not the same length
        A :class:`~MDAnalysis.core.Residue` in
        one of the strands not a valid nucleic acid
    ValueError
        If a given residue pair from the provided strands returns an empty
        :class:`~MDAnalysis.core.groups.AtomGroup` when selecting the atom
        pairs used in the distance calculations


    .. versionadded:: 2.7.0
    """

    def __init__(
        self,
        strand1: ResidueGroup,
        strand2: ResidueGroup,
        o2_name: str = "O2",
        c2_name: str = "C2",
        g_name: str = "G",
        a_name: str = "A",
        u_name: str = "U",
        t_name: str = "T",
        c_name: str = "C",
        **kwargs,
    ) -> None:

        selections: Tuple[List[mda.AtomGroup], List[mda.AtomGroup]] = (
            self.select_strand_atoms(
                strand1,
                strand2,
                c2_name,
                o2_name,
                g_name=g_name,
                a_name=a_name,
                t_name=t_name,
                u_name=u_name,
                c_name=c_name,
            )
        )

        super(MinorPairDist, self).__init__(
            selections[0], selections[1], **kwargs
        )


class MajorPairDist(NucPairDist):
    r"""Minor-Pair base pair distance for
    selected residues over a trajectory.

    Takes two :class:`~MDAnalysis.core.groups.ResidueGroup` objects and
    calculates the Major-groove hydrogen bond length between the nitrogen
    and oxygen atoms over the trajectory. Bases are matched by their index
    in the two :class:`~MDAnalysis.core.groups.ResidueGroup`
    provided as arguments.

    Parameters
    ----------
    strand1: List[Residue]
        First list of bases
    strand2: List[Residue]
        Second list of bases
    o6_name: str (optional)
        Name of Oxygen 6 of nucleic acids;
        by default assigned to "O6"
    n4_name: str (optional)
        Name of Nitrogen 4 of nucleic acids;
        by default assigned to "N4"
    g_name: str (optional)
        Name of Guanine in topology;
        by default assigned to "G"
    a_name: str (optional)
        Name of Adenine in topology;
        by default assigned to "A"
    u_name: str (optional)
        Name of Uracil in topology;
        by default assigned to "U"
    t_name: str (optional)
        Name of Thymine in topology;
        by default assigned to "T"
    c_name: str (optional)
        Name of Cytosine in topology;
        by default assigned to "C"
    **kwargs:
        arguments for :class:`~MDAnalysis.analysis.base.AnalysisBase`

    Attributes
    ----------
    results.distances: numpy.ndarray
        Distances are stored in a 2d numpy array with axis 0 (first index)
        indexing the trajectory frame and axis 1 (second index) selecting the
        Residue pair.
    times: numpy.ndarray
        Simulation times for analysis.

    Raises
    ------
    ValueError
        A :class:`~MDAnalysis.core.Residue`
        in one of the strands not a valid nucleic acid
    ValueError
        If a given residue pair from the provided strands returns an empty
        :class:`~MDAnalysis.core.groups.AtomGroup` when selecting the atom
        pairs used in the distance calculations
    ValueError
        if the selections given are not the same length


    .. versionadded:: 2.7.0
    """

    def __init__(
        self,
        strand1: ResidueGroup,
        strand2: ResidueGroup,
        n4_name: str = "N4",
        o6_name: str = "O6",
        g_name: str = "G",
        a_name: str = "A",
        u_name: str = "U",
        t_name: str = "T",
        c_name: str = "C",
        **kwargs,
    ) -> None:

        selections: Tuple[List[mda.AtomGroup], List[mda.AtomGroup]] = (
            self.select_strand_atoms(
                strand1,
                strand2,
                o6_name,
                n4_name,
                g_name=g_name,
                a_name=a_name,
                t_name=t_name,
                u_name=u_name,
                c_name=c_name,
            )
        )

        super(MajorPairDist, self).__init__(
            selections[0], selections[1], **kwargs
        )