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

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


import unittest
import numpy
import math
from PyAutoDock.Tests.test_scorer import ScorerTest, WeightedDataBase
from PyAutoDock.scorer import WeightedMultiTerm
from PyAutoDock.electrostatics import Electrostatics



class EstatFactorTest(ScorerTest):
    def test_332(self):
        q = 1.60217733E-19  # (Coloumbs) charge on electron
        eps0 = 8.854E-12    # (C**2/Joule*meter) vacuum permittivity
        J_per_cal = 4.18400 # Joules per calorie
        avo = 6.02214199E23 # Avogadro's number

        factor = (q*avo*q*1E10 )/(4.0*math.pi*eps0*J_per_cal*1000)
        
        self.assertFloatEqual(332.071, factor)
# EstatFactorTest



class ElectrostaticsTest(ScorerTest):

    def setUp(self):
        self.setup_ind_ind() # set up MolecularSystem


    def test_init(self):
        """instantiate the Electrostatics class."""
        es = Electrostatics()
        es.set_molecular_system(self.ms)
        self.assertEqual(es.__class__, Electrostatics)


    dddp_dict = {
        # distance : value (now un-weighted by 0.1146 term)
        0.00 : (332.0*(0.1146/28.254785)),
        0.01 : (332.0*(0.1146/27.685275))}

    def test_epsilon(self):
        """Test distance dependent dielectric permittivity function: get_dddp.
        """
        es = Electrostatics()
        es.set_molecular_system(self.ms)
        for dist, dddp_value in self.dddp_dict.items():
            value = es.get_dddp(dist)
            self.assertFloatEquals( value, dddp_value)


    def test_estat_dims(self):
        "Make sure the pairwise matrix is of the correct shape"
        es = Electrostatics()
        es.set_molecular_system(self.ms)
        array = es.get_score_array()
        self.assertEquals(len(array), len(self.ind.allAtoms))
        self.assertEquals(len(array[0]), len(self.ind.allAtoms))
        # if subsetA is ind and subsetB is ind,
        # then array[a][b] should work where a in A and b in B.


    def Xtest_estat_nocharge(self):
        """Test dealing with missing charges from the molecular system.
        """
        # ValueError is not the correct Error to raise
        self.assertRaises(ValueError, Electrostatics,
                          self.ms, self.crn, self.ind)

# ElectrostaticsTest




class WeightedElectrostaticsTest(WeightedDataBase):
    def test_estat_vs_data(self):
        """Test Electrostatics against Autogrid3 data (weight in test)
        """
        # get the result
        es = Electrostatics()
        es.set_molecular_system(self.ms)
        result = numpy.add.reduce(es.get_score_array())

        # check answers
        self._get_estat_data()
        for res, weighted_data in zip(result, self.estat_data):
            # when comparing from file, round to four digits (like in file)
            self.assertFloatEquals((res*self.estat_weight),
                                   weighted_data, digits=4)


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

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



if __name__ == '__main__':

    test_cases = [
        'EstatFactorTest',
        'ElectrostaticsTest',          #  1 second
        'WeightedElectrostaticsTest',
        ]
    
    unittest.main( argv=([__name__ ,'-v'] + test_cases) )

