#! /usr/bin/env python
## gp_regression_stan.py

import stan
import numpy as np
import matplotlib.pyplot as plt

def reg_fn(t): return(10+5*np.sin(t)+t**2/5.0)

# Simulate data
gen = np.random.default_rng(seed=0)
n = 40
m = 50
T = 10
x = np.linspace(start=0, stop=T, num=n)
y = [gen.normal(loc=reg_fn(x_i)) for x_i in x]
sm_data = {'n':n, 'x':x, 'y':y, 'a':1, 'b':0.5, 'm':m}

# Initialise stan object
with open('gp_regression.stan','r',newline='') as f:
    sm = stan.build(f.read(),sm_data,random_seed=1)

# Select the number of MCMC chains and iterations, then sample
chains, samples, burn = 1, 10000, 1000
fit=sm.sample(num_chains=chains, num_samples=samples, num_warmup=burn, save_warmup=False)

# Plot regression function and posterior for rho
fig,axs=plt.subplots(1,2,figsize=(10,4),constrained_layout=True)
fig.canvas.manager.set_window_title('GP regression posterior')
f = np.mean(fit['fn_vals'], axis=1)
grid = np.linspace(start=0, stop=T, num=m)
true_f = [reg_fn(x_i) for x_i in grid]
r = fit['rho'][0]
axs[0].plot(grid,f)
axs[0].plot(grid,true_f, color='c', lw=2, linestyle='--')
axs[0].scatter(x,y, color='black')
axs[0].set_title('Posterior mean regression function')
axs[0].set_xlabel(r'$x$')
h = axs[1].hist(r,200, density=True);
axs[1].set_title('Approximate posterior density of '+r'$\rho$')
axs[1].set_xlabel(r'$\rho$')
plt.show()