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

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


import unittest
import numpy
import math

from PyAutoDock.Tests.test_scorer import ScorerTest
from PyAutoDock.Tests.test_scorer import WeightedDataBase

from PyAutoDock.scorer import WeightedMultiTerm
from PyAutoDock.desolvation import Desolvation


class DesolvationTest(ScorerTest):
    def setUp(self):
        self.setup_hsg1_ind() # set up hsg1-ind MolecularSystem

    def test_get_score_array_dims(self):
        ds = Desolvation()
        ds.set_molecular_system(self.ms)
        array = ds.get_score_array()
        self.assertEquals( array.shape, (len(self.hsg1.allAtoms),
                                         len(self.ind.allAtoms)))
# DesolvationTest



class WeightedDesolvationTest(WeightedDataBase):
    def test_dsolv_vs_data(self):
        """Test Desolvation against Autogrid3 data (weight in test)
        """
        # get the result
        ds = Desolvation()
        ds.set_molecular_system(self.ms)
        result = numpy.add.reduce(ds.get_score_array())

        # check answers
        self._get_dsolv_data()
        for res, weighted_data in zip(result, self.dsolv_data):
            # when comparing from file, round to four digits (like in file)
            self.assertFloatEquals((res*self.dsolv_weight),
                                   weighted_data, digits=4)
##         for (res,
##              weighted_data,
##              element) in zip(result,
##                              self.dsolv_data,
##                              self.ind.allAtoms.autodock_element):
##             print "e,r,d,q, diff: ", element, (res*self.dsolv_weight), weighted_data,
##             try:
##                 # check to see if answers are off by a constant factor
##                 # meaning there's a piece of the term missing or wrong
##                 q = ((res*self.dsolv_weight)/weighted_data)
##                 diff = ((res*self.dsolv_weight) - weighted_data)
##                 print q, diff
##             except:
##                 print


    def test_wmt_dsolv_vs_data(self):
        """Test Desolvation against Autogrid3 data (weight in WMT)
        """
        # construct the wmt scorer
        wmt = WeightedMultiTerm()
        wmt.set_molecular_system(self.ms)
        wmt.add_term( Desolvation(), self.dsolv_weight) 

        # get result
        result = numpy.add.reduce(wmt.get_score_array())
        
        # check answers
        self._get_dsolv_data()
        for res, weighted_data in zip(result, self.dsolv_data):
            # when comparing from file, round to four digits (like in file)
            self.assertFloatEquals(res, weighted_data, digits=4)
# WeightedDesolvationTest


if __name__ == '__main__':
    test_cases = [
        'DesolvationTest',
        'WeightedDesolvationTest',
        ]
    unittest.main( argv=([__name__ ,'-v'] + test_cases) )

