#
# $Id: test_GridMap.py,v 1.8 2005/11/08 21:27:28 rhuey Exp $
#

import random
import types
import unittest

from MolKit.molecule import Atom
from MolKit import Read

from PyAutoDock.AutoGrid import GridMap
from PyAutoDock.MolecularSystem import MolecularSystem


#
# Unit Tests
#

class GridMapTest(unittest.TestCase):
    npts = [5, 5, 5]
    spacing = 1.0
    center = [0., 0., 0.]

    def setUp(self):
        element = 'C'
        name = element + str(0)
        
        at = Atom(name=name, chemicalElement=element)

        at._charges = {'gridmap': 1.0}
        at.chargeSet = 'gridmap'
        at.number = 1
        at._coords = [[0.,0.,0.]]
        at.conformation = 0

        #these 2 would change between maps:
        at.autodock_element = element
        at.AtVol = 12.77  #for carbon map
        
        self.atom = at
    
    def tearDown(self):
        pass


class GridMapConstructorTest(GridMapTest):
    def test_constructor(self):
        """ test GridMap instance is created with defaults
        """
        gm = GridMap(self.atom)
        self.assertEqual(gm.__class__, GridMap)


    def test_constructor_with_params(self):
        """test GripMap instantiation with parameters.
        """
        gm = GridMap(self.atom, self.npts, self.spacing, self.center)
        self.assertEquals( gm.npts,  self.npts)
        self.assertEquals( gm.spacing,  self.spacing)
        self.assertEquals( gm.center,  self.center)
        
# end GridMapTest


class GridMapGeneratorTest(GridMapTest):
    def test_get_entity_list_length(self):
        gm = GridMap(self.atom, self.npts, self.spacing, self.center)
        ent_gen = gm.get_entity_list()
        num_npts = self.npts[0]*self.npts[1]*self.npts[2]
        self.assertEqual( len(list(ent_gen)), num_npts)


    def test_simple1(self):
        gm = GridMap(self.atom, [1,1,1], 1.0, self.center)
        gm_gen = gm.get_entity_list()
        self.assertEqual(gm_gen.next().coords, self.center)


    def test_simple2(self):
        gm = GridMap(self.atom, [3,3,3], 1.0, self.center)
        gm_gen = gm.get_entity_list()
        self.assertEqual(gm_gen.next().coords, [-1., -1., -1.])
        self.assertEqual(gm_gen.next().coords, [ 0., -1., -1.])
        self.assertEqual(gm_gen.next().coords, [ 1., -1., -1.])
        self.assertEqual(gm_gen.next().coords, [-1.,  0., -1.])
        self.assertEqual(gm_gen.next().coords, [ 0.,  0., -1.])
        self.assertEqual(gm_gen.next().coords, [ 1.,  0., -1.])
        self.assertEqual(gm_gen.next().coords, [-1.,  1., -1.])
        self.assertEqual(gm_gen.next().coords, [ 0.,  1., -1.])
        self.assertEqual(gm_gen.next().coords, [ 1.,  1., -1.])

        self.assertEqual(gm_gen.next().coords, [-1., -1.,  0.])
        self.assertEqual(gm_gen.next().coords, [ 0., -1.,  0.])
        self.assertEqual(gm_gen.next().coords, [ 1., -1.,  0.])
        self.assertEqual(gm_gen.next().coords, [-1.,  0.,  0.])
        self.assertEqual(gm_gen.next().coords, [ 0.,  0.,  0.])
        self.assertEqual(gm_gen.next().coords, [ 1.,  0.,  0.])
        self.assertEqual(gm_gen.next().coords, [-1.,  1.,  0.])
        self.assertEqual(gm_gen.next().coords, [ 0.,  1.,  0.])
        self.assertEqual(gm_gen.next().coords, [ 1.,  1.,  0.])

        self.assertEqual(gm_gen.next().coords, [-1., -1.,  1.])
        self.assertEqual(gm_gen.next().coords, [ 0., -1.,  1.])
        self.assertEqual(gm_gen.next().coords, [ 1., -1.,  1.])
        self.assertEqual(gm_gen.next().coords, [-1.,  0.,  1.])
        self.assertEqual(gm_gen.next().coords, [ 0.,  0.,  1.])
        self.assertEqual(gm_gen.next().coords, [ 1.,  0.,  1.])
        self.assertEqual(gm_gen.next().coords, [-1.,  1.,  1.])
        self.assertEqual(gm_gen.next().coords, [ 0.,  1.,  1.])
        self.assertEqual(gm_gen.next().coords, [ 1.,  1.,  1.])
        
        self.assertRaises( StopIteration, gm_gen.next)


    def test_wierd(self):
        """shows that 2x2x2 and 3x3x3 GridMaps are identical
        """
        xdim = 2*random.randint(1,10) # pick a number from 1 to 10...
        ydim = 2*random.randint(1,10)
        zdim = 2*random.randint(1,10)
        even_npts = [xdim, ydim, zdim]
        odd_npts = [xdim+1, ydim+1, zdim+1]
        # even_npts = [2,2,2]
        # odd_npts = [3,3,3]
        spacing = random.random()
        center = [random.random(),random.random(),random.random()]
        
        gm222 = GridMap(self.atom, even_npts, spacing, center)
        gm_gen222 = gm222.get_entity_list()
        gm333 = GridMap(self.atom, odd_npts,  spacing, center)
        gm_gen333 = gm333.get_entity_list()
        try:
            while 1:
                self.assertEqual( gm_gen222.next().coords,
                                  gm_gen333.next().coords)
            self.assertEqual(0,1) # just in case 
        except StopIteration:
            pass
# GridMapGeneratorTest



class GridMapMolecularSystemTest(GridMapTest):

    def test_call_twice(self):
        # create GridMap...
        map = GridMap(self.atom, npts=[3,3,3])
        # ... and a molecule...
        mol = Read('Data/hsg1_sm.pdbqs')[0]
        
        # ... and a  molecular system
        ms = MolecularSystem()
        mol_ix = ms.add_entities(mol.allAtoms)
        map_ix = ms.add_entities(map)
        
        self.assertEqual( 27, len(list(ms.get_entities(1) )))
        self.assertEqual( 27, len(list(ms.get_entities(1) )))


    def test_ms2(self):
        # create GridMap
        map = GridMap(self.atom, npts=[3,3,3])

        m = Read('Data/hsg1_sm.pdbqs')[0]
        
        # create molecular system
        ms = MolecularSystem(m.allAtoms)
        ms.add_entities(map )
        
        # don't make a copy of your iterator!!
        grid_map_entities = ms.get_entities(1)
        
        # check type
        self.assertEqual( type(grid_map_entities), types.GeneratorType)
        # check length
        self.assertEqual( 27, len(list(grid_map_entities)))

        array = []
        for ax, a in enumerate(ms.get_entities(0)):
            array.append([])
            row = array[ax]
            for bx, b in enumerate(ms.get_entities(1)):
                row.append(b)
        self.assertEqual( 27, len(row))

# GridMapMolecularSystemTest



class GridMapTestEntityAccess(GridMapTest):
    
    def test_get_entity(self):
        """test grid_map.get_entity"""
        # create GridMap
        map = GridMap(self.atom, npts=[3,3,3])
        self.assertEqual( self.atom, map.get_entity())

    def test_change_entity_attribute(self):
        # create GridMap with atom
        map = GridMap(self.atom, npts=[3,3,3])

        # change an attribute in the supplied entity
        self.atom.chemicalElement = 'A'

        # does the map.get_entity-returned  entity show it?
        self.assertEqual( self.atom.chemicalElement,
                          map.get_entity().chemicalElement)

# GridMapTestEntityAccess




    

if __name__ == '__main__':
    
    test_cases = [
        'GridMapConstructorTest',
        'GridMapGeneratorTest',
        'GridMapMolecularSystemTest',
        'GridMapTestEntityAccess'
        # <new testcase here>
        ]

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



