#!/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
# pmx is Copyright (C) 2013-2022 by Vytautas Gapsys
#
#                        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 AND VYTAUTAS GAPSYS DISCLAIM 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 *
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='''
Take any itp or top and repartition mass:
triple hydrogen mass and remove the excess mass from heavy atoms
''')

    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')

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

    return args


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

    top_file = args.intop
    top_file_ext = top_file.split('.')[-1]
    outfile = args.outfile

    # 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 = TopolBase(top_file)

    topol.repartition_mass( )

    # write topology
    topol.write( outfile )


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


if __name__ == '__main__':
    entry_point()
