#  begin by clearing memory
rm(list=ls())
#  Load libraries
library(cmdstanr)
library(posterior)
library(readxl)
# Set directory
setwd("/Users/mwatson/Dropbox/TVExtreme/ReplicationFiles/Stan/WeatherDamages")

# --- Run for raw damages then for normalized damages 
# Specify model input
stan_dir <- "/Users/mwatson/Dropbox/TVExtreme/ReplicationFiles/Stan/Stan_Code"
application_dir = "/Users/mwatson/Dropbox/TVExtreme/ReplicationFiles/Stan/WeatherDamages"
model_str <- "PGP_Static_Model"

# ==== Damages Model ... un-normalized ====
# application_str = "Disasters"
# data_str <- "Disasters"
# data_file <- file.path(application_dir, paste0(data_str,"_y.xlsx"))
# y <- read_excel(data_file,col_names = FALSE)
# data_file <- file.path(application_dir, paste0(data_str,"_tau.xlsx"))
# tau <- read_excel(data_file,col_names = FALSE)
# data_file <- file.path(application_dir, paste0(data_str,"_nobs.xlsx"))
# nobs <- read_excel(data_file,col_names = FALSE)

# # ------ Damages Model ... normalized ------
application_str = "Disasters_Normalized"
data_str <- "Disasters"
data_file <- file.path(application_dir, paste0(data_str,"_y_norm.xlsx"))
y <- read_excel(data_file,col_names = FALSE)
data_file <- file.path(application_dir, paste0(data_str,"_tau_norm.xlsx"))
tau <- read_excel(data_file,col_names = FALSE)
data_file <- file.path(application_dir, paste0(data_str,"_nobs.xlsx"))
nobs <- read_excel(data_file,col_names = FALSE)

# Process data and get dimensions
# Convert to matrix
y <- as.matrix(y)
y <- t(y)
nobs <- as.matrix(nobs) 
nobs <- c(nobs)
# Convert tau to a vector
tau <- as.matrix(tau) 
tau <- c(tau)
#  Get dimensions of cs_data
T <- dim(y)[1]

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

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

#  Inittial Values of parameters and seed value
trans_xi_level_init <- 0.7
ln_s_level_init <- -0.05
m_level_init <- 0.6

init_list_1 <- list(
    trans_xi_level = trans_xi_level_init,
    ln_s_level = ln_s_level_init,
    m_level = m_level_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
)
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
)
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
)
seed_value <- 87123

#  Carry Out Analysis

#  Model file
model_file <- file.path(stan_dir, paste0(model_str, ".stan"))

# Input Files
init_file <- file.path(application_dir, "Parm.init.json")
data_file <- file.path(application_dir, paste0(data_str,".json"))

# Output Files
draws_file <- file.path(application_dir, paste0("CSV/",application_str,".",model_str,".Draws."))
output_file <- file.path(application_dir, paste0("Output/",application_str,".",model_str))

# Compile model
mod <- cmdstan_model(model_file)

# MCMC Draws model
fit <- mod$sample(
    data = data_list,
    iter_warmup = 5000,
    iter_sampling = 10000,
    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,".diagonstics.txt"))
fit$diagnostic_summary()
# fit summary
sink(paste0(output_file,".summary.txt"))
fit$print(max_rows = 1000)
sink(file = NULL)

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

tmp <- subset_draws(draws_arr, c("xi"))
xi_draws <- merge_chains(tmp)
write.csv(xi_draws, paste0(draws_file, "xi.csv"))
# #
# tmp <- subset_draws(draws_arr, c("psi"))
# psi_draws <- merge_chains(tmp)
# write.csv(psi_draws, paste0(draws_file, "psi.csv"))
# #
# tmp <- subset_draws(draws_arr, c("lambda"))
# lambda_draws <- merge_chains(tmp)
# write.csv(lambda_draws, paste0(draws_file, "lambda.csv"))




