############################################################################# 
# Computes the effective sample size using the algorithm in Section 2.3 of
# the paper (http://arxiv.org/pdf/1011.0175v1.pdf) by Madeline Thompson.
# The algorithm is taken from earlier work on 'Initial Sequence Estimators'
# by multiple authors. 
# 
# The 'magic' argument is isnpired by the Appendix in the NUTS paper by
# Gelman & Hoffmann (2013).
# 
# Args: 
#   chain - matrix object with each sample (possibly multivariate) as a row
#   mu - vector of means (must of the same length as ncols(chain)
#   adj - Set to TRUE to enable Initial Convex Sequence Estimator (see Section
#         2.3 of the paper by Thompson, referred above).
#   magic - cutoff used in Appendix of NUTS paper by Gelman & Hoffmann (2013)
# Returns:
#   effective sample sizes for the time series in each col of 'chain'
############################################################################# 
ess <- function(chain, mu=NULL, adj=TRUE, magic=0.05) 
{
  dims <- dim(chain);
  M <- dims[1];
  K <- dims[2];

  # If M is of 1 dimensions
  if (is.null(M)) {
    M <- length(chain);
  }
  if (is.null(K) || is.na(K)) {
    K <- 1;
  }

  if (is.null(mu)) {
    if (K != 1) {
      mu = colMeans(chain); 
    }
    else {
      mu = mean(chain);
    }
  }

  x <- array();
  for (i in 1:K) {
    if(K != 1) { 
      sigma <- sqrt(rhof(chain[,i], M, 0, mu[i], 1)); # Center first term at 1
    }
    else {
      sigma <- sqrt(rhof(chain, M, 0, mu[i], 1)); # Center first term at 1
    }

    sum <- 0;
    prev <- 0; # storage for previous
    for (s in 1:(M-1)) {
      if (K != 1) {
        rho <- rhof(chain[,i], M, s, mu[i], sigma);
      }
      else {
        rho <- rhof(chain, M, s, mu[i], sigma);
      }

      # Break if less than magic number, or adjacent is negative
      if (!adj && rho < magic) {
        break;
      }
      else if (adj && (prev + rho) <= 0) {
        break;
      }

      else {
        #sum <- sum + (1 - s/M)*rho; # NUTS paper
        sum <- sum + rho; # Thompson version
        prev <- rho;
      }
    }
    x[i] <- (M/(1+2*sum));
  }

  return(x);
}

rhof <- function(chain, M, s, mu, sigma) {
  # Can optimize by saving the chain shifted by mu instead of doing it everytime...
  a <- sum((chain[1:(M-s)] - mu)*(chain[(s+1):M ] - mu)/(sigma^2*(M-s)));
  return(a);
}