# Import packages
%matplotlib inline
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import pymc as pm
from pymc.gp.cov_funs import matern
import pymc.gp as gp

# Import MMP data
mmp = pd.read_csv('mmp.csv')
# Import TropWater data
tropwater = pd.read_csv('tropwater.csv')

mmpTN = np.array([mmp.Total_Nitrogen.values,mmp.LATITUDE.values,mmp.LONGITUDE.values]).T

tropwaterTN = np.array([tropwater.TN_uM.values,tropwater.LATITUDE.values,tropwater.LONGITUDE.values]).T

TN = np.r_[mmpTN,tropwaterTN]

diff_degree = pm.Uniform('diff_degree',0.1,4,value=1)

scale = pm.Uniform('scale',0,10,value=1.0)

amp = pm.Exponential('amp',1,value=1)

# Prior covariance function
@pm.deterministic
def C(amp=amp, scale=scale, diff_degree=diff_degree):
    return gp.Covariance(matern.euclidean, 
                                 diff_degree = diff_degree, 
                                 amp = amp, scale = scale)

# Overall mean
mu = pm.Normal('mu', 0, 1e-5, value=0)

# Mean function, which is constant
def constant(x, val):
    return np.zeros(x.shape[:-1], dtype=float) + val

# Wrap mean function for GP submodel
@pm.deterministic
def M(mu=mu):
    return gp.Mean(constant, val=mu)

# Lat/lon for observations
obs_coords = TN.T[1:].T
# Unique lat/lon for points on the mesh
mesh = np.array(list(set([tuple(x) for x in obs_coords])))

# GP prior
TN_prior = gp.GPSubmodel('TN_prior', M, C, mesh)

# Observation SD prior
sd_obs = pm.Uniform('sd_obs', 0, 1000, value=2)
tau_obs = pm.Lambda('tau_obs', lambda sd=sd_obs: sd**-2)

Yi = pm.Normal('Yi', mu=TN_prior.f(obs_coords), tau=tau_obs, value=TN.T[0], observed=True)

M = pm.MCMC(locals())

%pdb

M.sample(10000, 5000)

pm.Matplot.summary_plot([M.amp, M.diff_degree, M.scale])

n = 5000
m = 100
xmin, ymin = TN[:, 1:].min(0)
xmax, ymax = TN[:, 1:].max(0)

xplot = np.linspace(xmin, xmax, m)
yplot = np.linspace(ymin, ymax, m)
dplot = np.dstack(np.meshgrid(xplot, yplot))

TN[:, 1:].min(0)

Msurf = np.zeros(dplot.shape[:2])
E2surf = np.zeros(dplot.shape[:2])

from pymc.gp import point_eval

# Get E[v] and E[v**2] over the entire posterior
for i in xrange(n):
    
    # Reset all variables to their values at frame i of the trace
    M.remember(0, i)
    # Evaluate the observed mean
    Msurf_i, Vsurf_i = point_eval(M.TN_prior.M_obs.value, M.TN_prior.C_obs.value, dplot)
    Msurf += Msurf_i/n
    # Evaluate the observed covariance with one argument
    E2surf += (Vsurf_i + Msurf_i**2)/n

# Get the posterior variance and standard deviation
Vsurf = E2surf - Msurf**2
SDsurf = np.sqrt(Vsurf)

TN[:, 1][::-1]

TN[:, 1]

plt.imshow?

ymin, ymax, xmax, xmin

Msurf.shape

# Plot mean and standard deviation surfaces
plt.imshow(Msurf.T, extent=[ymin, ymax, xmin, xmax],interpolation='nearest', origin='lower')
plt.plot(TN[:, 2], TN[:, 1], 'r.', markersize=4)
plt.axis([ymin, ymax, xmin, xmax])
plt.title('Posterior predictive mean surface')
plt.colorbar()

plt.imshow(SDsurf.T, extent=[ymin, ymax, xmin, xmax],interpolation='nearest', origin='lower')
plt.plot(TN[:, 2], TN[:, 1],'r.',markersize=4)
plt.axis([ymin, ymax, xmin, xmax])
plt.title('Posterior predictive standard deviation surface')
plt.colorbar()