# -*- 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
#

"""MOL2 file format --- :mod:`MDAnalysis.coordinates.MOL2`
========================================================

Classes to work with Tripos_ molecule structure format (MOL2_) coordinate and
topology files. Used, for instance, by the DOCK_ docking code.


Example for working with mol2 files
-----------------------------------

To open a mol2, remove all hydrogens and save as a new file, use the following::

  u = Universe("Molecule.mol2")
  gr = u.select_atoms("not name H*")
  print(len(u.atoms), len(gr))
  gr.write("Molecule_noh.mol2")

.. _MOL2: http://www.tripos.com/data/support/mol2.pdf
.. _Tripos: http://www.tripos.com/
.. _DOCK: http://dock.compbio.ucsf.edu/


See Also
--------

rdkit: rdkit_ is an open source cheminformatics toolkit with more mol2
       functionality


.. _rdkit: http://www.rdkit.org/docs/GettingStartedInPython.html


Classes
-------

.. autoclass:: MOL2Reader
   :members:

.. autoclass:: MOL2Writer
   :members:


Notes
-----

* The MDAnalysis :class:`MOL2Reader` and :class:`MOL2Writer` only handle the
  MOLECULE, SUBSTRUCTURE, ATOM, and BOND record types. Other records are not
  currently read or preserved on writing.
* As the CRYSIN record type is not parsed / written, MOL2 systems always have
  dimensions set to ``None`` and dimensionless MOL2 files are written.


MOL2 format notes
-----------------

* MOL2 Format Specification:  (http://www.tripos.com/data/support/mol2.pdf)
* Example file (http://www.tripos.com/mol2/mol2_format3.html)::

    #    Name: benzene
    #    Creating user name: tom
    #    Creation time: Wed Dec 28 00:18:30 1988

    #    Modifying user name: tom
    #    Modification time: Wed Dec 28 00:18:30 1988

    @<TRIPOS>MOLECULE
    benzene
    12 12 1  0   0
    SMALL
    NO_CHARGES


    @<TRIPOS>ATOM
    1   C1  1.207   2.091   0.000   C.ar    1   BENZENE 0.000
    2   C2  2.414   1.394   0.000   C.ar    1   BENZENE 0.000
    3   C3  2.414   0.000   0.000   C.ar    1   BENZENE 0.000
    4   C4  1.207   -0.697  0.000   C.ar    1   BENZENE 0.000
    5   C5  0.000   0.000   0.000   C.ar    1   BENZENE 0.000
    6   C6  0.000   1.394   0.000   C.ar    1   BENZENE 0.000
    7   H1  1.207   3.175   0.000   H   1   BENZENE 0.000
    8   H2  3.353   1.936   0.000   H   1   BENZENE 0.000
    9   H3  3.353   -0.542  0.000   H   1   BENZENE 0.000
    10  H4  1.207   -1.781  0.000   H   1   BENZENE 0.000
    11  H5  -0.939  -0.542  0.000   H   1   BENZENE 0.000
    12  H6  -0.939  1.936   0.000   H   1   BENZENE 0.000
    @<TRIPOS>BOND
    1   1   2   ar
    2   1   6   ar
    3   2   3   ar
    4   3   4   ar
    5   4   5   ar
    6   5   6   ar
    7   1   7   1
    8   2   8   1
    9   3   9   1
    10  4   10  1
    11  5   11  1
    12  6   12  1
   @<TRIPOS>SUBSTRUCTURE
    1   BENZENE 1   PERM    0   ****    ****    0   ROOT

"""
import numpy as np

from . import base
from ..lib import util
from MDAnalysis.lib.util import store_init_arguments


class MOL2Reader(base.ReaderBase):
    """Reader for MOL2 structure format.

    .. versionchanged:: 0.11.0
       Frames now 0-based instead of 1-based.
       MOL2 now reuses the same Timestep object for every frame,
       previously created a new instance of Timestep each frame.
    .. versionchanged:: 0.20.0
       Allows for comments at top of file.
       Ignores status bit strings
    .. versionchanged:: 2.0.0
       Bonds attribute is not added if no bonds are present in MOL2 file
    .. versionchanged:: 2.2.0
       Read MOL2 files with optional columns omitted.
    """

    format = "MOL2"
    units = {"time": None, "length": "Angstrom"}

    @store_init_arguments
    def __init__(self, filename, **kwargs):
        """Read coordinates from `filename`.


        Parameters
        ----------
        filename : str or NamedStream
             name of the mol2 file or stream
        """
        super(MOL2Reader, self).__init__(filename, **kwargs)

        self.n_atoms = None

        blocks = []

        with util.openany(filename) as f:
            for i, line in enumerate(f):
                # found new molecules
                if "@<TRIPOS>MOLECULE" in line:
                    blocks.append({"start_line": i, "lines": []})
                if len(blocks):
                    blocks[-1]["lines"].append(line)
        self.n_frames = len(blocks)
        self.frames = blocks

        sections, coords = self.parse_block(blocks[0])
        self.n_atoms = len(coords)

        self.ts = self._Timestep(self.n_atoms, **self._ts_kwargs)

        self.ts = self._read_frame(0)

    def parse_block(self, block):
        sections = {}
        cursor = None
        for line in block["lines"]:
            if "@<TRIPOS>" in line:
                cursor = line.split("@<TRIPOS>")[1].strip().lower()
                sections[cursor] = []
                continue
            elif line.startswith("#") or line == "\n":
                continue
            sections[cursor].append(line)

        atom_lines = sections["atom"]
        if not len(atom_lines):
            raise Exception(
                "The mol2 (starting at line {0}) block has no atoms"
                "".format(block["start_line"])
            )
        elif self.n_atoms is None:
            # First time round, remember the number of atoms
            self.n_atoms = len(atom_lines)
        elif len(atom_lines) != self.n_atoms:
            raise ValueError(
                "MOL2Reader assumes that the number of atoms remains unchanged"
                " between frames; the current "
                "frame has {0}, the next frame has {1} atoms"
                "".format(self.n_atoms, len(atom_lines))
            )

        coords = np.zeros((self.n_atoms, 3), dtype=np.float32)
        for i, a in enumerate(atom_lines):
            aid, name, x, y, z, atom_type = a.split()[:6]

            # x, y, z = float(x), float(y), float(z)
            coords[i, :] = x, y, z

        return sections, coords

    def _read_next_timestep(self, ts=None):
        if ts is None:
            ts = self.ts
        else:
            # TODO: cleanup _read_frame() to use a "free" Timestep
            raise NotImplementedError("MOL2Reader cannot assign to a timestep")
        frame = self.frame + 1
        return self._read_frame(frame)

    def _read_frame(self, frame):
        try:
            block = self.frames[frame]
        except IndexError:
            errmsg = (
                f"Invalid frame {frame} for trajectory with length "
                f"{len(self)}"
            )
            raise IOError(errmsg) from None

        sections, coords = self.parse_block(block)

        for sect in ["molecule", "substructure"]:
            try:
                self.ts.data[sect] = sections[sect]
            except KeyError:
                pass

        self.ts.positions = np.array(coords, dtype=np.float32)

        if self.convert_units:
            # in-place !
            self.convert_pos_from_native(self.ts._pos)
        self.ts.frame = frame

        return self.ts

    def _reopen(self):
        # Make frame think it's before start, so calling next
        # reads first frame
        self.ts.frame = -1


class MOL2Writer(base.WriterBase):
    """mol2 format writer

    Write a file in the Tripos_ molecule structure format (MOL2_).

    Note
    ----
    :class:`MOL2Writer` can only be used to write out previously loaded MOL2 files.
    If you're trying to convert, for example, a PDB file to MOL you should
    use other tools, like rdkit_.

    Here is an example how to use rdkit_ to convert a PDB to MOL::

      from rdkit import Chem
      mol = Chem.MolFromPDBFile("molecule.pdb", removeHs=False)
      Chem.MolToMolFile(mol, "molecule.mol" )

    MOL2 writer is currently not available for rdkit master. It requires SYBYL
    atomtype generation.
    See page 7 for list of SYBYL atomtypes
    (http://tripos.com/tripos_resources/fileroot/pdfs/mol2_format2.pdf).


    .. _rdkit: http://www.rdkit.org/docs/GettingStartedInPython.html

    .. versionchanged:: 0.11.0
       Frames now 0-based instead of 1-based

    """

    format = "MOL2"
    multiframe = True
    units = {"time": None, "length": "Angstrom"}

    def __init__(self, filename, n_atoms=None, convert_units=True):
        """Create a new MOL2Writer

        Parameters
        ----------
        filename: str
            name of output file
        convert_units: bool (optional)
            units are converted to the MDAnalysis base format; [``True``]
        """
        self.filename = filename
        self.convert_units = (
            convert_units  # convert length and time to base units
        )

        self.frames_written = 0

        self.file = util.anyopen(self.filename, "w")  # open file on init

    def close(self):
        self.file.close()

    def encode_block(self, obj):
        """
        Parameters
        ----------
        obj : AtomGroup or Universe
        """
        # Issue 2717
        try:
            obj = obj.atoms
        except AttributeError:
            errmsg = "Input obj is neither an AtomGroup or Universe"
            raise TypeError(errmsg) from None

        traj = obj.universe.trajectory
        ts = traj.ts

        try:
            molecule = ts.data["molecule"]
        except KeyError:
            errmsg = "MOL2Writer cannot currently write non MOL2 data"
            raise NotImplementedError(errmsg) from None

        # Need to remap atom indices to 1 based in this selection
        mapping = {a: i for i, a in enumerate(obj.atoms, start=1)}

        # only write bonds if the Bonds attribute exists (Issue #3057)
        if hasattr(obj, "bonds"):
            # Grab only bonds between atoms in the obj
            # ie none that extend out of it
            bondgroup = obj.intra_bonds
            bonds = sorted((b[0], b[1], b.order) for b in bondgroup)
            bond_lines = ["@<TRIPOS>BOND"]
            bls = [
                "{0:>5} {1:>5} {2:>5} {3:>2}".format(
                    bid, mapping[atom1], mapping[atom2], order
                )
                for bid, (atom1, atom2, order) in enumerate(bonds, start=1)
            ]
            bond_lines.extend(bls)
            bond_lines.append("\n")
            bond_lines = "\n".join(bond_lines)
        else:
            bondgroup = []
            bond_lines = ""

        atom_lines = ["@<TRIPOS>ATOM"]
        atom_lines.extend(
            "{0:>4} {1:>4} {2:>13.4f} {3:>9.4f} {4:>9.4f}"
            " {5:>4} {6} {7} {8:>7.4f}"
            "".format(
                mapping[a],
                a.name,
                a.position[0],
                a.position[1],
                a.position[2],
                a.type,
                a.resid,
                a.resname,
                a.charge,
            )
            for a in obj.atoms
        )
        atom_lines.append("\n")
        atom_lines = "\n".join(atom_lines)

        try:
            substructure = ["@<TRIPOS>SUBSTRUCTURE\n"] + ts.data[
                "substructure"
            ]
        except KeyError:
            substructure = ""

        check_sums = molecule[1].split()
        check_sums[0], check_sums[1] = str(len(obj.atoms)), str(len(bondgroup))

        # prevent behavior change between repeated calls
        # see gh-2678
        molecule_0_store = molecule[0]
        molecule_1_store = molecule[1]

        molecule[1] = "{0}\n".format(" ".join(check_sums))
        molecule.insert(0, "@<TRIPOS>MOLECULE\n")

        return_val = (
            "".join(molecule) + atom_lines + bond_lines + "".join(substructure)
        )

        molecule[0] = molecule_0_store
        molecule[1] = molecule_1_store
        return return_val

    def _write_next_frame(self, obj):
        """Write a new frame to the MOL2 file.

        Parameters
        ----------
        obj : AtomGroup or Universe


        .. versionchanged:: 1.0.0
            Renamed from `write_next_timestep` to `_write_next_frame`.
        """
        block = self.encode_block(obj)
        self.file.writelines(block)