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

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


import numpy
import math
import random
import string
import time
import unittest


from PyAutoDock.MolecularSystem import MolecularSystem

from MolKit import Read
from MolKit.molecule import Atom

from PyAutoDock.scorer import Distance

from PyAutoDock.AutoDockScorer import AutoDock305Scorer
#from PyAutoDock.AutoDockScorer import AutoDock4Scorer

from PyAutoDock.scorer import WeightedMultiTerm

from PyAutoDock.vanDerWaals import VanDerWaals
from PyAutoDock.vanDerWaals import HydrogenBonding
from PyAutoDock.vanDerWaals import NewHydrogenBonding
from PyAutoDock.vanDerWaals import NewVanDerWaals
from PyAutoDock.vanDerWaals import NewHydrogenBonding12_10
from PyAutoDock.vanDerWaals import HBondVectorCalculator

from PyAutoDock.electrostatics import Electrostatics
from PyAutoDock.desolvation import Desolvation
from PyAutoDock.desolvation import NewDesolvation



##################################################################

# UNIT TESTS

##################################################################


class ScorerTest(unittest.TestCase):
    from PyAutoDock.desolvation import Desolvation
    
    err_epsilon = 0.000001
    def assertFloatEquals(self, float1, float2, digits=None):
        if digits is not None:
            float1 = round(float1, digits)
            float2 = round(float2, digits)
        difference = abs(float2 - float1)
        if difference == 0.0:
            self.assertEquals(float1, float2)
        else:
            try:
                eps = abs((float1/float2) - 1.0)
            except ZeroDivisionError:
                eps = abs((float2/float1) - 1.0)
            self.assertEquals((eps < self.err_epsilon), True,
                              msg="%f != %f; eps=%f" % (float1, float2, eps))
    # 1 method, 2 names
    assertFloatEqual = assertFloatEquals

    
    def setup_4pt_3pt(self):
        self.four_pdbqs = Read('Data/4pt.pdbqs')
        self.three_pdbqs = Read('Data/3pt.pdbqs')
        self.ms = MolecularSystem()
        self.ms.add_entities(self.four_pdbqs.chains.residues.atoms)
        self.ms.add_entities(self.three_pdbqs.chains.residues.atoms)

    
    def setup_4pt_3pt_v4(self):
        self.four_pdbqt = Read('Data/4pt.pdbqt')
        self.three_pdbqt = Read('Data/3pt.pdbqt')
        self.ms = MolecularSystem()
        self.ms.add_entities(self.four_pdbqt.chains.residues.atoms)
        self.ms.add_entities(self.three_pdbqt.chains.residues.atoms)


    def setup_crn_ind(self):
        # create a complete MolecularSystem
        self.crn = Read('Data/1crn.pdb')
        self.ind = Read('Data/ind.out.pdbq')
        self.ind[0].buildBondsByDistance()
        self.ms = MolecularSystem()
        self.ms.add_entities(self.crn.chains.residues.atoms)
        self.ms.add_entities(self.ind.chains.residues.atoms)


    def setup_ind_ind(self):
        # create a complete MolecularSystem
        #warning this indinavir has different coordinates
        self.ind = Read('Data/ind.out.pdbq')
        self.ind[0].buildBondsByDistance()
        self.ms = MolecularSystem()
        self.ms.add_entities(self.ind.chains.residues.atoms)
        self.ms.add_entities(self.ind.chains.residues.atoms)


    def setup_hsg1_ind(self):
        from MolKit import molecule, protein
        #turn off bhtreeFlag to make rvector test pass
        molecule.bhtreeFlag = 0
        protein.bhtreeFlag = 0
        self.hsg1 = Read('Data/hsg1.pdbqs')
        self.hsg1[0].buildBondsByDistance()
        self.ind = Read('Data/docked_ind.pdbq')
        self.ms = MolecularSystem()
        self.ms.add_entities(self.hsg1.chains.residues.atoms)
        self.ms.add_entities(self.ind.chains.residues.atoms)


#    def setup_hsg1_ind_v4(self):
#        from MolKit import molecule, protein
#        #turn off bhtreeFlag to make rvector test pass
#        molecule.bhtreeFlag = 0
#        protein.bhtreeFlag = 0
#        self.hsg1 = Read('Data/hsg1.pdbqt')
#        self.hsg1[0].buildBondsByDistance()
#        self.ind = Read('Data/ind.pdbqt')
#        self.ms = MolecularSystem()
#        self.ms.add_entities(self.hsg1.chains.residues.atoms)
#        self.ms.add_entities(self.ind.chains.residues.atoms)

# ScorerTest



class DistanceTest(ScorerTest):
    def test_distance_init(self):
        """Test initialization of the Distance scorer.
        """
        self.setup_crn_ind() # set up MolecularSystem
        ds = Distance()
        ds.set_molecular_system(self.ms)
        self.assertEquals(ds.__class__, Distance)


    def test_distance_score(self):
        """Test get_score method of the Distance scorer.
        """
        self.setup_crn_ind() # set up MolecularSystem
        # get result to be tested
        ds = Distance()
        ds.set_molecular_system(self.ms)
        result = ds.get_score_array()
        # do local calculation
        a1 = self.crn.chains.residues.atoms[0].coords
        a2 = self.ind.chains.residues.atoms[0].coords
        my_dist = math.sqrt( (a2[0]-a1[0])**2 +
                             (a2[1]-a1[1])**2 +
                             (a2[2]-a1[2])**2)
        # compare result to be tested to local calculation
        self.assertFloatEquals( result[0][0], my_dist)


# DistanceTest



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

        # get the weighted data
        data_file = open('Data/evaluator_results')
        self.lines = data_file.readlines()
        data_file.close()

    def _get_estat_data(self):
        self.estat_data = [float(string.split(l)[2]) for l in self.lines[1:]]
    def _get_hbond_data(self):
        self.hbond_data = [float(string.split(l)[3]) for l in self.lines[1:]]
    def _get_vdw_data(self):
        self.vdw_data   = [float(string.split(l)[4]) for l in self.lines[1:]]
    def _get_dsolv_data(self):
        self.dsolv_data = [float(string.split(l)[5]) for l in self.lines[1:]]
        
    err_epsilon = 0.07
    err_epsilon = 0.25 # for vdw

    def test_ad305_vs_data(self):
        """Test AutoDock305Scorer against Autogrid3 data (weight in scorer)
        """
        # get the result
        ad305 = AutoDock305Scorer()
        ad305.set_molecular_system(self.ms)
        result = numpy.add.reduce(ad305.get_score_array())
        # check answers
        self._get_estat_data()
        self._get_hbond_data()
        self._get_vdw_data()
        self._get_dsolv_data()
        for res,e,h,v,d in zip(result,self.estat_data, self.hbond_data,
                           self.vdw_data, self.dsolv_data):
            data = e+h+v+d
            # when comparing from file, round to four digits (like in file)
            self.assertFloatEquals(res, data, digits=4)
        


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

#        # get the weighted data
#        data_file = open('Data/evaluator_results_v4')
#        self.lines = data_file.readlines()
#        data_file.close()

#    def _get_estat_data(self):
#        self.estat_data = [float(string.split(l)[2]) for l in self.lines[1:]]
#    def _get_hbond_data(self):
#        self.hbond_data = [float(string.split(l)[3]) for l in self.lines[1:]]
#    def _get_vdw_data(self):
#        self.vdw_data   = [float(string.split(l)[4]) for l in self.lines[1:]]
#    def _get_dsolv_data(self):
#        self.dsolv_data = [float(string.split(l)[5]) for l in self.lines[1:]]
#        
#    err_epsilon = 0.07
#    err_epsilon = 0.25 # for vdw

#    def test_ad4_vs_data(self):
#        """Test AutoDock4Scorer against Autogrid4 data (weight in scorer)
#        """
#        # get the result
#        ad4 = AutoDock4Scorer()
#        ad4.set_molecular_system(self.ms)
#        result = numpy.add.reduce(ad4.get_score_array())
#        # check answers
#        self._get_estat_data()
#        self._get_hbond_data()
#        self._get_vdw_data()
#        self._get_dsolv_data()
#        for res,e,h,v,d in zip(result,self.estat_data, self.hbond_data,
#                           self.vdw_data, self.dsolv_data):
#            data = e+h+v+d
#            # when comparing from file, round to four digits (like in file)
#            self.assertFloatEquals(res, data, digits=4)
#        


class WeightedDataBase(ScorerTest):
    """This tests the new terms against old (weighted) data in
    the file Tests/Data/evaluator_results

    Done: Electrostatic, Desolvation, VanDerWaals, HydrogenBonding
    """
    estat_weight = 0.1146 # Autogrid3 weight
    dsolv_weight = 0.1711 # Autogrid3 weight
    hbond_weight = 0.0656 # Autogrid3 weight
    vdw_weight = 0.1485 # Autogrid3 weight

    err_epsilon = 0.07
    err_epsilon = 0.25 # for vdw
    def setUp(self):
        self.setup_hsg1_ind() # set up hsg1-ind MolecularSystem

        # get the weighted data
        data_file = open('Data/evaluator_results')
        self.lines = data_file.readlines()
        data_file.close()

    def _get_estat_data(self):
        self.estat_data = [float(string.split(l)[2]) for l in self.lines[1:]]
    def _get_hbond_data(self):
        self.hbond_data = [float(string.split(l)[3]) for l in self.lines[1:]]
    def _get_vdw_data(self):
        self.vdw_data   = [float(string.split(l)[4]) for l in self.lines[1:]]
    def _get_dsolv_data(self):
        self.dsolv_data = [float(string.split(l)[5]) for l in self.lines[1:]]



class WeightedDataTest(WeightedDataBase):
    def test_estat_dsolv_hbond_vdw_vs_data(self):
        # construct the wmt scorer
        wmt = WeightedMultiTerm()
        wmt.set_molecular_system(self.ms)
        wmt.add_term( Electrostatics(), self.estat_weight)
        wmt.add_term( Desolvation(), self.dsolv_weight)
        wmt.add_term( VanDerWaals(), self.vdw_weight)
        wmt.add_term( HydrogenBonding(), self.hbond_weight)

        # get result
        result = numpy.add.reduce(wmt.get_score_array())
        
        # check answers
        self._get_estat_data()
        self._get_dsolv_data()
        self._get_hbond_data()
        self._get_vdw_data()
        for res, es, ds, hb, vdw in zip(result, self.estat_data,
                                        self.dsolv_data, self.hbond_data,
                                        self.vdw_data):
            data = es + ds + hb + vdw
            # when comparing from file, round to four digits (like in file)
            self.assertFloatEquals(res, data, digits=4)


    def test_breakdown(self):
        # construct the wmt scorer
        es = WeightedMultiTerm()
        es.set_molecular_system(self.ms)
        es.add_term( Electrostatics(), self.estat_weight)
        
        ds = WeightedMultiTerm()
        ds.set_molecular_system(self.ms)
        ds.add_term( Desolvation(), self.dsolv_weight)
        
        vdw = WeightedMultiTerm()
        vdw.set_molecular_system(self.ms)
        vdw.add_term( VanDerWaals(), self.vdw_weight) 

        hb = WeightedMultiTerm()
        hb.set_molecular_system(self.ms)
        hb.add_term( HydrogenBonding(), self.hbond_weight)
        

        # get results
        es_array = es.get_score_array()
        es_list = numpy.add.reduce(es_array)
        self._get_estat_data()
        ctr = 0
        for res, data in zip(es_list, self.estat_data):
            self.assertFloatEquals(res, data, digits=4)
            ctr=ctr+1
            
        ds_array = ds.get_score_array()
        ds_list = numpy.add.reduce(ds_array)
        self._get_dsolv_data()
        for res, data in zip(ds_list, self.dsolv_data):
            self.assertFloatEquals(res, data, digits=4)

        hb_array = hb.get_score_array()
        hb_list = numpy.add.reduce(hb_array)
        self._get_hbond_data()
        for res, data in zip(hb_list, self.hbond_data):
            self.assertFloatEquals(res, data, digits=4)

        vdw_array = vdw.get_score_array()
        vdw_list = numpy.add.reduce(vdw_array)
        self._get_vdw_data()                         
        for res, data in zip(vdw_list, self.vdw_data):
            self.assertFloatEquals(res, data, digits=4)


    def test_get_score_per_term(self):
        """Compare WMT get_score_per_term results with empirical data"""
        wmt = WeightedMultiTerm()
        wmt.set_molecular_system(self.ms)
        wmt.add_term( Electrostatics(), self.estat_weight)
        wmt.add_term( Desolvation(), self.dsolv_weight)
        wmt.add_term( VanDerWaals(), self.vdw_weight)
        wmt.add_term( HydrogenBonding(), self.hbond_weight)

        # get results
        scorer_result = wmt.get_score_per_term()

        #construct lists of empirical data: self.estat_data etc
        self._get_estat_data()
        self._get_dsolv_data()
        self._get_hbond_data()
        self._get_vdw_data()

        ag3_data = [numpy.add.reduce(self.estat_data),
                numpy.add.reduce(self.dsolv_data),
                numpy.add.reduce(self.vdw_data),
                numpy.add.reduce(self.hbond_data)]

        # check answers
        for res, data in zip(scorer_result,ag3_data): 
            self.assertFloatEquals(res, data, digits=4)
# WeightedDataTest






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

    def test_init(self):
        """Test initialization of WeightedMultiTerm class"""
        wmt = WeightedMultiTerm()
        wmt.set_molecular_system(self.ms)
        self.assertEqual( wmt.__class__, WeightedMultiTerm)


    def test_add_term(self):
        """Test adding terms to WeightedMultiTerm class
        """
        wmt = WeightedMultiTerm()
        wmt.set_molecular_system(self.ms)
        wmt.add_term(Electrostatics(), 1.0)
        self.assertEquals( wmt.terms[-1][0].__class__, Electrostatics)


    def test_estat_wmt_compare_score(self):
        """Compare Estat (or any term) and WMT results"""
        # get EStat scorer
        es = Electrostatics()
        es.set_molecular_system(self.ms)
        # get WMT scorer
        wmt = WeightedMultiTerm()
        wmt.set_molecular_system(self.ms)
        wmt.add_term(Electrostatics(), 1.0) # weight = 1.0
        # compare results
        self.assertFloatEquals(es.get_score(), wmt.get_score())

        
    def test_estat_wmt_compare_score2(self):
        """Compare Estat and WMT results (two weighted estat terms)"""
        # get EStat array
        es = Electrostatics()
        es.set_molecular_system(self.ms)
        # get WMT array
        wmt = WeightedMultiTerm()
        wmt.set_molecular_system(self.ms)
        wmt.add_term(Electrostatics(), 0.5)
        wmt.add_term(Electrostatics(), 0.5) # add 2nd term
        # compare results
        self.assertFloatEquals(es.get_score(), wmt.get_score())
# WeightedMultiTermTest



class SevenPointTest(ScorerTest):
    def setUp(self):
        self.setup_4pt_3pt() # set up small MolecularSystem
        

    def read_estat_data(self):
        # read the data and construct 
        data_file = open('Data/4pt3pt_estat_data')
        lines = data_file.readlines()
        self.estat_data = []
        for l in lines:
            self.estat_data.append(map(float, string.split(l)))


    def read_dsolv_data(self):
        data_file = open('Data/4pt3pt_dsolv_data')
        lines = data_file.readlines()
        # read in 3x4 array
        d = []
        for l in lines[:3]:
            d.append(map(float, string.split(l)))
        self.threexfour_data = d
        # read in the 4x3 array
        d = []
        for l in lines[-4:]:
            d.append(map(float, string.split(l)))
        self.fourxthree_data = d

    def test_distance_symmetry(self):
        """Test that Distance(x,y).get_score() = Distance(y,x).get_score()
        """
        dc1 = Distance()
        dc1.set_molecular_system(self.ms)
        dc2 = Distance()
        dc2.set_molecular_system(self.ms)
        self.assertFloatEquals(dc1.get_score(), dc2.get_score())


    def test_estat(self):
        """Compare Electrostatics with test data file"""
        self.read_estat_data()
        # calculate estat array
        es = Electrostatics()
        es.set_molecular_system(self.ms)
        res_array = es.get_score_array()
        # compare with data from file
        res_sum = numpy.add.reduce(numpy.add.reduce(res_array))
        data_sum = numpy.add.reduce(numpy.add.reduce(self.estat_data))
        self.assertFloatEquals( res_sum, data_sum)
        # check i,j values
        for i in xrange(len(res_array)): # rows
            for j in xrange(len(res_array[0])): # cols
                self.assertFloatEquals(res_array[i][j], self.estat_data[j][i])


    def test_dsolv_4x3(self):
        """Compare Desolvation result (4x3) with test data file"""
        self.read_dsolv_data()
        data = self.fourxthree_data
#3/1        data = self.threexfour_data
        # calculate hydrophobic array
        dsolv = Desolvation()
        dsolv.set_molecular_system(self.ms)
        res_array = dsolv.get_score_array()
        
        # compare with data from file
        res_sum = numpy.add.reduce(numpy.add.reduce(res_array))
        data_sum = numpy.add.reduce(numpy.add.reduce(data))
        self.assertFloatEquals( res_sum, data_sum)
        # check i,j values
        for i in xrange(len(res_array)): # rows
            for j in xrange(len(res_array[0])): # cols
                self.assertFloatEquals(res_array[i][j], data[i][j])


#    def test_dsolv_3x4(self):
#        """Compare Desolvation result (3x4) with test data file"""
#        self.read_dsolv_data()
#        data = self.threexfour_data
#        # calculate hydrophobic array
#        hps = Desolvation()
#        hps.set_molecular_system(self.ms)
#        res_array = hps.get_score_array()
#        
#        # compare with data from file
#        res_sum = numpy.add.reduce(numpy.add.reduce(res_array))
#        data_sum = numpy.add.reduce(numpy.add.reduce(data))
#        ####self.assertFloatEquals( res_sum, data_sum)
#        # check i,j values
#        print "dsolv_3x4"
#        print "res_array.shape=", numpy.array(res_array).shape
#        print "self.threexfour.shape=", numpy.array(self.threexfour_data).shape
#        
#        for i in xrange(len(res_array)): # rows
#            for j in xrange(len(res_array[0])): # cols
#                print i,',',j,':',res_array[i][j], '-', data[i][j]
#                ####self.assertFloatEquals(res_array[i][j], data[i][j])


    def test_dsolv_symmetry(self):
        """Test symmetry of desolvation model."""
        pass
    

# SevenPointTest





if __name__ == '__main__':

    test_cases = [
        'DistanceTest',                #  7 seconds
        'WMTClassTest',                # 17 seconds
        'SevenPointTest',              #  0 seconds
        'WeightedDataTest',            # 11 seconds
        'AutoDock305Test',
        #'AutoDock4Test',
        ]
    
    unittest.main( argv=([__name__ ,'-v'] + test_cases) )

