#
#
#
#
# $Id: test_prepareReceptor.py,v 1.3 2011/09/07 23:02:25 annao Exp $
#

import unittest, os, sys
from string import split, find, strip
from MolKit import Read

class BaseTests(unittest.TestCase):

    def test_default(self):
        """
        test default options 
        options
          [-o pdbqs filename]  (output file)
          [-v]    verbose output          
          [-A]  type(s) of repairs to make:\n\t\t bonds_hydrogens, bonds, hydrogens, None
          [-G]  add Gasteiger charges        
          [-C]  do not add charges        
          [-U]  cleanup by merging/removing:\n\t\t nphs_lps, nphs, lps, waters, nonstdres, 
          [-M]  interactive           automatic 
        """
        testfile = "test_new.pdbqs"
        cmd = "set PYTHONPATH=%s"%sys.path[0]
        cmd = cmd +"; ../prepare_receptor.py -r hsg1.pdb -A hydrogens -U nphs_waters -o " + testfile
        os.system(cmd)
        fptr = open(testfile)
        lines = fptr.readlines()
        std_ptr = open("std_hsg1.pdbqs")
        std_lines = std_ptr.readlines()
        ctr = 0
        badlines = []
        for l, std_l in zip(lines, std_lines):
            if find(l, 'ATOM')==0 and find(std_l,'ATOM')==0:
                ind = find(l, '-0.000')
                if ind>-1:
                    l = l[:ind] + ' ' + l[ind+1:]
                ind = find(std_l, '-0.000')
                if ind>-1:
                    std_l = std_l[:ind] + ' ' + l[ind+1:]
                self.assertEqual(l, std_l)



    def test_donot_remove_waters(self):
        """
        test donot remove_waters
        options
           -U nphs.... 
        """
        testfile = "test_keep_waters.pdbqs" 
        cmd = "set PYTHONPATH = %s"%sys.path[0] 
        cmd = cmd + "; ../prepare_receptor.py -r hsg1.pdb -A hydrogens -U nphs -o " + testfile
        #cmd = "$PYTHONPATH/AutoDockTools/Utilities/prepare_receptor.py -r hsg1.pdb -A hydrogens -U nphs -o " + testfile
        os.system(cmd)
        fptr = open(testfile)
        lines = fptr.readlines()
        std_ptr = open("std_with_waters.pdbqs")
        std_lines = std_ptr.readlines()
        for l, std_l in zip(lines, std_lines):
            if find(l, 'ATOM')==0 and find(std_l,'ATOM')==0:
                ind = find(l, '-0.000')
                if ind>-1:
                    l = l[:ind] + ' ' + l[ind+1:]
                ind = find(std_l, '-0.000')
                if ind>-1:
                    std_l = std_l[:ind] + ' ' + l[ind+1:]
                self.assertEqual(strip(l), strip(std_l))


    def test_remove_non_std_chains(self):
        """
        12/29: need better testmolecule for this one
        test remove_non_std_chains
        options
           -U 
        """
        testfile = "test_no_nonstd_chains.pdbqs"
        cmd = "set PYTHONPATH=%s"%sys.path[0]
        cmd = cmd + "; ../prepare_receptor.py -r hsg1.pdb -U nphs_nonstdres -A checkhydrogens -o " + testfile
        #cmd = "$PYTHONPATH/AutoDockTools/Utilities/prepare_receptor.py -r hsg1.pdb -A hydrogens -U nphs_nonstdres -o " + testfile
        os.system(cmd)
        fptr = open(testfile)
        lines = fptr.readlines()
        std_ptr = open("std_hsg1.pdbqs")
        std_lines = std_ptr.readlines()
        for l, std_l in zip(lines, std_lines):
            if find(l, 'ATOM')==0 and find(std_l,'ATOM')==0:
                ind = find(l, '-0.000')
                if ind>-1:
                    l = l[:ind] + ' ' + l[ind+1:]
                ind = find(std_l, '-0.000')
                if ind>-1:
                    std_l = std_l[:ind] + ' ' + l[ind+1:]
                self.assertEqual(strip(l), strip(std_l))



#    def test_do_not_add_polar_hydrogens(self):
#        """
#        THIS IS NOT SUPPORTED!!!
#        test do_not_add_polar_hydrogens
#        options
#           -A  repairs==""
#        """
#        testfile = "test_nopolarhydrogens.pdbqs"
#        cmd = "set PYTHONPATH=%s"%sys.path[0]
#        cmd = cmd + "; ../prepare_receptor.py -r hsg1.pdbqs -A '' -o " + testfile
#        #cmd = "$PYTHONPATH/AutoDockTools/Utilities/prepare_receptor.py -r hsg1.pdb -o " + testfile
#        os.system(cmd)
#        sleep(5)
#        fptr = open(testfile)
#        lines = fptr.readlines()
#        std_ptr = open("std_no_polar_hydrogens.pdbqs")
#        std_lines = std_ptr.readlines()
#        for l, std_l in zip(lines, std_lines):
#            if find(l, 'ATOM')==0 and find(std_l,'ATOM')==0:
#                ind = find(l, '-0.000')
#                if ind>-1:
#                    l = l[:ind] + ' ' + l[ind+1:]
#                self.assertEqual(strip(l), strip(std_l))


    def test_add_gasteiger_charges(self):
        """
        test add_gasteiger_charges
        options
           -G 
        """
        testfile = "test_gasteiger.pdbqs" 
        cmd = "set PYTHONPATH=%s"%sys.path[0]
        cmd = cmd + "; ../prepare_receptor.py -r hsg1.pdbqs  -G -o " + testfile
        #cmd = cmd + "; ../prepare_receptor.py -r hsg1.pdbqs -U waters -G -o " + testfile
        #cmd = "$PYTHONPATH/AutoDockTools/Utilities/prepare_receptor.py -r hsg1.pdb -A hydrogens -U nphs_waters -G -o " + testfile
        os.system(cmd)
        fptr = open(testfile)
        lines = fptr.readlines()
        std_ptr = open("std_gasteiger.pdbqs")
        std_lines = std_ptr.readlines()
        for l, std_l in zip(lines, std_lines):
            if find(l, 'ATOM')==0 and find(std_l,'ATOM')==0:
                ind = find(l, '-0.000')
                if ind>-1:
                    l = l[:ind] + ' ' + l[ind+1:]
                ind = find(std_l, '-0.000')
                if ind>-1:
                    std_l = std_l[:ind] + ' ' + l[ind+1:]
                #self.assertEqual(strip(l), strip(std_l))
                line1 = strip(l).split()
                line2 = strip(std_l).split()
                for s1, s2 in zip(line1, line2):
                    if not s1.isalnum(): self.assertAlmostEqual(float(s1), float(s2), places=2)
                    else: self.assertEqual(s1, s2)
                        

    def test_add_no_charges(self):
        """
        test add_no_charges
        options
           -C 
        """
        testfile = "test_add_no_charges.pdbqs"
        cmd = "set PYTHONPATH=%s"%sys.path[0]
        cmd = cmd + "; ../prepare_receptor.py -r std_hsg1.pdbqs  -A None -C -o " + testfile
        #cmd = "$PYTHONPATH/AutoDockTools/Utilities/prepare_receptor.py -r std_hsg1.pdbqs -C -o " + testfile
        os.system(cmd)
        fptr = open(testfile)
        lines = fptr.readlines()
        std_ptr = open("std_hsg1.pdbqs")
        std_lines = std_ptr.readlines()
        for l, std_l in zip(lines, std_lines):
            if find(l, 'ATOM')==0 and find(std_l,'ATOM')==0:
                ind = find(l, '-0.000')
                if ind>-1:
                    l = l[:ind] + ' ' + l[ind+1:]
                ind = find(std_l, '-0.000')
                if ind>-1:
                    std_l = std_l[:ind] + ' ' + l[ind+1:]
                self.assertEqual(strip(l), strip(std_l))


    def ztest_do_not_add_bonds(self):
        """
        test do_not_add_bonds
        options
           -A None
        @@need a better file to test for this one...
        """
        
        testfile = "test_do_not_add_bonds.pdbqs"
        cmd = "set PYTHONPATH=%s"%sys.path[0]
        cmd = cmd + "; ../prepare_receptor.py -r std_hsg1.pdbqs -A checkhydrogens -o " + testfile
        os.system(cmd)
        fptr = open(testfile)
        lines = fptr.readlines()
        std_ptr = open("std_hsg1.pdbqs")
        std_lines = std_ptr.readlines()
        for l, std_l in zip(lines, std_lines):
            if find(l, 'ATOM')==0 and find(std_l,'ATOM')==0:
                ind = find(l, '-0.000')
                if ind>-1:
                    l = l[:ind] + ' ' + l[ind+1:]
                ind = find(std_l, '-0.000')
                if ind>-1:
                    std_l = std_l[:ind] + ' ' + l[ind+1:]
                self.assertEqual(l, std_l)




if __name__ == '__main__':
    unittest.main()
