#
#
# A Gillet 08/2008
#

"""Unit test for statetocoords.py
"""

from AutoDockTools.Conformation import Conformation
from MolKit import Read
from mglutil.math.statetocoords import StateToCoords
import unittest, math
import numpy
from UserList import UserList

class StateToCoordsTest(unittest.TestCase):
    def setUp(self):
        self.decimals = 2 # for rounding; 7 is SciPy default.
        # value obtain from first run id 1 result of statecoords.xml
        self.origin =  numpy.array([ -1.413100, 0.115500, 0.685400])
        self.translation = numpy.array([1.144057, -0.003141, -7.605415])
        self.quaternion = numpy.array([ 0.284704, 0.255287, 0.535881, -0.752731])
        self.axisangle = numpy.array([0.432465, 0.387780, 0.814003, -82.345039])
        self.torsions = numpy.array([57.82, -175.79, 180.00, -162.36, 81.42, 100.35] )
        
        # value obtain from first run id 1 result of statecoords.dlg 
        self.expected_coords = numpy.array([(2.191,  -0.823,  -5.475),
                                            (1.363,  -0.143,  -9.679),
                                            (3.370,   1.839, -11.573)])
        
        self.lig = Read('statecoord_lig.pdbqt')[0]
        self.lig.buildBondsByDistance()    

    def test_statetocoords_from_axisangle(self):
        """ test getting molecule coordinate by using statetocoords with 
            axis/angle information 
        """
        self.lig.allAtoms.addConformation(self.lig.allAtoms.coords)  # preserve initial, work in new slot
        # we add a new conformation to ligand so we conserver original coords
        self.lig.allAtoms.addConformation(self.lig.allAtoms.coords)
        conf_index = len(self.lig.allAtoms[0]._coords) -1
        stoc = StateToCoords(self.lig, self.origin,conf_index)
        self.lig.stoc = stoc
        state= Conformation(self.lig, self.origin,self.translation,self.axisangle , self.torsions)
        self.lig.allAtoms.setConformation(conf_index)
        # get coordinate from state
        coords = state.getCoords()
        c0 = self.lig.allAtoms.coords[0]
        c16 = self.lig.allAtoms.coords[15]
        c22 = self.lig.allAtoms.coords[21]
        self.assertArrayEqual(self.expected_coords[0], c0)
        self.assertArrayEqual(self.expected_coords[1], c16)
        self.assertArrayEqual(self.expected_coords[2], c22)

    def test_statetocoords_from_quaternion(self):
        """ test getting molecule coordinate by using statetocoords with 
            quaternion information 
        """
        self.lig.allAtoms.addConformation(self.lig.allAtoms.coords)  # preserve initial, work in new slot
        # we add a new conformation to ligand so we conserver original coords
        self.lig.allAtoms.addConformation(self.lig.allAtoms.coords)
        conf_index = len(self.lig.allAtoms[0]._coords) -1
        stoc = StateToCoords(self.lig, self.origin,conf_index)
        self.lig.stoc = stoc
        quaternion= (self.quaternion[3],self.quaternion[:3])
        state= Conformation(self.lig, self.origin,self.translation,quaternion , self.torsions)
        self.lig.allAtoms.setConformation(conf_index)
        # get coordinate from state
        coords = state.getCoords()
        c0 = self.lig.allAtoms.coords[0]
        c16 = self.lig.allAtoms.coords[15]
        c22 = self.lig.allAtoms.coords[21]
        self.assertArrayEqual(self.expected_coords[0], c0)
        self.assertArrayEqual(self.expected_coords[1], c16)
        self.assertArrayEqual(self.expected_coords[2], c22)   
             
    def tearDown(self):
        pass       

    def assertArrayEqual(self, a1, a2, decimals=None):
        """Round the arrays according to decimals and compare

        Use self.decimals if decimals is not supplied"""
        if decimals:
            d = decimals
        else:
            d = self.decimals

        for point in xrange(len(a1)):
            self.assertAlmostEqual(a1[point],a2[point],d)

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