# Write file from within Sage worksheet --

open('/tmp/8schools.stan','w').write("""
data {
  int<lower=0> J; // number of schools
  real y[J]; // estimated treatment effects
  real<lower=0> sigma[J]; // s.e. of effect estimates
}
parameters {
  real mu;
  real<lower=0> tau;
  real eta[J];
}
transformed parameters {
  real theta[J];
  for (j in 1:J)
    theta[j] <- mu + tau * eta[j];
}
model {
  eta ~ normal(0, 1);
  y ~ normal(theta, sigma);
}""")

%r

library(rstan)

%r

# WARNING: This takes a while to evaluate.  It's compiling C++ code,
# running it, etc., behind the scenes.  Be patient.

schools_dat <- list(J = 8,
                    y = c(28,  8, -3,  7, -1,  1, 18, 12),
                    sigma = c(15, 10, 16, 11,  9, 11, 10, 18))

fit <- stan(file = '/tmp/8schools.stan', data = schools_dat,
            iter = 1000, chains = 4)

%r

#
# This is a demo of using model_code argument since
# we can use this file directly or put the string in R
# directly.
#
schools_code <- paste(readLines('/tmp/8schools.stan'), collapse = '\n')
fit1 <- stan(model_code = schools_code, data = schools_dat,
             iter = 1000, chains = 4)

%r

fit2 <- stan(fit = fit1, data = schools_dat, iter = 10000, chains = 4)

%r

print(fit2)
plot(fit2)

%r

la <- extract(fit2, permuted = TRUE) # return a list of arrays
mu <- la$mu

### return an array of three dimensions: iterations, chains, parameters
a <- extract(fit2, permuted = FALSE)

### use S3 functions as.array (or as.matrix) on stanfit objects
a2 <- as.array(fit2)
m <- as.matrix(fit2)

%r

print(fit, digits = 1)