#!/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
import os
import shutil
import commands
from glob import glob
from pmx.forcefield import MDP
import argparse


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' % err_file
        fp = open(err_file, 'w')
        print >>fp, s
        print >>fp, out
        fp.close()
        sys.exit(1)
    else:
        print "%-90s" % s, ': ok'


def make_dir_tree(skip=4):

    # catch all run?.? dirs
    dirs = glob('run?.?')
    for d in dirs:
        print '\tPreparing run %s' % d
        os.chdir(d)
        os.mkdir('morphes')
        cmd = 'echo 0| trjconv -f traj.trr -s topol.tpr -skip %d -b 2001 -sep -o morphes/frame.gro' % skip
        run_command(make_dir_tree, cmd)
        os.chdir('morphes')

        # put each file into single directory
        file_list = glob('frame*.gro')
        for f in file_list:
            name = f.split('.')[0]
            os.mkdir(name)
            shutil.move(f, name)
        os.chdir('..')
        os.chdir('..')


def prepare_mdp_files(template_mdp, sw_time, sc_alpha, sc_sigma):

    mdp = MDP().read(template_mdp)

    # make 0->1 file
    mdp['free-energy'] = 'yes'
    mdp['init-lambda'] = 0
    # time is in ps -> convert to steps
    nsteps = sw_time*500
    mdp['nsteps'] = int(nsteps)
    # delta lambda = 1./nsteps
    delta_lambda = 1./float(nsteps)
    mdp['delta-lambda'] = delta_lambda
    mdp['sc-alpha'] = sc_alpha
    mdp['sc-sigma'] = sc_sigma

    fp = open('crooks_TI_runA.mdp', 'w')
    print >>fp, mdp
    fp.close()

    # make 1->0 file
    mdp['init-lambda'] = 1
    mdp['delta-lambda'] = -delta_lambda

    fp = open('crooks_TI_runB.mdp', 'w')
    print >>fp, mdp
    fp.close()


def make_run_input_files():

    dirs = glob('run?.?')
    for d in dirs:
        print '\n\tPreparing run input files %s' % d
        mdp_file = None
        if d.split('.')[0][-1] == 'A':  # 0->1
            mdp_file = 'crooks_TI_runA.mdp'
        elif d.split('.')[0][-1] == 'B':  # 1->0
            mdp_file = 'crooks_TI_runB.mdp'
        dir_list = glob(os.path.join(d, 'morphes')+'/frame*')
        for f in dir_list:
            fname = os.path.basename(f)+'.gro'
            gro_file = os.path.join(f, fname)
            cmd = 'grompp -f %s -c %s -o %s/topol.tpr' % (mdp_file, gro_file, f)
            run_command(make_run_input_files, cmd)
            os.system('rm mdout.mdp')


def parse_options():
    parser = argparse.ArgumentParser(description='Script for setting up plain FGTI runs...')

    # TODO: add following options
    # 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")
    # 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"),

    parser.add_argument('-d',
                        metavar='dir',
                        dest='dir',
                        type=str,
                        help='directory with equlibrated states',
                        default='.')
    parser.add_argument('-mdp',
                        metavar='mdp',
                        dest='mdp',
                        type=str,
                        help='template TI mdp file',
                        default='TI_template.mdp')
    parser.add_argument('-sw_time',
                        metavar='sw_time',
                        dest='sw_time',
                        type=float,
                        help='switching time to use for FGTI runs',
                        default=50.)
    parser.add_argument('-sc_alpha',
                        metavar='sc_alpha',
                        dest='sc_alpha',
                        type=float,
                        help='soft-core alpha',
                        default=0.3)
    parser.add_argument('-sc_sigma',
                        metavar='sc_sigma',
                        dest='sc_sigma',
                        type=float,
                        help='soft-core sigma',
                        default=0.25)
    parser.add_argument('-skip',
                        metavar='skip',
                        dest='skip',
                        type=int,
                        help='skip # frames',
                        default=4)

    args = parser.parse_args()
    return args


def main(args):

    here = os.getcwd()

    run_dir = args.dir
    mdp_file = args.mdp
    sw_time = args.sw_time
    sc_alpha = args.sc_alpha
    sc_sigma = args.sc_sigma

    print '\n\t Preparing FGTI runs in directory..: %s' % run_dir
    print '\t Template mdp file to use............: %s' % mdp_file
    print '\t Switching time to use...............: %8d ps' % int(sw_time)
    print '\t Soft-core alpha to use..............: %8.3f' % sc_alpha
    print '\t Soft-core sigma to use..............: %8.3f' % sc_sigma
    print '\n'


    os.chdir(run_dir)

    print '\t Preparing mdp input files........... '

    prepare_mdp_files(mdp_file, sw_time, sc_alpha, sc_sigma)


    print '\t Preparing directory tree............ '
    make_dir_tree(skip=args.skip)
    make_run_input_files()
    os.chdir(here)
    print '\n\t............... DONE .................\n'


if __name__ == '__main__':
    args = parse_options()
    main(args)
