#'
#' 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))
    theta <- rnorm(G, mu, tau)
    family <- get(family, mode = "function", envir = parent.frame())
    inv_link <- family()$linkinv
    alpha <- inv_link(rep(theta, times=Ng))
    likelihood <- family()$family
    if (likelihood == "binomial") {
        r <- rbinom(N, 1, alpha)
        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, alpha)
        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, alpha, 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),
         draw_theta=theta)
}

#'
#'
#' 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
    draw_theta <- instance$draw_theta
    Ng <- length(draw_theta)

    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]")
    params_group <- paste0("theta[", 1:Ng, "]")

    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]
    post_group <- as.matrix(fit)[,params_group]
    S <- nrow(post)
    ## thin down to 1023 draws so that we get 1024 bins
    idx  <- round(seq(1, S, length=1024-1))
    post <- post[idx,]
    post_group <- post_group[idx,]
    colnames(post) <- c("mu", "tau")
    unlist(list(rank=c(colSums(sweep(post, 2, draw) < 0), colSums(sweep(post_group, 2, draw_theta) < 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
    }
}

