## to use
##
## markers <- read.marker.data("filename.csv")
## result <- analyze.data(markers, model="full", DIC=FALSE)
##           model argument defaults to "full", alternatives are "f0" and "t0"
##           DIC argument defaults to FALSE, alternative is TRUE
## print.summary(result)

count.genos <- function(x) {
  genos <- numeric(3)
  genos[1] <- sum(x==0, na.rm=TRUE)
  genos[2] <- sum(x==1, na.rm=TRUE)
  genos[3] <- sum(x==2, na.rm=TRUE)
  genos
}

read.marker.data <- function(filename) {
  markers <- read.csv(filename, header=TRUE, na.strings=".")

  n.pops <-length(unique(markers$pop))
  n.loci <- ncol(markers)-1

  locus <- colnames(markers)

  n <- array(dim=c(n.pops, n.loci, 3))
  N <- matrix(nrow=n.pops, ncol=n.loci)
  for (i in 1:n.pops) {
    pop.n <- unique(markers$pop)[i]
    ## x has one row for each individual in the population
    ## each locus is in a different column
    x <- subset(markers, pop==pop.n)
    for (j in 1:n.loci) {
      ## +1 because indiv is in column 1
      n[i, j, ] <- count.genos(x[,j+1])
      N[i,j] <- sum(n[i,j,])
    }
  }
  list(n.pops=n.pops,
       n.loci=n.loci,
       n=n,
       N=N)
}

get.llike <- function(n, p, f) {
  x <- numeric(3)
  x[1] <- p^2 + f*p*(1.0 - p)
  x[2] <- 2.0*p*(1.0 - p)*(1.0 - f)
  x[3] <- (1.0 - p)^2 + f*p*(1.0 - p)
  prob <- dmultinom(n, prob=x, log=TRUE)
  prob
}

report.DIC <- function(fit, n) {
  n.pops <- dim(n)[1]
  n.loci <- dim(n)[2]
  f <- fit$BUGSoutput$sims.list$f
  p <- fit$BUGSoutput$sims.list$p
  n.samp <- dim(p)[1]

  ## calculate Dbar
  ##
  llike <- 0.0
  for (k in 1:n.samp) {
    for (i in 1:n.pops) {
      if (n.loci == 1) {
        llike <- llike + get.llike(n[i,1,], p[k,i,1], f[k])
      } else {
        for (j in 1:n.loci) {
          llike <- llike + get.llike(n[i,j,], p[k,i,j], f[k])
        }
      }
    }
  }
  Dbar <- -2.0*llike/n.samp

  ## calculate Dhat
  ##
  f <- fit$BUGSoutput$mean$f
  p <- fit$BUGSoutput$mean$p
  llike <- 0.0
  for (i in 1:n.pops) {
    if (n.loci == 1) {
      llike <- llike + get.llike(n[i,1,], p[i], f)
    } else {
      for (j in 1:n.loci) {
        llike <- llike + get.llike(n[i,j,], p[i,j], f)
      }
    }
  }
  Dhat <- -2.0*llike
  pD <- Dbar - Dhat

  list(Dbar=Dbar,
       Dhat=Dhat,
       pD=pD,
       DIC=Dbar+pD)
}

## markers - a list containing n, N, n.pops, and n.loc
##   n - an n.pops x n.loci x 3 array of genotype counts
##   N - an n.pops x n.loci array of sample sizes
##   n.pops - number of populations in the sample
##   n.loci - number of loci scored
##
analyze.data <- function(markers,
                         n.sample=25000,
                         n.burnin=5000,
                         n.thin=25,
                         n.chains=1,
                         model="full",
                         DIC=FALSE)
{
  require(R2jags)

  n <- markers$n
  N <- markers$N
  n.pops <- markers$n.pops
  n.loci <- markers$n.loci

  n.sample <- n.sample
  n.thin <- n.thin
  n.chains <- n.chains
  n.iter   <- n.sample + n.burnin

  jags.data <- c("n",
                 "N",
                 "n.pops",
                 "n.loci")
  if (DIC) {
    jags.params <- c("f",
                     "theta",
                     "p")
  } else {
    jags.params <- c("f",
                     "theta")
  }

  if (model=="full") {
    model.file <- "f-statistics.jags"
  } else if (model=="f0") {
    model.file <- "f-statistics-f0.jags"
  } else if (model=="t0") {
    model.file <- "f-statistics-t0.jags"
  } else {
    stop(paste("Unrecognized model option: ", model, sep=""))
  }
  fit <- jags(data=jags.data,
              inits=NULL,
              parameters.to.save=jags.params,
              model.file=model.file,
              n.chains=n.chains,
              n.iter=n.iter,
              n.burnin=n.burnin,
              n.thin=n.thin)
  if (DIC) {
    dic <- report.DIC(fit, markers$n)
    cat("Dbar: ", round(dic$Dbar, 1), "\n", sep="")
    cat("Dhat: ", round(dic$Dhat, 1), "\n", sep="")
    cat("pD:   ", round(dic$pD, 1), "\n", sep="")
    cat("DIC:  ", round(dic$DIC, 1), "\n", sep="")
  }
  fit
}

print.summary <- function(fit, digits=3) {
  old.width <- options(width=180)
  print(fit, digits=digits)
  options(width=old.width$width)
}