#'
#' Utilities for SBC validation
#'
library(devtools)
devtools::load_all("../..")
library(rstan)

#'
#' Simulates a draw from the prior and fake data for it. This will be
#' the data generating step in the simulation. The function recieves
#' the problem data, job specifics and hyper-parameters which
#' determine the scenario (choices for the prior).
#'
simulate_fake <- function(data, job, family, mean_mu, sd_mu, sd_tau, samp_sd, ...) {
    G <- length(unique(data$group))
    N <- length(data$group)
    rl <- rle(data$group)
    Ng <- rl$lengths
    mu <- rnorm(1, mean_mu, sd_mu)
    tau <- abs(rnorm(1, 0, sd_tau))
    alpha_g <- rnorm(G, mu, tau)
    family <- get(family, mode = "function", envir = parent.frame())
    inv_link <- family()$linkinv
    theta <- inv_link(rep(alpha_g, times=Ng))
    likelihood <- family()$family
    if (likelihood == "binomial") {
        r <- rbinom(N, 1, theta)
        y <- as.numeric(tapply(r, data$group, sum))
        fake <- data.frame(r=y, nr=Ng-y, group=1:G)
    }
    if (likelihood == "poisson") {
        count <- rpois(N, theta)
        y <- as.numeric(tapply(count, data$group, sum))
        fake <- data.frame(y=y, n=Ng, group=1:G)
    }
    if (likelihood == "gaussian") {
        y_i <- rnorm(N, theta, samp_sd)
        y <- as.numeric(tapply(y_i, data$group, mean))
        fake <- data.frame(y=y, y_se=samp_sd/sqrt(Ng), group=1:G)
    }
    list(fake=fake,
         draw=c(mu=mu, tau=tau))
}

#'
#'
#' Procedure to fit each fake data set using our fitting
#' procedure. This method obtains the problem data, job details and an
#' **instance** of the scenario as generated by `simulate_fake`.
#'

fit_rbest <- function(data, job, instance, ...) {
    fake <- instance$fake
    draw <- instance$draw

    pars <- job$pars$prob.pars
    prior_mean_mu <- pars$mean_mu
    prior_sd_mu <- pars$sd_mu
    prior_sd_tau <- pars$sd_tau
    samp_sd <- pars$samp_sd
    family <- pars$family

    model <- switch(family,
                    binomial=cbind(r, nr) ~ 1 | group,
                    poisson=y ~ 1 + offset(log(n)) | group,
                    gaussian=cbind(y, y_se) ~ 1 | group)

    fit <- gMAP(model, data=fake, family=family,
                tau.dist="HalfNormal",
                tau.prior=cbind(0, prior_sd_tau),
                beta.prior=cbind(c(prior_mean_mu),c(prior_sd_mu)),
                chains=2)

    params <- c("beta[1]", "tau[1]")

    sampler_params <- rstan::get_sampler_params(fit$fit, inc_warmup=FALSE)
    n_divergent <- sum(sapply(sampler_params, function(x) sum(x[,'divergent__'])) )

    fit_sum <- rstan::summary(fit$fit)$summary
    min_Neff <- ceiling(min(fit_sum[params, "n_eff"], na.rm=TRUE))

    post <- as.matrix(fit)[,params]
    S <- nrow(post)
    ## thin down to 1023 draws so that we get 1024 bins
    post <- post[round(seq(1, S, length=1024-1)),]
    colnames(post) <- c("mu", "tau")
    list(rank=colSums(sweep(post, 2, draw) < 0), min_Neff=min_Neff, n_divergent=n_divergent)
}

scale_ranks <- function(Nbins, scale=1) {
    ## scale must evenly divide the total number of bins
    assert_that(round(Nbins/scale) == Nbins/scale)
    breaks <- (0:(Nbins/scale))
    Nbreaks <- length(breaks)
    function(scen) {
        vars <- grep("^rank.", names(scen), value=TRUE)
        res <- lapply(vars, function(v) hist(ceiling((scen[[v]]+1)/scale), breaks=breaks, plot=FALSE, include.lowest=FALSE)$counts)
        names(res) <- gsub("^rank", "count", vars)
        res$rank <- breaks[-Nbreaks]
        res <- as.data.frame(do.call(cbind, res))
        res
    }
}

