#!/usr/bin/env python
#
#

__author__ = "Randall J. Radmer"
__version__ = "1.0"
__doc__ = """Load AMBER prep files."""


import sys, os
import getopt
#

class Parser:
    """Prep file object"""
    def __init__(self, filename):
        self.dPrepResidueTitles={}
        self.dPrepAtomNameToIndex={}
        self.dPrepProps={}
        self.dPrepZMat={}
        self.dPrepConnected={}
        self.dPrepImpropers={}
        self.mainAtoms={}

        dummList={}

        fIn = open(filename)
        mode='startLine1'
        for line0 in fIn:
            line=line0.rstrip()
            if mode=='startLine1':
                IDBGEN, IREST, ITYPF = line.split()
                idbgen=int(IDBGEN)
                irest=int(IREST)
                self.itypf=int(ITYPF)
                mode='startLine2'
            elif mode=='startLine2':
                NAMDBF=line
                mode='newResLine1'
            elif mode=='newResLine1':
                if line=='STOP':
                    break
                currentResidueTitle=line
                mode='newResLine2'
            elif mode=='newResLine2':
                NAMF=line
                mode='newResLine3'
            elif mode=='newResLine3':
                NAMRES, INTX, KFORM = line.split()
                if 'int' != INTX.lower():
                    raise Exception, 'Error: INTX not equal to "INT" -- only internal coords currently supported.'
                self.dPrepResidueTitles[NAMRES]=currentResidueTitle
                self.dPrepAtomNameToIndex[NAMRES]={}
                self.dPrepProps[NAMRES]={}
                self.dPrepZMat[NAMRES]={}
                self.dPrepConnected[NAMRES]=[]
                self.dPrepImpropers[NAMRES]=[]
                dummList[NAMRES]=[]
                mode='newResLine4'
            elif mode=='newResLine4':
                IFIXC, IOMIT, ISYMDU, IPOS = line.split()
                if 'corr' != IFIXC[:4].lower():
                    raise Exception, 'Error: IFIXC not equal to "CORR" -- only internal coords currently supported.'
                mode='newResLine5'
            elif mode=='newResLine5':
                CUT = float(line)
                if CUT>0:
                    raise Exception, 'Error: CUT is greater than zero -- use explicit loop closing.'
                mode='readAtoms'
            elif mode=='readAtoms':
                if line=='':
                    mode='postReadAtoms'
                    continue
                try:
                    (I, IGRAPH, ISYMBL, ITREE,
                     NA, NB, NC, R, THETA, PHI, CHG) = line.split()
                    chg=float(CHG)
                except ValueError:
                    (I, IGRAPH, ISYMBL, ITREE,
                     NA, NB, NC, R, THETA, PHI) = line.split()
                    if IGRAPH=='DUMM' and ISYMBL=='DU':
                        chg=0.0
                    else:
                        chg=None

                i=int(I)
                self.dPrepAtomNameToIndex[NAMRES][IGRAPH]=i
                self.dPrepProps[NAMRES][i]=(IGRAPH, ISYMBL, ITREE, chg)

                na=int(NA)
                nb=int(NB)
                nc=int(NC)
                r=float(R)
                theta=float(THETA)
                phi=float(PHI)
                self.dPrepZMat[NAMRES][(i, na, nb, nc)]=(r, theta, phi, chg)

                if ITREE=='M' and IGRAPH!='DUMM':
                    if NAMRES not in self.mainAtoms:
                        self.mainAtoms[NAMRES]=[i]
                    else:
                        self.mainAtoms[NAMRES].append(i)

                if IGRAPH=='DUMM':
                    dummList[NAMRES].append(i)

                if i  not in dummList[NAMRES] and \
                   na not in dummList[NAMRES] and \
                   i>0 and na > 0:
                    self.dPrepConnected[NAMRES].append((min(i,na), max(i,na)))

            elif mode=='postReadAtoms' and line=='':
                continue
            elif mode=='postReadAtoms' and line=='IMPROPER':
                mode='readImproper'
                continue
            elif mode=='postReadAtoms' and line=='LOOP':
                mode='readLoop'
                continue
            elif mode=='postReadAtoms' and line=='CHARGE':
                mode='readCharge'
                continue
            elif mode=='postReadAtoms' and line=='DONE':
                mode='newResLine1'
                continue
            elif mode=='postReadAtoms':
                raise Exception, 'Error: Bad postReadAtoms mode at line: %s' % line
            elif mode=='readImproper':
                if line=='':
                    mode='postReadAtoms'
                    continue
                NA, NB, NC, ND = line.split()
                try:
                    na=self.dPrepAtomNameToIndex[NAMRES][NA]
                except KeyError:
                    na=NA
                try:
                    nb=self.dPrepAtomNameToIndex[NAMRES][NB]
                except KeyError:
                    nb=NB
                try:
                    nc=self.dPrepAtomNameToIndex[NAMRES][NC]
                except KeyError:
                    nc=NC
                try:
                    nd=self.dPrepAtomNameToIndex[NAMRES][ND]
                except KeyError:
                    nd=ND
                item = (na, nb, nc, nd)
                self.dPrepImpropers[NAMRES].append(item)
            elif mode=='readLoop':
                if line=='':
                    mode='postReadAtoms'
                    continue
                NA, NB, = line.split()
                na=self.dPrepAtomNameToIndex[NAMRES][NA]
                nb=self.dPrepAtomNameToIndex[NAMRES][NB]
                self.dPrepConnected[NAMRES].append((min(na,nb), max(na,nb)))
            elif mode=='readCharge':
                if line=='':
                    mode='postReadAtoms'
                    continue
                atomIndexList=sorted(self.dPrepProps[NAMRES].keys())
                for CHG in line.split():
                    chg=float(CHG)
                    for i in atomIndexList:
                        (IGRAPH, ISYMBL, ITREE,
                         refChg) = self.dPrepProps[NAMRES][i]
                        if refChg is None:
                            self.dPrepProps[NAMRES][i] = (IGRAPH, ISYMBL, ITREE, chg)
                            break
            else:
                raise Exception, 'Error: Bad "mode" value: %s' % mode


    def getMainAtoms(self, resName):
        return self.mainAtoms[resName]

    def getNumAtomsByResName(self, resName):
        return len(self.dPrepProps[resName])

    def getConnectedAtomPairs(self, resName):
        return self.dPrepConnected[resName]

    def getAtomPropertiesByIndex(self, resName, atomIndex):
        return self.dPrepProps[resName][atomIndex]

    def getAtomPropertiesByName(self, resName, atomName):
        try:
            atomIndex=self.dPrepAtomNameToIndex[resName][atomName]
        except KeyError:
            raise Exception, 'Error: Unknown residue and atom combinations: %s -- %s' \
                 % (resName, atomName)
        (IGRAPH, ISYMBL, ITREE, chg) = self.dPrepProps[resName][atomIndex]
        return (atomIndex, ISYMBL, ITREE, chg)


def parseCommandLine():
    opts, args_proper = getopt.getopt(sys.argv[1:], 'h')
    for option, parameter in opts:
        if option=='-h': usageError()
    return (args_proper)

def main():
    args_proper = parseCommandLine()
    try:
        filename = args_proper[0]
    except IndexError:
        usageError()

    return Parser(filename)

def usageError():
    print 'usage: %s inputPrepFilename' \
         % os.path.basename(sys.argv[0])
    sys.exit(1)

if __name__=='__main__':
    main()