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

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

def rename_HA( a ):
    if (a.name=='HA2'):
        a.name = 'HA1'
    if (a.name=='HA3'):
        a.name = 'HA2'

def rename_HB( a ):
    if (a.name=='HB2'):
        a.name = 'HB1'
    if (a.name=='HB3'):
        a.name = 'HB2'

def rename_HG( a ):
    if (a.name=='HG2'):
        a.name = 'HG1'
    if (a.name=='HG3'):
        a.name = 'HG2'
    if (a.name=='HG12'):
        a.name = 'HG11'
    if (a.name=='HG13'):
        a.name = 'HG12'
    if (a.name=='2HG1'):
        a.name = '1HG1'
    if (a.name=='3HG1'):
        a.name = '2HG1'

def rename_HD( a ):
    if (a.name=='HD2'):
        a.name = 'HD1'
    if (a.name=='HD3'):
        a.name = 'HD2'
    if a.resname=='ILE':
        if a.name=='1HD1':
            a.name = 'HD1'
        if a.name=='2HD1':
            a.name = 'HD2'
        if a.name=='3HD1':
            a.name = 'HD3'

def rename_HE( a ):
    if (a.name=='HE2') and ('HI' not in a.resname):
        a.name = 'HE1'
    if (a.name=='HE3'):
        a.name = 'HE2'

def rename_HH( a ):
    if (a.name=='2HH1'):
        a.name = 'HH12'
    if (a.name=='1HH1'):
        a.name = 'HH11'
    if (a.name=='1HH2'):
        a.name = 'HH21'
    if (a.name=='2HH2'):
        a.name = 'HH22'

def rename_HE_TRP( a ):
    if (a.name=='HE2'):
        a.name = 'HE3'

def rename_CT3( a ):
    if (a.name=='CA'):
        a.name = 'CH3'
    if (a.name=='HA1') or (a.name=='1HA'):
        a.name = 'HH31'
    if (a.name=='HA2') or (a.name=='2HA'):
        a.name = 'HH32'
    if (a.name=='HA3') or (a.name=='3HA'):
        a.name = 'HH33'


def rename_hydrogens(m,ff,bMaestro=False):
    for a in m.atoms:
	# Maestro
        if bMaestro:
            if( a.resname=='ALA'):
                continue
            elif( a.resname=='ARG'):
                rename_HB( a )
                rename_HG( a )
                rename_HD( a )
                rename_HH( a )
            elif( a.resname=='ASN'):
                rename_HB( a )
            elif( a.resname=='ASP'):
                rename_HB( a )
            elif( a.resname=='ASPP'):
                rename_HB( a )
            elif( a.resname=='ASH'):
                rename_HB( a )
            elif( 'CY' in a.resname):
                rename_HB( a )
            elif( a.resname=='GLY'):
                rename_HA( a )
            elif( a.resname=='GLU'):
                rename_HB( a )
                rename_HG( a )
            elif( a.resname=='GLN'):
                rename_HB( a )
                rename_HG( a )
            elif( 'HI' in a.resname or 'HS' in a.resname ):
                rename_HB( a )
            elif( a.resname=='ILE'):
                rename_HG( a )
                rename_HD( a )
            elif( a.resname=='LEU'):
                rename_HB( a )
            elif( a.resname=='LYS'):
                rename_HB( a )
                rename_HG( a )
                rename_HD( a )
                rename_HE( a )
            elif( a.resname=='LYN' or a.resname=='LSN'):
                rename_HB( a )
                rename_HG( a )
                rename_HD( a )
                rename_HE( a )
            elif( a.resname=='MET'):
                rename_HB( a )
                rename_HG( a )
            elif( a.resname=='PHE'):
                rename_HB( a )
            elif( a.resname=='PRO'):
                rename_HB( a )
                rename_HG( a )
                rename_HD( a )
            elif( a.resname=='SER'):
                rename_HB( a )
#            elif( a.resname=='THR'):
#                continue
            elif( a.resname=='TRP'):
                rename_HB( a )
            elif( a.resname=='TYR'):
                rename_HB( a )
            elif( a.resname=='CT3'):
                rename_CT3( a )
#            elif( a.resname=='VAL'):
#                continue
#            elif( a.resname=='ACE'):
#            elif( a.resname=='NME'):
            if (a.resname=='HOH') and (a.name=='H1'):
                a.name = 'HW1'
            if (a.resname=='HOH') and (a.name=='H2'):
                a.name = 'HW2'
	# H and HN
        if (a.name=='H2' or a.name=='H11' or a.name=='H1') and bMaestro:
            if ff.startswith('charmm'):
                a.name = 'HN'
            else:
                a.name = 'H'
	# H and HN
        if (a.name=='H') and ff.startswith('charmm'):
            a.name = 'HN'
        if (a.name=='H1' or a.name=='1H') and ff.startswith('charmm') and a.resname=='PRO': # terminal proline in charmm
            a.name = 'HN1'
        if (a.name=='H2' or a.name=='2H') and ff.startswith('charmm') and a.resname=='PRO': # terminal proline in charmm
            a.name = 'HN2'
        if (a.name=='HN') and (ff.startswith('opls') or ff.startswith('amber')):
            a.name = 'H'
	# SER
        if (a.resname=='SER') and (a.name=='HG') and ff.startswith('charmm'):
            a.name = 'HG1'
        if (a.resname=='SER') and (a.name=='HG1') and (ff.startswith('opls') or ff.startswith('amber')):
            a.name = 'HG'
	# CYS
        if (a.resname=='CYS') and (a.name=='HG') and ff.startswith('charmm'):
            a.name = 'HG1'
        if (a.resname=='CYS') and (a.name=='HG1') and (ff.startswith('opls') or ff.startswith('amber')):
            a.name = 'HG'
	# ACE
        if (a.resname=='ACE') and (a.name=='H1' or a.name=='1H' or a.name=='H'):
            a.name = 'HH31'
        if (a.resname=='ACE') and (a.name=='H2' or a.name=='2H' or a.name=='H'):
            a.name = 'HH32'
        if (a.resname=='ACE') and (a.name=='H3' or a.name=='3H' or a.name=='H'):
            a.name = 'HH33'
	# NME
        if (a.resname=='NME') and (a.name=='CA' or a.name=='C'):
            a.name = 'CH3'
        if (a.resname=='NME') and (a.name=='HA1' or a.name=='1HA'):
            a.name = 'HH31'
        if (a.resname=='NME') and (a.name=='HA2' or a.name=='2HA'):
            a.name = 'HH32'
        if (a.resname=='NME') and (a.name=='HA3' or a.name=='3HA'):
            a.name = 'HH33'
    


def main(argv):

   options=[
   Option( "-ft", "string", "charmm" , "force field type (charmm, amber99sb, amber99sb*-ildn, oplsaa"),
   Option( "-maestro", "bool", False , "conversion from Maestro prepared structure"),
#   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 = ('Rename hydrogens (and some other atoms) to match the force fields, ',
                '',
                )

    
   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']
   bMaestro = False
   if cmdl['-maestro']==True:
       bMaestro = True

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

   rename_hydrogens(m,ff,bMaestro)

   m.write(oname,bPDBTER=True)

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

