#  begin by clearing memory
rm(list=ls())
#  Load libraries
library(cmdstanr)
library(posterior)
library(readxl)
# Set directory
setwd("/Users/mwatson/Dropbox/TVExtreme/CitySize/CmdStanR")
#
# Specify model input
# --------------- City size Model -------------
stan_dir <- "/Users/mwatson/Dropbox/TVExtreme/ReplicationFiles/Stan/Stan_Code"
application_str = "CitySize"
application_dir = "/Users/mwatson/Dropbox/TVExtreme/ReplicationFiles/Stan/CitySize"
data_str <- "CitySize_Data_30"

#  Model names
# Test versus constant parameters
# model_str_1 <- "GEV_4RW_v_constant_Model_1"
# model_str_2 <- "GEV_4RW_v_constant_Model_2"
# Test versus alpha RW
# model_str_1 <- "GEV_4RW_v_alpha_Model_1"
# model_str_2 <- "GEV_4RW_v_alpha_Model_2"
# Test versus xi, sigma, mu RWs
model_str_1 <- "GEV_4RW_v_xi_sigma_mu_Model_1"
model_str_2 <- "GEV_4RW_v_xi_sigma_mu_Model_2"

#  Model files
model_file_1 <- file.path(stan_dir, paste0(model_str_1, ".stan"))
draws_file_1 <- file.path(application_dir, paste0("CSV/",application_str,".",model_str_1,".Draws."))
output_file_1 <- file.path(application_dir, paste0("Output/",application_str,".",model_str_1))

model_file_2 <- file.path(stan_dir, paste0(model_str_2, ".stan"))
draws_file_2 <- file.path(application_dir, paste0("CSV/",application_str,".",model_str_2,".Draws."))
output_file_2 <- file.path(application_dir, paste0("Output/",application_str,".",model_str_2))

# Read Data From Excel File
data_file <- file.path(application_dir, paste0(data_str,".xlsx"))
cs_data <- read_excel(data_file,col_names = FALSE)

# Process data and get dimensions
# Convert to matrix
y <- as.matrix(cs_data)
#  Get dimensions of cs_data
T <- dim(y)[1]
k <- dim(y)[2]

# Parameter values
#  xi bounds
x_min <- -0.9
x_max <- 2.1

# Scale gamma parameters
sg_xi <- 1.0
sg_alpha <- 0.5
sg_s <- 0.5

# Create list of data inputs for Stan
data_list <- list(
    T = T,
    nobs = k,
    xi_min = x_min,
    xi_max = x_max,
    sg_xi = sg_xi,
    sg_alpha = sg_alpha,
    sg_s = sg_s,
    y = t(y)
)

#  Inittial Values of parameters and seed value
trans_xi_level_init <- 0.0
ln_s_level_init <- 0.0
m_level_init <- 1.5
g_xi_init <- 0.2
g_alpha_init <- 0.2
g_s_init <- 0.2
g_m_init <- 0.2

init_list_1 <- list(
    trans_xi_level = trans_xi_level_init,
    ln_s_level = ln_s_level_init,
    m_level = m_level_init,
    g_xi = g_xi_init,
    g_alpha = g_alpha_init,
    g_s = g_s_init,
    g_m = g_m_init
)
init_list_2 <- list(
    trans_xi_level = trans_xi_level_init + 0.1,
    ln_s_level = ln_s_level_init + 0.1,
    m_level = m_level_init + 0.1,
    g_xi = g_xi_init,
    g_alpha = g_alpha_init,
    g_s = g_s_init,
    g_m = g_m_init
)
init_list_3 <- list(
    trans_xi_level = trans_xi_level_init - 0.1,
    ln_s_level = ln_s_level_init - 0.1,
    m_level = m_level_init - 0.1,
    g_xi = g_xi_init,
    g_alpha = g_alpha_init,
    g_s = g_s_init,
    g_m = g_m_init
)
init_list_4 <- list(
    trans_xi_level = trans_xi_level_init + 0.0,
    ln_s_level = ln_s_level_init + 0.1,
    m_level = m_level_init + 0.1,
    g_xi = g_xi_init,
    g_alpha = g_alpha_init,
    g_s = g_s_init,
    g_m = g_m_init
)
seed_value <- 87123


# Results for Model 1
# Compile model
mod1 <- cmdstan_model(model_file_1)

# MCMC Draws model
fit1 <- mod1$sample(
    data = data_list,
    iter_warmup = 5000,
    iter_sampling = 25000,
    init=list(
        init_list_1, # chain 1
        init_list_2, # chain 2
        init_list_3, # chain 3
        init_list_4  # chain 4
    ),
    refresh = 5000,
    chains = 4,
    parallel_chains = 4,
    seed = seed_value
) 

# fit diagnostics
sink(paste0(output_file_1,".diagonstics.txt"))
fit1$diagnostic_summary()
sink(file = NULL)

# Save Draws of key parameters
draws_arr <- fit1$draws()

tmp <- subset_draws(draws_arr, c("xi_1"))
xi_1_draws <- merge_chains(tmp)
write.csv(xi_1_draws, paste0(draws_file_1, "xi_1.csv"))
#
tmp <- subset_draws(draws_arr, c("sigma_1"))
sigma_1_draws <- merge_chains(tmp)
write.csv(sigma_1_draws, paste0(draws_file_1, "sigma_1.csv"))
#
tmp <- subset_draws(draws_arr, c("mu_1"))
mu_1_draws <- merge_chains(tmp)
write.csv(mu_1_draws, paste0(draws_file_1, "mu_1.csv"))

tmp <- subset_draws(draws_arr, c("xi_2"))
xi_2_draws <- merge_chains(tmp)
write.csv(xi_2_draws, paste0(draws_file_1, "xi_2.csv"))
#
tmp <- subset_draws(draws_arr, c("sigma_2"))
sigma_2_draws <- merge_chains(tmp)
write.csv(sigma_2_draws, paste0(draws_file_1, "sigma_2.csv"))
#
tmp <- subset_draws(draws_arr, c("mu_2"))
mu_2_draws <- merge_chains(tmp)
write.csv(mu_2_draws, paste0(draws_file_1, "mu_2.csv"))


# Results for Model 2
# Compile model
mod2 <- cmdstan_model(model_file_2)

# MCMC Draws model
fit2 <- mod2$sample(
    data = data_list,
    iter_warmup = 5000,
    iter_sampling = 25000,
    init=list(
        init_list_1, # chain 1
        init_list_2, # chain 2
        init_list_3, # chain 3
        init_list_4  # chain 4
    ),
    refresh = 5000,
    chains = 4,
    parallel_chains = 4,
    seed = seed_value
) 
# fit diagnostics
sink(paste0(output_file_2,".diagonstics.txt"))
fit2$diagnostic_summary()
sink(file = NULL)

# Save Draws of key parameters
draws_arr <- fit2$draws()

tmp <- subset_draws(draws_arr, c("xi_1"))
xi_1_draws <- merge_chains(tmp)
write.csv(xi_1_draws, paste0(draws_file_2, "xi_1.csv"))
#
tmp <- subset_draws(draws_arr, c("sigma_1"))
sigma_1_draws <- merge_chains(tmp)
write.csv(sigma_1_draws, paste0(draws_file_2, "sigma_1.csv"))
#
tmp <- subset_draws(draws_arr, c("mu_1"))
mu_1_draws <- merge_chains(tmp)
write.csv(mu_1_draws, paste0(draws_file_2, "mu_1.csv"))

tmp <- subset_draws(draws_arr, c("xi_2"))
xi_2_draws <- merge_chains(tmp)
write.csv(xi_2_draws, paste0(draws_file_2, "xi_2.csv"))
#
tmp <- subset_draws(draws_arr, c("sigma_2"))
sigma_2_draws <- merge_chains(tmp)
write.csv(sigma_2_draws, paste0(draws_file_2, "sigma_2.csv"))
#
tmp <- subset_draws(draws_arr, c("mu_2"))
mu_2_draws <- merge_chains(tmp)
write.csv(mu_2_draws, paste0(draws_file_2, "mu_2.csv"))
