library(stringr)
library(reshape2)
#goal here is to take a given stanfit object (titled fit here)
#and extract, mark up, and melt a given set of parameters so that
#the data can be processed easily with ggplot2

extract_nums <- function(str)
{
    #https://stackoverflow.com/a/12727871
    str_extract_all(str,"\\(?[0-9]+\\)?")[[1]]
}
remap <- function(mapping) {
    function(index) {
        mapping[index]
    }
}



parameter_list <- c(
    `1`="Elapsed Duration",
    `2`="asinh(Number Brands)",
    `3`="asinh(High SDI)",
    `4`="asinh(High-Medium SDI)",
    `5`="asinh(Medium SDI)",
    `6`="asinh(Low-Medium SDI)",
    `7`="asinh(Low SDI)"
)
group_list <- c(
    `1`="Infections & Parasites",
    `2`="Neoplasms",
    `3`="Blood & Immune system",
    `4`="Endocrine, Nutritional, and Metabolic",
    `5`="Mental & Behavioral",
    `6`="Nervous System",
    `7`="Eye and Adnexa",
    `8`="Ear and Mastoid",
    `9`="Circulatory",
    `10`="Respiratory",
    `11`="Digestive",
    `12`="Skin & Subcutaneaous tissue",
    `13`="Musculoskeletal",
    `14`="Genitourinary",
    `15`="Pregancy, Childbirth, & Puerperium",
    `16`="Perinatal Period",
    `17`="Congential",
    `18`="Symptoms, Signs etc.",
    `19`="Injury etc.",
    `20`="External Causes",
    `21`="Contact with Healthcare",
    `22`="Special Purposes"
)

param <- "beta"
param_cols <- grep(param,fit@sim$fnames_oi)
param_col_names <- fit@sim$fnames_oi[param_cols]



#TODO: get it to collect all the data.

itt = 1
#extract data for iteration
data <- as.data.frame(fit@sim$samples[[itt]][beta_cols])
#add col number
data <- cbind(id=rownames(data),data)
#drop warmups
sample_start <- fit@sim$warmup2[itt]+1
sample_size <- fit@sim$n_save[itt]
data <- data[sample_start:sample_size,]
#add chain
data["chain"] <- itt


#melt it
id.vars <- c("id","chain")
data <- melt(data,id.vars)

#extract group and parameter identifiers
a <- t(as.data.frame(mapply(extract_nums,as.character(data$variable))))
a <- as.data.frame(a)
#add group and parameter identifiers back in
data["group"] <- a[1]
data["group"] <- as.integer(data$group)
data["parameter"] <- a[2]




ggplot(subset(data,parameter <=2), aes(x=value)) +
    stat_density() +
    xlim(-2,3) +
    facet_grid(group ~ parameter,
               labeller = labeller(group=group_list,parameter=parameter_list)
            ) +
    scale_y_discrete(labels = NULL) +
    theme(strip.text.y = element_text(angle = 0))