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

from PyAutoDock.MolecularSystem import MolecularSystem
import unittest
import numpy
import random

from MolKit import Read
from MolKit.molecule import Atom

class MolecularSystemTest(unittest.TestCase):
    def setUp(self):
        pass

    def tearDown(self):
        pass

##     def test_fail(self):
##         self.assertEqual(0, 1)


class MolecularSystemInitTest(MolecularSystemTest):
    def test_init_none(self):
        """Make sure instances are MolecularSystems
        """
        ms = MolecularSystem()
        self.assertEquals(ms.__class__, MolecularSystem)


    def test_init_1crn(self):
        """test instantiation with a single AtomSet
        """
        crn = Read('Data/1crn.pdb')
        atoms = crn.findType(Atom)
        ms = MolecularSystem(atoms)
        self.assertEquals(ms.__class__, MolecularSystem)
        self.assertEquals(len(ms.entity_sets), 1)
        self.assertEquals(len(ms.get_entities(0)), len(atoms))



class MolecularSystemAddTest(MolecularSystemTest):
    def test_add_none(self):
        """test adding an AtomSet to empty MolecularSystem
        """
        crn = Read('Data/1crn.pdb')
        ms = MolecularSystem()
        old_len = len(ms.entity_sets)
        ms.add_entities(crn.allAtoms)
        self.assertEquals(len(ms.entity_sets), old_len + 1)


    def test_add_1crn(self):
        """test adding an AtomSet to the single AtomSet MolecularSystem
        """
        crn = Read('Data/1crn.pdb')
        ms = MolecularSystem(crn.allAtoms)
        old_len = len(ms.entity_sets)
        ind = Read('Data/ind.out.pdbq')
        ms.add_entities(ind.allAtoms)
        self.assertEquals(len(ms.entity_sets), old_len + 1)
        

class MolecularSystemEnumTest(MolecularSystemTest):
    def setUp(self):
        # create a complete MolecularSystem
        self.crn = Read('Data/1crn.pdb')
        self.ind = Read('Data/ind.out.pdbq')
        self.ind[0].buildBondsByDistance()
        # the index of an AtomSet passed to the __init__ is 0
        self.ms = MolecularSystem(self.crn.allAtoms)
        self.crn_ix = 0
        # MolecularSystem.add returns the index of the added AtomSet
        self.ind_ix = self.ms.add_entities(self.ind.allAtoms)
        
    def tearDown(self):
        pass

    
    def test_interface_coords(self):
        """Test call to MolecularSystem.get_coords()
        """
#        self.assertRaises(DeprecationWarning,
#                          self.ms.get_coords, self.crn_ix )
        
        # make the call
        coords = self.ms.get_coords(self.crn_ix)
        # check length
        self.assertEquals(len(coords), len(self.crn.chains.residues.atoms))
        # check values
        self.assertEquals(coords[0],
                          self.crn.chains.residues.atoms[0].coords)


    def test_interface_coords_subset(self):
        """Test call to MolecularSystem.get_coords() with subset
        """
#        self.assertRaises(DeprecationWarning,
#                          self.ms.get_coords, self.ind_ix )
        # make the call
        coords = self.ms.get_coords(self.ind_ix)

        # check length
        self.assertEquals(len(coords), len(self.ind.chains.residues.atoms))
        # check values
        self.assertEquals(coords[0],
                          self.ind.chains.residues.atoms[0].coords)


class BugTests(MolecularSystemTest):
    def setUp(self):
        # create a complete MolecularSystem
        self.ind = Read('Data/ind.out.pdbq')
        self.ind[0].buildBondsByDistance()
        self.ms = MolecularSystem(self.ind.allAtoms)
        self.ind_ix = 0

    def test_autodock_element_type(self):
        """test that get_element returns the autodock_element
        
        This test was 
        """
        elms = self.ms.get_entities(self.ind_ix).autodock_element
        for elm, result in zip(self.ind.allAtoms.autodock_element, elms):
            self.assertEquals(elm, result)



class InstantiationTest(MolecularSystemTest):
    # test adding a list of atom_sets, like a MoleculeSet!!
    pass


class RequiredAttributeTest(MolecularSystemTest):
    def test(self):
        """Check required attributes of ind.out.pdbq
        """
        ind = Read('Data/ind.out.pdbq')
        ind[0].buildBondsByDistance()
        ms = MolecularSystem(ind.allAtoms)
        atom_set_ix = 0
        attr_list = ['coords', 'charge', 'bonds']
        ms.check_required_attributes( atom_set_ix, attr_list )

    
    def test_fail_with_zero_length_attr(self):
        """Check zero_length required_attribute raises exception
        """
        ind = Read('Data/ind.out.pdbq')
        #ind[0].buildBondsByDistance()
        ms = MolecularSystem(ind.allAtoms)
        atom_set_ix = 0
        attr_list = ['coords', 'charge', 'bonds']
        self.assertRaises(ValueError, ms.check_required_attributes, atom_set_ix, attr_list )


    def test_fail_with_missing_attr(self):
        """Check missing required_attribute raises exception
        """
        ind = Read('Data/ind.out.pdbq')
        ind[0].buildBondsByDistance()
        ms = MolecularSystem(ind.allAtoms)
        atom_set_ix = 0
        attr_list = ['coords', 'charge', 'bonds', 'AtVol']
        self.assertRaises(AttributeError, ms.check_required_attributes, atom_set_ix, attr_list )



class AtomSetConfigurationTest(MolecularSystemTest):
    def test_init(self):
        """test that the initial configuration is (None, None).
        """
        ms = MolecularSystem()
        self.assertEqual( ms.configuration, (None, None))


    def test_init_with_atom_set(self):
        """instantiation with an atom_set makes configuration (0, 0).
        """
        ind = Read('Data/ind.out.pdbq')
        ms = MolecularSystem(ind.allAtoms)
        atom_set_ix = 0
        self.assertEqual( ms.configuration, (atom_set_ix, atom_set_ix))


    def test_add_first_atom_set(self):
        """test configuration after adding atom_set to empty MS
        """
        ms = MolecularSystem()
        self.assertEqual( ms.configuration, (None, None))
        ind = Read('Data/ind.out.pdbq')
        atom_set_ix = ms.add_entities(ind.allAtoms)
        self.assertEqual( ms.configuration, (atom_set_ix, atom_set_ix))


    def test_add_second_atom_set(self):
        """test configuration after adding atom_set to non-empty MS
        """
        ms = MolecularSystem()
        self.assertEqual( ms.configuration, (None, None))
        ind = Read('Data/ind.out.pdbq')
        atom_set1_ix = ms.add_entities(ind.allAtoms)
        self.assertEqual( ms.configuration, (atom_set1_ix, atom_set1_ix))
        atom_set2_ix = ms.add_entities(ind.allAtoms)
        self.assertEqual( ms.configuration, (atom_set1_ix, atom_set2_ix))


    def test_set_configuration(self):
        """test modifying a fully configured MS: (ix, jx) -> (jx, ix)
        """
        # set up the MolecularSystem
        ms = MolecularSystem()
        ind = Read('Data/ind.out.pdbq')
        atom_set1_ix = ms.add_entities(ind.allAtoms) #0
        atom_set2_ix = ms.add_entities(ind.allAtoms) #1
        
        # check that the configuration from the add method is correct
        self.assertEqual( ms.configuration, (atom_set1_ix, atom_set2_ix))
        
        # reset the configuration
        configuration = (atom_set2_ix, atom_set1_ix)
        ms.set_configuration( ix = atom_set2_ix, jx = atom_set1_ix )
        
        # check that the configuration by the add method is correct
        self.assertEqual( ms.configuration, configuration )

    
##     def test_bad_configuration(self):
##         # ix and jx must be valid atomSetNums !!
##         configuration = (ix, jx)
##         self.assertRaises( ms.set_configuration, configuration )
        
# AtomSetConfigurationTest



class SetCoordsTest(MolecularSystemTest):
    def test_set_coords(self):
        """ Check setting coords to new values
        """
        # set up the MolecularSystem
        ms = MolecularSystem()
        ind = Read('Data/ind.out.pdbq')
        atom_set1_ix = ms.add_entities(ind.allAtoms) #0

        # modify the coords of atom_set1
        coords = ms.get_coords(atom_set1_ix)
        new_coords = numpy.array(coords) + 1.0
        ms.set_coords(atom_set1_ix, new_coords)
        for ms_coord, local_coord  in zip( ms.get_coords(atom_set1_ix),
                                           new_coords ):
            self.assertEqual( ms_coord[0], local_coord[0])
            self.assertEqual( ms_coord[1], local_coord[1])
            self.assertEqual( ms_coord[2], local_coord[2])


    def test_set_coords_cleared_dist_mats(self):
        """ Check setting coords clears corresponding distmats
        """
        # set up the MolecularSystem
        ms = MolecularSystem()
        n3pt = Read('Data/3pt.pdbq')
        atom_set1_ix = ms.add_entities(n3pt.allAtoms) #0
        n4pt = Read('Data/4pt.pdbq')
        atom_set2_ix = ms.add_entities(n4pt.allAtoms) #0

        #process
        mat = ms.get_dist_mat( atom_set1_ix, atom_set2_ix)
        coords = ms.get_coords(atom_set1_ix)
        new_coords = numpy.array(coords) + 1.0
        ms.set_coords(atom_set1_ix, new_coords)

        #check
        self.assertEqual(ms._dist_mats[atom_set1_ix][atom_set2_ix], None)
        

# SetCoordsTest


class GetDistanceTest(MolecularSystemTest):
    def test_single_distance(self):
        """ Check calculating internal distances creates 0 diagonal
        """
        # setup
        n3pt = Read('Data/3pt.pdbq')
        ms = MolecularSystem()
        n3ix = ms.add_entities(n3pt.allAtoms)
        # process
        d33 = ms.get_dist_mat( n3ix, n3ix)
        # check
        self.assertEqual( 3, len(d33) )
        self.assertEqual( 3, len(d33[0]) )

        # check self diagonal is 0.0
        for rx, row in enumerate(d33):
            for cx, col in enumerate(row):
                if rx == cx:
                    self.assertEqual( d33[rx][cx], 0.0)
        


    def test_symmetry(self):
        """ Check dims of distmats and of transposed distmats
        """
        n3pt = Read('Data/3pt.pdbq')
        n4pt = Read('Data/4pt.pdbq')
        
        ms = MolecularSystem()
        # add the molecules
        n3ix = ms.add_entities(n3pt.allAtoms)
        n4ix = ms.add_entities(n4pt.allAtoms)

        d33 = ms.get_dist_mat( n3ix, n3ix)
        d34 = ms.get_dist_mat( n3ix, n4ix)
        d43 = ms.get_dist_mat( n4ix, n3ix)
        d44 = ms.get_dist_mat( n4ix, n4ix)

        self.assertEqual( 3, len(d33) )
        self.assertEqual( 3, len(d33[0]) )
        self.assertEqual( 3, len(d34) )
        self.assertEqual( 4, len(d34[0]) )
        self.assertEqual( 4, len(d43) )
        self.assertEqual( 3, len(d43[0]) )
        self.assertEqual( 4, len(d44) )
        self.assertEqual( 4, len(d44[0]) )
# GetDistanceTest
    

class CheckDistanceCutoffTest(MolecularSystemTest):
    def test_check_distance_cutoff(self):
        """check distmat for cutoff distance(s)
        """
        # setup
        n3pt = Read('Data/3pt.pdbq')
        ms = MolecularSystem()
        n3ix = ms.add_entities(n3pt.allAtoms)
        n4pt = Read('Data/4pt.pdbq')
        n4ix = ms.add_entities(n4pt.allAtoms)

        # process
        d34 = ms.check_distance_cutoff( n3ix, n4ix, 2.1)
        d34_ok = ms.check_distance_cutoff( n3ix, n4ix, 2.0)

        # check
        self.assertEqual( d34, 2.0)
        self.assertEqual( d34_ok, None)
        #check for the side effect of storing calculated mat
        self.assertEqual( len(ms._dist_mats[0][1]), 3)


def _get_coords(res):
    return res.childByName.get('CA', res.atoms[0]).coords


def _set_coords(res, residueCoordList):
    res.atoms.updateCoords(residueCoordList)
    

class ResidueEntityTest(MolecularSystemTest):

    def test_use_residue_entities(self):
        """check mechanism for using coords_access_func
        """
        hsg1 = Read('Data/hsg1.pdbqs')
        ms = MolecularSystem(hsg1.chains.residues, _get_coords)
        hsg_ix = 0
        ind = Read('Data/ind.out.pdbq')

        #check ind
        ind_ix = ms.add_entities(ind.chains.residues, _get_coords)
        self.assertEqual(ms.get_coords(ind_ix)[0], ind.allAtoms[0].coords)

        #check hsg1.residue[0]
        hsg1_0_crds = hsg1.chains.residues[0].childByName['CA'].coords
        self.assertEqual(ms.get_coords(hsg_ix)[0], hsg1_0_crds)

        
    def test_atoms_vs_residues(self):
        """check scoring atoms vs residues
        """
        hsg1 = Read('Data/hsg1.pdbqs')
        ms = MolecularSystem(hsg1.chains.residues, _get_coords)
        hsg_ix = 0
        ind = Read('Data/ind.out.pdbq')
        ind_ix = ms.add_entities(ind.allAtoms)

        #check dist shape
        dist = ms.get_dist_mat(hsg_ix, ind_ix)
        self.assertEqual(len(dist), len(ms.get_entities(hsg_ix)))
        self.assertEqual(len(dist), len(hsg1.chains.residues))

        self.assertEqual(len(dist[0]), len(ms.get_entities(ind_ix)))
        self.assertEqual(len(dist[0]), len(ind.allAtoms))


    def test_setting_coords(self):
        """ Check setting coords to new values
        """
        # set up the MolecularSystem
        hsg1 = Read('Data/hsg1.pdbqs')
        ms = MolecularSystem(hsg1.chains.residues, _get_coords)
        hsg_ix = 0
        ind = Read('Data/ind.out.pdbq')
        ind_ix = ms.add_entities(ind.allAtoms)

        old_coords = ms.get_coords(hsg_ix)[:]

        # modify the coords of hsg1
        new_coords = []
        for r in hsg1.chains.residues:
            new_coords.append((numpy.array(r.atoms.coords) + numpy.array([1.,0.,0.])).tolist())

        ms.set_coords(hsg_ix, new_coords, _set_coords)

        modified_coords = ms.get_coords(hsg_ix)
        self.assertEqual( old_coords[0][0], modified_coords[0][0] -1.0)


#
if __name__ == '__main__':
    unittest.main( argv=['__main__', '-v'] )
        
            
