library(bayesplot)
available_mcmc(pattern = "_nuts_")
library(ggplot2)

library(rstan)
#Resources: https://github.com/stan-dev/rstan/wiki/RStan-Getting-Started

#save unchanged models instead of recompiling
rstan_options(auto_write = TRUE)
#allow for multithreaded sampling
options(mc.cores = parallel::detectCores())

#test installation, shouldn't get any errors
#example(stan_model, package = "rstan", run.dontrun = TRUE)

################ Pull data from database ######################
library(RPostgreSQL)

driver <- dbDriver("PostgreSQL")

con <- dbConnect(
    driver,
    user='root',
    password='root',
    dbname='aact_db',
    host='will-office'
    )

query <- dbSendQuery(
    con,
    "select * from formatted_data_with_planned_enrollment;"
    )
df <- fetch(query)
df <- na.omit(df)

query2 <-dbSendQuery(con,"select count(*) from \"DiseaseBurden\".icd10_categories ic ;")
n_categories <- fetch(query2)

################ Format Data ###########################


categories <- df["category_id"]

x <- df["elapsed_duration"]
x["log_n_brands"] <- log(df$n_brands + 1)
x["h_sdi_val"] <- log(df$h_sdi_val + 1)
x["hm_sdi_val"] <- log(df$hm_sdi_val + 1)
x["m_sdi_val"] <- log(df$m_sdi_val + 1)
x["lm_sdi_val"] <- log(df$lm_sdi_val + 1)
x["l_sdi_val"] <- log(df$l_sdi_val + 1)
x["enrollment"] <- df["current_enrollment"] / df["planned_enrollment"]
#square terms
#x["enrollment^2"] <- x["enrollment"]^2
#x["elapsed_duration^2"] <- x["elapsed_duration"]^2
#x["n_brands^2"] <- x["n_brands"]^2
#break out
x["status_NYR"] <- ifelse(df["current_status"]=="Not yet recruiting",1,0)
x["status_Rec"] <- ifelse(df["current_status"]=="Recruiting",1,0)
x["status_ANR"] <- ifelse(df["current_status"]=="Active, not recruiting",1,0)
x["status_EBI"] <- ifelse(df["current_status"]=="Enrolling by invitation",1,0)

y <- ifelse(df["final_status"]=="Terminated",1,0)

################################# DATA EXPLORATION ############################

#Plot terminated vs completed
#Plot duration for terminated vs completed
#Plot different times of

################################# FIT #########################################
#setup data (named list)
trials_data <- list(
    D = ncol(x),#
    N = nrow(x),
#    L = n_categories$count,
    y = as.vector(y),
#    ll = as.vector(categories$category_id),
    x = as.matrix(x)
)

#fit <- stan(file='Logistic.stan', data = trials_data) #attempt at simple model
fit <- stan(file='Hierarchal_Logistic.stan', data = trials_data)

################################# ANALYZE #####################################
print(fit)

color_scheme_set("darkgray")
#trace plots
plot(fit, pars=c("mu"), plotfun="trace")
plot(fit, pars=c("sigma"), plotfun="trace")

#other diagnostics
logpost <- log_posterior(fit)
nuts_prmts <- nuts_params(fit)
posterior <- as.array(fit)


mcmc_parcoord(posterior,pars=c(
    "mu[1]",
    "mu[2]",
    "mu[3]",
    "mu[4]",
    "mu[5]",
    "mu[6]",
    "mu[7]",
    "mu[8]",
    "mu[9]",
    "mu[10]",
    "mu[11]",
    "mu[12]"
    ),
    np=nuts_prmts
)

mcmc_pairs(
    posterior,
    np = nuts_prmts,
    pars=c(
        "mu[1]",
        "mu[2]",
        "sigma[1]",
        "sigma[2]"
    ),
    off_diag_args = list(size = 0.75)
)

mcmc_parcoord(posterior,pars=c(
    "sigma[1]",
    "sigma[2]",
    "sigma[3]",
    "sigma[4]",
    "sigma[5]",
    "sigma[6]",
    "sigma[7]",
    "sigma[8]",
    "sigma[9]",
    "sigma[10]",
    "sigma[11]",
    "sigma[12]"
),
np=nuts_prmts
)
# check model estimates
#mu
plot(fit, pars=c("mu"), show_density =TRUE, ci_level=0.8)
plot(fit, pars=c("mu"), plotfun = "hist")

#sigma
plot(fit, pars=c("sigma"), show_density =TRUE, ci_level=0.8)
plot(fit, pars=c("sigma"), plotfun = "hist")
