# 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,shutil
import commands
from glob import glob
from pmx import *
from pmx.ndx import *
from pmx.library import _one_letter
from pmx.forcefield import MDP


ff_path = os.path.join(os.environ.get('GMXDATA'),'top/amber99sbmut.ff')


def run_command( func, string ):
    s = func.__name__+'(): '+ string
    err_file = func.__name__+'_ERROR.log'
    status, out = commands.getstatusoutput( string )
    if status != 0:
        print >>sys.stderr, 'ERROR in %s ' % func.__name__
        print >>sys.stderr, 'Output written to %s'  % os.path.abspath(os.path.join('.',err_file))
        fp = open(err_file,'w')
        print >>fp, s
        print >>fp, out
        fp.close()
        sys.exit(1)
    else:
        print("%-90s" % s, ': ok')

chain_ids = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'

class MD:

    def __init__(self, **kwargs):

        self.ff = 'amber99sbmut'
        self.water = 'tip3p'
        self.conc = .15
        self.box_type = 'triclinic'
        self.box_size = 1.2
        self.vsite = False
        self.princ = False
        self.md_in_conf = 'md_in.pdb'

        for key, val in kwargs.items():
            setattr(self,key,val)

    def setup(self):
        self.generate_topology( self.pdb_in)
        self.renumber_pdb( )
        self.generate_topology( self.pdb_in)
        self.generate_sim_box( )
        self.fill_sim_box(  )
        self.tpr_from_box()
        self.add_ions(  )
        self.make_em_tpr()
        self.run_em()
        self.pdb_from_em()
        self.generate_topology('min.pdb')
        self.compact_repr( 'min.pdb', self.md_in_conf)
        self.clean_backups()

    def write_residue_map( self, model ):
        fp = open('residue_map.txt','w')
        for r in model.residues:
            if r.chain_id == ' ':
                chain = '_'
                r.old_chain_id = '_'
            else:
                chain = r.chain_id
            print >>fp, '%d|%s|%s -> %d|%s|%s' %( r.orig_id, r.resname, r.old_chain_id, r.id, r.resname, chain)
        fp.close()

    def renumber_pdb(self ):
        m = Model('gmx.pdb')
        for i, chain in enumerate(m.chains):
            if chain.id != ' ':
                for r in chain.residues:
                    r.old_chain_id = chain.id
                chain.set_chain_id( chain_ids[i] )
        self.write_residue_map( m )
        m.write('start.pdb')
        self.pdb_in = 'start.pdb'

    def __str__(self):
        s = '< MD (%s) > ' % self.pdb_in
        return s

    def clean_backups(self):
        files = glob('#*#')
        for f in files:
            os.unlink(f)

    def generate_topology(self, pdb_in):

        run = 'pdb2gmx -f %s -water %s -ff %s -o gmx.pdb ' % ( pdb_in, self.water, self.ff )
        if self.vsite:
            run+=' -vsite hydrogens'
        run_command( self.generate_topology, run )

    def generate_sim_box(self ):
        if self.princ:
            run = 'echo 4 | editconf -f gmx.pdb -o tmp.pdb -c -princ '
            run_command( run_editconf, run)
            run = 'editconf -f tmp.pdb -o ed.pdb -d %g -bt %s ' % (self.box_size, self.box_type)
            run_command( generate_sim_box, run)
        else:
            run = 'editconf -f gmx.pdb -o ed.pdb -d %g -bt %s ' % (self.box_size, self.box_type)
            run_command( self.generate_sim_box, run)

    def fill_sim_box(self ):
        if self.water == 'spce': water = 'spc216'
        else: water = self.water
        run = 'genbox -cp ed.pdb -cs %s -p -o box.pdb' % os.path.join(ff_path,water)
        run_command( self.fill_sim_box, run)

    def tpr_from_box(self):
        run = 'grompp -f ~/mdp/em.mdp -c box.pdb'
        run_command( self.tpr_from_box, run)

    def get_solvent_index(self ):
        run = 'echo q | make_ndx -f topol.tpr -o tmp.ndx'
        run_command( self.get_solvent_index, run)
        ndx = IndexFile("tmp.ndx")
        return ndx.names.index('SOL')

    def add_ions(self ):
        idx = self.get_solvent_index()
        run = 'echo %d | genion  -conc %g -neutral -p -o ion.pdb' % (idx, self.conc)
        run_command( self.add_ions, run)

    def make_em_tpr(self):
        run = 'grompp -f ~/mdp/em.mdp -c ion.pdb'
        run_command( self.make_em_tpr, run)

    def run_em(self):
        run = 'mdrun -v -c em.pdb'
        run_command( self.run_em, run)

    def pdb_from_em(self):
        run = 'echo 0| trjconv -f em.pdb -s topol.tpr -o min.pdb'
        run_command( self.pdb_from_em, run)

    def compact_repr(self, pdb_in, pdb_out ):
        run = 'echo 0 | trjconv -f %s -s topol.tpr -ur compact -pbc mol -o %s' % (pdb_in, pdb_out)
        run_command( self.compact_repr, run)



class FreeEnergyMD(MD):

    def __init__(self, mutation_file, **kwargs):

        MD.__init__( self )

        for key, val in kwargs.items():
            setattr(self,key,val)
        self.mutation_tags = []
        self.read_mutation_file( mutation_file )
        self.mutations = []
        if self.mutation_tags:
            self.mutations_from_tags()
        self.runs = []
        self.is_single_chain = False

    def read_mutation_file(self, mut_file ):
        print()'\n\t\t\tReading mutation file: %s\n' % mut_file)
        l = open(mut_file).readlines()
        count = 1
        for line in l:
            entr = line.strip()
            if entr:
                print('\t\t\t (%d) ->  %s' % (count, entr))
                self.mutation_tags.append( entr )
                count+=1


    def setup(self):
        self.read_pdb()
        print('\n\n')
        for m in self.mutations:
            self.setup_mutation_run(m)

    def read_pdb( self ):
        self.model = Model(self.md_in_conf)
        if len(self.model.chains) == 1:
            self.is_single_chain = True

    def mutations_from_tags(self):
        for t in self.mutation_tags:
            mut = t.split()
            new_mut = []
            for m in mut:
                resid, resn, chain = m.split('|')
                resid = int(resid)
                new_mut.append( ( resid, resn, chain ) )
            self.mutations.append( new_mut )

    def setup_mutation_run(self, mutation ):
        muts = []
        affected_chains = []
        for m in mutation:
            resid, resn, chain = m
            residue = self.model.residues[resid-1]
            resname = _one_letter[residue.resname]
            muts.append( chain+'.'+resname+str(resid)+resn )
            if chain not in affected_chains:
                affected_chains.append( chain )
        name = '-'.join(muts)
        if str(self.__class__).split('.')[1] == 'DiscreteTI':
            name+='.dti'
        elif str(self.__class__).split('.')[1] == 'CrooksMD':
            name+='.crooks'
        print('\n\t\t\tPreparing mutation run -> %s \n' % name)
        if os.path.isdir( name ):
            shutil.rmtree(name)
        os.mkdir(name)
        self.runs.append( name )
        shutil.copy(self.md_in_conf, name)
        script_file = os.path.join(name,'mutations.txt')
        fp = open(script_file,'w')
        for m in mutation:
            resid, resn, chain = m
            print>>fp, resid, resn
        fp.close()
        self.make_mutation(name, affected_chains)


    def make_mutation(self, path, affected_chains ):
        os.chdir(path)
        run = '~/software/pmx/scripts/mutate_beta45.py -f %s -script mutations.txt -o mut.pdb' % self.md_in_conf
        run_command( self.make_mutation, run )
        self.generate_topology( 'mut.pdb')
        if self.is_single_chain:
            itp_file = 'topol_Protein.itp'
            run = '~/software/pmx/scripts/make_bstate_beta45.py -itp %s' % itp_file
            run_command( self.make_mutation, run )
            shutil.move(itp_file, itp_file+'.save')
            shutil.move('newtop.itp', itp_file)
        else:
            for chain in affected_chains:
                print('Applying changes to topology of chain %s' % chain)
                itp_file = 'topol_Protein_chain_%s.itp' % chain
                run = '~/software/pmx/scripts/make_bstate_beta45.py -itp %s' % itp_file
                run_command( self.make_mutation, run )
                shutil.move(itp_file, itp_file+'.save')
                shutil.move('newtop.itp', itp_file)
        os.chdir('..')


class CrooksMD( FreeEnergyMD ):

    def __init__(self, mutation_file, **kwargs):
        FreeEnergyMD.__init__(self, mutation_file)
        for key, val in kwargs.items():
            setattr(self,key,val)


    def do_crooks( self, mdp_file, min_mdp_file, time, nruns ):
        mdp = MDP().read( mdp_file)
        min_mdp = MDP().read( min_mdp_file)
        self.prepare_min_mdp_files(min_mdp, time )
        self.prepare_eq_mdp_files(mdp, time )
        for r in self.runs:
            print('\n\t\t\tSetting up Crooks run -> %s\n' % r)

            self.minimize_states( r )
            self.setup_equilibration_runs( r, nruns )

    def minimize_states( self, path ):

        os.chdir(path)
        run = 'grompp -f ../crooks_minA.mdp -c gmx.pdb -o minA.tpr'
        run_command( self.minimize_states, run )
        run = 'grompp -f ../crooks_minB.mdp -c gmx.pdb -o minB.tpr'
        run_command( self.minimize_states, run )
        print('\n\t\t\tRunning energy minimization on %s ( state A ) \n' % ( path ))
        run = 'mdrun -v -c emA.pdb -s minA.tpr'
        run_command( self.minimize_states, run )
        print('\n\t\t\tRunning energy minimization on %s ( state B ) \n' % ( path ))
        run = 'mdrun -v -c emB.pdb -s minB.tpr'
        run_command( self.minimize_states, run )
        os.chdir('..')

    def prepare_eq_mdp_files( self, mdp, time ):
        nsteps = 1000*time/.002
        mdp['nsteps'] = nsteps
        mdp['init-lambda'] = 0
        fp = open('crooks_eqA.mdp','w')
        print >>fp, mdp
        fp.close()
        mdp['init-lambda'] = 1
        fp = open('crooks_eqB.mdp','w')
        print >>fp, mdp
        fp.close()

    def prepare_min_mdp_files( self, mdp, time ):
        mdp['init-lambda'] = 0
        fp = open('crooks_minA.mdp','w')
        print >>fp, mdp
        fp.close()
        mdp['init-lambda'] = 1
        fp = open('crooks_minB.mdp','w')
        print >>fp, mdp
        fp.close()


    def setup_equilibration_runs( self, path, nruns ):
        os.chdir(path)
        for i in range(nruns):
            print('\n\t\t\tPreparing run input file for  %s ( run %d ) \n' % ( path, i ))
            os.mkdir('runA.%d' % i)
            os.mkdir('runB.%d' % i)
            run = 'grompp -f ../crooks_eqA.mdp -c emA.pdb -o runA.%d/topol.tpr' % i
            run_command( self.setup_equilibration_runs, run )
            run = 'grompp -f ../crooks_eqB.mdp -c emB.pdb -o runB.%d/topol.tpr' % i
            run_command( self.setup_equilibration_runs, run )
            self.clean_backups()
        os.chdir('..')


class DiscreteTI( FreeEnergyMD ):

    def __init__(self, mutation_file, **kwargs):
        FreeEnergyMD.__init__(self, mutation_file)
        for key, val in kwargs.items():
            setattr(self,key,val)

    def read_lambda_steps( self, filename ):
        l = open(filename).readlines()
        self.lambda_steps = []
        for line in l:
            entr = line.strip()
            if entr:
                self.lambda_steps.append( float(entr) )

    def prepare_dti_mdp_files( self, mdp_file, time ):
        mdp = MDP().read(mdp_file)
        nsteps = 1000*time/.002
        mdp['nsteps'] = nsteps
        for lda in self.lambda_steps:
            fp = open('dti_%4.3f.mdp' % round(lda,3),'w' )
            mdp['init-lambda'] = lda
            print >>fp, mdp
            fp.close()

    def prepare_dti_min_mdp_files( self, mdp_file):
        mdp = MDP().read(mdp_file)
        for lda in self.lambda_steps:
            fp = open('dti_min_%4.3f.mdp' % round(lda,3),'w' )
            mdp['init-lambda'] = lda
            print >>fp, mdp
            fp.close()

    def setup_runs( self, path ):
        os.chdir( path )
        for lda in self.lambda_steps:
            min_mdp = 'dti_min_%4.3f.mdp' % round(lda,3)
            run_mdp = 'dti_%4.3f.mdp' % round(lda,3)
            run_dir = 'run_%4.3f' % round(lda,3)
            os.mkdir( run_dir )
            print('\n\t\t\tRunning energy minimization on %s/%s \n' % ( path, run_dir ))
            run = 'grompp -f ../%s -c gmx.pdb -o %s/em.tpr' % (min_mdp, run_dir )
            run_command( self.setup_runs, run )
            os.chdir( run_dir )
            run = 'mdrun -v -c em.pdb -s em.tpr'
            run_command( self.setup_runs, run )
            os.chdir('..')
            print('\n\t\t\tPreparing run input file for  %s/%s \n' % ( path, run_dir ))
            run = 'grompp -f ../%s -c %s/em.pdb -o %s/topol.tpr' % (run_mdp, run_dir, run_dir )
            run_command( self.setup_runs, run )
            self.clean_backups()
        os.chdir( '..')

    def do_dti( self, mdp, min_mdp, time ):
        self.prepare_dti_min_mdp_files( min_mdp)
        self.prepare_dti_mdp_files( mdp, time)
        for r in self.runs:
            print('\n\t\t\tSetting up Discrete TI run -> %s\n' % r)
            self.setup_runs( r )



def main(argv):

    version = "1.0"

    options = [
        Option( "-water", "string", "tip3p", "water model"),
        Option( "-box_type", "string", "triclinic", "box geometry (triclinic, octahedron, dodecahedron)"),
        Option( "-box_size", "float", 1.2, "distance from solute to box"),
        Option( "-conc", "float", 0.15, "ion concentration"),
        Option( "-vsite", "bool", False, "use virtual sites"),
        Option( "-fe.dti", "bool", False, "do discrete TI setup"),
        Option( "-dti_run_time", "float", 10., "Simulation time [ns] at each lambda point"),
        Option( "-fe.crooks", "bool", False, "do Crooks setup"),
        Option( "-n_crooks_runs", "int", 1, "setup # crooks runs for each mutation"),
        Option( "-crooks_run_time", "float", 50., "Simulation time [ns] for crooks equilibrium runs"),
        Option( "-skip_md_setup", "bool", False, "skip md setup and use -f as starting configuration for free energy runs"),
        ]

    files = [
        FileOption("-f", "r",["pdb"],"protein", "input pdb file"),
        FileOption("-m", "r/o",["txt"],"mutations", "mutations to make"),
        FileOption("-crooks_mdp", "r/o",["mdp"],"template", "template run mdp file for crooks equilibrium runs"),
        FileOption("-dti_mdp", "r/o",["mdp"],"template", "template run mdp file for discrete TI calculations"),
        FileOption("-min_mdp", "r/o",["mdp"],"em", "template minimization mdp file ( for TI or Crooks )"),
        FileOption("-lambda_steps", "r/o",["txt"],"lambda_steps", "text file with lambda steps for DTI runs"),
        ]



    help_text = ("Script for setting up plain MD runs",
                 )


    cmdl = Commandline( argv, options = options,
                        fileoptions = files,
                        program_desc = help_text,
                        check_for_existing_files = False, version = version)


    # make some basic checks
    bFreeEnergy = False
    if cmdl['-fe.dti'] or cmdl['-fe.crooks']:
        bFreeEnergy = True
        mutation_file = cmdl['-m']
        try:
            open(mutation_file).readlines()
        except:
            print >>sys.stderr, 'Error: Cannot open %s' % mutation_file
            sys.exit(1)
        if cmdl['-fe.dti']: # require lambda steps
            try:
                open(cmdl['-lambda_steps']).readlines()
            except:
                print >>sys.stderr, 'Error: Cannot open %s' % cmdl['-lambda_steps']
                sys.exit(1)



    pdb_in = cmdl['-f']
    water = cmdl['-water']
    box_type = cmdl['-box_type']
    box_size = cmdl['-box_size']
    conc = cmdl['-conc']
    free_energy_start_pdb = None
    bVsite = cmdl['-vsite']
    if bVsite and bFreeEnergy:
        print >>sys.stderr, 'Error: Cannot use virtual sites in free energy calculations !'
        sys.exit(1)


    if not cmdl['-skip_md_setup']:
        md = MD( pdb_in = pdb_in, water = water, conc = conc, box_type = box_type, box_size = box_size )
        md.setup()
        free_energy_start_pdb = md.md_in_conf
    else:
        if not bFreeEnergy:
            print('Nothing to do........... (no MD, no Free Energy)')
            sys.exit()
    if bFreeEnergy:
        if not free_energy_start_pdb:
            free_energy_start_pdb = pdb_in
        if cmdl['-fe.crooks']:
            crooksMD = CrooksMD( mutation_file, md_in_conf = free_energy_start_pdb)
            crooksMD.setup()
            crooksMD.do_crooks( cmdl['-crooks_mdp'], cmdl['-min_mdp'], cmdl['-crooks_run_time'], cmdl['-n_crooks_runs'] )

        if cmdl['-fe.dti']:
            dti = DiscreteTI( mutation_file, md_in_conf = free_energy_start_pdb )
            dti.setup()
            dti.read_lambda_steps( cmdl['-lambda_steps'] )
            dti.do_dti(cmdl['-dti_mdp'], cmdl['-min_mdp'], cmdl['-dti_run_time'] )


    print('\n\t\t\t ....... DONE .......... \n')



##     n_crooks_runs = cmdl['-n_crooks_runs']
##     if cmdl.opt['-m'].is_set:
##         mut_file  = cmdl['-m']
##     else:
##         mutations = []

##     if cmdl.opt['-mdp'].is_set:
##         mdp = MDP().read( cmdl['-mdp'] )
##     if cmdl.opt['-min_mdp'].is_set:
##         min_mdp = MDP().read( cmdl['-min_mdp'] )
##     if cmdl.opt['-lambda_steps'].is_set:
##         lambda_file = cmdl['-lambda_steps']


##     md = MD( pdb_in = pdb_in, water = water, conc = conc, box_type = box_type, box_size = box_size )
## #    md.setup()

##     dti = DiscreteTI( mut_file, md_in_conf = md.md_in_conf )
##     dti.setup()
##     dti.read_lambda_steps( lambda_file )
##     dti.do_dti(mdp, min_mdp, 10. )



##     crooksMD = CrooksMD( mut_file, md_in_conf = md.md_in_conf)
##     crooksMD.setup()
##     crooksMD.do_crooks(mdp, 20, cmdl['-n_crooks_runs'])

#    fe_md = FreeEnergyMD( mutations, md_in_conf = md.md_in_conf )
#    fe_md.setup()


##     generate_topology(pdb_in, water, 'amber99sbmut')
##     generate_sim_box( box_type, box_size )
##     fill_sim_box( water )
##     tpr_from_box()
##     add_ions( conc )
##     make_em_tpr()
##     run_em()
##     pdb_from_em()
##     generate_topology('min.pdb', water, 'amber99sbmut')
##     compact_repr( 'min.pdb')
##     clean_backups()



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


