#!/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 os
import argparse
from pmx.forcefield import Topology
from pmx.utils import ff_selection
from pmx.utils import multiple_replace
from pmx.alchemy import gen_hybrid_top, write_split_top
from pmx.scripts.cli import check_unknown_cmd


def _change_outfile_format(filename, ext):
    head, tail = os.path.split(filename)
    name, ex = os.path.splitext(tail)
    new_name = os.path.join(head, name+'.'+ext)
    return new_name


# =============
# Input Options
# =============
def parse_options():
    parser = argparse.ArgumentParser(description='''
This script fills in the B state to a topology file (itp or top) according to
the hybrid residues present in the file. If you provide a top file with
include statemets, by default the script will run through the included itp
files too; this can turned off using the --norecursive flag.

You need to use this script after having mutated a structure file with pmx
mutate, and after having passed that mutated structure through pdb2gmx.
''')

    parser.add_argument('-p',
                        metavar='topol',
                        dest='intop',
                        type=str,
                        help='Input topology file (itp or top). '
                        'Default is "topol.top"',
                        default='topol.top')
    parser.add_argument('-o',
                        metavar='outfile',
                        dest='outfile',
                        type=str,
                        help='Output topology file. '
                        'Default is "pmxtop.top"',
                        default='pmxtop.top')
    parser.add_argument('-ff',
                        metavar='ff',
                        dest='ff',
                        type=str.lower,
                        help='Force field to use. If -p is a top file, it is '
                        'not necessary to specify the forcefield, as it will '
                        'be determined automatically. If -p is an itp file, '
                        'then -ff is needed, and if not provided a list of '
                        'available ff will be shown.',
                        default=None)
    parser.add_argument('--split',
                        dest='split',
                        help='Write separate topologies for the vdW and charge'
                        ' transformations.',
                        default=False,
                        action='store_true')
    parser.add_argument('--scale_mass',
                        dest='scale_mass',
                        help='Scale the masses of dummy atoms by multiplying '
                        'with scale_mass factor',
                        default=0.33, #TODO: need a function to check for oscillation period: need to ensure 2*pi*sqrt(m1*m2/((m1+m2)*k)) > 10*dt
                        action='store')
    parser.add_argument('--scale_dih',
                        dest='scale_dih',
                        help='Scale the dihedrals that have a dummy. ',
                        default=1.0,
                        action='store')
    parser.add_argument('--norecursive',
                        dest='recursive',
                        help='Whether to fill the B states also for all itp '
                        'files included in the provided topology file. '
                        'Default is True. This flag sets it to False.',
                        default=True,
                        action='store_false')

    args, unknown = parser.parse_known_args()
    check_unknown_cmd(unknown)

    # ff selection is required if file provided is an itp
    if args.intop.split('.')[-1] == 'itp' and args.ff is None:
        args.ff = ff_selection()

    return args


# ====
# Main
# ====
def main(args):

    top_file = args.intop
    top_file_ext = top_file.split('.')[-1]
    outfile = args.outfile
    ff = args.ff
    scale_mass = float(args.scale_mass)
    recursive = args.recursive
    scaleDih = float(args.scale_dih)

    # if input is itp but output is else, rename output
    if top_file_ext == 'itp' and outfile.split('.')[-1] != 'itp':
        outfile = _change_outfile_format(outfile, 'itp')
        print('log_> Setting outfile name to %s' % outfile)

    # load topology file
    topol = Topology(top_file, ff=ff, version='new')

    # fill the B states
    pmxtop, pmxitps = gen_hybrid_top(topol=topol, recursive=recursive,
                                     verbose=True, scaleDih=scaleDih)

    # get the base path for the .itp file output
    basepath = os.path.abspath(outfile).rsplit('/',1)[0]

    # write hybrid itps if present
    replace = {}
    if len(pmxitps) > 0:
        for pmxitp in pmxitps:
            itp_fn = os.path.basename(pmxitp.filename)
            out_fn = '{0}/pmx_{1}'.format(basepath,itp_fn)
            # store old/new itp names for replacement in top file
            replace[itp_fn] = 'pmx_{0}'.format(itp_fn) #out_fn
            print('\nlog_> Writing itp file "%s""' % out_fn)
            pmxitp.write(out_fn, scale_mass=scale_mass)

    # write hybrid topology
    print('\nlog_> Writing topology file "%s""' % outfile)
    pmxtop.write(outfile, scale_mass=scale_mass)

    # if we modified itp files included in top, fix top file
    if len(pmxitps) > 0:
        multiple_replace(outfile, replace, isfile=True)

    # separated topologies
    if args.split is True:
        write_split_top(pmxtop=pmxtop, outfile=outfile, scale_mass=scale_mass,
                        verbose=True)

    print('')
    print('b-states filled...........')
    print('')


def entry_point():
    args = parse_options()
    main(args)


if __name__ == '__main__':
    entry_point()
