from pymc.gp import *
from pymc.gp.cov_funs import matern
from pymc import *
from numpy import *
from numpy.testing import *
from numpy.random import normal


# class test_MCMC(TestCase):
#     def test(self):

x = arange(-1., 1., .1)


def make_model():
    # Prior parameters of C
    # Matern seems to be segfaulting...
    # diff_degree = Lognormal('diff_degree', mu=1.4, tau=100, verbose=0)
    diff_degree = Uniform('diff_degree', .2, 3)
    amp = Lognormal('amp', mu=.4, tau=1., verbose=0)
    scale = Lognormal('scale', mu=.5, tau=1., verbose=0)

    # The deterministic C is valued as a Covariance object.
    @deterministic(verbose=0)
    def C(eval_fun=matern.euclidean,
          diff_degree=diff_degree, amp=amp, scale=scale):
        return Covariance(
            eval_fun, diff_degree=diff_degree, amp=amp, scale=scale)

    # Prior parameters of M
    a = Normal('a', mu=1., tau=1., verbose=0)
    b = Normal('b', mu=.5, tau=1., verbose=0)
    c = Normal('c', mu=2., tau=1., verbose=0)

    # The mean M is valued as a Mean object.
    def linfun(x, a, b, c):
        return a * x ** 2 + b * x + c

    @deterministic(verbose=0)
    def M(eval_fun=linfun, a=a, b=b, c=c):
        return Mean(eval_fun, a=a, b=b, c=c)

    # The GP itself
    fmesh = array([-.5, .5])
    submod = GPSubmodel('submodel', M, C, fmesh)

    # Observation precision
    V = Gamma('V', alpha=3., beta=.002 / 3., verbose=0)

    # The data d is just array-valued. It's normally distributed about
    # GP.f(obs_x).
    @observed(verbose=0)
    def d(value=array([3.1, 2.9]), mu=submod.f, V=V, verbose=0):
        """
        Data
        """
        mu_eval = mu(fmesh)
        return flib.normal(value, mu_eval, 1. / V)

    return locals()

GPSampler = MCMC(make_model())
GPSampler.use_step_method(
    gp.GPEvaluationGibbs,
    GPSampler.submod,
    GPSampler.V,
    GPSampler.d)
GPSampler.assign_step_methods()
GPSampler.sample(iter=500, burn=0, thin=10, progress_bar=0)
#
# if __name__ == '__main__':
#     import unittest
#     unittest.main()
