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


"""
PDBQT structure files in MDAnalysis --- :mod:`MDAnalysis.coordinates.PDBQT`
===========================================================================

MDAnalysis reads coordinates from PDBQT_ files and additional optional
data such as B-factors, partial charge and AutoDock_ atom types.  It
is also possible to substitute a PDBQT file for a PSF file in order to
define the list of atoms (but no connectivity information will be
available in this case).

.. _PDBQT:
   https://autodock.scripps.edu/wp-content/uploads/sites/56/2021/10/AutoDock4.2.6_UserGuide.pdf
.. _AutoDock:
   http://autodock.scripps.edu/
"""

import os
import errno
import itertools
import numpy as np
import warnings

from ..lib import util
from . import base


class PDBQTReader(base.SingleFrameReaderBase):
    """PDBQTReader that reads a PDBQT-formatted file, no frills.

    Records read:
     - CRYST1 for unitcell A,B,C, alpha,beta,gamm
     - ATOM. HETATM for x,y,z

    Original `PDB format documentation`_ with `AutoDOCK extensions`_


    .. _PDB format documentation:
       http://www.wwpdb.org/documentation/file-format-content/format32/v3.2.html
    .. _AutoDOCK extensions:
       http://autodock.scripps.edu/faqs-help/faq/what-is-the-format-of-a-pdbqt-file

    =============  ============  ===========  =============================================
    COLUMNS        DATA  TYPE    FIELD        DEFINITION
    =============  ============  ===========  =============================================
    1 -  6         Record name   "CRYST1"
    7 - 15         Real(9.3)     a              a (Angstroms).
    16 - 24        Real(9.3)     b              b (Angstroms).
    25 - 33        Real(9.3)     c              c (Angstroms).
    34 - 40        Real(7.2)     alpha          alpha (degrees).
    41 - 47        Real(7.2)     beta           beta (degrees).
    48 - 54        Real(7.2)     gamma          gamma (degrees).

    1 -  6         Record name   "ATOM  "
    7 - 11         Integer       serial       Atom  serial number.
    13 - 16        Atom          name         Atom name.
    17             Character     altLoc       Alternate location indicator. IGNORED
    18 - 21        Residue name  resName      Residue name.
    22             Character     chainID      Chain identifier.
    23 - 26        Integer       resSeq       Residue sequence number.
    27             AChar         iCode        Code for insertion of residues. IGNORED
    31 - 38        Real(8.3)     x            Orthogonal coordinates for X in Angstroms.
    39 - 46        Real(8.3)     y            Orthogonal coordinates for Y in Angstroms.
    47 - 54        Real(8.3)     z            Orthogonal coordinates for Z in Angstroms.
    55 - 60        Real(6.2)     occupancy    Occupancy.
    61 - 66        Real(6.2)     tempFactor   Temperature  factor.
    67 - 70        LString(4)    footnote     Usually blank. IGNORED.
    71 - 76        Real(6.4)     partialChrg  Gasteiger PEOE partial charge *q*.
    79 - 80        LString(2)    atomType     AutoDOCK atom type *t*.
    =============  ============  ===========  =============================================

    We ignore torsion notation and just pull the partial charge and atom type columns::

         COMPND    NSC7810
         REMARK  3 active torsions:
         REMARK  status: ('A' for Active; 'I' for Inactive)
         REMARK    1  A    between atoms: A7_7  and  C22_23
         REMARK    2  A    between atoms: A9_9  and  A11_11
         REMARK    3  A    between atoms: A17_17  and  C21_21
         ROOT
         123456789.123456789.123456789.123456789.123456789.123456789.123456789.123456789. (column reference)
         ATOM      1  A1  INH I           1.054   3.021   1.101  0.00  0.00     0.002 A
         ATOM      2  A2  INH I           1.150   1.704   0.764  0.00  0.00     0.012 A
         ATOM      3  A3  INH I          -0.006   0.975   0.431  0.00  0.00    -0.024 A
         ATOM      4  A4  INH I           0.070  -0.385   0.081  0.00  0.00     0.012 A
         ATOM      5  A5  INH I          -1.062  -1.073  -0.238  0.00  0.00     0.002 A
         ATOM      6  A6  INH I          -2.306  -0.456  -0.226  0.00  0.00     0.019 A
         ATOM      7  A7  INH I          -2.426   0.885   0.114  0.00  0.00     0.052 A
         ATOM      8  A8  INH I          -1.265   1.621   0.449  0.00  0.00     0.002 A
         ATOM      9  A9  INH I          -1.339   2.986   0.801  0.00  0.00    -0.013 A
         ATOM     10  A10 INH I          -0.176   3.667   1.128  0.00  0.00     0.013 A
         ENDROOT
         BRANCH   9  11
         ATOM     11  A11 INH I          -2.644   3.682   0.827  0.00  0.00    -0.013 A
         ATOM     12  A16 INH I          -3.007   4.557  -0.220  0.00  0.00     0.002 A
         ATOM     13  A12 INH I          -3.522   3.485   1.882  0.00  0.00     0.013 A
         ATOM     14  A15 INH I          -4.262   5.209  -0.177  0.00  0.00    -0.024 A
         ATOM     15  A17 INH I          -2.144   4.784  -1.319  0.00  0.00     0.052 A
         ATOM     16  A14 INH I          -5.122   4.981   0.910  0.00  0.00     0.012 A
         ATOM     17  A20 INH I          -4.627   6.077  -1.222  0.00  0.00     0.012 A
         ATOM     18  A13 INH I          -4.749   4.135   1.912  0.00  0.00     0.002 A
         ATOM     19  A19 INH I          -3.777   6.285  -2.267  0.00  0.00     0.002 A
         ATOM     20  A18 INH I          -2.543   5.650  -2.328  0.00  0.00     0.019 A
         BRANCH  15  21
         ATOM     21  C21 INH I          -0.834   4.113  -1.388  0.00  0.00     0.210 C
         ATOM     22  O1  INH I          -0.774   2.915  -1.581  0.00  0.00    -0.644 OA
         ATOM     23  O3  INH I           0.298   4.828  -1.237  0.00  0.00    -0.644 OA
         ENDBRANCH  15  21
         ENDBRANCH   9  11
         BRANCH   7  24
         ATOM     24  C22 INH I          -3.749   1.535   0.125  0.00  0.00     0.210 C
         ATOM     25  O2  INH I          -4.019   2.378  -0.708  0.00  0.00    -0.644 OA
         ATOM     26  O4  INH I          -4.659   1.196   1.059  0.00  0.00    -0.644 OA
         ENDBRANCH   7  24
         TORSDOF 3
         123456789.123456789.123456789.123456789.123456789.123456789.123456789.123456789. (column reference)

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

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

    def _read_first_frame(self):
        coords = []
        unitcell = np.zeros(6, dtype=np.float32)
        with util.openany(self.filename) as pdbfile:
            for line in pdbfile:
                # Should only break at the 'END' of a model definition
                # and prevent premature exit for a torsion termination
                # , eg, ENDBRANCH
                if line.startswith("END\n"):
                    break
                if line.startswith("CRYST1"):
                    # lengths
                    x, y, z = np.float32(
                        (line[6:15], line[15:24], line[24:33])
                    )
                    # angles
                    A, B, G = np.float32(
                        (line[33:40], line[40:47], line[47:54])
                    )
                    unitcell[:] = x, y, z, A, B, G
                if line.startswith(("ATOM", "HETATM")):
                    # convert all entries at the end once for optimal speed
                    coords.append([line[30:38], line[38:46], line[46:54]])
        self.n_atoms = len(coords)
        self.ts = self._Timestep.from_coordinates(coords, **self._ts_kwargs)
        self.ts.dimensions = unitcell
        self.ts.frame = 0  # 0-based frame number
        if self.convert_units:
            # in-place !
            self.convert_pos_from_native(self.ts._pos)
            if self.ts.dimensions is not None:
                self.convert_pos_from_native(self.ts.dimensions[:3])

    def Writer(self, filename, **kwargs):
        """Returns a permissive (simple) PDBQTWriter for *filename*.

        Parameters
        ----------
        filename : str
            filename of the output PDBQT file

        Returns
        -------
        :class:`PDBQTWriter`

        """
        return PDBQTWriter(filename, **kwargs)


class PDBQTWriter(base.WriterBase):
    """PDBQT writer that implements a subset of the PDB_ 3.2 standard and the PDBQT_ spec.

    .. _PDB: http://www.wwpdb.org/documentation/file-format-content/format32/v3.2.html
    .. _PDBQT: https://autodock.scripps.edu/wp-content/uploads/sites/56/2021/10/AutoDock4.2.6_UserGuide.pdf

    .. versionchanged:: 2.6.0
       Files are now written in `wt` mode, and keep extensions, allowing
       for files to be written under compressed formats
    """

    fmt = {
        "ATOM": (
            "ATOM  {serial:5d} {name:<4.4s} {resName:<4.4s}"
            "{chainID:1.1s}{resSeq:4d}{iCode:1.1s}"
            "   {pos[0]:8.3f}{pos[1]:8.3f}{pos[2]:8.3f}{occupancy:6.2f}"
            "{tempFactor:6.2f}    {charge:< 1.3f} {element:<2.2s}\n"
        ),
        "REMARK": "REMARK     {0}\n",
        "TITLE": "TITLE     {0}\n",
        "CRYST1": (
            "CRYST1{box[0]:9.3f}{box[1]:9.3f}{box[2]:9.3f}"
            "{ang[0]:7.2f}{ang[1]:7.2f}{ang[2]:7.2f} "
            "{spacegroup:<11s}{zvalue:4d}\n"
        ),
    }
    format = "PDBQT"
    units = {"time": None, "length": "Angstrom"}
    pdb_coor_limits = {"min": -999.9995, "max": 9999.9995}

    def __init__(self, filename, **kwargs):
        self.filename = util.filename(filename, ext="pdbqt", keep=True)
        self.pdb = util.anyopen(self.filename, "wt")

    def close(self):
        self.pdb.close()

    def write(self, selection, frame=None):
        """Write selection at current trajectory frame to file.

        Parameters
        ----------
        selection : AtomGroup
            The selection to be written
        frame : int (optional)
            optionally move to frame index `frame` before writing; the default
            is to write the current trajectory frame

        Note
        ----
        The last letter of the
        :attr:`~MDAnalysis.core.groups.Atom.segid` is used as the PDB
        chainID.


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

        """
        try:
            u = selection.universe
        except AttributeError:
            errmsg = "Input obj is neither an AtomGroup or Universe"
            raise TypeError(errmsg) from None

        if frame is not None:
            u.trajectory[frame]  # advance to frame
        else:
            try:
                frame = u.trajectory.ts.frame
            except AttributeError:
                frame = 0  # should catch cases when we are analyzing a single PDB (?)

        atoms = selection.atoms  # make sure to use atoms (Issue 46)
        coor = (
            atoms.positions
        )  # can write from selection == Universe (Issue 49)

        # Check attributes
        attrs = {}
        missing_topology = []
        for attr, dflt in (
            ("altLocs", " "),
            ("charges", 0.0),
            ("icodes", " "),
            ("names", "X"),
            ("occupancies", 1.0),
            ("resids", 1),
            ("resnames", "UNK"),
            ("tempfactors", 0.0),
            ("types", "  "),
        ):
            try:
                attrs[attr] = getattr(atoms, attr)
            except AttributeError:
                attrs[attr] = itertools.cycle((dflt,))
                missing_topology.append(attr)
        # Order of preference: chainids -> segids -> blank string
        try:
            attrs["chainids"] = atoms.chainids
        except AttributeError:
            try:
                attrs["chainids"] = atoms.segids
            except AttributeError:
                attrs["chainids"] = itertools.cycle((" ",))
                missing_topology.append("chainids")
        if missing_topology:
            warnings.warn(
                "Supplied AtomGroup was missing the following attributes: "
                "{miss}. These will be written with default values. "
                "".format(miss=", ".join(missing_topology))
            )

        # check if any coordinates are illegal (coordinates are already
        # in Angstroem per package default)
        if not self.has_valid_coordinates(self.pdb_coor_limits, coor):
            self.close()
            try:
                os.remove(self.filename)
            except OSError as err:
                if err.errno == errno.ENOENT:
                    pass
            raise ValueError(
                "PDB files must have coordinate values between {0:.3f}"
                " and {1:.3f} Angstroem: No file was written."
                "".format(
                    self.pdb_coor_limits["min"], self.pdb_coor_limits["max"]
                )
            )

        # Write title record
        # http://www.wwpdb.org/documentation/file-format-content/format32/sect2.html
        line = "FRAME " + str(frame) + " FROM " + str(u.trajectory.filename)
        self.pdb.write(self.fmt["TITLE"].format(line))

        # Write CRYST1 record
        # http://www.wwpdb.org/documentation/file-format-content/format32/sect8.html
        box = self.convert_dimensions_to_unitcell(u.trajectory.ts)
        self.pdb.write(
            self.fmt["CRYST1"].format(
                box=box[:3], ang=box[3:], spacegroup="P 1", zvalue=1
            )
        )

        # Write atom records
        # http://www.wwpdb.org/documentation/file-format-content/format32/sect9.html
        for serial, (
            pos,
            name,
            resname,
            chainid,
            resid,
            icode,
            occupancy,
            tempfactor,
            charge,
            element,
        ) in enumerate(
            zip(
                coor,
                attrs["names"],
                attrs["resnames"],
                attrs["chainids"],
                attrs["resids"],
                attrs["icodes"],
                attrs["occupancies"],
                attrs["tempfactors"],
                attrs["charges"],
                attrs["types"],
            ),
            start=1,
        ):
            serial = util.ltruncate_int(serial, 5)  # check for overflow here?
            resid = util.ltruncate_int(resid, 4)
            name = name[:4]
            if len(name) < 4:
                name = " " + name  # customary to start in column 14
            chainid = chainid.strip()[-1:]  # take the last character

            self.pdb.write(
                self.fmt["ATOM"].format(
                    serial=serial,
                    name=name,
                    resName=resname,
                    chainID=chainid,
                    resSeq=resid,
                    iCode=icode,
                    pos=pos,
                    occupancy=occupancy,
                    tempFactor=tempfactor,
                    charge=charge,
                    element=element,
                )
            )

        self.close()