addAnno <- function(dmrs, outputLoc = "nearestLocation", featureLocForDistance="TSS", 
                    bindingRegion=c(-2000, 2000), organism = "hg38"){

    library(GenomicRanges)
    library(ChIPpeakAnno)
    library(org.Hs.eg.db)

    dmrs = GRanges(dmrs)

    if(organism == "hg38"){   

        library(TxDb.Hsapiens.UCSC.hg38.knownGene)

        annoData <- toGRanges(TxDb.Hsapiens.UCSC.hg38.knownGene)

    } 

    if(organism == "hg19"){

        library(TxDb.Hsapiens.UCSC.hg19.knownGene)

        annoData <- toGRanges(TxDb.Hsapiens.UCSC.hg19.knownGene)

    }

    seqlevelsStyle(dmrs) <- seqlevelsStyle(annoData)

    anno_dmrs <- annotatePeakInBatch(dmrs, AnnotationData = annoData, 
                                    output = outputLoc, 
                                    FeatureLocForDistance = featureLocForDistance,
                                    bindingRegion = bindingRegion)



    anno_dmrs$symbol <- xget(anno_dmrs$feature, org.Hs.egSYMBOL)

    return(anno_dmrs)

}


main <- function(input, output, params, log) {

    # Log

    out <- file(log$out, open = "wt")

    err <- file(log$err, open = "wt")

    sink(out, type = "output")

    sink(err, type = "message")

    # Script

    library(minfi)
    library(DMRcate)
    library(rtracklayer)

    dmrs <- readRDS(input$rds)

    # params
    outputLoc <- params$output # "nearestLocation"
    featureLocForDistance <- params$featureLocForDistance # "TSS"
    bindingRegion <- params$bindingRegion  # c(-2000, 2000)
    organism <- params$organism

    # output 
    save <- output$csv

    # run annotation
    dmrs = addAnno(dmrs, outputLoc, featureLocForDistance, bindingRegion, organism)

    # save output

    write.csv(as.data.frame(dmrs), save)

    rtracklayer::export(dmrs, output$bed) 

    saveRDS(dmrs, file = output$rds)


}

main(snakemake@input, snakemake@output, snakemake@params, snakemake@log)

plotBoxplots <- function(beta, phenodata, fill, group = "Sample_Group"){

  t_beta = t(beta)
  t_beta = as.data.frame(t_beta, drop=FALSE)
  t_beta$Sample_name = phenodata$Sample_Name
  t_beta$Group = as.character(phenodata[[group]])
  melt_beta = melt(t_beta)

  colnames(melt_beta) = c("Sample","Group", "CpG", "B.Val" )

  bp <- ggplot(melt_beta, aes(x=Sample, y=B.Val, fill = Group)) + 
    geom_boxplot()+
    labs(title="B vals Samples Pre Norm",x="Sample", y = "B.Val") + 
    theme_classic() + coord_flip() + scale_fill_manual(values=fill)

  return(bp)


}

main <- function(input, output, params, log) {

  # Log

  out <- file(log$out, open = "wt")

  err <- file(log$err, open = "wt")

  sink(out, type = "output")

  sink(err, type = "message")

  # Script

  library(minfi)
  library(ggplot2)
  library(reshape2)

  GRset <- readRDS(input$rds)

  beta <- getBeta(GRset)

  phenodata <- pData(GRset)

  fill <- unlist(params$fill)

  # TODO See BW-31
  bp <- plotBoxplots(beta, phenodata, fill, group = "status")

  ggsave(output$pdf, plot = bp)


}

main(snakemake@input, snakemake@output, snakemake@params, snakemake@log)

densityHeatmapWrap <- function(GRset, col, samples_names, fill, cluster_columns = TRUE, clustering_distance_columns = "ks"){

  names(fill) = pData(GRset)[, col]

  ha1 = HeatmapAnnotation(group = pData(GRset)[, col], col = list(group = fill))

  mat <- as.matrix(getBeta(GRset))

  colnames(mat) <- pData(GRset)[, samples_names]

  p <- densityHeatmap(mat, top_annotation = ha1, 
                      cluster_columns = cluster_columns, clustering_distance_columns = "ks", 
                      ylab = "Beta-values")

  return(p)

}


main <- function(input, output, params, log) {

  # Log

  out <- file(log$out, open = "wt")

  err <- file(log$err, open = "wt")

  sink(out, type = "output")

  sink(err, type = "message")

  # Script

  library(minfi)
  library(ComplexHeatmap)

  GRset <- readRDS(input$rds)

  samples_names <- params$samples_names

  col <- params$group

  fillList <- params$fill

  # Check if names supplied are in GRSet metadata
  stopifnot(names(fillList) %in% pData(GRset)[, col])

  fill <- sapply(pData(GRset)[, col], function(x, fillList){as.character(fillList[x])}, fillList = fillList)

  # ensure character vector in correct order for plotting
  stopifnot(names(fill) == pData(GRset)[, col])

  cluster_columns <- params$cluster_columns

  clustering_distance_columns <- params$clustering_distance_columns

  # Plot

  pdf(output$pdf)

  p <- densityHeatmapWrap(GRset, col, samples_names, fill, cluster_columns = params$cluster_columns, 
                          clustering_distance_columns = params$clustering_distance_columns)

  draw(p)

  dev.off()

}

main(snakemake@input, snakemake@output, snakemake@params, snakemake@log)

modelMatrix <- function(data) {

  # Get phenotype data

  names <- colnames(data)

  # Set condition factor

  if ("condition" %in% names) {

    condition <- factor(data$condition)

    n.condition <- nlevels(condition)

    is.condition <- n.condition > 1

  } else {

    is.condition <- FALSE

  }

  # Set batch factor

  if ("batch" %in% names) {

    batch <- factor(data$batch)

    n.batch <- nlevels(batch)

    is.batch <- n.batch > 1

  } else {

    is.batch <- FALSE

  }

  # Set block factor

  if ("block" %in% names) {

    block <- factor(data$block)

    n.block <- nlevels(block)

    is.block <- n.block > 1

  } else {

    is.block <- FALSE

  }

  # Construct design matrix

  if (is.condition & !is.batch & !is.block) {

    design <- model.matrix(~ 0 + condition)

  }

  if (is.condition & is.batch & !is.block) {

    design <- model.matrix(~ 0 + condition + batch)

  }

  if (is.condition & !is.batch & is.block) {

    design <- model.matrix(~ 0 + condition + block)

  }

  if (is.condition & is.batch & is.block) {

    design <- model.matrix(~ 0 + condition + batch + block)

  }

  # Rename condition coefficients

  which.condition <- seq_len(n.condition)

  colnames(design)[which.condition] <- levels(condition)

  # Required for DMRcate if using numeric coef
  # colnames(design)[1] <- "(Intercept)"

  # Return design matrix

  design

}

# Makes sure sample tsv is in the same order as methylation object coldata

checkOrder <- function(sampledata, rds, colname = "Sample_Name") {

  sampledata <- sampledata[match(colData(rds)[[colname]],sampledata$sample),]

  return(sampledata)

}


main <- function(input, output, params, log, config) {

  # Log

  out <- file(log$out, open = "wt")

  err <- file(log$err, open = "wt")

  sink(out, type = "output")

  sink(err, type = "message")

  # Script

  library(DMRcate)
  library(rtracklayer)
  library(minfi)
  library(limma)

  GRset <- readRDS(input$rds)

  data <- read.table(params$samples, header = T)

  data <- checkOrder(data, GRset)

  mod <- modelMatrix(data)

  #### Run DMRCate 

  # Already ratio converted so no need to run below:
  # GRsetRatio <- ratioConvert(GRset)

  GRsetRatio <- GRset

  arraytype = params$arraytype

  # cpg.annotate expects array string as 450K but config specifies as HM450
  if(arraytype == "HM450"){

    arraytype = "450K"

  }

  analysis.type = params$analysistype

  # Changing to limma style contrast matrix makes easier to specify comparison explicitly in config
  # See BW-23
  con <- params$contrast

  # limma style contrast matrix
  cont.matrix <- makeContrasts(contrasts=con, levels=mod)

  # must be a column name in cont.matrix
  # Should correspond to contrast given as param con
  coef <- colnames(cont.matrix)[1]
  stopifnot(coef == con)

  fdr = params$fdr

  # See BW-23 Added cont.matrix = cont.matrix
  # Allowed us to remove intercept from desing matrix
  # Otherwise design matrix must have intercept and coef must be numeric
  cpg.annotation <- cpg.annotate("array", GRsetRatio, arraytype = arraytype, cont.matrix = cont.matrix,
                               analysis.type = analysis.type, design = mod, coef = coef,
                               contrasts = TRUE, fdr = fdr)


  write.csv(as.data.frame(cpg.annotation@ranges), file = output$csv, quote = F)

  saveRDS(cpg.annotation, file = output$rds)

}

main(snakemake@input, snakemake@output, snakemake@params, snakemake@log, snakemake@config)

liftover <- function(results.ranges, chainfile){

  # hg19ToHg38.over.chain from https://hgdownload.soe.ucsc.edu/goldenPath/hg19/liftOver/

  ch <- import.chain(chainfile)

  results <- rtracklayer::liftOver(results.ranges, ch)

  results <- unlist(results)

  return(results)
}

main <- function(input, output, params, log, config) {

  # Log

  out <- file(log$out, open = "wt")

  err <- file(log$err, open = "wt")

  sink(out, type = "output")

  sink(err, type = "message")

  # Script

  library(DMRcate)
  library(rtracklayer)
  library(minfi)
  library(ExperimentHub)

  cpg.annotation <- readRDS(input$rds)

  # Below are default values
  # fdr cutoff default applied and advised by authors
  # see opt pcutoff = "fdr"
  dmrcoutput <- dmrcate(cpg.annotation, lambda=1000, C=2)

  results.ranges <- extractRanges(dmrcoutput, genome = "hg19")
  #results.ranges <- dmrcateExtractRanges(dmrcoutput, genome = "hg19")

  seqlevelsStyle(results.ranges) <- "UCSC"

  # chainfile <- params$chainfile
  chainfile <- input$chainfile

  if(params$liftover){

    results <- liftover(results.ranges, chainfile)

  } else {

    results <- results.ranges

  }

  # order by fdr of choice

  fdr <- params$fdr

  results = results[order(mcols(results)[[fdr]]),]

  write.csv(as.data.frame(results), file = output$csv, quote = F)

  saveRDS(results, file = output$rds)

}

main(snakemake@input, snakemake@output, snakemake@params, snakemake@log, snakemake@config)

plotPCA <- function(x, ...) {

  UseMethod("plotPCA")

}

plotPCA.GenomicRatioSet <- function(x, col, fill) {

  mat <- getBeta(x)

  var <- matrixStats::rowVars(mat)

  num <- min(500, length(var))

  ind <- order(var, decreasing = TRUE)[seq_len(num)]

  pca <- prcomp(t(mat[ind, ]))

  pct <- (pca$sdev ^ 2) / sum(pca$sdev ^ 2)

  dat <- data.frame(PC1 = pca$x[,1], PC2 = pca$x[,2], group = pData(x)[, col])

  p <- ggplot(dat, aes_string(x = "PC1", y = "PC2", color = "group")) + 
    geom_point(size = 3) + 
    xlab(paste0("PC1: ", round(pct[1] * 100), "% variance")) + 
    ylab(paste0("PC2: ", round(pct[2] * 100), "% variance")) + 
    coord_fixed() + scale_colour_manual(values=fill)

  return(p)
}

main <- function(input, output, params, log) {

  # Log function

  out <- file(log$out, open = "wt")

  err <- file(log$err, open = "wt")

  sink(out, type = "output")

  sink(err, type = "message")

  # Script function

  library(ggplot2)

  library(minfi)

  x <- readRDS(input$rds)

  # convert fill to named vector as preserves names
  fill = unlist(params$fill)

  p <- plotPCA(x, col = params$group, fill = fill)

  ggsave(output$pdf, plot = p)

}

main(snakemake@input, snakemake@output, snakemake@params, snakemake@log)

prepRes <- function(results, yaxis = "meandiff", symbol = "overlapping.genes"){

  results <- as.data.frame(mcols(results)[,c(symbol, yaxis)])

  # Change colnames (expects data parsed as symbol, enrichment)

  colnames(results) <- c("symbol","score")

  # Rank by value

  results <- results[order(-results$score),]

  results$rank <- rep(1:length(results$score))

  return(results)
}

#### Mean diff rank plots 
# ranked_data assumed to be 2 col df with gene symbol and enrichment score titled symbol and score 
plotRank <- function(ranked_data, save, genes_of_interest = NULL, 
                     textsize = 0.8, pad = 2, pointsize = 1, highlight = 30, 
                     title = "Rank Plot", yaxis = "meandiff", number_show = 10, num_score_thres = 0.5){

  library(plyr)
  library(ggrepel)
  library(ggplot2)

  # Define gene names of interest if none provided

  if (is.null(genes_of_interest)){

    genes_of_interest <- c(head(ranked_data$symbol, (number_show/2)), tail(ranked_data$symbol, (number_show/2)))

    genes_of_interest <- genes_of_interest[!is.na(genes_of_interest)]
  }

  # Define state column for plotting of colour points and labelled points

  ranked_data$state = rep(FALSE,length(rownames(ranked_data)))

  ranked_data$state = ifelse(ranked_data$rank %in% 1:highlight | ranked_data$rank %in% (length(ranked_data$rank)-highlight):length(ranked_data$rank), "Yes", "No")

  # This will also pick up DMRs with associated genes that are not unique (multiple peaks per gene if they are within the top hits)
  ranked_data$state = ifelse(ranked_data$symbol %in% genes_of_interest & abs(ranked_data$score) > num_score_thres, "Highlight", ranked_data$state)

  # Write out underlying CSV file of rank in case required downstream

  write.table(as.data.frame(ranked_data), file = save, quote = F, sep = "\t", row.names = FALSE)

  # Ggplot obj

  num_highlight <- round(min(number_show, nrow(ranked_data[ranked_data$state == "Highlight",]))/2)
  highlight_df <- rbind.data.frame(head(ranked_data[ranked_data$state == "Highlight",], num_highlight), tail(ranked_data[ranked_data$state == "Highlight",], num_highlight))

  p = ggplot(ranked_data, aes(rank, score)) +
        geom_point(aes(col=state), size = pointsize, alpha = 0.5) + 
        geom_point(data = subset(ranked_data, state %in% c('Yes')), aes(col = state), size = pointsize, alpha = 0.05) + 
        geom_point(data = subset(ranked_data, state %in% c("Highlight")), aes(col = state), size = pointsize, alpha = 0.5) +
        scale_color_manual(values = c("#f05b5b", "#828282", "#5b88f0"), guide = FALSE) +
        geom_text_repel(data = highlight_df, aes(label = symbol), size = textsize, point.padding = pad) + 
        geom_rect(data = head(ranked_data, 1), aes(ymin=-num_score_thres, ymax=num_score_thres, xmin=-Inf, xmax=Inf), alpha= 0.5) +
        geom_hline(yintercept = 0, linetype="solid", color = "black", size=0.5, alpha= 0.5) + 
        geom_hline(yintercept = -num_score_thres, linetype="dashed", color = "black", size=0.5, alpha= 0.5) + 
        geom_hline(yintercept = num_score_thres, linetype="dashed", color = "black", size=0.5, alpha= 0.5) + 
        ggtitle(title) + 
        labs(y = yaxis, x = "Rank") + 
        theme_classic() + 
        theme(plot.title = element_text(hjust = 0.5, size=rel(1.6)) ,
              axis.text.x = element_text(angle = 45, vjust = 1, hjust=1, size=rel(1.6)), 
              axis.title.x = element_text(size=rel(1.6)),
              axis.text.y = element_text(size=rel(1.6)),
              axis.title.y = element_text(size=rel(1.6)))

  return(p)

}

main <- function(input, output, params, log) {

  # Log

  out <- file(log$out, open = "wt")

  err <- file(log$err, open = "wt")

  sink(out, type = "output")

  sink(err, type = "message")

  # Script

  library(minfi)
  library(DMRcate)

  results <- readRDS(input$rds)

  # Define params for rank plot function
  # note dmrcate results file for ranking is names meandiff (in the past it was meanbetafc)
  save <- output$tsv
  genes_of_interest <- params$genes_of_interest
  width <- params$width
  height <- params$height
  textsize <- params$textsize
  pad <- params$pad
  pointsize <- params$pointsize
  highlight <- params$highlight
  title <- params$title
  yaxis <- params$yaxis
  number_show <- params$number_show
  num_score_thres <- params$num_score_thres

  symbol <- params$symbol # overlapping.genes

  # Run prep results table 
  ranked_data <- prepRes(results, yaxis, symbol = symbol)

  # Plot rank plot
  p <- plotRank(ranked_data, save, genes_of_interest,
                     textsize, pad, pointsize, highlight, 
                     title, yaxis, number_show, num_score_thres)

  ggsave(output$pdf, plot = p, width = width,  height = height)

}

main(snakemake@input, snakemake@output, snakemake@params, snakemake@log)