#!/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 map ligand atoms for hybrid structure/topology generation.
"""

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
from pmx.ndx import *
from pmx.scripts.cli import check_unknown_cmd
from pmx.ligand_alchemy import *
import time
import warnings
from rdkit import Chem


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

    parser = argparse.ArgumentParser(description='''
Provided two structures find atoms to be morphed.

''', 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('-o1',
                        metavar='pairs1.dat',
                        dest='o1',
                        type=str,
                        help='Output pairs: column1:mol1, column2:mol2. '
                        'Default is "pairs1.dat"',
                        default='pairs1.dat')
    parser.add_argument('-o2',
                        metavar='pairs2.dat',
                        dest='o2',
                        type=str,
                        help='Output pairs: column1:mol2, column2:mol1. '
                        'Default is "pairs2.dat"',
                        default='pairs2.dat')
    parser.add_argument('-opdb1',
                        metavar='out_pdb1.pdb',
                        dest='opdb1',
                        type=str,
                        help='Optional output: superimposed structure 1. ')
    parser.add_argument('-opdb2',
                        metavar='out_pdb2.pdb',
                        dest='opdb2',
                        type=str,
                        help='Optional output: superimposed structure 2. ')
    parser.add_argument('-opdbm1',
                        metavar='out_pdb_morphe1.pdb',
                        dest='opdbm1',
                        type=str,
                        help='Optional output: morphable atoms in structure 1. ')
    parser.add_argument('-opdbm2',
                        metavar='out_pdb_morphe2.pdb',
                        dest='opdbm2',
                        type=str,
                        help='Optional output: morphable atoms in structure 2. ')
    parser.add_argument('-score',
                        metavar='out_score.dat',
                        dest='score',
                        type=str,
                        help='Optional output: score of the morph. '
                        'Default is "out_score.dat"',
                        default='out_score.dat')
    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('-log',
                        metavar='mapping.log',
                        dest='log',
                        type=str,
                        help='Output: log file. '
                        'Default is "mapping.log"',
                        default='mapping.log')
    parser.add_argument('--no-alignment',
                        dest='alignment',
                        help='Should the alignment method be disabled (default enabled)',
                        action='store_false')
    parser.add_argument('--no-mcs',
                        dest='mcs',
                        help='Should the MCS method be disabled (default enabled)',
                        action='store_false')
    parser.add_argument('--no-H2H',
                        dest='H2H',
                        help='Should non-polar hydrogens be discarded from morphing into any other hydrogen (default True)',
                        action='store_false')
    parser.add_argument('--H2Hpolar',
                        dest='H2Hpolar',
                        help='Should polar hydrogens be morphed into polar hydrogens (default False)',
                        action='store_true')
    parser.add_argument('--H2Heavy',
                        dest='H2Heavy',
                        help='Should hydrogen be morphed into a heavy atom (default False)',
                        action='store_true')
    parser.add_argument('--RingsOnly',
                        dest='RingsOnly',
                        help='Should rings only be used in the MCS search and alignemnt (default False)',
                        action='store_true')
    parser.add_argument('--dMCS',
                        dest='dMCS',
                        help='Should the distance criterium be also applied in the MCS based search (default False)',
                        action='store_true')
    parser.add_argument('--swap',
                        dest='swap',
                        help='Try swapping the molecule order which would be a cross-check and require double execution time (default False)',
                        action='store_true')
    parser.add_argument('--no-chirality',
                        dest='chirality',
                        help='Perform chirality check for MCS mapping (default True)',
                        action='store_false')
    parser.add_argument('--d',
                        metavar='0.05',
                        dest='d',
                        type=float,
                        help='Distance (nm) between atoms to consider them morphable for alignment approach (default 0.05 nm).',
                        default=0.05)
    parser.add_argument('--timeout',
                        metavar='10',
                        dest='timeout',
                        type=int,
                        help='Maximum time (s) for an MCS search (default 10 s).',
                        default=10)

    parser.set_defaults(alignment=True,mcs=True,H2H=True,H2Hpolar=False,H2Heavy=False,
                        RingsOnly=False,dMCS=False,swap=False,chirality=True)
    args, unknown = parser.parse_known_args()
    check_unknown_cmd(unknown)

    return args

def entry_point():
    args = parse_options()
    if args.H2Heavy==True:
        args.H2H=True
    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
    bH2H = args.H2H
    bH2Hpolar = args.H2Hpolar
    bH2Heavy = args.H2Heavy
    bChiral = args.chirality
    bSwap = args.swap
    bRingsOnly = args.RingsOnly
    d = args.d
    timeout = args.timeout
    bdMCS = args.dMCS 

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

#####################################
   # identify the methods to use 
    bAlignment = args.alignment
    bMCS = args.mcs
    if bMCS==False and bAlignment==False:
        doLog(logfile,"No method (alignment, mcs) was selected.",commandName='atomMapping')
        sys.exit(0)
    doLog(logfile,"Morphable atoms will be identified using the following methods:",commandName='atomMapping')
    doLog(logfile,"Alignment: {0}".format(bAlignment),commandName='atomMapping')
    doLog(logfile,"MCS: {0}".format(bMCS),commandName='atomMapping')

#####################################
    # read index
    read_from_idx1 = False
    read_from_idx2 = False
    if args.n1 is not None:
        read_from_idx1 = True
        idx1 = IndexFile(args.n1)
    if args.n2 is not None:
        read_from_idx2 = True
        idx2 = IndexFile(args.n2)

#####################################
    # reformat PDB to read two letter atoms properly (why does it have to be so painful?)
    if not os.path.exists( args.i1 ):
        doLog(logfile,'Input pdb1 not found. Exiting...',commandName='atomMapping')
        sys.exit(0)
    if not os.path.exists( args.i2 ):
        doLog(logfile,'Input pdb2 not found. Exiting...',commandName='atomMapping')
        sys.exit(0)
    pid = os.getpid()
    pdbName1,atomNameID1,sigmaHoleID1 = reformatPDB(args.i1,1,pid)
    pdbName2,atomNameID2,sigmaHoleID2 = reformatPDB(args.i2,2,pid)
    mol1 = Chem.MolFromPDBFile(pdbName1,removeHs=False,sanitize=False)
    mol2 = Chem.MolFromPDBFile(pdbName2,removeHs=False,sanitize=False)
    try:
        Chem.SanitizeMol(mol1)
    except:
        pass
    try:
        Chem.SanitizeMol(mol2)
    except:
        pass
    os.remove(pdbName1)
    os.remove(pdbName2)

#####################################
    ## prepare molecules for MCS ####
    molForMcs1 = cp.deepcopy(mol1)
    molForMcs2 = cp.deepcopy(mol2)
    try:
        rdmolops.AssignAtomChiralTagsFromStructure(mol1)
        rdmolops.AssignAtomChiralTagsFromStructure(mol2)
        molForMcs1 = cp.deepcopy(mol1)
        molForMcs2 = cp.deepcopy(mol2)
        rdmolops.AssignStereochemistry(mol1)
        rdmolops.AssignStereochemistry(mol2)
    except:
        doLog(logfile,"Chirality not assigned",commandName='atomMapping')

#####################################
    # deal with the -RingsOnly flag
    bYesRings = checkRingsOnlyFlag(mol1,mol2)
    if(bRingsOnly==True):
        if( bYesRings==True ):
            bRingsOnly=True
        else:
            doLog(logfile,"-RingsOnly flag is unset, because one (or both) molecule has no rings",commandName='atomMapping')
    
    n1 = []
    n2 = []

#########################
    # make molecule copies #
    molcp1 = cp.deepcopy(mol1)
    molcp2 = cp.deepcopy(mol2)

#########################
    # create an object of LigandsAtomMapping class
    doLog(logfile,"Initializing the main LigandsAtomMapping object",commandName='atomMapping')
    ligMap = LigandAtomMapping(mol1=mol1, mol2=mol2, molForMcs1=molForMcs1, molForMcs2=molForMcs2, bH2H=bH2H, bH2Hpolar=bH2Hpolar, bH2Heavy=bH2Heavy, bdMCS=bdMCS, bRingsOnly=bRingsOnly, d=d, bChiral=bChiral, sigmaHoleID1=sigmaHoleID1, sigmaHoleID2=sigmaHoleID2, timeout=timeout, logfile=logfile, bElements=False, bCarbonize=False, commandName='atomMapping' )



#########################
    # mcs
    bMCSfailed = False
    n1mcs = []
    n2mcs = []
    if(bMCS==True):
        doLog(logfile,"The topology matching approach will be used (MCS)",commandName='atomMapping')
        doLog(logfile,"fmcs module: Copyright (c) 2012 Andrew Dalke Scientific AB",commandName='atomMapping')
#########################
        # mcs: only rings
        if(bRingsOnly==True):
            doLog(logfile,"MCS: rings only...",commandName='atomMapping')
            ligMap.bElements = False
            ligMap.bCarbonize = False
            ligMap.bRingsOnly = True
            n1mcs,n2mcs = ligMap.mcs()
        else:
#########################
            # mcs: all atoms, match elements, modified molecules
            n1A = []
            n2A = []
            ligMap.bElements = True
            ligMap.bCarbonize = True
            ligMap.bRingsOnly = False
            doLog(logfile,"MCS: trying to run an MCS using all atoms: modified molecules, match elements...",commandName='atomMapping')
            n1A,n2A = ligMap.mcs()
            doLog(logfile,"MCS: all atoms, modified molecules, size of mapping: {0}".format(len(n1A)),commandName='atomMapping')

#########################
            # mcs: all atoms, match elements, unmodified molecules
            ligMap.bElements = True
            ligMap.bCarbonize = False
            ligMap.bRingsOnly = False
            doLog(logfile,"MCS: trying to run an MCS using all atoms: unmodified molecules, match elements...",commandName='atomMapping')
            n1B,n2B = ligMap.mcs()
            doLog(logfile,"MCS: all atoms, unmodified molecules, size of mapping: {0}".format(len(n1B)),commandName='atomMapping')

#########################
            # mcs: select the best mapping
            n1mcs,n2mcs = ligMap._compare_mappings_by_size( mol1,mol2, n1A, n2A, n1B, n2B )

#########################
            # mcs: rings only
            if( bYesRings==True ):
                doLog(logfile,"MCS: trying to run an MCS using rings only...",commandName='atomMapping')
                ligMap.bElements = False
                ligMap.bCarbonize = False
                ligMap.bRingsOnly = True
                n1C,n2C = ligMap.mcs()
                n1mcs,n2mcs = ligMap._compare_mappings_by_size( mol1,mol2, n1mcs, n2mcs, n1C, n2C ) 

        doLog(logfile,"MCS: using variant with the mapping size of {0}".format(len(n1mcs)),commandName='atomMapping')

        if len(n1mcs)==0:
             bMCSfailed = True

#########################
    # alignment
    n1align = []
    n2align = []
    if( (bAlignment==True) or (bMCSfailed==True) ):
        if bMCSfailed==True:
            doLog(logfile,"The MCS approach did not find any mapping",commandName='atomMapping')
        doLog(logfile,"The alignment approach will be used",commandName='atomMapping')
        doLog(logfile,"Tosco, P., Balle, T. & Shiri, F. Open3DALIGN: an open-source software aimed at unsupervised ligand alignment. J Comput Aided Mol Des 25:777-83 (2011)",commandName='atomMapping')
        doLog(logfile,"Alignment is based on atom logP contributions: S. A. Wildman and G. M. Crippen JCICS _39_ 868-873 (1999)",commandName='atomMapping')

#########################
    # alignment: rings only
        if(bRingsOnly==True):
            ligMap.bRingsOnly = True
            n1align,n2align,pyO3A = ligMap.alignment()
            doLog(logfile,"Alignment: size of mapping: {0}".format(len(n1align)),commandName='atomMapping')
        else:
#########################
            # alignment: all atoms
            doLog(logfile,"Alignment: trying to align all atoms...",commandName='atomMapping')
            ligMap.bRingsOnly = False
            n1align,n2align,pyO3A = ligMap.alignment()
            doLog(logfile,"Alignment: size of mapping: {0}".format(len(n1align)),commandName='atomMapping')
#########################
            # alignment: rings only
            if( bYesRings==True ):
                doLog(logfile,"Alignment: trying to align rings only...",commandName='atomMapping')
                ligMap.bRingsOnly = True
                mol1 = cp.deepcopy(molcp1)
                mol2 = cp.deepcopy(molcp2)
                n1Balign,n2Balign,pyO3A = ligMap.alignment()
                doLog(logfile,"Size of mapping: {0}".format(len(n1Balign)),commandName='atomMapping')

#########################
                # alignment: select the best mapping
                if(len(n1align)<=len(n1Balign)):
                    doLog(logfile,"Alignment: using ring only alignment result.",commandName='atomMapping')
                    n1align = n1Balign
                    n2align = n2Balign
                else:
                    doLog(logfile,"Alignment: using all atom alignment result.",commandName='atomMapping')

#########################
    # select the best mapping: mcs vs alignment
    if len(n1align)>=len(n1mcs):
        doLog(logfile,"The final result is based on the Alignment.",commandName='atomMapping')
        n1 = n1align
        n2 = n2align
    else:
        doLog(logfile,"The final result is based on the MCS.\n",commandName='atomMapping')
        n1 = n1mcs
        n2 = n2mcs

#########################
    # swap TODO


       
#########################
    # check
    if( len(n1) != len(n2) ):
        doLog(logfile,"ERROR: something went wrong.",commandName='atomMapping')
        doLog(logfile,"ERROR: Number of the morphable atoms in the ligands does not match.",commandName='atomMapping')
        sys.exit(1)

#########################
    # calculate score
    score = ligMap._calc_score(n1,n2,mol1,mol2)

#########################
    # report results
    doLog(logfile,"Final results:",commandName='atomMapping')
    if( bH2H==True or bH2Heavy==True ):
        doLog(logfile,"Atoms considered in mol1: {0}".format(mol1.GetNumAtoms()),commandName='atomMapping')
        doLog(logfile,"Atoms considered in mol2: {0}".format(mol2.GetNumAtoms()),commandName='atomMapping')
    else:
        doLog(logfile,"Atoms considered in mol1: {0}".format(mol1.GetNumHeavyAtoms()),commandName='atomMapping')
        doLog(logfile,"Atoms considered in mol2: {0}".format(mol2.GetNumHeavyAtoms()),commandName='atomMapping')
    doLog(logfile,"Morphable atoms in both molecules: {0} {1}".format(len(n1),len(n2)),commandName='atomMapping')
    doLog(logfile,"Dissimilarity (distance) score: %.4f\n" % score,commandName='atomMapping')
    if args.score!=None:
        fp = open(args.score,'w')
        fp.write("Score: %.4f\n" % score)
        fp.close()

#########################
    # restore atom names
    restoreAtomNames(mol1,atomNameID1)
    restoreAtomNames(mol2,atomNameID2)

#########################
    # output pdb
    if args.opdb1 != None:
        Chem.MolToPDBFile(molcp1,args.opdb1)
    if args.opdb2 != None:
        try:
            Chem.rdMolAlign.AlignMol(mol2,molcp1,atomMap=list(zip(n2,n1)))
        except:
            doLog(logfile,"PDB output: cannot superimpose -opdb2 structure. Maybe no morphable atoms have been found",commandName='atomMapping')
        Chem.MolToPDBFile(mol2,args.opdb2)
    if args.opdbm1!=None:
        mol = submol_by_index(molcp1,n1)
        Chem.MolToPDBFile(mol,args.opdbm1)
    if args.opdbm2!=None:
        mol = submol_by_index(mol2,n2)
        Chem.MolToPDBFile(mol,args.opdbm2)

#########################
    # output pairs
    ligMap._write_pairs(n1,n2,args.o1)
    ligMap._write_pairs(n2,n1,args.o2)


if __name__ == '__main__':
    entry_point()

