#
#
#
#
# $Id: test_prepareLigand.py,v 1.5 2012/09/14 19:01:10 rhuey Exp $
#

import unittest, os, sys
from string import split, find
from MolKit import Read
#import MolKit.molecule
#import MolKit.protein
#MolKit.molecule.bhtreeFlag = 0
#MolKit.protein.bhtreeFlag = 0

class BaseTests(unittest.TestCase):

    def test_default(self):
        """
          l:vo:d:A:CKU:B:R:M
          Required parameter
           [-l]    ligand_filename
          Optional parameters:         defaults vs  options
           -o outputfilename         stem.pdbq     filename
           -d dictionary to write entry listing atom_type list, rbonds and zero_charge atoms list
           -A repair type                  ''      bonds_hydrogens;bonds;hydrogens
           -C do not add charges
           -K add Kollman charges         False      True  
           -U cleanup:merge type(s)    nphs_lps    nphs,lps,'' 
           -B type(s) of rotatable bnds to allow   amide_backbone_guanidinium,
                                                   amide_backbone,amide_guanidinium, 
                                                   backbone_guanidinium, 
                                                   amide, backbone,guanidinium,
                                                   ''
           -B set root                    auto     index 
           -M mode                    automatic    interactive

        """
        cmd = 'set PYTHONPATH=%s'%sys.path[0]
        cmd = cmd + "; ../prepare_ligand.py -l ind.pdb -o test_default.pdbq"
        os.system(cmd)
        fptr = open("test_default.pdbq")
        lines = fptr.readlines()
        fptr.close()
        std_ptr = open("std_ind.pdbq")
        std_lines = std_ptr.readlines()
        std_ptr.close()
        ctr = 0
        for l, std_l in zip(lines, std_lines):
            if find(l, 'REMARK')==-1:
                self.assertEqual(l, std_l)
                ctr+=1



    def test_allow_amide_torsions(self):
        """
        test allow amide torsions
        options
           -B amide
        """
        cmd = 'set PYTHONPATH=%s'%sys.path[0]
        cmd = cmd+"; ../prepare_ligand.py -l ind.pdb -B amide -o test_allow_amide.pdbq"
        os.system(cmd)
        fptr = open("test_allow_amide.pdbq")
        lines = fptr.readlines()
        fptr.close()
        std_ptr = open("std_ind_with_amide.pdbq")
        std_lines = std_ptr.readlines()
        std_ptr.close()
        ctr = 0
        for l, std_l in zip(lines, std_lines):
            if find(l, 'REMARK')==-1:
                self.assertEqual(l, std_l)
                ctr+=1
        #print 'found ', ctr, ' identical output lines'


    def test_set_root_to_5(self):
        """
        test specifying 5th atom as root
        options
           -l 5  set root to 5th atom
        """
        cmd = 'set PYTHONPATH=%s'%sys.path[0]
        cmd = cmd + "; ../prepare_ligand.py -l ind.pdb -R 5 -o test_set_root_to_5.pdbq"
        os.system(cmd)
        fptr = open("test_set_root_to_5.pdbq")
        lines = fptr.readlines()
        fptr.close()
        std_ptr = open("std_ind_setroot_to_5.pdbq")
        std_lines = std_ptr.readlines()
        std_ptr.close()
        ctr = 0
        for l, std_l in zip(lines, std_lines):
            if find(l, 'REMARK')==-1:
                self.assertEqual(l, std_l)
                ctr+=1
        #print 'found ', ctr, ' identical output lines'


    def test_inactivate_all_torsions(self):
        """
        test inactivate_all_torsions 
        options
           -Z 
        """
        cmd = 'set PYTHONPATH=%s'%sys.path[0]
        cmd = cmd+"; ../prepare_ligand.py -l ind.pdb -Z -o test_inactivate_all_torsions.pdbq"
        os.system(cmd)
        fptr = open("test_inactivate_all_torsions.pdbq")
        lines = fptr.readlines()
        fptr.close()
        std_ptr = open("std_ind_inactivate_all_torsions.pdbq")
        std_lines = std_ptr.readlines()
        std_ptr.close()
        ctr = 0
        for l, std_l in zip(lines, std_lines):
            if find(l, 'REMARK')==-1:
                self.assertEqual(l, std_l)
                ctr+=1
        #print 'found ', ctr, ' identical output lines'


    def wtest_setroot_method(self):
        #after autoroot?
        pass


    def wtest_toggle_torsion(self):
        pass

    def wtest_toggle_guanidinium_torsions(self):
        pass

    def wtest_toggle_amide_torsions(self):
        pass

    def wtest_toggle_ppbb_torsions(self):
        pass

    def wtest_toggle_all_torsions(self):
        pass


    def wtest_change_aromatic_criteria(self):
        #use mtx
        pass

    def wtest_limit_torsions(self):
        #use both types
        pass


    def wtest_set_carbon_names(self):
        #????
        pass


    def wtest_write_reference_file(self):
        #same torTree but different coords(?)
        #need to match up atoms....
        pass



class LigandFiles(unittest.TestCase):
    """
    Check 14 ligands in AutoDockTools/Tests/ligandfiles
    """
    def setUp(self):
        self.datadir = "ligandfiles"
        #self.files = [ 'cholesterol.pdb']
        #self.files = [ 'hxw1_rit.pdb']
        #self.files = [ 'pieceCRN.pdb']
        #self.files = [ 'fivPIECE.pdb']
        #self.files = [ 'ld.pdb']
        #self.files = [ 'txp.pdb']
        #self.files = [ 'faicar.pdbq']
        #self.files = [ 'ligandPROB.pdbq']
        #self.files = [ 'lov.pdbq']
        #self.files = [ 'h2.pdbq']
        #self.files = [ 'adenosine.mol2']
        #self.files = [ 'db0000.mol2'] #CHECK HYDROGENS
        #self.files = [ 'mtx.mol2']
        #self.files = [ 'xa.mol2']
        self.files = [ 'cholesterol.pdb', 'hxw1_rit.pdb', 'pieceCRN.pdb',
                       'fivPIECE.pdb', 'ld.pdb', 'txp.pdb', 'faicar.pdbq',
                       'ligandPROB.pdbq', 'lov.pdbq', 'h2.pdbq', 'adenosine.mol2',
                       'db0000.mol2', 'mtx.mol2', 'xa.mol2']
        self.hDict = {}
        self.hDict['cholesterol'] = 0
        self.hDict['hxw1_rit'] = 0
        self.hDict['pieceCRN'] = 1
        self.hDict['fivPIECE'] = 0
        self.hDict['ld'] = 0
        self.hDict['txp'] = 0
        self.hDict['faicar'] = 0
        self.hDict['ligandPROB'] = 0
        self.hDict['lov'] = 0
        self.hDict['h2'] = 0
        self.hDict['adenosine'] = 0
        self.hDict['db0000'] = 0
        self.hDict['mtx'] = 0
        self.hDict['xa'] = 0



    def test_compare_ligand_files(self):
        """
        Test whether the written outputfile contents match stds
        """
        for f in self.files:
            stem = os.path.splitext(f)[0] 
            #print "\n\ntesting ", stem
            inputfilename = os.path.join(self.datadir, f)
            #print "inputfilename=", inputfilename
            outputfilename = "test_pl_" + stem + ".pdbq"
            #print "outputfilename=", outputfilename
            stdfilename = os.path.join(self.datadir, stem +".out.pdbq")
            #print "stdfilename=", stdfilename
            if stem in ['fivPIECE', 'ld']:
                cmd = 'set PYTHONPATH=%s'%sys.path[0]
                cmd = cmd+";../prepare_ligand.py -l " + inputfilename + " -K -o " + outputfilename
            elif stem in ['db0000', 'faicar', 'ligandPROB', 'lov', 'h2', 'adenosine', 'mtx', 'xa']:
                #these already had charges:
                cmd = 'set PYTHONPATH=%s'%sys.path[0]
                cmd = cmd+"; ../prepare_ligand.py -l " + inputfilename + " -C -o " + outputfilename 
                #cmd = "$PYTHONPATH/AutoDockTools/Utilities/prepare_ligand.py -l " + inputfilename + " -C -o " + outputfilename 
            elif stem in ['pieceCRN']:
                cmd = 'set PYTHONPATH=%s'%sys.path[0]
                cmd = cmd+";../prepare_ligand.py -l " + inputfilename + " -K  -o " + outputfilename
            elif self.hDict[stem]:
                cmd = 'set PYTHONPATH=%s'%sys.path[0]
                cmd = cmd+";../prepare_ligand.py -l " + inputfilename + " -A hydrogens -o " + outputfilename
            else:
                cmd = 'set PYTHONPATH=%s'%sys.path[0]
                cmd = cmd+";../prepare_ligand.py -l " + inputfilename + " -o " + outputfilename
            os.system(cmd)
            fptr = open(outputfilename)
            lines = fptr.readlines()
            fptr.close()
            std_ptr = open(stdfilename)
            std_lines = std_ptr.readlines()
            std_ptr.close()
            ctr = 0
            ok = 1
            numatoms = 0
            std_numatoms = 0
            numbranches = 0
            std_numbranches = 0
            ctroot = 0
            numroot = 0
            std_ctroot = 0
            std_numroot = 0
            numaromCs = 0
            std_numaromCs = 0
            match = {}
            match['atoms'] = True
            match['number_tors'] = True
            match['number_branches'] = True
            match['TORSDOF'] = True
            match['root'] = True
            match['tors_atomnames'] = True
            match['aromaticCs'] = True
            for l, std_l in zip(lines, std_lines):
                if find(l, "active torsions")>-1:
                    testnum = int(split(l)[1])
                    stdnum = int(split(std_l)[1])
                    if testnum!=stdnum:
                        #print stem, ' FAILED active torsions'
                        match['number_tors'] = (testnum, stdnum)
                    #self.assertEqual(testnum, stdnum)
                elif find(l, 'REMARK')==-1 and ok:
                    try:
                        self.assertEqual(l, std_l)
                        ctr+=1
                    except:
                        ok = 0
                        print stem, " ERROR: lines did not match "
                        print l, std_l

            for l in lines:
                if find(l, "HETATM")==0 or find(l, "ATOM")==0:
                    numatoms += 1
                    ll = split(l)
                    if ll[1][0]=='A':
                        numaromCs += 1
                elif find(l, "BRANCH")==0:
                    numbranches += 1
                elif find(l, "ROOT")==0:
                    ctroot = 1
                elif find(l, "ENDROOT")==0:
                    ctroot = 0
                elif find(l, "TORSDOF")==0:
                    testTORSDOF = int(split(l)[1])
                    #self.assertEqual(testnum, stdnum)
                if ctroot:
                    numroot += 1


            for std_l in std_lines:
                if find(std_l, "HETATM")==0 or find(std_l, "ATOM")==0:
                    std_numatoms += 1
                    ll = split(std_l)
                    if ll[1][0]=='A':
                        std_numaromCs += 1
                elif find(std_l, "BRANCH")==0:
                    std_numbranches += 1
                elif find(std_l, "ROOT")==0:
                    std_ctroot = 1
                elif find(std_l, "ENDROOT")==0:
                    std_ctroot = 0
                elif find(std_l, "TORSDOF")==0:
                    stdTORSDOF = int(split(std_l)[1])
                    if testTORSDOF!=stdTORSDOF:
                        match['TORSDOF'] = (testTORSDOF, stdTORSDOF)
                if std_ctroot:
                    std_numroot += 1

            if numatoms!=std_numatoms:
                #print stem, ' FAILED numatoms'
                match['atoms'] = ( numatoms, std_numatoms)
            #self.assertEqual(numatoms, std_numatoms) 
            if numbranches!=std_numbranches:
                #print stem, ' FAILED numbranches'
                match['number_branches'] = ( numbranches, std_numbranches)
            #self.assertEqual(numbranches, std_numbranches) 
            if numroot!=std_numroot:
                #print stem, ' FAILED numroot'
                match['root'] = ( numroot, std_numroot)
            #self.assertEqual(numroot, std_numroot) 

            #????match['tors_atomnames'] = True
            if numaromCs!=std_numaromCs:
                #print stem, ' FAILED numaromCs'
                match['numaromCs'] = ( numaromCs, std_numaromCs)

            #print stem, ' new vs std:',
            if ok: 
                pass
                #print ' passed'
            else:
                ct = 0
                firstE = 1
                for k, v in match.items():
                    #print k, v
                    if v != True:
                        if firstE:
                            firstE = 0
                            print
                        print "FAILED ", k, ':', v[0], ' - ', v[1]
                        ct += 1
                if ct==0:
                    print " differed in lines only"
            #print 'found ', ctr, ' identical output lines'

    

class LigandMolecules(unittest.TestCase):
    """
    Check 14 ligands built from files in AutoDockTools/Tests/ligandfiles
    """
    def setUp(self):
        self.datadir = "ligandfiles"
        #self.files = [ 'cholesterol.pdb']
        #self.files = [ 'hxw1_rit.pdb']
        #self.files = [ 'pieceCRN.pdb']
        #self.files = [ 'fivPIECE.pdb']
        #self.files = [ 'ld.pdb']
        #self.files = [ 'txp.pdb']
        #self.files = [ 'faicar.pdbq']
        #self.files = [ 'ligandPROB.pdbq']
        #self.files = [ 'lov.pdbq']
        #self.files = [ 'h2.pdbq']
        #self.files = [ 'adenosine.mol2']
        #self.files = [ 'db0000.mol2'] #CHECK HYDROGENS
        #self.files = [ 'mtx.mol2']
        #self.files = [ 'xa.mol2']
        self.files = [ 'cholesterol.pdb', 'hxw1_rit.pdb', 'pieceCRN.pdb',
                       'txp.pdb', 'faicar.pdbq',
                       'lov.pdbq', 'h2.pdbq', 'adenosine.mol2',
                       'db0000.mol2']
        self.hDict = {}
        self.hDict['cholesterol'] = 0
        self.hDict['hxw1_rit'] = 0
        self.hDict['pieceCRN'] = 1
        self.hDict['fivPIECE'] = 0
        self.hDict['ld'] = 0
        self.hDict['txp'] = 0
        self.hDict['faicar'] = 0
        self.hDict['ligandPROB'] = 0
        self.hDict['lov'] = 0
        self.hDict['h2'] = 0
        self.hDict['adenosine'] = 0
        self.hDict['db0000'] = 0
        self.hDict['mtx'] = 0
        self.hDict['xa'] = 0


    def test_ligands(self):
        """
        Test whether the names of atoms in rotatable bonds match stds
        """
        for f in self.files:
            stem = os.path.splitext(f)[0] 
            #print "\n\ntesting ", stem
            inputfilename = os.path.join(self.datadir, f)
            #print "inputfilename=", inputfilename
            outputfilename = "test_pl_" + stem + ".pdbq"
            #print "outputfilename=", outputfilename
            stdfilename = os.path.join(self.datadir, stem +".out.pdbq")
            #print "stdfilename=", stdfilename
            if stem in ['fivPIECE', 'ld']:
                cmd = 'set PYTHONPATH=%s'%sys.path[0]
                cmd = cmd+";../prepare_ligand.py -l " + inputfilename + " -K -o " + outputfilename
            elif stem in ['db0000', 'faicar', 'ligandPROB', 'lov', 'h2', 'adenosine', 'mtx', 'xa']:
                #these already had charges:
                cmd = 'set PYTHONPATH=%s'%sys.path[0]
                cmd = cmd+";../prepare_ligand.py -l " + inputfilename + " -C -o " + outputfilename 
            elif stem in ['pieceCRN']:
                cmd = 'set PYTHONPATH=%s'%sys.path[0]
                cmd = cmd+";../prepare_ligand.py -l " + inputfilename + " -K  -o " + outputfilename
            elif self.hDict[stem]:
                cmd = 'set PYTHONPATH=%s'%sys.path[0]
                cmd = cmd+";../prepare_ligand.py -l " + inputfilename + " -o " + outputfilename
            else:
                cmd = 'set PYTHONPATH=%s'%sys.path[0]
                cmd = cmd+";../prepare_ligand.py -l " + inputfilename + " -o " + outputfilename
            os.system(cmd)
            lig = Read(outputfilename)[0]
            std_lig = Read(stdfilename)[0]
            #self.assertEqual(len(lig.torTree.torsionMap),len(std_lig.torTree.torsionMap) 
            #self.assertEqual(lig.ROOT.name, std_lig.ROOT.name)
            #check that the names of the atoms in rotatable bnds match
            lig_list = []
            for n in lig.torTree.torsionMap:
                bnd = n.bond
                ats = lig.allAtoms
                lig_list.append((ats[bnd[0]].name, ats[bnd[1]].name))
            for n in std_lig.torTree.torsionMap:
                bnd = n.bond
                std_ats = std_lig.allAtoms
                std_names = (std_ats[bnd[0]].name, std_ats[bnd[1]].name)
                #print "std_names=", std_names
                rev_std_names = (std_ats[bnd[1]].name, std_ats[bnd[0]].name)
                #print "rev_std_names=", rev_std_names
                #self.assertEqual( (std_names in lig_list or rev_std_names in lig_list), True)
            del lig
            del std_lig




if __name__ == '__main__':
    test_cases = [
        'BaseTests',
        'LigandFiles',
        ]
    
    unittest.main( argv=([__name__ ,] + test_cases) )
    #unittest.main()

