############################################################################
#
# Authors: Ruth Huey, William Lindstrom
#
# Copyright: A. Olson TSRI 2006
#
#############################################################################

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

import unittest
import numpy
import math

from MolKit import Read
from MolKit.molecule import Atom

from PyAutoDock.MolecularSystem import MolecularSystem
from PyAutoDock.trilinterp_scorer import TrilinterpScorer, TrilinterpScorer_AD3

scorer = None
scorer2 = None
scorer3 = None

class BaseTests(unittest.TestCase):
    """ tests using maps whose values are sums of coords at pts
"""
    global scorer


    def setup_for_coords_tests(self):
        global scorer
        self.molecule = Read("Data/test.pdbqt")[0]
        self.ms = MolecularSystem()
        self.ms.add_entities(self.molecule.allAtoms)
        scorer = self.scorer = TrilinterpScorer('Data/test3', ['C'])
        self.scorer.set_molecular_system(self.ms)


    def setUp(self):
        #if not hasattr(self, 'molecule'):
        global scorer
        if not scorer:
            #print 'creating scorer'
            self.setup_for_coords_tests()
            #print 'scorer is created'
        else:
            self.scorer = scorer
            #print "scorer already set up!"


    def tearDown(self):
        pass
        #print 'in tearDown'


    def test_constructor(self):
        """Test creating TrilinterpScorer based on sum of coordinates"""
        tls = self.scorer
        self.assertEqual(tls.spacing, .25)
        self.assertEqual(len(tls.map_data['C']), 81)
        self.assertEqual(len(tls.map_data['e']), 81)
        self.assertEqual(len(tls.map_data['d']), 81)
        self.assertEqual(tls.npts[0], 81)
        self.assertEqual(tls.npts[1], 81)
        self.assertEqual(tls.npts[2], 81)
        self.assertEqual(tls.cen[0], 0)
        self.assertEqual(tls.cen[1], 0)
        self.assertEqual(tls.cen[2], 0)
        self.assertEqual(tls.x_low, -10)
        self.assertEqual(tls.y_low, -10)
        self.assertEqual(tls.z_low, -10)
        self.assertEqual(tls.x_high, 10)
        self.assertEqual(tls.y_high, 10)
        self.assertEqual(tls.z_high, 10)


    def test_get_score(self):
        """Test get_score method for sums of coordinates"""
        self.assertEqual(round(self.scorer.get_score(), 3), 23.203)


    def test_get_score_array(self):
        """Test get_score_array method for sums of coordinates"""
        tls = self.scorer
        score_array = tls.get_score_array()
        self.assertEqual(len(score_array), len(tls.ms.get_entities(0)))
        self.assertEqual(round(tls.get_score(),3), round(numpy.add.reduce(score_array), 3))



class AutoDockDlgComparisonTests(unittest.TestCase):
    """ compare python scorer and autodock4 results for same maps and same ligand
"""
    global scorer2
    
    #data extracted from
    #/mgl/works/rhuey/toplevel_flex_receptors/1met_variations/1MET_A/1a30_1met/1a30_1met.010.dlg
    #calculated on bluefish 5/04/2006
    data = [    -0.20 -0.12, -0.01 +0.27, +0.10 -0.12, -0.05 -0.25, -0.08 +0.16, -0.08 -0.36,
                -0.51 +0.41, -0.12 -0.22, +0.02 +0.29, +0.12 -0.18, -0.29 +0.09, -0.12 -0.27,
                -0.48 +0.24, -0.03 +0.29, -0.23 -0.08, -0.31 -0.01, -0.36 -0.03, -0.50 +0.00,
                -0.41 +0.01, -0.13 -0.16, -0.51 -0.02, +0.07 +0.46, +0.13 +0.40, -0.03 -0.38,
                -0.38 -0.70, -0.36 -0.64, -0.20 -0.00, -0.08 +0.10, -0.01 -0.37, -0.01 +0.71,
                +0.01 +0.29, +0.00 +0.00, +0.00 +0.00, -0.25 -0.02, -0.08 +0.08, +0.02 -0.42,
                +0.02 +0.35, -0.34 -0.36, -0.41 -0.32, -0.06 +0.11, +0.04 -0.10, +0.04 -0.06,
                +0.06 -0.05, +0.00 +0.00, +0.00 +0.00, -0.31 -0.06, -0.24 +0.17, +0.18 -0.94,
                -0.18 +0.38, -0.31 -0.34, -0.40 -0.26, -0.05 +0.35, -0.42 -0.35, +0.04 -0.16,
                +0.09 -0.12]
        

    def setup_for_coords_tests(self):
        global scorer2
        mol = self.molecule = Read("Data/1a30_docked.pdbqt")[0]
        mol.buildBondsByDistance()
        #self.molecule = Read("Data/1a30_small.pdbqt")[0]
        self.ms = MolecularSystem()
        ats = mol.chains.residues[-1].atoms[:2] + mol.chains.residues[-2].atoms[:2] 
        for atom in ats:
            atom.ignore = True
        self.ms.add_entities(mol.allAtoms)
        scorer2 = self.scorer = TrilinterpScorer('Data/1met_rigid', ['C', 'N', 'HD', 'OA'], atoms_to_ignore = ats)
        self.scorer.set_molecular_system(self.ms)


    def setUp(self):
        #if not hasattr(self, 'molecule'):
        global scorer2
        if not scorer2:
            #print 'creating scorer'
            self.setup_for_coords_tests()
            #print 'scorer is created'
        else:
            self.scorer = scorer2
            #print "scorer already set up!"


    def tearDown(self):
        pass
        #print 'in tearDown'


    def test_constructor(self):
        """Test creating TrilinterpScorer for comparison with autodock4 results"""
        tls = self.scorer
        self.assertEqual(tls.spacing, .375)
        self.assertEqual(len(tls.map_data['C']), 81)
        self.assertEqual(len(tls.map_data['N']), 81)
        self.assertEqual(len(tls.map_data['HD']), 81)
        self.assertEqual(len(tls.map_data['OA']), 81)
        self.assertEqual(len(tls.map_data['e']), 81)
        self.assertEqual(len(tls.map_data['d']), 81)
        self.assertEqual(tls.npts[0], 81)
        self.assertEqual(tls.npts[1], 81)
        self.assertEqual(tls.npts[2], 81)
        self.assertEqual(tls.cen[0], 0.148)
        self.assertEqual(tls.cen[1], 0.225)
        self.assertEqual(tls.cen[2], -0.086)
        self.assertEqual(tls.x_low, -14.852)
        self.assertEqual(tls.y_low, -14.775)
        self.assertEqual(tls.z_low, -15.086)
        self.assertEqual(tls.x_high, 15.148)
        self.assertEqual(tls.y_high, 15.225)
        self.assertEqual(tls.z_high, 14.914)


    def test_get_score(self):
        """Test get_score method for comparison with autodock4 results"""
        #value for 1a30_small.pdbqt
        #self.assertEqual(round(self.scorer.get_score(), 3), -0.928)
        #this is for whole, docked molecule
        self.assertEqual(round(self.scorer.get_score(), 3), -9.891)


    def test_get_score_array(self):
        """Test get_score_array method with autodock4 results"""
        tls = self.scorer
        score_array = tls.get_score_array()
        self.assertEqual(len(score_array), len(tls.ms.get_entities(0)))
        self.assertEqual(round(tls.get_score(),3), round(numpy.add.reduce(score_array), 3))


    def test_get_score_array(self):
        """Test compare individual get_score_array results with dlg values"""
        tls = self.scorer
        score_array = tls.get_score_array()
        atoms = tls.ms.get_entities(0)
        for i in range(len(self.data)):
            #for sa, val in zip(score_array, self.data):
            sa = round(score_array[i], 3)
            val = round(self.data[i], 3)
            delta = round(score_array[i]-self.data[i], 3)
            if abs(delta)>0.009:
                print "% 3s % 6.4f % 6.4f % 6.4f" %(atoms[i].name, sa, val, delta)
            self.failUnless(abs(delta) < 0.01)



class AutoDock3DlgComparisonTests(unittest.TestCase):
    """ compare python scorer and autodock3 results for same maps and same ligand
"""
    global scorer3
    
    #data extracted from
    #/mgl/work4/rhuey/TutorialDetails/tutorial/Results/ind.dlg
    #calculated on saul 12/03/2002
    data = [-0.2322 , -0.4299 , -0.4219 , -0.4400 , -0.0317 , -0.1908 , -0.0511 ,
            -0.5202 , -0.5009 , -0.6585 , -0.6529 , -0.4660 , -0.4813 , -0.4741 ,
            -0.4537 ,  0.0114 , -0.0886 , -0.4789 , -0.5394 , -0.5576 , -0.5088 ,
            -0.5926 , -0.5250 , -0.4889 , -0.3922 , -0.3898 ,  0.0724 , -0.3801 ,
             0.5569 , -0.9632 , -0.3209 , -0.3615 , -0.3444 , -0.3817 , -0.2976 ,
            -0.3250 , -0.4661 , -0.3075 , -0.3046 , -0.2732 , -0.6905 , -0.1646 ,
            -0.1873 , -0.5121 , -0.5990 , -0.5547 , -0.5868 , -0.3334 , -0.3641 ]
 
        

    def setup_for_coords_tests(self):
        global scorer3
        mol = self.molecule = Read("Data/dlg_coords_ind.pdbq")[0]
        mol.buildBondsByDistance()
        self.ms = MolecularSystem()
        self.ms.add_entities(mol.allAtoms)
        scorer3 = self.scorer = TrilinterpScorer_AD3('Data/test_hsg1', ['C', 'N', 'H', 'O', 'A'])
        self.scorer.set_molecular_system(self.ms)


    def setUp(self):
        #if not hasattr(self, 'molecule'):
        global scorer3
        if not scorer3:
            #print 'creating scorer'
            self.setup_for_coords_tests()
            #print 'scorer is created'
        else:
            self.scorer = scorer3
            #print "scorer already set up!"


    def tearDown(self):
        pass
        #print 'in tearDown'


    def test_constructor(self):
        """Test creating TrilinterpScorer for comparison with autodock3 results"""
        tls = self.scorer
        self.assertEqual(tls.spacing, .375)
        self.assertEqual(len(tls.map_data['C']), 61)
        self.assertEqual(len(tls.map_data['N']), 61)
        self.assertEqual(len(tls.map_data['H']), 61)
        self.assertEqual(len(tls.map_data['A']), 61)
        self.assertEqual(len(tls.map_data['O']), 61)
        self.assertEqual(len(tls.map_data['e']), 61)
        self.assertEqual(tls.npts[0], 61)
        self.assertEqual(tls.npts[1], 61)
        self.assertEqual(tls.npts[2], 61)
        self.assertEqual(tls.cen[0], 2.5)
        self.assertEqual(tls.cen[1], 6.5)
        self.assertEqual(tls.cen[2], -7.5)
        self.assertEqual(tls.x_low, -8.75)
        self.assertEqual(tls.y_low, -4.75)
        self.assertEqual(tls.z_low, -18.75)
        self.assertEqual(tls.x_high, 13.75)
        self.assertEqual(tls.y_high, 17.75)
        self.assertEqual(tls.z_high, 3.75)


    def test_get_score(self):
        """Test get_score method for comparison with autodock3 results"""
        #this is for whole, docked molecule
        score = self.scorer.get_score()
        delta = abs(round(score, 3))- 18.67
        #check that the absolute value of the difference is less than one percent
        self.failUnless(abs(delta/score)*100 < 1)


    def test_get_score_array(self):
        """Test get_score_array method with autodock3 results"""
        tls = self.scorer
        score_array = tls.get_score_array()
        self.assertEqual(len(score_array), len(tls.ms.get_entities(0)))
        self.assertEqual(round(tls.get_score(),3), round(numpy.add.reduce(score_array), 3))


    def test_get_score_array(self):
        """Test compare individual get_score_array results with dlg values"""
        tls = self.scorer
        score_array = tls.get_score_array()
        atoms = tls.ms.get_entities(0)
        for i in range(len(self.data)):
            #for sa, val in zip(score_array, self.data):
            sa = round(score_array[i], 3)
            val = round(self.data[i], 3)
            delta = round(score_array[i]-self.data[i], 3)
            if abs(delta)>0.009:
                print "% 3s % 6.4f % 6.4f % 6.4f" %(atoms[i].name, sa, val, delta)
            self.failUnless(abs(delta) < 0.01)


    def test_compare_with_atom_properties(self):
        """Test compare individual get_score_array results with ligand
        temperatureFactor and occupancy values [these are the fields where
        autodock writes out the vdw+hbond and the estat values]"""
        tls = self.scorer
        score_array = tls.get_score_array()
        atoms = tls.ms.get_entities(0)
        for i in range(len(atoms)):
            at = atoms[i]
            delta = score_array[i] - (at.temperatureFactor + at.occupancy)
            self.failUnless(abs(delta) < 0.015)


if __name__ == '__main__':

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

