edgeR <- function(counts, group, ylim = NULL, xlim = NULL){
  d <- DGEList(counts)
  d <- suppressWarnings(edgeR::calcNormFactors(d))
  design <- model.matrix(~group)
  d <- estimateDisp(d, design)
  plotBCV(d, ylim = ylim, main = 'edgeR', xlim = xlim)
  fit <- glmFit(d,design)
  lrt <- glmLRT(fit, coef = 2)
  pval <- lrt$table$PValue
  padj <- p.adjust(pval, "BH")
  cbind(pval = pval, padj = padj)
}

edger_ngenes <- function(counts, group, ylim = NULL, xlim = NULL){
  d <- DGEList(counts)
  d <- suppressWarnings(edgeR::calcNormFactors(d))
  cdr <- colSums(counts > 0)
  design <- model.matrix(~ group + cdr)
  d <- estimateDisp(d, design)
  plotBCV(d, ylim = ylim, main = 'edgeR', xlim = xlim)
  fit <- glmFit(d,design)
  lrt <- glmLRT(fit, coef = 2)
  pval <- lrt$table$PValue
  padj <- p.adjust(pval, "BH")
  cbind(pval = pval, padj = padj)
}

DESeq2 <- function(counts, group, ylim = NULL, xlim = NULL){
  colData <- data.frame(group = group)
  dse <- DESeqDataSetFromMatrix(countData=counts, colData=colData, design=~group)
  colData(dse)$group <- as.factor(colData(dse)$group)
  dse <- DESeq2::estimateSizeFactors(dse, type="poscounts")
  dse <- estimateDispersions(dse, minmu=1e-3)
  dse <- nbinomWaldTest(dse, minmu=1e-3, betaPrior=TRUE)
  rr <- results(dse)
  cbind(pval = rr$pvalue, padj = rr$padj)
}

limma <- function(counts, group, ylim = NULL, xlim = NULL){
    design <- model.matrix(~ group)
    nf <- suppressWarnings(edgeR::calcNormFactors(counts))
    y <- voom(counts, design, plot = FALSE, lib.size = colSums(counts) * nf)
    fit <- lmFit(y, design)
    fit <- eBayes(fit)
    tt <- topTable(fit, coef = 2, n = nrow(counts), sort.by = "none")
    pval <- tt$P.Value
    padj <- tt$adj.P.Val
    cbind(pval = pval, padj = padj)
}

scde <- function(counts, group, ylim=NULL, xlim=NULL){
  counts = matrix(as.integer(counts),nrow=nrow(counts),ncol=ncol(counts))
  if(is.null(colnames(counts))) colnames(counts)=paste0("sample",1:ncol(counts))
  require(scde)

  # calculate error models
  o.ifm <- scde.error.models(counts = counts, groups = group, n.cores = 1, threshold.segmentation = TRUE, save.crossfit.plots = FALSE, save.model.plots = FALSE, verbose = 0)
  # estimate gene expression prior
  o.prior <- scde.expression.prior(models = o.ifm, counts = counts, length.out = 400, show.plot = FALSE)
  # calculate differential expression
  ediff <- scde.expression.difference(o.ifm, counts, o.prior, groups  =  group, n.randomizations  =  150, n.cores  =  1, verbose  =  0)
  lfc <- ediff$mle
  pval=(1-pnorm(abs(ediff$Z)))*2
  padj=(1-pnorm(abs(ediff$cZ)))*2
  out = cbind(pval,padj,lfc)
  return(out)
}

### copied code from FPR_mocks.Rmd on September 14, 2017.
MAST <- function(counts, group, ylim = NULL, xlim = NULL){
  require(MAST)
  tpm <- counts*1e6/colSums(counts)
  tpm <- log2(tpm+1)
  sca <- FromMatrix(tpm,  cData=data.frame(group=group))
  #sca <- FromMatrix(counts,  cData=data.frame(group=group))

  # Adaptive thresholding from MAST vignette
  freq_expressed <- 0.2
  thres <- thresholdSCRNACountMatrix(assay(sca), nbins = 10, min_per_bin = 50, conditions = group)
  #par(mfrow=c(5,4))
  #plot(thres)
  assays(sca) <- list(thresh=thres$counts_threshold, tpm=assay(sca))
  expressed_genes <- freq(sca) > freq_expressed
  sca <- sca[expressed_genes,]

  ngeneson <- apply(counts,2,function(x) mean(x>0))
  CD <- colData(sca)
  CD$ngeneson <- ngeneson
  CD$cngeneson <- CD$ngeneson-mean(ngeneson)
  colData(sca) <- CD
  ## differential expression
  fit <- zlm(~ cngeneson + group , sca = sca)
  lrFit <- lrTest(fit, 'group')
  pval <- lrFit[, 'hurdle', 'Pr(>Chisq)']
  padj <- p.adjust(pval, method = "BH")

  ### MAST filtered the genes, so make a list that is consistent with the original count matrix.
  pvalAll = vector(length=nrow(counts))
  pvalAll[] = 1
  names(pvalAll)=rownames(counts)
  pvalAll[match(names(pval),names(pvalAll))] = pval

  padjAll = vector(length=nrow(counts))
  padjAll[] = 1
  names(padjAll)=rownames(counts)
  padjAll[match(names(padj),names(padjAll))] = padj

  out = cbind(pval = pvalAll, padj = padjAll, logfc = NA)
  return(out)
}

NODES <- function(counts, group, xlim, ylim){
  require(NODES)
  g=ifelse(group==0,"A","B")
  colnames(counts)=g
  normCounts=pQ(counts)
  res=NODES::NODES(data=normCounts,group=colnames(normCounts))
  pval=vector(length=nrow(counts))
  names(pval)=rownames(counts)
  pval[rownames(normCounts)]=res$Fisher
  pval[is.na(pval)]=1
  padj=p.adjust(pval,"BH")
  lfc=NA
  out=cbind(pval,padj,lfc)
  return(out)
}

metagenomeSeq <- function(counts, group, xlim, ylim){
  require(metagenomeSeq)
  design <- model.matrix(~group)
  pheno <- AnnotatedDataFrame(data.frame(group=group))
  rownames(pheno) <- colnames(counts)
  p <- cumNormStatFast(counts)
  dat <- newMRexperiment(counts=counts, phenoData=pheno, featureData = NULL, libSize = colSums(counts), normFactors = metagenomeSeq::calcNormFactors(counts, p=p))
  fit <- fitZig(dat,design)
  lfc <- fit$eb$coefficients[,"group1"]
  pval <- fit$eb$p.value[,"group1"]
  padj <- p.adjust(pval)
  out <- cbind(pval,padj,lfc)
  return(out)
}

Seurat <- function(counts, group, xlim=NULL, ylim=NULL){
    require(Seurat)
    seur = CreateSeuratObject(counts, project=paste0("cell",as.character(group)), display.progress = FALSE)
    res <- FindMarkers(seur, ident.1 = "cell0", ident.2 = "cell1", print.bar=FALSE)
    pval = res$p_val[match(rownames(counts),rownames(res))]
    padj = p.adjust(pval,"BH")
    out=cbind(pval,padj)
    return(out)
}

#https://github.com/miaozhun/DEsingle downloaded on November 20, 2017
DEsingle <- function(counts, group, xlim=NULL, ylim=NULL){
  require(DEsingle)
  res = DEsingle(counts=counts, group=group)
  pval = res[,"pvalue"]
  padj = res[,"pvalue.adj.FDR"]
  out=cbind(pval,padj)
  return(out)
}

zingeR_edgeR <- function(counts, group, ylim = NULL, xlim = NULL){
  #require(zingeR)
  d <- DGEList(counts)
  d <- suppressWarnings(edgeR::calcNormFactors(d))
  design <- model.matrix(~ group)
  weights <- zingeR::zeroWeightsLS(counts = d$counts, design = design, maxit = 300,
                           normalization = "TMM", verbose = F)
  d$weights <- weights
  d <- estimateDisp(d, design)
  plotBCV(d, ylim = ylim, main = 'zingeR', xlim = xlim)
  fit <- glmFit(d,design)
  lrt <- zinbwave::glmWeightedF(fit, coef = 2, independentFiltering = TRUE)
  cbind(pval = lrt$table$PValue, padj =lrt$table$padjFilter)
}

zingeR_DESeq2 <- function(counts, group, ylim = NULL, xlim = NULL){
  #require(zingeR)
  colData <- data.frame(group = group)
  design <- model.matrix(~ group)
  dse <- DESeqDataSetFromMatrix(countData = counts, colData = colData,
                                design = ~group)
  weights <- zingeR::zeroWeightsLS(counts = counts, design = design, maxit = 300,
                           normalization = "DESeq2_poscounts", colData = colData,
                           designFormula = ~group, verbose = F)
  assays(dse)[["weights"]] <- weights
  dse <- DESeq2::estimateSizeFactors(dse, type="poscounts")
  dse <- estimateDispersions(dse)
  dse <- nbinomWaldTest(dse, betaPrior = TRUE, useT = TRUE,
                        df = rowSums(weights) - 2)
  rr <- results(dse)
  cbind(pval = rr$pvalue, padj = rr$padj)
}

zinbwave_edgeR <- function(counts, group, zinb, ylim = NULL, xlim = NULL, main = 'ZINB-WaVE'){
  d=DGEList(counts)
  d=suppressWarnings(edgeR::calcNormFactors(d))
  design=model.matrix(~group)
  weights <- computeObservationalWeights(zinb, d$counts)
  d$weights <- weights
  d=estimateDisp(d, design)
  plotBCV(d, ylim = ylim, main = main)
  fit=glmFit(d,design)
  lrt=glmWeightedF(fit,coef=2, independentFiltering = TRUE)
  cbind(pval = lrt$table$PValue, padj =lrt$table$padjFilter)
}

zinbwave_DESeq2 <- function(counts, group, zinb){
  colData=data.frame(group=group)
  design=model.matrix(~group)
  dse=DESeqDataSetFromMatrix(countData=counts, colData=colData, design=~group)
  weights <- computeObservationalWeights(zinb, counts(dse))
  assays(dse)[["weights"]]=weights
  dse = DESeq2::estimateSizeFactors(dse, type="poscounts")
  dse = estimateDispersions(dse, minmu=1e-3)
  dse <- nbinomWaldTest(dse, minmu=1e-3, betaPrior=TRUE, useT=TRUE, df=rowSums(weights)-2)
  res = results(dse)
  cbind(pval = res$pvalue, padj = res$padj)
}

zinbwave_limma <- function(counts, group, zinb){
  design <- model.matrix(~group)
  nf <- edgeR::calcNormFactors(counts)
  zeroWeights <- computeObservationalWeights(zinb, counts)
  y <- voom(counts, design, plot=FALSE, lib.size = colSums(counts)*nf,
            weights = zeroWeights)
  y$weights <- y$weights * zeroWeights
  fit <- lmFit(y, design, weights=y$weights)
  fit$df.residual <- rowSums(zeroWeights) - ncol(design)
  fit <- eBayes(fit)
  tt <- topTable(fit,coef=2,n=nrow(counts), sort.by = "none")
  pval <- tt$P.Value
  baseMean = unname(rowMeans(sweep(counts,2,nf,FUN="*")))
  hlp <- pvalueAdjustment_kvdb(baseMean=baseMean, pValue=pval)
  padj <- hlp$padj
  cbind(pval = pval, padj = padj)
}

def wald_test(counts, group):
  test = de.test.wald(
    data=counts.values,
    formula_loc="~ 1 + group",
    factor_loc_totest="group",
    sample_description=pd.DataFrame().assign(group=group),
    batch_size=100,
    training_strategy="DEFAULT",
    dtype="float64",
    gene_names=counts.index
  )
  return pd.DataFrame().assign(
    pval=test.pval,
    padj=test.qval
  )

def lrt_test(counts, group):
  test = de.test.lrt(
    data=counts.values,
    # TODO check what's correct. This is even inconsistent within diffxpy.
    full_formula_scale="~ 1",
    full_formula_loc="~ 1 + group",
    reduced_formula_scale="~ 1",
    reduced_formula_loc="~ 1",
    sample_description=pd.DataFrame().assign(group=group),
    batch_size=100,
    training_strategy="DEFAULT",
    dtype="float64",
    gene_names=counts.index
  )

def wald_test_counts(counts, group):
  n_genes = np.sum(counts > 0, axis=1)
  test = de.test.wald(
    data=counts.values,
    formula_loc="~ group + n_genes",
    factor_loc_totest=["group"],
    coef_to_test="group",
    sample_description=pd.DataFrame().assign(group=pd.Categorical(group), n_genes=np.float64(n_genes.values)),
    batch_size=100,
    training_strategy="DEFAULT",
    dtype="float64",
    gene_names=counts.index
  )
  return pd.DataFrame().assign(
    pval=test.pval,
    padj=test.qval
  )

# def test_scqtl(counts, group):
#     umi = counts.values.T
#     num_samples = umi.shape[0]
#     size_factors = np.ones(len(group))
#     test = de.test.wald(
#         data=umi.astype(np.float64),
#         formula_loc="~ 1 + group",
#         factor_loc_totest="group",
#         sample_description=pd.DataFrame().assign(group=group),
#         batch_size=100,
#         training_strategy="DEFAULT",
#         dtype="float64",
#         size_factors=size_factors,
#         gene_names=counts.index
#     )
#     design = test.model_estim.input_data.design_loc
#     size_factor = size_factors.reshape(1, num_samples).T
#     onehot = np.ones((num_samples, 1))
#     init = scqtl.tf.fit(umi=umi.astype(np.float32),
#          onehot=onehot.astype(np.float32),
#          design=design.astype(np.float32),
#          size_factor=size_factor.astype(np.float32),
#          learning_rate=1e-3,
#          max_epochs=8000,
#           verbose=True
#     )
#     res = scqtl.tf.fit(umi=umi.astype(np.float32),
#          onehot=onehot.astype(np.float32),
#          design=design.astype(np.float32),
#          size_factor=size_factor.astype(np.float32),
#          learning_rate=1e-3,
#          max_epochs=8000,
#         warm_start=init[:3],
#           verbose=True
#     )
#     wald_zinb = wald_test(res[0][0, :], np.maximum(res[1][0, :], 1e-299))
#     return wald_zinb

import warnings
import sys
from multiprocessing import Pool
warnings.filterwarnings('ignore')


def _fit_model(data, model_class):
    model = model_class.from_formula(
        "gene ~ 1 + group",
         data,
         missing='raise'
    )
    try:
        res = model.fit(
            method='nm',
            disp=False,
            maxiter=5000
        )
        test = res.wald_test("group[T.1]=0")
        return test.pvalue
    except:
        return np.nan


def fit_models(count_matrix, group, model_class):
    p = Pool(16)
    assert len(group) == count_matrix.shape[1]
    _pvals = []
    pvals = []
    for i in range(count_matrix.shape[0]):
        if not i % 200:
            sys.stdout.write("model {}\r".format(i))
            sys.stdout.flush()
        data = pd.DataFrame().assign(gene=count_matrix.values[i, :]).assign(group=group)
        _pvals.append(p.apply_async(_fit_model, (data, model_class)))
    for i, r in enumerate(_pvals):
        if not i % 200:
            sys.stdout.write("gene: {}\r".format(i))
            sys.stdout.flush()
        pvals.append(r.get())
    p.close()
    return np.array(pvals)

diffxpy_wald = function(counts, group, ylim = NULL, xlim=NULL) {
  res = py$wald_test(as.data.frame(t(counts)), group)
}

diffxpy_lrt = function(counts, group, ylim = NULL, xlim=NULL) {
  res = py$lrt_test(as.data.frame(t(counts)), group)
}

diffxpy_wald_counts = function(counts, group, ylim = NULL, xlim=NULL) {
  res = py$wald_test_counts(as.data.frame(t(counts)), group)
}

# scqtl_wald = function(counts, group, ylim=NULL, xlim=NULL) {
#   pval = py$test_scqtl(as.data.frame(counts), group)
#   data.frame(pval=pval, padj=p.adjust(pval, method="BH"))
# }

statsmodels_nb = function(counts, group, ylim=NULL, xlim=NULL) {
  pval = py$fit_models(as.data.frame(counts), group, py$sm$NegativeBinomial)
  data.frame(pval=pval, padj=p.adjust(pval, method="BH"))
}

statsmodels_zip = function(counts, group, ylim=NULL, xlim=NULL) {
  pval = py$fit_models(as.data.frame(counts), group, py$sm$ZeroInflatedPoisson)
  data.frame(pval=pval, padj=p.adjust(pval, method="BH"))
}

statsmodels_zinb = function(counts, group, ylim=NULL, xlim=NULL) {
  pval = py$fit_models(as.data.frame(counts), group, py$sm$ZeroInflatedNegativeBinomialP)
  data.frame(pval=pval, padj=p.adjust(pval, method="BH"))
}