library(latex2exp)

index_seeds = c(1:25)
it_num = 1

args = commandArgs(TRUE)
dgm = as.logical(as.numeric(args[1])) # fit the data generating model (dgm) or the approx. model

# Parameter initialization -----------------------------------------------------
if(dgm) {
    par_index = list()
    par_index$alpha = 1:12
    par_index$zeta = 13:16
    par_index$diag_R = 17:20
    par_index$init = 21:22
    
    true_par = rep(0, max(do.call('c', par_index)))
    true_par[par_index$alpha] = c( 50,  -5,   5,
                                  100,  10, -10,
                                  100, -10,  10,
                                   50,   5,  -5)
    true_par[par_index$zeta] = c(-2, -2, -1.5, -1.5)
    true_par[par_index$diag_R] = c(log(4), log(4), log(4), log(4))
    true_par[par_index$init] = c(0, 0)
    
    labels = c(TeX(r'(Model A $\alpha_{1,1}:$ state 1 intercept for $y_1$ )'),
               TeX(r'(Model A $\alpha_{1,2}:$ state 2 slope for $y_1$ )'),
               TeX(r'(Model A $\alpha_{1,3}:$ state 3 slope for $y_1$ )'),
               TeX(r'(Model A $\alpha_{2,1}:$ state 1 intercept for $y_2$ )'),
               TeX(r'(Model A $\alpha_{2,2}:$ state 2 slope for $y_2$ )'),
               TeX(r'(Model A $\alpha_{2,3}:$ state 3 slope for $y_2$ )'),
               TeX(r'(Model A $\alpha_{3,1}:$ state 1 intercept for $y_3$ )'),
               TeX(r'(Model A $\alpha_{3,2}:$ state 2 slope for $y_3$ )'),
               TeX(r'(Model A $\alpha_{3,3}:$ state 3 slope for $y_3$ )'),
               TeX(r'(Model A $\alpha_{4,1}:$ state 1 intercept for $y_4$ )'),
               TeX(r'(Model A $\alpha_{4,2}:$ state 2 slope for $y_4$ )'),
               TeX(r'(Model A $\alpha_{4,3}:$ state 3 slope for $y_4$ )'),
               "logit baseline 1 -> 2", "logit baseline 2 -> 3",
               "logit baseline 3 -> 1", "logit baseline 3 -> 2",
               "log R(1,1)", "log R(2,2)", "log R(3,3)", "log R(4,4)",
               "logit init S2", "logit init S3")
    
    labels_short = c(TeX(r'($\alpha_{1,1}$ (Model A) )'),
                     TeX(r'($\alpha_{1,2}$ (Model A) )'),
                     TeX(r'($\alpha_{1,3}$ (Model A) )'),
                     TeX(r'($\alpha_{2,1}$ (Model A) )'),
                     TeX(r'($\alpha_{2,2}$ (Model A) )'),
                     TeX(r'($\alpha_{2,3}$ (Model A) )'),
                     TeX(r'($\alpha_{3,1}$ (Model A) )'),
                     TeX(r'($\alpha_{3,2}$ (Model A) )'),
                     TeX(r'($\alpha_{3,3}$ (Model A) )'),
                     TeX(r'($\alpha_{4,1}$ (Model A) )'),
                     TeX(r'($\alpha_{4,2}$ (Model A) )'),
                     TeX(r'($\alpha_{4,3}$ (Model A) )'),
                     "logit baseline 1 -> 2", "logit baseline 2 -> 3",
                     "logit baseline 3 -> 1", "logit baseline 3 -> 2",
                     "log R(1,1)", "log R(2,2)", "log R(3,3)", "log R(4,4)",
                     "logit init S2", "logit init S3", 
                     "log G(1,1)", "log G(2,2)", "log G(3,3)", "log G(4,4)")
} else {
    par_index = list()
    par_index$alpha = 1:8
    par_index$zeta = 9:12
    par_index$diag_R = 13:16
    par_index$init = 17:18
    par_index$diag_G = 19:22
    
    true_par = rep(0, max(do.call('c', par_index)))
    true_par[par_index$alpha] = c( -5,   5,
                                   10, -10,
                                  -10,  10,
                                    5,  -5)
    true_par[par_index$zeta] = c(-2, -2, -1.5, -1.5)
    true_par[par_index$diag_R] = c(log(4), log(4), log(4), log(4))
    true_par[par_index$init] = c(0, 0)
    true_par[par_index$diag_G] = c(log(4), log(4), log(4), log(4))
    
    labels = c(TeX(r'(Model B $\alpha^{*}_{1,1}:$ state 2 slope for $y_1$ )'),
               TeX(r'(Model B $\alpha^{*}_{1,2}:$ state 3 slope for $y_1$ )'),
               TeX(r'(Model B $\alpha^{*}_{2,1}:$ state 2 slope for $y_2$ )'),
               TeX(r'(Model B $\alpha^{*}_{2,2}:$ state 3 slope for $y_2$ )'),
               TeX(r'(Model B $\alpha^{*}_{3,1}:$ state 2 slope for $y_3$ )'),
               TeX(r'(Model B $\alpha^{*}_{3,2}:$ state 3 slope for $y_3$ )'),
               TeX(r'(Model B $\alpha^{*}_{4,1}:$ state 2 slope for $y_4$ )'),
               TeX(r'(Model B $\alpha^{*}_{4,2}:$ state 3 slope for $y_4$ )'),
               "logit baseline 1 -> 2", "logit baseline 2 -> 3",
               "logit baseline 3 -> 1", "logit baseline 3 -> 2",
               "log R(1,1)", "log R(2,2)", "log R(3,3)", "log R(4,4)",
               "logit init S2", "logit init S3", 
               "log G(1,1)", "log G(2,2)", "log G(3,3)", "log G(4,4)")
    
    labels_short = c(TeX(r'($\alpha^{*}_{1,1}$ (Model B) )'),
                     TeX(r'($\alpha^{*}_{1,2}$ (Model B) )'),
                     TeX(r'($\alpha^{*}_{2,1}$ (Model B) )'),
                     TeX(r'($\alpha^{*}_{2,2}$ (Model B) )'),
                     TeX(r'($\alpha^{*}_{3,1}$ (Model B) )'),
                     TeX(r'($\alpha^{*}_{3,2}$ (Model B) )'),
                     TeX(r'($\alpha^{*}_{4,1}$ (Model B) )'),
                     TeX(r'($\alpha^{*}_{4,2}$ (Model B) )'),
                     "logit baseline 1 -> 2", "logit baseline 2 -> 3",
                     "logit baseline 3 -> 1", "logit baseline 3 -> 2",
                     "log R(1,1)", "log R(2,2)", "log R(3,3)", "log R(4,4)",
                     "logit init S2", "logit init S3", 
                     "log G(1,1)", "log G(2,2)", "log G(3,3)", "log G(4,4)")
}

# Estimate the initial state probabilities
init_prob_mat = matrix(nrow = length(index_seeds), ncol = 3)
c_s = 1
for(seed in index_seeds) {
    load(paste0('Data/data_format', seed, '.rda'))
    first_ind = c(0, which(diff(data_format[,"id"]) != 0)) + 1
    init_state = data_format[first_ind, "state"]
    init_prob_mat[c_s, ] = c(sum(init_state == 1), sum(init_state == 2),
                             sum(init_state == 3)) / length(init_state)
    c_s = c_s + 1
}

init_par_est = colMeans(init_prob_mat)
true_par[par_index$init[1]] = log(init_par_est[2] / (1 - sum(init_par_est[2:3])))
true_par[par_index$init[2]] = log(init_par_est[3] / (1 - sum(init_par_est[2:3])))

# -----------------------------------------------------------------------------
# Create mcmc trace plots and histograms
# -----------------------------------------------------------------------------
chain_list = vector(mode = "list", length = length(index_seeds))

ind = 0
post_means_mat = NULL
covg = NULL

for(seed in index_seeds){
    
    it_seq = 1:it_num
    covg_val = FALSE
    
    for(it in it_seq) {
        file_name = paste0('Model_out/mcmc_out_',seed, '_it_', it, '_', as.numeric(dgm), '.rda')
        if(file.exists(file_name)) {
            load(file_name)
            print(paste0(seed, ": ", file_name))
            print("accept")
            print(mcmc_out$accept)
            
            par_index = mcmc_out$par_index
            
            if(it == 1) {
                ind = ind + 1
                chain_list[[ind]] = mcmc_out$chain[500:1000, ]
            } else {
                chain_list[[ind]] = rbind(chain_list[[ind]], mcmc_out$chain)
            }
            
            rm(mcmc_out)
            
            covg_val = TRUE
        } else {
            print(paste0("Missing! ", file_name))
        }
        
    }
    
    if(covg_val) {
        post_means_temp = matrix(colMeans(chain_list[[ind]]), nrow = 1)
        post_means_mat = rbind(post_means_mat, post_means_temp)
        
        covg_low = apply(chain_list[[ind]], 2, function(x){quantile(x, prob=.025)})
        covg_high = apply(chain_list[[ind]], 2, function(x){quantile(x, prob=.975)})
        covg_temp = as.numeric(true_par <= covg_high & true_par >= covg_low)
        covg = rbind(covg, covg_temp)    
    }
}

stacked_chains = do.call( rbind, chain_list)

pdf_title = paste0('Plots/trace_plot_', as.numeric(dgm), '.pdf')
pdf(pdf_title)
par(mfrow=c(4, 2), mar = c(3, 3, 2, 1), oma = c(0, 0, 1, 0)) # bot, left, top, right
lab_ind = 0
for(s in names(par_index)){
    temp_par = par_index[[s]]
    for(r in temp_par){
        # lab_ind = lab_ind + 1
        lab_ind = r
        parMean = round( mean(stacked_chains[,r], na.rm = T), 4)
        parMedian = round( median(stacked_chains[,r], na.rm = T), 4)
        upper = quantile( stacked_chains[,r], prob=.975, na.rm = T)
        lower = quantile( stacked_chains[,r], prob=.025, na.rm = T)
        
        title_color = "black"
        
        y_limit = range(stacked_chains[,r])
        plot( NULL, ylab=NA, main=labels[lab_ind], xlim=c(1,nrow(chain_list[[1]])),
              ylim=y_limit, xlab = paste0("95% CI: [", round(lower, 4),
                                          ", ", round(upper, 4), "]"),
              col.main = title_color)
        
        for(seed in 1:length(chain_list)) { lines( chain_list[[seed]][,r], type='l', col=seed) }
        
        x_label = paste0('Mean =',toString(parMean),
                         ' Median =',toString(parMedian),
                         ' True =', round(true_par[r], 3)) 
        
        hist( stacked_chains[,r], breaks=sqrt(nrow(stacked_chains)), ylab=NA, freq=FALSE,
              xlab=x_label, main = "")
        abline( v=upper, col='red', lwd=2, lty=2)
        abline( v=lower, col='purple', lwd=2, lty=2)
        abline( v=true_par[r], col='green', lwd=2, lty=2)
    }   
}
dev.off()

pdf(paste0('Plots/box_plot_', as.numeric(dgm), '.pdf'))
par(mfrow=c(4, 3), 
    mar = c(3, 3, 2, 1), 
    oma = c(0, 0, 1, 0), 
    mgp = c(0.5, 1, 0), # default is c(3, 1, 0)
    cex = 0.9) # bot, left, top, right
lab_ind = 0
for(s in names(par_index)){
    temp_par = par_index[[s]]
    
    temp=1
    for(r in temp_par){
        temp = temp + 1
        
        lab_ind = r
        boxplot(post_means_mat[,r], main = labels_short[r],
                xlab = paste0("true value = ", true_par[r]),
                ylab = "")
        abline(h = true_par[r], col = 'red')
        if(s == "alpha") {
          if(!dgm){
            if(temp%%3 == 0){
              plot.new()
              temp=1
            }
          }
        }
    }
    
    par(mfrow=c(4, 4))
}

dev.off()

