"""
    Routines for PDB2PQR

    This module contains the protein object used in PDB2PQR and associated
    methods
    
    ----------------------------
   
    PDB2PQR -- An automated pipeline for the setup, execution, and analysis of
    Poisson-Boltzmann electrostatics calculations

    Copyright (c) 2002-2011, Jens Erik Nielsen, University College Dublin; 
    Nathan A. Baker, Battelle Memorial Institute, Developed at the Pacific 
    Northwest National Laboratory, operated by Battelle Memorial Institute, 
    Pacific Northwest Division for the U.S. Department Energy.; 
    Paul Czodrowski & Gerhard Klebe, University of Marburg.

	All rights reserved.

	Redistribution and use in source and binary forms, with or without modification, 
	are permitted provided that the following conditions are met:

		* Redistributions of source code must retain the above copyright notice, 
		  this list of conditions and the following disclaimer.
		* Redistributions in binary form must reproduce the above copyright notice, 
		  this list of conditions and the following disclaimer in the documentation 
		  and/or other materials provided with the distribution.
        * Neither the names of University College Dublin, Battelle Memorial Institute,
          Pacific Northwest National Laboratory, US Department of Energy, or University
          of Marburg nor the names of its contributors may be used to endorse or promote
          products derived from this software without specific prior written permission.

	THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 
	ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 
	WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. 
	IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, 
	INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, 
	BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 
	DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF 
	LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE 
	OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED 
	OF THE POSSIBILITY OF SUCH DAMAGE.

    ----------------------------

"""

__date__ = "13 May 2008"
__author__ = "Todd Dolinsky, Yong Huang"

from pdb import *
from structures import *
from aa import *
from na import *


class Protein:
    """
    Protein class

    The protein class represents the parsed PDB, and provides a
    hierarchy of information - each Protein contains a list of Chain
    objects as provided in the PDB file.  Each Chain then contains its
    associated list of Residue objects, and each Residue contains a list
    of Atom objects, completing the hierarchy.
    """

    def __init__(self, pdblist, definition):
        """
        Initialize using parsed PDB file

        Parameters
            pdblist: List of Classes of PDB lines as created
        """
        self.chainmap = {}
        self.chains = []
        self.residues = []
        self.referencemap = definition.map
        self.patchmap = definition.patches

        chainDict = {}
        previousAtom = None
        residue = []
        numModels = 0
        numChains = 1
        count = 0

        for record in pdblist:  # Find number of chains
            if isinstance(record, TER):
                numChains += 1

        for record in pdblist:
            if isinstance(record, ATOM) or isinstance(record, HETATM):

                if (
                    record.chainID == ""
                    and numChains > 1
                    and record.resName not in ["WAT", "HOH"]
                ):
                    # Assign a chain ID
                    record.chainID = string.ascii_uppercase[count]

                chainID = record.chainID
                resSeq = record.resSeq
                resName = record.resName
                iCode = record.iCode

                if previousAtom == None:
                    previousAtom = record

                if chainID not in chainDict:
                    myChain = Chain(chainID)
                    chainDict[chainID] = myChain

                if (
                    resSeq != previousAtom.resSeq
                    or iCode != previousAtom.iCode
                    or chainID != previousAtom.chainID
                ):
                    myResidue = self.createResidue(residue, previousAtom.resName)
                    chainDict[previousAtom.chainID].addResidue(myResidue)
                    residue = []

                residue.append(record)
                previousAtom = record

            elif isinstance(record, END):
                myResidue = self.createResidue(residue, previousAtom.resName)
                chainDict[previousAtom.chainID].addResidue(myResidue)
                residue = []

            elif isinstance(record, MODEL):
                numModels += 1
                if residue == []:
                    continue
                if numModels > 1:
                    myResidue = self.createResidue(residue, previousAtom.resName)
                    chainDict[previousAtom.chainID].addResidue(myResidue)
                    break

            elif isinstance(record, TER):
                count += 1

        if residue != [] and numModels <= 1:
            myResidue = self.createResidue(residue, previousAtom.resName)
            chainDict[previousAtom.chainID].addResidue(myResidue)

        # Keep a map for accessing chains via chainID

        self.chainmap = chainDict.copy()

        # Make a list for sequential ordering of chains

        if chainDict.has_key(""):
            chainDict["ZZ"] = chainDict[""]
            del chainDict[""]

        keys = chainDict.keys()
        keys.sort()

        for key in keys:
            self.chains.append(chainDict[key])

        for chain in self.chains:
            # if chain.numResidues() == 1:
            # We cannot support Amino Acid chains with only one residue-
            # It is unclear whether they are Nterm, Cterm, or both.
            # residue = chain.residues[0]
            # if isinstance(residue, Amino):
            # raise ValueError, "Unable to support amino acid chains of only one residue (%s)" % residue

            for residue in chain.getResidues():
                self.residues.append(residue)

    def createResidue(self, residue, resname):
        """
        Create a residue object.  If the resname is a known residue
        type, try to make that specific object, otherwise just make
        a standard residue object.

        Parameters
            residue:  A list of atoms (list)
            resname:  The name of the residue (string)

        Returns:
            residue:  The residue object (Residue)
        """
        try:
            refobj = self.referencemap[resname]
            if refobj.name != resname:  # Patched!
                obj = "%s(residue, refobj)" % refobj.name
                residue = eval(obj)
                residue.reference = refobj
            else:
                obj = "%s(residue, refobj)" % resname
                residue = eval(obj)
        except (KeyError, NameError):
            residue = Residue(residue)
        return residue

    def printAtoms(self, atomlist, chainflag=False, pdbfile=False):
        """
        Get the text for the entire protein
        Parameters
            atomlist:  The list of atoms to include (list)
            chainflag: Flag whether to print chainid or not -
                          Defaults to False
        Returns
            text:      The list of (stringed) atoms (list)
        """
        self.reSerialize()
        text = []
        currentchainID = None
        for atom in atomlist:
            # Print the "TER" records between chains
            if currentchainID == None:
                currentchainID = atom.chainID
            elif atom.chainID != currentchainID:
                currentchainID = atom.chainID
                text.append("TER\n")

            if not atom.resName:
                continue

            if pdbfile == True:
                text.append("%s\n" % atom.getPDBString())
            else:
                text.append("%s\n" % atom.getPQRString(chainflag=chainflag))
        text.append("TER\nEND")
        return text

    def createHTMLTypeMap(self, definition, outfilename):
        """
        Create an HTML typemap file at the desired location. If a
        type cannot be found for an atom a blank is listed.

        Parameters
            definition: The definition objects.
            outfilename:  The name of the file to write (string)
        """
        from forcefield import Forcefield
        from aconf import STYLESHEET

        # Cache the initial atom numbers
        numcache = {}
        for atom in self.getAtoms():
            numcache[atom] = atom.serial
        self.reSerialize()

        amberff = Forcefield("amber", definition, None)
        charmmff = Forcefield("charmm", definition, None)

        file = open(outfilename, "w")
        file.write("<HTML>\n")
        file.write("<HEAD>\n")
        file.write("<TITLE>PQR Typemap (beta)</TITLE>\n")
        file.write('<link rel="stylesheet" href="%s" type="text/css">\n' % STYLESHEET)
        file.write("</HEAD>\n")
        file.write("<BODY>\n")
        file.write(
            "<H3>This is a developmental page including the atom type for the atoms in the PQR file.</H3><P>\n"
        )
        file.write("<TABLE CELLSPACING=2 CELLPADDING=2 BORDER=1>\n")
        file.write(
            "<tr><th>Atom Number</th><th>Atom Name</th><th>Residue Name</th><th>Chain ID</th><th>AMBER Atom Type</th><th>CHARMM Atom Type</th></tr>\n"
        )

        for atom in self.getAtoms():
            if isinstance(atom.residue, (Amino, WAT, Nucleic)):
                resname = atom.residue.ffname
            else:
                resname = atom.residue.name

            ambergroup = amberff.getGroup(resname, atom.name)
            charmmgroup = charmmff.getGroup(resname, atom.name)

            file.write(
                "<tr><td>%s</td><td>%s</td><td>%s</td><td>%s</td><td>%s</td><td>%s</td></tr>\n"
                % (
                    atom.serial,
                    atom.name,
                    resname,
                    atom.chainID,
                    ambergroup,
                    charmmgroup,
                )
            )

        file.write("</table>\n")
        file.write("</BODY></HTML>\n")
        file.close()

        # Return the original numbers back
        for atom in self.getAtoms():
            atom.serial = numcache[atom]

        del numcache
        del amberff
        del charmmff

    def reSerialize(self):
        """
        Generate new serial numbers for atoms in the protein
        """
        count = 1
        for atom in self.getAtoms():
            atom.set("serial", count)
            count += 1

    def getResidues(self):
        """
        Return the list of residues in the entire protein
        """
        return self.residues

    def numResidues(self):
        """
        Get the number of residues for the entire protein (including
        multiple chains)

        Returns
            count:  Number of residues in the protein (int)
        """
        return len(self.getResidues())

    def numAtoms(self):
        """
        Get the number of atoms for the entire protein(including
        multiple chains)
        """
        return len(self.getAtoms())

    def getAtoms(self):
        """
        Return all Atom objects in list format

        Returns
            atomlist:  List of Atom objects in the protein (list)
        """

        atomlist = []
        for chain in self.chains:
            for atom in chain.getAtoms():
                atomlist.append(atom)
        return atomlist

    def getCharge(self):
        """
        Get the total charge on the protein
        NOTE:  Since the misslist is used to identify incorrect
               charge assignments, this routine does not list the
               3 and 5 termini of nucleic acid chains as having
               non-integer charge even though they are (correctly)
               non-integer.
        Returns:
            misslist: List of residues with non-integer
                      charges (list)
            charge:   The total charge on the protein (float)
        """
        charge = 0.0
        misslist = []
        for chain in self.chains:
            for residue in chain.get("residues"):
                rescharge = residue.getCharge()
                charge += rescharge
                if isinstance(residue, Nucleic):
                    if residue.is3term or residue.is5term:
                        continue
                if float("%i" % rescharge) != rescharge:
                    misslist.append(residue)
        return misslist, charge

    def getChains(self):
        """
        Get the chains object

        Returns
            chains: The list of chains in the protein (chain)
        """
        return self.chains

    def getSummary(self):
        output = []
        for chain in self.chains:
            output.append(chain.getSummary())
        return " ".join(output)
