from numpy.testing import *
import numpy as np
import pymc as pm
from types import UnboundMethodType

submod = pm.gp.GPSubmodel('x5',pm.gp.Mean(lambda x:0*x),pm.gp.FullRankCovariance(pm.gp.cov_funs.exponential.euclidean, amp=1, scale=1),np.linspace(-1,1,21))
x = [pm.MvNormalCov('x0',np.zeros(5),np.eye(5)),
    pm.Gamma('x1',4,4,size=3),
    pm.Gamma('x2',2,2),
    pm.Binomial('x3',100,.4),
    pm.Bernoulli('x4',.5),
    submod.f]

do_not_implement_methods = ['iadd','isub','imul','itruediv','ifloordiv','imod','ipow','ilshift','irshift','iand','ixor','ior']
uni_methods = ['neg','pos','abs','invert','index']
rl_bin_methods = ['div', 'truediv', 'floordiv', 'mod', 'divmod', 'pow', 'lshift', 'rshift', 'and', 'xor', 'or']+['add', 'mul', 'sub']
bin_methods = ['call','getitem']

class test_special_methods(TestCase):
    def test_call(self):
        for y in x[:2]:
            assert_equal(y[0].value, y.value[0])
    def test_getitem(self):
        assert_equal(x[5](3).value, x[5].value(3))

for dnim in do_not_implement_methods:
    def meth(self, dnim=dnim):
        for y in x:
            if hasattr(y, '__%s__'%dnim):
                raise AssertionError, 'Method %s implemented in class %s'%(dnim, y.__class__)
    setattr(test_special_methods, 'test_'+dnim, UnboundMethodType(meth, None, test_special_methods))

for dnim in uni_methods:
    def meth(self, dnim=dnim):
        # These methods only work on integer-valued variables.
        if dnim in ['index','invert']:
            testvars = [x[3]]
        # All the others work on all numeric ar ndarray variables.
        else:
            testvars = x[:5]
        for y in testvars:
            assert_equal(getattr(y,'__%s__'%dnim)().value, getattr(y.value,'__%s__'%dnim)())
    setattr(test_special_methods, 'test_'+dnim, UnboundMethodType(meth, None, test_special_methods))

for dnim in rl_bin_methods:
    def meth(self, dnim=dnim):
        # These only work for boolean-valued variables.
        if dnim in ['or','xor','and']:
            testvars = [x[4]]
            for y in testvars:
                assert_equal(getattr(y,'__%s__'%dnim)(True).value, getattr(y.value,'__%s__'%dnim)(True))
        # These only work for integers and booleans.
        elif dnim in ['lshift','rshift']:
            testvars = x[3:5]
            for y in testvars:
                assert_equal(getattr(y,'__%s__'%dnim)(2).value, getattr(y.value,'__%s__'%dnim)(2))
        # These should work for all numeric or ndarray variables
        else:
            testvars = x[:4]
            for y in testvars:
                assert_equal(getattr(y,'__%s__'%dnim)(3.0).value, getattr(y.value,'__%s__'%dnim)(3.0))
    setattr(test_special_methods, 'test_'+dnim, UnboundMethodType(meth, None, test_special_methods))


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