########################
### Begin RStan code ###
########################

## Fit the Time-varying coefficients model:

setwd("C:/WQR/")
library(rstan) 
rm(list=ls(all=TRUE))

# Create the object ‘data’, containing the standardized wine quality 
# ratings and growing season temperature values:

data <- list(
N=90,
temp=c(-0.745864,0.270905,-0.542232,-0.560997,-0.699698,-0.769275,-0.259649,-0.837022,-0.015517,-0.349978,-0.990825,-1.604208,-1.015239,0.443452,0.512419,-0.394532,-1.488245,-0.277959,-1.444301,-1.410733,-0.426128,-1.211813,0.405774,1.330109,0.775508,1.60741,-0.854111,1.572564,-0.600824,1.302798,0.230446,-0.957562,0.215188,-0.212959,-1.720781,0.252418,-1.404324,-0.883102,-0.566645,0.62472,-0.566645,0.382724,-0.752796,-1.404324,0.326879,-1.795241,-0.529415,-0.845872,-0.845872,-0.47357,-0.343264,0.271033,-1.851086,0.531645,-0.585261,0.345494,1.015637,-1.395475,-0.864487,-0.5108,-0.659721,0.047652,0.531645,0.777303,-0.51965,-0.138498,-0.715566,0.568875,0.53317,1.617118,1.797776,1.014111,0.921951,-0.23902,0.524199,0.74758,0.058821,1.529413,0.617274,1.492183,1.101266,0.66195,0.140728,2.970221,1.052867,1.68578,2.169772,0.755026,0.140728,1.406554),
WR_obs=c(0.358684,0.836929,0.119561,0.119561,0.597806,-0.358684,0.836929,-0.358684,1.076051,1.315174,-0.836929,-2.510786,-1.076051,-0.836929,0.597806,-1.315174,-2.271664,0.597806,-0.597806,-0.597806,-0.358684,-2.989031,-0.836929,0.597806,-1.076051,1.076051,-1.076051,0.836929,0.119561,0.836929,0.358684,-1.554296,0.597806,0.836929,-1.315174,0.836929,-1.315174,-0.597806,-0.597806,1.076051,-0.597806,1.315174,0.597806,-2.749909,0.597806,-2.749909,0.597806,-0.119561,-2.032541,-0.597806,0.836929,0.597806,-1.076051,-0.597806,-0.597806,0.597806,0.358684,-0.836929,0.597806,0.358684,-0.358684,0.358684,1.076051,0.597806,-0.597806,0.836929,0.836929,-0.119561,0.597806,1.076051,1.076051,-0.358684,-0.597806,-0.119561,0.119561,0.597806,0.836929,0.119561,0.597806,0.358684,1.076051,0.597806,0.358684,0.836929,0.597806,1.315174,0.597806,0.358684,0.597806,1.076051))

# Create the object ‘TVC_wine_code’, which includes the code for fitting the 
# time-varying coefficients model that will be translated by 
# RStan to C++ code:

TVC_wine_code <- '

data {
	int<lower=0> N;
	real temp[N]; // Standardised average growing season temperature
	real WR_obs[N]; // Standardised Tastet & Lawton wine Quality Ratings 
	}

parameters {
	// Bounded time series of latent Wine Quality Ratings
	real<lower=-1.85108645079154, upper=2.97022104705866> WR_exp[N];

	// Location parameters
	real alpha[N];
	real l_alpha;
	real beta_lin[N];
	real l_beta_lin;
	real beta_sq[N];
	real l_beta_sq;
	real l_sigma;

	// Scale (standard deviation) parameters
	real<lower=0> sigma[N];
	real<lower=0> rho_alpha;
	real<lower=0> rho_beta_lin;
	real<lower=0> rho_beta_sq;
	real<lower=0> rho_sigma;
	}
	
transformed parameters {
	real rho2_alpha;
	real rho2_beta_lin;
	real rho2_beta_sq;
	real sigma2[N];
	
	// Scale (variance) parameters
	rho2_alpha <- pow(rho_alpha, 2);
	rho2_beta_lin <- pow(rho_beta_lin, 2);
	rho2_beta_sq <- pow(rho_beta_sq, 2);
	for(n in 1:N)
	sigma2[n] <- pow(sigma[n], 2);
	
	}
	
model {

	// Model for the observed Wine Quality Ratings
	for(n in 2:N)
	WR_obs[n] ~ normal(WR_exp[n], 0.0937*exp(-0.627*WR_obs[n]));

	// Random walk with drift model for the locally-linear trend coefficient
	for(n in 2:N)
	alpha[n] ~ normal(l_alpha + alpha[n-1], rho_alpha);

	// Random walk with drift model for the linear temperature coefficient
	for(n in 2:N)
	beta_lin[n] ~ normal(l_beta_lin + beta_lin[n-1], rho_beta_lin);	
	
	// Random walk with drift model for the quadratic temperature coefficient	
	for(n in 2:N)
	beta_sq[n] ~ normal(l_beta_sq + beta_sq[n-1], rho_beta_sq);	

	// Stochastic volatility component
	for(n in 2:N)
	sigma[n] ~ normal(l_sigma + sigma[n-1], rho_sigma);

	// Time-varying coefficients model for the latent Wine Quality Ratings
	for(n in 2:N)
	WR_exp[n] ~ normal(alpha[n] + beta_lin[n] * temp[n] + beta_sq[n] * pow(temp[n], 2), sigma[n]);

	// Initial latent wine quality rating (1920)
	WR_exp[1] ~ normal(alpha[1] + beta_lin[1] * temp[1] + beta_sq[1] * pow(temp[1], 2), sigma[1]);
	
	// Initial stochastic value
	sigma[1] ~ uniform(0, 2);
	}

generated quantities {
	real log_lik[N];
	real Predict[N];
	real Resids[N];
	real mean_alpha;
	real mean_beta_lin;
	real mean_beta_sq;
	real mean_sigma2;
	real WR_hat[N];
	real Opt_temp[N];
	real R2[N];

	// Time-averaged estimates for the time-varying coefficients
	mean_alpha <- mean(alpha);
	mean_beta_lin <- mean(beta_lin);
	mean_beta_sq <- mean(beta_sq);
	mean_sigma2 <- mean(sigma2);

	for (n in 1:N){
	
	
	// Posterior log-likelihood for the time-varying coefficients model
	log_lik[n] <- normal_log(WR_exp[n], alpha[n] + beta_lin[n] * temp[n] + beta_sq[n] * pow(temp[n], 2), sigma[n]);

	// Predicted and residual values
	Predict[n] <- alpha[n] + beta_lin[n] * temp[n] + beta_sq[n] * pow(temp[n], 2);
	Resids[n] <- WR_exp[n] - Predict[n];

	// Posterior Predicted Datasets
	WR_hat[n] <- normal_rng(alpha[n] + beta_lin[n] * temp[n] + beta_sq[n] * pow(temp[n], 2), sigma[n]);

	// Time-varying estimate for optimum growing season temperature
	Opt_temp[n] <- (0.895331125133887*(-beta_lin[n]/(2*beta_sq[n]))) + 17.1740019813739;
	
	// Goodness-of-fit of the TVCs model
	R2[n] <- pow(beta_lin[n],2) + pow(beta_sq[n],2);

	}

}
'

# Call the stan function to fit the model ‘TVC_wine_code’ through HMC,
# and write samples and diagnostic files to the working directory:

fit <- stan(model_code = TVC_wine_code, 
		data = data, 
		seed = 666, 
		iter = 4000, 
		warmup = 3000, 
		refresh = 50, 
		chains = 3,
		sample_file = "TVC_HMC_samples", 
		diagnostic_file = "TVC_HMC_MCMCdiag")


# Calculate the Watanabe/Akaike, or Widely Applicable, Information Criterion,
# (WAIC) using the code written by Aki Vehtari and Andrew Gelman 
# (unpublished manuscript, 31 May 2014):


colVars <- function(stanfit) {
n <- dim(stanfit)[[1]]; 
c <- dim(stanfit)[[2]]; 
return(.colMeans(((stanfit - matrix(.colMeans(stanfit, n, c), nrow = n, ncol = c, byrow = TRUE)) ^ 2), n, c) * n / (n - 1))}

waic <- function(stanfit){
log_lik <- extract (stanfit, "log_lik")$log_lik
dim(log_lik) <- if (length(dim(log_lik))==1) c(length(log_lik),1) else
c(dim(log_lik)[1], prod(dim(log_lik)[2:length(dim(log_lik))]))
S <- nrow(log_lik)
n <- ncol(log_lik)
lpd <- log(colMeans(exp(log_lik)))
p_waic <- colVars(log_lik)
elpd_waic <- lpd - p_waic
WAIC <- -2*elpd_waic
loo_weights_raw <- 1/exp(log_lik-max(log_lik))
loo_weights_normalized <- loo_weights_raw/
matrix(colMeans(loo_weights_raw),nrow=S,ncol=n,byrow=TRUE)
loo_weights_regularized <- pmin (loo_weights_normalized, sqrt(S))
elpd_loo <- log(colMeans(exp(log_lik)*loo_weights_regularized)/
colMeans(loo_weights_regularized))
LOOCV <- -2*elpd_loo
p_loo <- lpd - elpd_loo
pointwise <- cbind(WAIC,LOOCV,lpd,p_waic,elpd_waic,p_loo,elpd_loo)
total <- colSums(pointwise)
se <- sqrt(n*colVars(pointwise))
return(list(WAIC=total["WAIC"], LOOCV=total["LOOCV"],elpd_waic=total["elpd_waic"],
p_waic=total["p_waic"], elpd_loo=total["elpd_loo"], p_loo=total["p_loo"],
pointwise=pointwise, total=total, "Standard Error"=se))
}

# Print the posterior parameter estimates along with the WAIC:

print(fit, probs = c(0.05, 0.10, 0.16, 0.50, 0.84, 0.90, 0.95), digits = 3)
print(waic(fit))

# Plot the posterior parameter estimates and the MC trace of each chain
# and parameter:

plot(fit, ask=TRUE)
trace <- traceplot(fit, inc_warmup = FALSE, ask=TRUE)

## End of the TVCs regression model fitting.

#############################
#############################
#############################

## Fit the Constant coefficients model:

setwd("C:/WQR/")
library(rstan) 
rm(list=ls(all=TRUE))

# Create the object ‘data’, containing the standardized wine quality 
# ratings, the growing season temperature values and the values for Year

data <- list(
N=90,
year=c(-1.703369,-1.665091,-1.626813,-1.588535,-1.550257,-1.511979,-1.473701,-1.435423,-1.397145,-1.358867,-1.320589,-1.282311,-1.244033,-1.205755,-1.167477,-1.1292,-1.090922,-1.052644,-1.014366,-0.976088,-0.93781,-0.899532,-0.861254,-0.822976,-0.784698,-0.74642,-0.708142,-0.669864,-0.631586,-0.593308,-0.55503,-0.516752,-0.478474,-0.440196,-0.401918,-0.363641,-0.325363,-0.287085,-0.248807,-0.210529,-0.172251,-0.133973,-0.095695,-0.057417,-0.019139,0.019139,0.057417,0.095695,0.133973,0.172251,0.210529,0.248807,0.287085,0.325363,0.363641,0.401918,0.440196,0.478474,0.516752,0.55503,0.593308,0.631586,0.669864,0.708142,0.74642,0.784698,0.822976,0.861254,0.899532,0.93781,0.976088,1.014366,1.052644,1.090922,1.1292,1.167477,1.205755,1.244033,1.282311,1.320589,1.358867,1.397145,1.435423,1.473701,1.511979,1.550257,1.588535,1.626813,1.665091,1.703369),
temp=c(-0.745864,0.270905,-0.542232,-0.560997,-0.699698,-0.769275,-0.259649,-0.837022,-0.015517,-0.349978,-0.990825,-1.604208,-1.015239,0.443452,0.512419,-0.394532,-1.488245,-0.277959,-1.444301,-1.410733,-0.426128,-1.211813,0.405774,1.330109,0.775508,1.60741,-0.854111,1.572564,-0.600824,1.302798,0.230446,-0.957562,0.215188,-0.212959,-1.720781,0.252418,-1.404324,-0.883102,-0.566645,0.62472,-0.566645,0.382724,-0.752796,-1.404324,0.326879,-1.795241,-0.529415,-0.845872,-0.845872,-0.47357,-0.343264,0.271033,-1.851086,0.531645,-0.585261,0.345494,1.015637,-1.395475,-0.864487,-0.5108,-0.659721,0.047652,0.531645,0.777303,-0.51965,-0.138498,-0.715566,0.568875,0.53317,1.617118,1.797776,1.014111,0.921951,-0.23902,0.524199,0.74758,0.058821,1.529413,0.617274,1.492183,1.101266,0.66195,0.140728,2.970221,1.052867,1.68578,2.169772,0.755026,0.140728,1.406554),
WR_obs=c(0.358684,0.836929,0.119561,0.119561,0.597806,-0.358684,0.836929,-0.358684,1.076051,1.315174,-0.836929,-2.510786,-1.076051,-0.836929,0.597806,-1.315174,-2.271664,0.597806,-0.597806,-0.597806,-0.358684,-2.989031,-0.836929,0.597806,-1.076051,1.076051,-1.076051,0.836929,0.119561,0.836929,0.358684,-1.554296,0.597806,0.836929,-1.315174,0.836929,-1.315174,-0.597806,-0.597806,1.076051,-0.597806,1.315174,0.597806,-2.749909,0.597806,-2.749909,0.597806,-0.119561,-2.032541,-0.597806,0.836929,0.597806,-1.076051,-0.597806,-0.597806,0.597806,0.358684,-0.836929,0.597806,0.358684,-0.358684,0.358684,1.076051,0.597806,-0.597806,0.836929,0.836929,-0.119561,0.597806,1.076051,1.076051,-0.358684,-0.597806,-0.119561,0.119561,0.597806,0.836929,0.119561,0.597806,0.358684,1.076051,0.597806,0.358684,0.836929,0.597806,1.315174,0.597806,0.358684,0.597806,1.076051))

# Create the object ‘CNT_wine_code’, which includes the code for fitting the 
# constant coefficients model that will be translated by 
# RStan to C++ code:


CNT_wine_code <- '

data {
	int<lower=0> N;
	real temp[N]; // Standardised average growing season temperature
	real WR_obs[N]; // Standardised Tastet & Lawton wine Quality Ratings 
	real year[N];  // Standardised Year
	}

parameters {
	// Bounded time series of latent Wine Quality Ratings
	real<lower=-1.85108645079154, upper=2.97022104705866> WR_exp[N];

	// Location parameters
	real I;
	real alpha;
	real beta_lin;
	real beta_sq;
	
	// Scale (standard deviation) parameter
	real<lower=0> sigma;
	}

transformed parameters {
	real sigma2;
	
	// Scale (variance) parameters
	sigma2 <- pow(sigma, 2);
	}
	
model {
	
	// Model for the observed Wine Quality Ratings	
	for(n in 2:N)
	WR_obs[n] ~ normal(WR_exp[n], 0.0937*exp(-0.627*WR_obs[n]));

	// Constant coefficients model for the latent Wine Quality Ratings
	for(n in 2:N)
	WR_exp[n] ~ normal(I + alpha * year[n] + beta_lin * temp[n] + beta_sq * pow(temp[n], 2), sigma);

	// Initial latent wine quality rating (1920)
	WR_exp[1] ~ normal(I + alpha * year[1] + beta_lin * temp[1] + beta_sq * pow(temp[1], 2), sigma);

	}
generated quantities {
	real log_lik[N];
	real Predict[N];
	real Resids[N];
	real WR_hat[N];
	real Opt_temp;
	real R2;
	
	for (n in 1:N){
	
	// Posterior log-likelihood for the time-varying coefficients model
	log_lik[n] <- normal_log(WR_exp[n], I + alpha * year[n] + beta_lin * temp[n] + beta_sq * pow(temp[n], 2), sigma);

	// Predicted and residual values
	Predict[n] <- I + alpha * year[n] + beta_lin * temp[n] + beta_sq * pow(temp[n], 2);
	Resids[n] <- WR_exp[n] - Predict[n];

	// Posterior Predicted Datasets
	WR_hat[n] <- normal_rng(I + alpha * year[n] + beta_lin * temp[n] + beta_sq * pow(temp[n], 2), sigma);

	}

	// Posterior estimate for optimum growing season temperature
	Opt_temp <- (0.895331125133887*(-beta_lin/(2*beta_sq))) + 17.1740019813739;

	// Goodness-of-fit of the TVCs model
	R2 <- pow(beta_lin,2) + pow(beta_sq,2);

	}
'

# Call the stan function to fit the model ‘CNT_wine_code’ through HMC,
# and write samples and diagnostic files to the working directory:

fit <- stan(model_code = CNT_wine_code, 
		data = data, 
		seed = 666, 
		iter = 4000, 
		warmup = 3000, 
		refresh = 50, 
		chains = 3,
		sample_file = "CNT_HMC_samples", 
		diagnostic_file = "CNT_HMC_MCMCdiag")


# Calculate the Watanabe/Akaike, or Widely Applicable, Information Criterion,
# (WAIC) using the code written by Aki Vehtari and Andrew Gelman 
# (unpublished manuscript, 31 May 2014):


colVars <- function(stanfit) {
n <- dim(stanfit)[[1]]; 
c <- dim(stanfit)[[2]]; 
return(.colMeans(((stanfit - matrix(.colMeans(stanfit, n, c), nrow = n, ncol = c, byrow = TRUE)) ^ 2), n, c) * n / (n - 1))}

waic <- function(stanfit){
log_lik <- extract (stanfit, "log_lik")$log_lik
dim(log_lik) <- if (length(dim(log_lik))==1) c(length(log_lik),1) else
c(dim(log_lik)[1], prod(dim(log_lik)[2:length(dim(log_lik))]))
S <- nrow(log_lik)
n <- ncol(log_lik)
lpd <- log(colMeans(exp(log_lik)))
p_waic <- colVars(log_lik)
elpd_waic <- lpd - p_waic
WAIC <- -2*elpd_waic
loo_weights_raw <- 1/exp(log_lik-max(log_lik))
loo_weights_normalized <- loo_weights_raw/
matrix(colMeans(loo_weights_raw),nrow=S,ncol=n,byrow=TRUE)
loo_weights_regularized <- pmin (loo_weights_normalized, sqrt(S))
elpd_loo <- log(colMeans(exp(log_lik)*loo_weights_regularized)/
colMeans(loo_weights_regularized))
LOOCV <- -2*elpd_loo
p_loo <- lpd - elpd_loo
pointwise <- cbind(WAIC,LOOCV,lpd,p_waic,elpd_waic,p_loo,elpd_loo)
total <- colSums(pointwise)
se <- sqrt(n*colVars(pointwise))
return(list(WAIC=total["WAIC"], LOOCV=total["LOOCV"],elpd_waic=total["elpd_waic"],
p_waic=total["p_waic"], elpd_loo=total["elpd_loo"], p_loo=total["p_loo"],
pointwise=pointwise, total=total, "Standard Error"=se))
}

# Print the posterior parameter estimates along with the WAIC:

print(fit, probs = c(0.05, 0.10, 0.16, 0.50, 0.84, 0.90, 0.95), digits = 3)
print(waic(fit))

# Plot the posterior parameter estimates and the MC trace of each chain
# and parameter:

plot(fit, ask=TRUE)
trace <- traceplot(fit, inc_warmup = FALSE, ask=TRUE)

## End of the constant regression model fitting.

########################
#### End RStan code ####
########################