# Purpose: Partner code for the RStan tutorial
# Date: February 17, 2022
# Author: Renny Doig

# Preliminaries ----------------------------------------------------------------

rm(list=ls())
library(tidyverse)
library(rstan)
library(extraDistr)


# Prepare data -----------------------------------------------------------------

Income <- read_csv("../Data/Income.csv")$Income

# enable parallel computing
options(mc.cores=5)

# set up data
n <- length(Income)
K <- 4

stan_data <- list(n=n, K=K, y=Income)


# Run stan ---------------------------------------------------------------------

# set up parallel cores
options(mc.cores=5)

# set up data
stan_data <- list(n=n, K=K, y=y)

# initial values
init_func <- function(){
  list(alpha = as.numeric(rdirichlet(1, rep(1, K))),
       mu = sort(rnorm(K, 70, 5)),
       sigma = 1 / sqrt(rgamma(1, 1, 1)))
}

fit <- stan(file = "gaussian_mixture.stan",
            data = stan_data,
            chains = 5,
            iter = 1e4,
            init = init_func,
            seed = 2)