#
# $Id: test_AutoGrid.py,v 1.11.10.1 2016/02/12 08:01:40 annao Exp $
#

import unittest
import os
from MolKit import Read

from PyAutoDock.AutoGrid import AutoGrid


#
# Unit Tests
#



class AutoGridTest(unittest.TestCase):
    def setUp(self):
        self.receptor = Read('Data/hsg1_sm.pdbqs')[0]
        self.receptor.buildBondsByDistance()


    def tearDown(self):
        pass

# AutoGridTest


class BasicAutoGridTest(AutoGridTest):

    def test_constructor(self):
        """ test that AutoGrid instance is created
        """
        ag = AutoGrid(self.receptor, center=[2.5, 6.5, -7.5])
        self.assertEqual(ag.__class__, AutoGrid)


    def test_atom_score_array(self):
        """test the AutoGrid estat_map score_array validity
        """
        ag = AutoGrid(self.receptor, atom_types = ['C'],
                      npts=[3,3,3],
                      center=[2.5, 6.5, -7.5])

        sa = ag.estat_map_scorer.get_score_array()
#        print sa
        self.assertEqual( sa.shape, (150,27))
        ag.write_maps()


    def test_write_maps_fld(self):
        """test the AutoGrid write_maps_fld
        """
        ag = AutoGrid(self.receptor, atom_types = ['C', 'A','N','O','S','H'],
                      npts=[3,3,3],
                      center=[2.5, 6.5, -7.5])
        filename = self.receptor.name + '.maps.fld'
        ag.write_maps_fld(filename)
        self.assertEqual( os.path.exists(filename), True)

        #compare with reference
        fptr = open("Data/ref.maps.fld")
        reflines = fptr.readlines()
        fptr.close()
        fptr = open(filename)
        testlines = fptr.readlines()
        fptr.close()
        for ref, test in zip(reflines, testlines):
            self.assertEqual(ref, test)
        


    def test_write_xyz_file(self):
        """test the AutoGrid write_xyz_file
        """
        ag = AutoGrid(self.receptor, atom_types = ['C', 'A','N','O','S','H'],
                      npts=[60,60,60],
                      center=[2.5, 6.5, -7.5])
        filename = self.receptor.name + '.maps.xyz'
        ag.write_xyz_file(filename)
        self.assertEqual( os.path.exists(filename), True)

        #compare with reference
        fptr = open("Data/ref.maps.xyz")
        reflines = fptr.readlines()
        fptr.close()
        fptr = open(filename)
        testlines = fptr.readlines()
        fptr.close()
        for ref, test in zip(reflines, testlines):
            self.assertEqual(ref, test)
        


class CompareMapsAutoGridTest(AutoGridTest):

    def write_map(self, atom_types):
        ag = AutoGrid(self.hsg1_sm, atom_types = atom_types,
                  npts=[2,2,2],
                  center=[2.5, 6.5, -7.5])
        ag.write_maps()
        

    def setUp(self):
        if not hasattr(self, 'hsg1_sm'):
            self.hsg1_sm = Read('Data/hsg1_sm.pdbqs')[0]
            self.hsg1_sm.buildBondsByDistance()


class CompareCMapsAutoGridTest(CompareMapsAutoGridTest):

    def test_compare_C_maps(self):
        """test AutoGrid C.map output with reference C.map
        """
        self.write_map(['C'])
        ref_filename = 'Data/hsg1_sm.C.map'
        fptr = open(ref_filename)
        reflines = fptr.readlines()
        fptr.close()
        map_filename = 'hsg1_sm.C.map'
        fptr = open(map_filename)
        testlines = fptr.readlines()
        fptr.close()
        #skip the 6 lines of header stuff...
        for ref, test in zip(reflines[6:], testlines[6:]):
            self.assertAlmostEqual(float(ref), float(test), 2)
            #truncate the numbers at 2 decimal places...
            #self.assertAlmostEqual(int(float(ref)/100)*100., int(float(test)/100)*100., 2)


class CompareAMapsAutoGridTest(CompareMapsAutoGridTest):

    def test_compare_A_maps(self):
        """test AutoGrid A.map output with reference A.map
        """
        self.write_map(['A'])
        ref_filename = 'Data/hsg1_sm.A.map'
        fptr = open(ref_filename)
        reflines = fptr.readlines()
        fptr.close()
        map_filename = 'hsg1_sm.A.map'
        fptr = open(map_filename)
        testlines = fptr.readlines()
        fptr.close()
        #skip the 6 lines of header stuff...
        for ref, test in zip(reflines[6:], testlines[6:]):
            self.assertAlmostEqual(float(ref), float(test), 2)
            #truncate the numbers at 2 decimal places...
            #self.assertAlmostEqual(int(float(ref)/100)*100., int(float(test)/100)*100., 2)



class CompareNMapsAutoGridTest(CompareMapsAutoGridTest):

    def test_compare_N_maps(self):
        """test AutoGrid N.map output with reference N.map
        """
        self.write_map(['N'])
        ref_filename = 'Data/hsg1_sm.N.map'
        fptr = open(ref_filename)
        reflines = fptr.readlines()
        fptr.close()
        map_filename = 'hsg1_sm.N.map'
        fptr = open(map_filename)
        testlines = fptr.readlines()
        fptr.close()
        #skip the 6 lines of header stuff...
        for ref, test in zip(reflines[6:], testlines[6:]):
            self.assertAlmostEqual(float(ref), float(test), 2)
            #truncate the numbers at 2 decimal places...
            #self.assertAlmostEqual(int(float(ref)/100)*100., int(float(test)/100)*100., 2)



class CompareOMapsAutoGridTest(CompareMapsAutoGridTest):
    def test_compare_O_maps(self):
        """test the AutoGrid O.map output with reference O.map
        """
        self.write_map(['O'])
        ref_filename = 'Data/hsg1_sm.O.map'
        fptr = open(ref_filename)
        reflines = fptr.readlines()
        fptr.close()
        map_filename = 'hsg1_sm.O.map'
        fptr = open(map_filename)
        testlines = fptr.readlines()
        fptr.close()
        #skip the 6 lines of header stuff...
        for ref, test in zip(reflines[6:], testlines[6:]):
            self.assertAlmostEqual(float(ref), float(test), 2)
            #truncate the numbers at 2 decimal places...
            #self.assertAlmostEqual(int(float(ref)/100)*100., int(float(test)/100)*100., 2)


class CompareHMapsAutoGridTest(CompareMapsAutoGridTest):
    def test_compare_H_maps(self):
        """test the AutoGrid H.map output with reference H.map
        """
        self.write_map(['H'])
        ref_filename = 'Data/hsg1_sm.H.map'
        fptr = open(ref_filename)
        reflines = fptr.readlines()
        fptr.close()
        map_filename = 'hsg1_sm.H.map'
        fptr = open(map_filename)
        testlines = fptr.readlines()
        fptr.close()
        #skip the 6 lines of header stuff...
        for ref, test in zip(reflines[6:], testlines[6:]):
            self.assertAlmostEqual(float(ref), float(test), 2)
            #truncate the numbers at 2 decimal places...
            #self.assertAlmostEqual(int(float(ref)/100)*100., int(float(test)/100)*100., 2)



class CompareEstatMapsAutoGridTest(CompareMapsAutoGridTest):
    def test_compare_e_maps(self):
        """test the AutoGrid estat output with reference estat output
        """
        #now compare the electrostatics maps...
        self.write_map(['C'])
        ref_filename = 'Data/hsg1_sm.e.map'
        fptr = open(ref_filename)
        reflines = fptr.readlines()
        fptr.close()
        map_filename = 'hsg1_sm.e.map'
        fptr = open(map_filename)
        testlines = fptr.readlines()
        fptr.close()
        for ref, test in zip(reflines[6:], testlines[6:]):
            self.assertAlmostEqual(float(ref), float(test), 2)
            #truncate the numbers at 2 decimal places...
            #self.assertAlmostEqual(int(float(ref)/100)*100., int(float(test)/100)*100., 2)




class CompareAllMapsAutoGridTest(CompareMapsAutoGridTest):
    def test_compare_all_maps(self):
        """compare all the AutoGrid output with reference output
        """
        #now compare all the maps...
        self.write_map(['C', 'A','O','N','H'])
        for t in ['C', 'A','O','N','H','e']:
            ref_filename = 'Data/hsg1_sm.' + t + '.map'
            fptr = open(ref_filename)
            reflines = fptr.readlines()
            fptr.close()
            map_filename = 'hsg1_sm.' + t + '.map'
            fptr = open(map_filename)
            testlines = fptr.readlines()
            fptr.close()
            print "comparing ", t, " maps"
            for ref, test in zip(reflines[6:], testlines[6:]):
                #truncate the numbers at 2 decimal places...
                self.assertAlmostEqual(float(ref), float(test), 2)


        
# end BasicAutoGridTest


if __name__ == '__main__':
    
    test_cases = [
        #'BasicAutoGridTest',
        'CompareCMapsAutoGridTest',  #PASS 8/4
        'CompareAMapsAutoGridTest',  #PASS 8/4
        'CompareNMapsAutoGridTest',
        'CompareOMapsAutoGridTest',
        'CompareHMapsAutoGridTest',
        'CompareEstatMapsAutoGridTest',  #PASS 8/4
        'CompareAllMapsAutoGridTest',
        # <new testcase here>
        ]

    unittest.main( argv=( [__name__, '-v'] + test_cases))



