#!/usr/bin/env python
# pmx  Copyright Notice
# ============================
#
# The pmx source code is copyrighted, but you can freely use and
# copy it as long as you don't change or remove any of the copyright
# notices.
#
# ----------------------------------------------------------------------
# pmx is Copyright (C) 2006-2013 by Daniel Seeliger
#
#                        All Rights Reserved
#
# Permission to use, copy, modify, distribute, and distribute modified
# versions of this software and its documentation for any purpose and
# without fee is hereby granted, provided that the above copyright
# notice appear in all copies and that both the copyright notice and
# this permission notice appear in supporting documentation, and that
# the name of Daniel Seeliger not be used in advertising or publicity
# pertaining to distribution of the software without specific, written
# prior permission.
#
# DANIEL SEELIGER DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS
# SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND
# FITNESS.  IN NO EVENT SHALL DANIEL SEELIGER BE LIABLE FOR ANY
# SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER
# RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF
# CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
# CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
# ----------------------------------------------------------------------

import sys,os
from pmx import *
from pmx.options import *
from pmx.parser import *
from pmx import library
from pmx.mutdb import *
from pmx.geometry import *

amber_termini_full = {
    'NALA':'ALA',
    'NGLY':'GLY',
    'NSER':'SER',
    'NTHR':'THR',
    'NLEU':'LEU',
    'NILE':'ILE',
    'NVAL':'VAL',
    'NASN':'ASN',
    'NGLN':'GLN',
    'NARG':'ARG',
    'NHID':'HID',
    'NHIE':'HIE',
    'NHIP':'HIP',
    'NTRP':'TRP',
    'NPHE':'PHE',
    'NTYR':'TYR',
    'NGLU':'GLU',
    'NASP':'ASP',
    'NLYP':'LYP',
    'NPRO':'PRO',
    'NCYS':'CYS',
    'NCYN':'CYN',
    'NCYX':'CYX',
    'NMET':'MET',
    'CALA':'ALA',
    'CGLY':'GLY',
    'CSER':'SER',
    'CTHR':'THR',
    'CLEU':'LEU',
    'CILE':'ILE',
    'CVAL':'VAL',
    'CASN':'ASN',
    'CGLN':'GLN',
    'CARG':'ARG',
    'CHID':'HID',
    'CHIE':'HIE',
    'CHIP':'HIP',
    'CTRP':'TRP',
    'CPHE':'PHE',
    'CTYR':'TYR',
    'CGLU':'GLU',
    'CASP':'ASP',
    'CLYP':'LYP',
    'CPRO':'PRO',
    'CCYS':'CYS',
    'CCYN':'CYN',
    'CCYX':'CYX',
    'CMET':'MET'
}

amber_termini_partial = {
    'NAL':'ALA',
    'NSE':'SER',
    'NTH':'THR',
    'NLE':'LEU',
    'NIL':'ILE',
    'NVA':'VAL',
    'NAR':'ARG',
    'NHI':'HIS',
    'NTR':'TRP',
    'NPH':'PHE',
    'NTY':'TYR',
    'NLY':'LYS',
    'NPR':'PRO',
    'NCY':'CYS',
    'CAL':'ALA',
    'CSE':'SER',
    'CTH':'THR',
    'CLE':'LEU',
    'CIL':'ILE',
    'CVA':'VAL',
    'CAR':'ARG',
    'CHI':'HIS',
    'CTR':'TRP',
    'CPH':'PHE',
    'CTY':'TYR',
    'CLY':'LYS',
    'CPR':'PRO',
    'CCY':'CYS',
    'CME':'MET',
}

def identify_nas_cas(r):
    newname = 'ASP'
    for a in r.atoms:
        if a.name=='ND2':
            return('ASN')
        if a.name=='OD2':
            return('ASP')
    return(newname)

def identify_gly_glu_gln(r):
    newname = 'GLY'
    for a in r.atoms:
        if a.name=='OE2':
            return('GLU')
        if a.name=='NE2':
            return('GLN')
    return(newname)

def identify_nmet_nme(r):
    newname = 'NME'
    for a in r.atoms:
        if 'S' in a.name:
            return('MET')
    return(newname)

def fix_amber_termini(m,ff):
    # get the termini
    termini = []
    for ch in m.chains:
        termini.append(ch.residues[0])
        termini.append(ch.residues[-1])

    # check (and fix) the names
    dict_full = amber_termini_full.keys()
    dict_part = amber_termini_partial.keys()
    for term in termini:
	# terminal name in the 4 letter dictionary
        if term.resname in dict_full:
            term.set_resname(amber_termini_full[term.resname])
            continue
	# terminal name in the 3 letter dictionary
        if term.resname in dict_part:
            term.set_resname(amber_termini_partial[term.resname])
            continue
	# other special cases
        if term.resname.startswith('NAS') or term.resname.startswith('CAS'):
            newname = identify_nas_cas(term)
            term.set_resname(newname)
            continue
        if term.resname.startswith('NCY') or term.resname.startswith('CCY'):
            newname = 'CYS'
            term.set_resname(newname)
            continue
        if term.resname.startswith('NGL') or term.resname.startswith('CGL'):
            newname = identify_gly_glu_gln(term)
            term.set_resname(newname)
            continue
        if term.resname.startswith('NHI') or term.resname.startswith('CHI'):
            newname = 'HIS'
            term.set_resname(newname)
            continue
        if term.resname.startswith('NLY') or term.resname.startswith('CLY'):
            newname = 'LYS'
            term.set_resname(newname)
            continue
        if term.resname.startswith('NME'):
            newname = identify_nmet_nme(term)
            term.set_resname(newname)
            continue

def rename_lowlevel(r,name):
    for a in r.atoms:
        a.resname = name

def rename_hie(r,ff):
    if ff.startswith('charmm'):
        newname = 'HSE'
    if ff.startswith('amber'):
        newname = 'HIE'
    if ff.startswith('opls'):
        newname = 'HISE'
    # actual rename
    rename_lowlevel(r,newname)

def rename_hid(r,ff):
    if ff.startswith('charmm'):
        newname = 'HSD'
    if ff.startswith('amber'):
        newname = 'HID'
    if ff.startswith('opls'):
        newname = 'HISD'
    # actual rename
    rename_lowlevel(r,newname)

def rename_hip(r,ff):
    if ff.startswith('charmm'):
        newname = 'HSP'
    if ff.startswith('amber'):
        newname = 'HIP'
    if ff.startswith('opls'):
        newname = 'HISH'
    # actual rename
    rename_lowlevel(r,newname)

def rename_his(r,ff):
    # identify histidine type
    bNd = False
    bNe = False
    for a in r.atoms:
        if a.name == 'HE2':
            bNe = True
        if a.name == 'HD1':
            bNd = True

    # no Hydrogen names
#    newname = 'HIE'
#    if ff.startswith('charmm'):
#	newname = 'HSE'
#    elif ff.startswith('opls'):
#	newname = 'HISE'

    # names with Hydrogens
    if (bNd==True) and (bNe==True):
        if ff.startswith('charmm'):
            newname = 'HSP'
        if ff.startswith('amber'):
            newname = 'HIP'
        if ff.startswith('opls'):
            newname = 'HISH'
        # actual rename
        rename_lowlevel(r,newname)
    elif (bNd==False) and (bNe==True):
        if ff.startswith('charmm'):
            newname = 'HSE'
        if ff.startswith('amber'):
            newname = 'HIE'
        if ff.startswith('opls'):
            newname = 'HISE'
        # actual rename
        rename_lowlevel(r,newname)
    elif (bNd==True) and (bNe==False):
        if ff.startswith('charmm'):
            newname = 'HSD'
        if ff.startswith('amber'):
            newname = 'HID'
        if ff.startswith('opls'):
            newname = 'HISD'
        # actual rename
        rename_lowlevel(r,newname)

def rename_lys(r,ff):
    # default names
    newname = 'LYS'
    if ff.startswith('amber') or ff.startswith('charmm'):
        newname = 'LYS'
    elif ff.startswith('opls'):
        newname = 'LYSH'

    # names with hydrogens
    bUncharged = False
    for a in r.atoms:
        if a.name == 'HZ2':
            bUncharged = True
        if a.name == 'HZ3':
            bUncharged = False
            break

    # if lysine neutral
    if bUncharged==True:
        if ff.startswith('amber'):
            newname = 'LYN'
        elif ff.startswith('charmm'):
            newname = 'LSN'
        elif ff.startswith('opls'):
            newname = 'LSN'

    # actual rename
    rename_lowlevel(r,newname)

def rename_asp(r,ff):
    # default names
    newname = 'ASP'

    # names with hydrogens
    bUncharged = False
    for a in r.atoms:
        if a.name == 'HD2':
            bUncharged = True
            break

    # if aspartate neutral
    if bUncharged==True:
        if ff.startswith('amber'):
            newname = 'ASH'
        elif ff.startswith('charmm'):
            newname = 'ASPP'
        elif ff.startswith('opls'):
            newname = 'ASPP'

    # actual rename
    rename_lowlevel(r,newname)

def rename_glu(r,ff):
    # default names
    newname = 'GLU'

    # names with hydrogens
    bUncharged = False
    for a in r.atoms:
        if a.name == 'HE2':
            bUncharged = True
            break

    # if glutamate neutral
    if bUncharged==True:
        if ff.startswith('amber'):
            newname = 'GLH'
        elif ff.startswith('charmm'):
            newname = 'GLUP'
        elif ff.startswith('opls'):
            newname = 'GLUP'

    # actual rename
    rename_lowlevel(r,newname)

def rename_cys(r,ff):
    # default names
    newname = 'CYS'
    if ff.startswith('opls'):
        newname = 'CYSH'

    # names with hydrogens
    bDisulfide = False
    for a in r.atoms:
        if a.name == 'HB1':
            bDisulfide = True
        if a.name.startswith('HG'):
            bDisulfide = False
            break

    # if disulfide
    if bDisulfide==True:
        if ff.startswith('amber'):
            newname = 'CYX'
        elif ff.startswith('charmm'):
            newname = 'CYS2'
        elif ff.startswith('opls'):
            newname = 'CYS2'

    # actual rename
    rename_lowlevel(r,newname)

def rename_cyx(r,ff):
    if ff.startswith('amber'):
        newname = 'CYX'
    elif ff.startswith('charmm'):
        newname = 'CYS2'
    elif ff.startswith('opls'):
        newname = 'CYS2'
    # actual rename
    rename_lowlevel(r,newname)

def rename_lsn(r,ff):
    if ff.startswith('amber'):
        newname = 'LYN'
    elif ff.startswith('charmm'):
        newname = 'LSN'
    elif ff.startswith('opls'):
        newname = 'LSN'
    # actual rename
    rename_lowlevel(r,newname)

def rename_aspp(r,ff):
    if ff.startswith('amber'):
        newname = 'ASH'
    elif ff.startswith('charmm'):
        newname = 'ASPP'
    elif ff.startswith('opls'):
        newname = 'ASPP'
    # actual rename
    rename_lowlevel(r,newname)

def rename_glup(r,ff):
    if ff.startswith('amber'):
        newname = 'GLH'
    elif ff.startswith('charmm'):
        newname = 'GLUP'
    elif ff.startswith('opls'):
        newname = 'GLUP'
    # actual rename
    rename_lowlevel(r,newname)

def rename_res(m,ff):
    for r in m.residues:
        if r.resname == 'HIS':
            rename_his(r,ff)
        if r.resname == 'HIE' or r.resname == 'HISE' or r.resname == 'HSE':
            rename_hie(r,ff)
        if r.resname == 'HID' or r.resname == 'HISD' or r.resname == 'HSD':
            rename_hid(r,ff)
        if r.resname == 'HIP' or r.resname == 'HISH' or r.resname == 'HSP':
            rename_hip(r,ff)
        if r.resname.startswith('LYS') or r.resname == 'LYP':
            rename_lys(r,ff)
        if r.resname == 'LSN' or r.resname == 'LYN':
            rename_lsn(r,ff)
        if r.resname == 'ASP':
            rename_asp(r,ff)
        if r.resname == 'ASH' or r.resname == 'ASPP':
            rename_aspp(r,ff)
        if r.resname =='GLU':
            rename_glu(r,ff)
        if r.resname == 'GLH' or r.resname == 'GLUP':
            rename_glup(r,ff)
        if r.resname == 'CYS' or r.resname == 'CYSH' or r.resname == 'CYN':
            rename_cys(r,ff)
        if r.resname == 'CYX' or r.resname == 'CYS2':
            rename_cyx(r,ff)

def rename_caps(m,ff):
# Amber: NME and ACE
# Charmm: CT3 and ACE
# OPLS: NAC and ACE
    for r in m.residues:
        if r.resname=='CT3': # charmm
            if ff.startswith('amber'):
                rename_lowlevel(r,'NME')
            elif ff.startswith('opls'):
                rename_lowlevel(r,'NAC')
            elif ff.startswith('maestro'):
                rename_lowlevel(r,'NMA')
        if r.resname=='NAC': # opls
            if ff.startswith('amber'):
                rename_lowlevel(r,'NME')
            elif ff.startswith('charmm'):
                rename_lowlevel(r,'CT3')
            elif ff.startswith('maestro'):
                rename_lowlevel(r,'NMA')
        if r.resname=='NME':
            if ff.startswith('opls'):
                rename_lowlevel(r,'NAC')
            elif ff.startswith('charmm'):
                rename_lowlevel(r,'CT3')
        if r.resname=='NMA': # maestro
            if ff.startswith('opls'):
                rename_lowlevel(r,'NAC')
            elif ff.startswith('charmm'):
                rename_lowlevel(r,'CT3')
            elif ff.startswith('amber'):
                rename_lowlevel(r,'NME')

def terminal_oxt(m):
    for a in m.atoms:
        if a.name=='OXT':
            a.name = 'O'
    
def main(argv):

   options=[
   Option( "-ft", "string", "charmm" , "force field type (charmm, amber99sb, amber99sb*-ildn, oplsaa"),
#   Option( "-align", "bool", True, "align side chains"),
        ]
    
   files = [
       FileOption("-f", "r",["pdb","gro"],"protein.pdb", "input structure file"),
       FileOption("-o", "w",["pdb","gro"],"out.pdb", "output structure file"),
       ]
    
   help_text = ('If hydrogens in the system are to be preserved, ',
                'this script would adjust residue names',
		'for pdb2gmx not to complain.',
		'Residues considered: His, Lys, Asp, Glu.',
                'Also the termini names will be fixed.'
                '',
                )

    
   cmdl = Commandline( argv, options = options,
                       fileoptions = files,
                       program_desc = help_text,
                       check_for_existing_files = False )
    
   iname = cmdl['-f']
   oname = cmdl['-o']
   ff = cmdl['-ft']

   m = Model(iname,bPDBTER=True,renumber_residues=False)

   # fix potential problems with the amber termini
   fix_amber_termini(m,ff)

   # rename residues properly
   rename_res(m,ff)

   # fix terminal OXT, if present
   terminal_oxt(m)

   # rename capping groups
   rename_caps(m,ff)

   m.write(oname,bPDBTER=True)

if __name__=='__main__':
    main(sys.argv)

