#!/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-2020 by Daniel Seeliger, 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 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.
# ----------------------------------------------------------------------

"""Program to build ligand hybrid topology.
"""

import sys
import argparse
from pmx import *
#from pmx.options import *
#from pmx.parser import *
from pmx.model import Model
from pmx.parser import read_and_format
from pmx.utils import get_ff_path, doLog
from pmx.scripts.cli import check_unknown_cmd
from pmx.forcefield import *
from pmx.ndx import *
from pmx.ligand_alchemy import *
import time
import warnings

# =============
# Input Options
# =============
def parse_options():

    parser = argparse.ArgumentParser(description='''
Provided two structures and topologies, build hybrid structure/topology.

''', formatter_class=argparse.RawTextHelpFormatter)

#    exclus = parser.add_mutually_exclusive_group()

    parser.add_argument('-i1',
                        metavar='lig1.pdb',
                        dest='i1',
                        type=str,
                        help='Input ligand structure 1. '
                        'Default is "lig1.pdb"',
                        default='lig1.pdb')
    parser.add_argument('-i2',
                        metavar='lig2.pdb',
                        dest='i2',
                        type=str,
                        help='Input ligand structure 2. '
                        'Default is "lig2.pdb"',
                        default='lig2.pdb')
    parser.add_argument('-itp1',
                        metavar='lig1.itp',
                        dest='itp1',
                        type=str,
                        help='Input ligand topology 1. '
                        'Default is "lig1.itp"',
                        default='lig1.itp')
    parser.add_argument('-itp2',
                        metavar='lig2.itp',
                        dest='itp2',
                        type=str,
                        help='Input ligand topology 2. '
                        'Default is "lig2.itp"',
                        default='lig2.itp')
    parser.add_argument('-pairs',
                        metavar='pairs.dat',
                        dest='pairs',
                        type=str,
                        help='Optional input: atom pair mapping. ')
#                        'Default is "lig2.itp"',
#                        default='lig2.itp')
    parser.add_argument('-n1',
                        metavar='scaffold1.ndx',
                        dest='n1',
                        type=str,
                        help='Optional input: index of atoms to consider for mol1 ')
    parser.add_argument('-n2',
                        metavar='scaffold2.ndx',
                        dest='n2',
                        type=str,
                        help='Optional input: index of atoms to consider for mol2 ')
    parser.add_argument('-oA',
                        metavar='mergedA.pdb',
                        dest='oA',
                        type=str,
                        help='Output: hybrid structure based on the ligand 1. '
                       'Default is "mergedA.pdb"',
                        default='mergedA.pdb')
    parser.add_argument('-oB',
                        metavar='mergedB.pdb',
                        dest='oB',
                        type=str,
                        help='Output: hybrid structure based on the ligand 2. '
                        'Default is "mergedB.pdb"',
                        default='mergedB.pdb')
    parser.add_argument('-oitp',
                        metavar='merged.itp',
                        dest='oitp',
                        type=str,
                        help='Output: hybrid topology. '
                        'Default is "merged.itp"',
                        default='merged.itp')
    parser.add_argument('-offitp',
                        metavar='ffmerged.itp',
                        dest='offitp',
                        type=str,
                        help='Output: atomtypes for hybrid topology. '
                        'Default is "ffmerged.itp"',
                        default='ffmerged.itp')
    parser.add_argument('-log',
                        metavar='hybrid.log',
                        dest='log',
                        type=str,
                        help='Output: log file. '
                        'Default is "hybrid.log"',
                        default='hybrid.log')
    parser.add_argument('--d',
                        metavar='0.05',
                        dest='d',
                        type=float,
                        help='Optional: if -pairs not provided, distance (nm) between atoms to consider them morphable for alignment approach (default 0.05 nm).',
                        default=0.05)
    parser.add_argument('--fit',
                        dest='bFit',
                        help='Fit mol2 onto mol1, only works if pairs.dat is provided',
                        action='store_true')
    parser.add_argument('--split',
                        dest='bSplit',
                        help='split the topology into separate transitions',
                        action='store_true')
    parser.add_argument('--scDUMm',
                        metavar='1.0',
                        dest='scDUMm',
                        type=float,
                        help='scale dummy masses using the counterpart atoms',
                        default=1.0)
    parser.add_argument('--scDUMa',
                        metavar='1.0',
                        dest='scDUMa',
                        type=float,
                        help='scale bonded dummy angle parameters',
                        default=1.0)
    parser.add_argument('--scDUMd',
                        metavar='1.0',
                        dest='scDUMd',
                        type=float,
                        help='scale bonded dummy dihedral parameters',
                        default=1.0)
    parser.add_argument('--deAng',
                        dest='bDeAng',
                        help='decouple angles composed of 1 dummy and 2 non-dummies',
                        action='store_true')

    parser.set_defaults(bFit=False,bSplit=False,bDeAng=False)
    args, unknown = parser.parse_known_args()
    check_unknown_cmd(unknown)

    return args

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


def main(args):
    """Run the main script.

    Parameters
    ----------
    args : argparse.Namespace
        The command line arguments
    """

#####################################
    # start timing
    stime = time.time()

#####################################
    # deal with the flags
    bFit = args.bFit
    bSplit = args.bSplit
    bDeAng = args.bDeAng
    d = args.d
    scDUMm = args.scDUMm
    scDUMa = args.scDUMa
    scDUMd = args.scDUMd

#####################################
    # log file
    logfile = open(args.log,'w')

#####################################
    # if pairs.dat is used, ndx will be discarded
    if args.pairs!=None and (args.n1!=None or args.n2!=None):
        doLog(logfile,'pairs.dat is provided, thus discarding index files')
        args.n1 = None
        args.n2 = None
    if args.n1!=None and args.n2==None:
        doLog(logfile,'Only one index file is provided, thus discarding index files')
        args.n1 = None
        args.n2 = None
    if args.n1==None and args.n2!=None:
        doLog(logfile,'Only one index file is provided, thus discarding index files')
        args.n1 = None
        args.n2 = None

#####################################
    # read structures
    doLog(logfile,'Reading ligand 1 from: "%s"' % args.i1)
    m1 = Model().read(args.i1)
    doLog(logfile,'Reading ligand 2 from: "%s"' % args.i2)
    m2 = Model().read(args.i2)

#####################################
    # read topologies
    doLog(logfile,'Reading topology 1 from: "%s"' % args.itp1)
    itp1 = TopolBase(args.itp1)
    doLog(logfile,'Reading topology 2 from: "%s"' % args.itp2)
    itp2 = TopolBase(args.itp2)

#####################################
    # read mapped atom pairs
    if args.pairs!=None:
        doLog(logfile,'Reading file with atom pairs: "%s"' % args.pairs)
        plst = readPairsFile(args.pairs)
    else:
        plst = None

#####################################
    # assign ff parameters
    doLog(logfile,"Assigning forcefield parameters....")
    assignFF(m1,itp1)
    assignFF(m2,itp2)

#####################################
    # read index
    grps = None
    if args.n1!=None and args.n2!=None:
        doLog(logfile,'Reading scaffold index file: "%s"' % args.n1)
        grp1 = IndexFile(args.n1).dic['scaffold']
        doLog(logfile,'Reading scaffold index file: "%s"' % args.n2)
        grp2 = IndexFile(args.n2).dic['scaffold']
        # now we add all atoms with bonds to scaffold atoms
        for b in itp1.bonds:
            if b[0] in grp1.ids and b[1] not in grp1.ids:
                grp1.ids.append(b[1])
                doLog(logfile,'Adding atom %s to scaffold 1' % m1.atoms[b[1]-1].name)
            elif b[1] in grp1.ids and b[0] not in grp1.ids:
                grp1.ids.append(b[0])
                doLog(logfile,'Adding atom %s to scaffold 1' % m1.atoms[b[0]-1].name)
        for b in itp2.bonds:
            if b[0] in grp2.ids and b[1] not in grp2.ids:
                grp2.ids.append(b[1])
                doLog(logfile,'Adding atom %s to scaffold 2' % m2.atoms[b[1]-1].name)
            elif b[1] in grp2.ids and b[0] not in grp2.ids:
                grp2.ids.append(b[0])
                doLog(logfile,'Adding atom %s to scaffold 2' % m2.atoms[b[0]-1].name)
        grps = [grp1.ids, grp2.ids]
    else:
        grps = None

#####################################
    # prepare model copies
    m3 = m1.copy() # m3 will contain all the atoms from m1, but with the coordinates of the matching atoms from m2
    m4 = m2.copy() # need to copy it when fitting

#####################################
    # fitting (optional)

#    if(bFit==True):
    superimposeStructures( args.i1, args.i2, m3, m2, plst, logfile )

#####################################
    # initialize the main object of the LigandHybridTopology class
    doLog(logfile,'Initializing the main object for LigandHybridTopology')
    hybridTop = LigandHybridTopology( m1=m1, m2=m2, m3=m3, m4=m4, itp1=itp1, itp2=itp2, bFit=bFit, scDUMm=scDUMm, scDUMa=scDUMa, scDUMd=scDUMd, bDeAng=bDeAng, pairList=plst, grps=grps, d=d, logfile=logfile )

#####################################
    # the main function: generate hybrid topology
    doLog(logfile,'Starting hybrid topology generation')
    hybridTop._makeHybridTop( )

#####################################
    # write hybrid topology
    doLog(logfile, 'Writing new itp file: "%s"' % args.oitp)
    hybridTop.newitp.write(args.oitp, target_qB = [hybridTop.qB] )
    doLog(logfile, 'Writing dummy forcefield file: "%s"' % args.offitp)
    hybridTop._write_ffitp( args.offitp, m1 )
    # write merged pdb
    if bFit==True:
        m1.write(args.oA)
        m1.write(args.oB)
    else:
        m1.write(args.oA)
        m3.write(args.oB)

#####################################
    # write split topology
    if bSplit==True: # write splitted topology
        hybridTop._write_split_itp( args.oitp )


#*******************************************************************************#
#*******************************************************************************#


    logfile.close()
if __name__ == '__main__':
    entry_point()
