#!/usr/bin/env python
import numpy as npy
from pymc.gp import Mean, Covariance, Realization, observe, plot_envelope, NearlyFullRankCovariance, FullRankCovariance
from pymc.gp.cov_funs import matern #, thinplate1d
import matplotlib
matplotlib.rcParams['axes.facecolor']=[1,1,1]

__all__ = ['surface_mean', 'M', 'C']
def surface_mean(x, val):
    """docstring for parabolic_fun"""
    return x * 0 + val

M1 = Mean(surface_mean, val = 0.)
M2 = Mean(surface_mean, val = 0.)
C1 = NearlyFullRankCovariance(eval_fun = matern.euclidean, diff_degree = 3.4, amp = .4, scale = 1.)
C2 = FullRankCovariance(eval_fun = matern.euclidean, diff_degree = 3.4, amp = .4, scale = 1.)


if __name__ == '__main__':
    import pylab as p
    p.close('all')
    x = p.linspace(-2,2)
    obs_x = p.array([-1., -0.5,  0., 0.5, 1])[:1]
    V = p.array([.1,.1,.1,.1,.1])[:1]
    data = p.array([-1, -0, 1, -0, 1])[:1]

    print C1(obs_x), C2(obs_x)
    print C1.cholesky(obs_x, nugget=V), C2.cholesky(obs_x, nugget=V)

    # p.figure(2)
    # for ox,v,d in zip(obs_x, V, data):
    #
    #     print "Observing at", ox, ":", v,",", d
    #     observe(M=M1, C=C1, obs_mesh=[ox], obs_V = [v], obs_vals = [d], cross_validate = True)
    # p.clf()
    # plot_envelope(M1,C1,mesh=x)
    # p.title('Sequential Observations')
    #
    # print "Observing all simultaneously"
    # observe(M=M2, C=C2, obs_mesh=obs_x[C1.obs_piv], obs_V = V[C1.obs_piv], obs_vals = data[C1.obs_piv], cross_validate = True)
    # p.figure(1)
    # plot_envelope(M2,C2,mesh=x)
    # p.title('Simultaneous Observations')
    #
    # print C1.Uo, C2.Uo