require(lumi)
require(lumiHumanIDMapping)
library(multtest)
library(limma)
library(splines)
library(sva)
library(glmnet)
library(glmnetcr)
library(GSVA)
library(combinat)
library(limma)

source("readData.r")
source("analyseData.r")
source("predictDiagnosis.r")

## ============================================================================================
##
## getDate
## - Convert from str on format ddMMMyyyy  to ISOdate
##
## ============================================================================================

getDate <- function(str) {
  year <- as.numeric(substr(str,6,9))
    day <- as.numeric(substr(str,1,2))

  month <- substr(str,3,5)
  if (month=="JAN")
    month <- 1
  else if (month=="FEB")
    month <- 2
  else if (month=="MAR")
    month <- 3
  else if (month=="APR")
    month <- 4
  else if (month=="MAY")
    month <- 5
  else if (month=="JUN")
    month <- 6
  else if (month=="JUL")
    month <- 7
  else if (month=="AUG")
    month <- 8
  else if (month=="SEP")
    month <- 9
  else if (month=="OCT")
    month <- 10
  else if (month=="NOV")
    month <- 11
  else if (month=="DEC")
    month <- 12

  ISOdate(year, month, day)
}

## ============================================================================================
##
## selectDataset
##
## ============================================================================================

selectDataset <- function(dataList, caseN, selVec) {

  data <- dataList$dataN1
  if (caseN=="N2")
    data <- dataList$dataN2
  else if (caseN=="N3")
    data <- dataList$dataN3

  data <- data[, selVec]
  followUpTime <- dataList$followUpTime[selVec]
  weights <- c()
  if (length(dataList$runNo)>0) {
    weights <- t(dataList$weight[dataList$runNo[selVec],])
    rownames(weights) <- rownames(data)
  }

  list(data=data,
       followUpTime=followUpTime,
       weights=weights)
}

## ============================================================================================
##
## computeFuncList
##
## ============================================================================================

computeMean <- function(data, weights=c()) {
  if (length(weights)==0) {
    (1/ncol(data))  * rowSums(data)
  } else {
    (1/rowSums(weights))  * rowSums(data*weights)
  }
}

computeVar <- function(data, weights=c()) {
  if (length(weights)==0) {
    mu <- matrix(rep(computeMean(data),ncol(data)), ncol=ncol(data))
    expr <- (data-mu)
    (1/(ncol(expr)-1)) * rowSums(expr*expr)
  } else {
    mu <- matrix(rep(computeMean(data, weights),ncol(data)), ncol=ncol(data))
    expr <- (data-mu) * weights
    (1/(ncol(expr)-1)) * rowSums(expr*expr)
  }
}

computeSd <- function(data, weights=c()) {
    sqrt(computeVar(data, weights))
}

###

computeFuncList <- function(caseStr) {

  ## caseStr=="Sd"
  commonNullDist <- TRUE
  sampleData <- function(data, curData, n, weights=c(), curWeights=c()) {
    ind <- sample(1:ncol(data), n)
    if (length(weights)==0) {
      list(data=data[, ind], weights=c())
    } else {
      list(data=data[, ind], weights=weights[, ind])
    }
  }
  computeSortedStat <- function(data, n, n0, weights=c()) {
    sort(as.numeric(computeSd(data, weights)))
  }
  computeNofExtreme <- function(expSd, nullDist) {
    colSums(nullDist<=matrix(rep(expSd,each=nSim), nrow=nSim))
  }
  xlab="genes ordered with respect to standard deviation"
  ylab="standard deviation"
  if (caseStr=="Exp1") {
    commonNullDist <- TRUE
    sampleData <- function(data, curData, n, weights=c(), curWeights=c()) {
      ind <- sample(1:ncol(data), n)
      if (length(weights)==0) {
        list(data=data[, ind], c())
      } else {
        list(data=data[, ind], weights=weights[, ind])
      }
    }
    computeSortedStat <- function(data, n, n0, weights=c()) {
      sort(abs(as.numeric(computeMean(data, weights))), decreasing=TRUE)
    }
    computeNofExtreme <- function(expSd, nullDist) {
      colSums(nullDist>=matrix(rep(expSd,each=nSim), nrow=nSim))
    }
    xlab="genes ordered with respect to absolute value of expectation"
    ylab="absolute value of expectation"
  } else if (caseStr=="Exp2") {
    commonNullDist <- FALSE
    sampleData <- function(data, curData, n, weights=c(), curWeights=c()) {
      mat <- matrix(rep(sample(c(-1,1), ncol(curData), replace=T), each=nrow(curData)),
                    nrow=nrow(curData))
      list(data=curData * mat, weights=curWeights)
    }
    computeSortedStat <- function(data, n, n0, weights=c()) {
      sort(abs(as.numeric(computeMean(data, weights))), decreasing=TRUE)
    }
    computeNofExtreme <- function(expSd, nullDist) {
      colSums(nullDist>=matrix(rep(expSd,each=nSim), nrow=nSim))
    }
    xlab="genes ordered with respect to absolute value of expectation"
    ylab="absolute value of expectation"
  } else if (caseStr=="ExpSd") {
    commonNullDist <- FALSE
    sampleData <- function(data, curData, n, weights=c(), curWeights=c()) {
      mat <- matrix(rep(sample(c(-1,1), ncol(curData), replace=T), each=nrow(curData)),
                    nrow=nrow(curData))
      list(data=curData * mat, weights=curWeights)
    }
    computeSortedStat <- function(data, n, n0, weights=c()) {
       sort(abs(as.numeric(computeMean(data, weights)/computeSd(data, weights))), decreasing=TRUE)
    }
    computeNofExtreme <- function(expSd, nullDist) {
      colSums(nullDist>=matrix(rep(expSd,each=nSim), nrow=nSim))
    }
    xlab="genes ordered with respect to absolute value of expectation/standard deviation"
    ylab="absolute value of expectation/standard deviation"
  } else if (caseStr=="Weight") {
    commonNullDist <- FALSE
    sampleData <- function(data, curData, n, weights=c(), curWeights=c()) {
      ind <- sample(1:ncol(curData))
      if (length(weights)==0) {
        list(data=curData[, ind], weights=c())
      } else {
        list(data=curData[, ind], weights=curWeights[, ind])
      }
    }
    computeSortedStat <- function(data, n, n0, weights=c()) {
      sort(abs(as.numeric((computeMean(data[,1:n], weights[,1:n])-
                           computeMean(data[,(n+1):(n+n0)], weights[,(n+1):(n+n0)]))/
                          sqrt(computeVar(data[,1:n], weights[,1:n])+
                               computeVar(data[,(n+1):(n+n0)], weights[,(n+1):(n+n0)])))),
           decreasing=TRUE)
    }
    computeNofExtreme <- function(expSd, nullDist) {
      colSums(nullDist>=matrix(rep(expSd,each=nSim), nrow=nSim))
    }
    xlab="genes ordered with respect to absolute value of weight"
    ylab="absolute value of weight"
  }

  list(commonNullDist=commonNullDist,
       sampleData=sampleData,
       computeSortedStat=computeSortedStat,
       computeNofExtreme=computeNofExtreme,
       xlab=xlab,
       ylab=ylab)
}

## ============================================================================================
##
## makeGeneList
##
## ============================================================================================

makeGeneList <- function(nofPeriods, periodLength, startStopTimes, selYear, filename) {

  caseN <- "N1"
  meanTimeForPeriod <- rep(NA, nofPeriods)
  for (p in 1:nofPeriods)
    meanTimeForPeriod[p] <- (startStopTimes[p+periodLength -1]+startStopTimes[p])/2

  days <- sort(c((1:7)*365, (0:7)*365+182))
  periodForDay <- rep(NA, length(days))
  periodForDay[days < mean(meanTimeForPeriod[1:2])] <- 1
  periodForDay[days > mean(meanTimeForPeriod[(length(meanTimeForPeriod)-1):(length(meanTimeForPeriod))])] <- nofPeriods
  for (p in 2:(nofPeriods-1))
    periodForDay[days >= mean(meanTimeForPeriod[(p-1):p]) & days <= mean(meanTimeForPeriod[(p):(p+1)])] <- p
  keep <- abs(days-meanTimeForPeriod[periodForDay])<60 ## Ikke avvik på mer enn 2 mnd.
  days <- days[keep]
  periodForYear <- periodForDay[keep]
  years <- round(days/365,1)

  curDataList <- selectDataset(dataList, caseN, dataList$withSpread)
  data <- curDataList$data
  followUpTime <- curDataList$followUpTime
  weights <- curDataList$weights

  curDataList0 <- selectDataset(dataList, caseN, dataList$withoutSpread)
  data0 <- curDataList0$data
  followUpTime0 <- curDataList0$followUpTime
  weights0 <- curDataList0$weights

  p <- periodForYear[years==selYear]
  curData <- data[, followUpTime>=startStopTimes[p] & followUpTime<=startStopTimes[p+periodLength-1]]
  curData0 <- data0[, followUpTime0>=startStopTimes[p] & followUpTime0<=startStopTimes[p+periodLength-1]]
  curWeights <- weights[, followUpTime>=startStopTimes[p] & followUpTime<=startStopTimes[p+periodLength-1]]
  curWeights0 <- weights0[, followUpTime0>=startStopTimes[p] & followUpTime0<=startStopTimes[p+periodLength-1]]

  genes <- rownames(curData)
  logFC <- as.numeric(computeMean(curData, curWeights)-computeMean(curData0, curWeights0))
  weightsForScore <-
    logFC/as.numeric(sqrt(computeVar(curData, curWeights) + computeVar(curData0, curWeights0)))

  res <- cbind(genes, round(weightsForScore,2), round(logFC,2))
  sortList <- sort.list(abs(weightsForScore), decreasing=T)
  res <- res[sortList,]
  colnames(res) <- c("gene", "weight", "logFC")

  write.table(res, filename, row.names=F, quote=F, sep="\t")
}

## ============================================================================================
##
## computeScore
##
## ============================================================================================

computeScore <- function(data, weights) {
  if (length(weights)!=1)
    colSums((matrix(rep(weights, ncol(data)), ncol=ncol(data)))*data, na.rm=TRUE)
  else
    weights*data
}

plotFollowUpTimesWithPredRes <- function(dataList, correctPredWithS, correctPredWithoutS, filename) {
  f <- dataList$followUpTime
  withS <- dataList$withSpread
  withoutS <- dataList$withoutSpread

  pdf(filename, height=3.25, width=15)
  plot(f[withS][!correctPredWithS], rep(0, length(f[withS][!correctPredWithS])), pch="|",
       xlim=rev(range(dataList$followUpTime)), ylim =c(-2,1), ylab="", yaxt="n",
       xlab="Follow up time (days)", cex=1.4, cex.axis=1.4, cex.lab=1.4)
  abline(v=-0)
  abline(v=365)
  abline(v=365*2)
  abline(v=365*3)
  abline(v=365*4)
  abline(v=365*5)
  abline(v=365*6)
  abline(v=365*7)

  ## Wrong result: red
  points(f[withS][!correctPredWithS], rep(0, length(f[withS][!correctPredWithS])), pch="|", cex=1.4, col=2)
  points(f[withoutS][!correctPredWithoutS], rep(-1.5, length(f[withoutS][!correctPredWithoutS])), pch="|", cex=1.4, col=2)

  ## Correct result: green
  points((f[withS][correctPredWithS]), rep(-0.1, length((f[withS])[correctPredWithS])), pch="|", cex=1.4, col=3)
  points((f[withoutS][correctPredWithoutS]), rep(-1.6, length((f[withoutS])[correctPredWithoutS])), pch="|", cex=1.4, col=3)

  f5 <- sort(f[withS])[(1:1000)*5]
  points(f5, rep(0.5, length(f5)), cex=1.4)

  f5 <- sort(f[withoutS])[(1:1000)*5]
  points(f5, rep(-1, length(f5)), cex=1.4)

  axis(2, at=c(-1.5,0), labels=c("without", "with"), xlab="", cex.axis=1.2)

  dev.off()
}


combinePlannedAndDelivered <- function(planned, delivered){
  planned$CP <- 1:356
  delivered$CP <- 1:356

  planned <- select(planned, -c(ID, Energy,GantryAngle,GantryDirection,BeamDose))  %>%
    separate(ControlPoints, into=c(paste0("B", 1:60), paste0("A", 1:60)), " ") %>%
    melt(id.vars="CP", variable.name = "MLCIndex", value.name="Planned")
  delivered <- select(delivered, -c(ID, Energy,GantryAngle,GantryDirection,BeamDose))  %>%
    separate(ControlPoints, into=c(paste0("B", 1:60), paste0("A", 1:60)), " ") %>%
    melt(id.vars="CP", variable.name = "MLCIndex", value.name="Delivered")

  dat <- data.frame("MLCIndex" = planned$MLCIndex, "CP"=planned$CP, "Planned"=planned$Planned, "Delivered"=delivered$Delivered, stringsAsFactors=FALSE)

  return(dat)
}


##########################################################################################################
# CDC Proportional Estimator Drawing
# Zhihang Dong
# Version 0.5
# 3/6/2019
##########################################################################################################

# Import just one library
library(dplyr)

# ====================================================================================================================
# @Function 1: rankProp gives you a smaller dataset grouped by the coordinate of the data points
# rounded to the closest 0.5 increment. For example, a datapoint P(36.229, 39.581) will be giving you
# (36.0, 39.5). These grouped data points naturally into several groups. We then calculate the proportion of
# time spent in each group as the percentage of the entire length of observation
#
# @Function 2: getFinalLvSet gives you the coordinate pair groups selected by the given levelSet
# Warning: when one group does not fulfill the levelSet requirement, the function will automatically append
# the second largest group, until the levelSet minimum has been fulfilled.
# Warning 2: levelSet in our definition has been inflated 100 times. For example, a lambda of 0.1 should be written
# as "10".
#
# @Function 3: propLvSet calculates our proportional estimator, which gives the first week, following when we have never
# observed any differences of the coordinate group produced by the getFinalLvSet algorithms.
#
# @Function 4: drawLevelSets gives you the line plot of our PropEstimator given a step size,
# which, by essence, is an increment value (5 or 10 recommended)
# There are two types of visualization to choose from:
# 1) Color Segments, which colors represent the amount of anchorLocations we have so far in our resulted dataset
# red = 1; orange = 2; green = 3-4; blue = 5-+inf
# 2) Annotaton of ANCH LOC numbers: a black text will appear above each step size \eta.
# =====================================================================================================================

rankProp <- function(dat){
  dat$long = round(dat$long*2, digit = 1)/2
  dat$lat = round(dat$lat*2, digit = 1)/2
  proptab<-dat %>%
    group_by(long, lat) %>%
    summarise(a_sum=sum(t.delta))
  proptab$prop_per = round((100*proptab$a_sum/sum(proptab$a_sum)),digits = 1)
  proptab = proptab[-c(3)]
  proptab = proptab[order(-proptab$prop_per),]
}

getFinalLvSet <- function(tab, levelSet){
  proptab = rankProp(tab)
  for (k in 1:nrow(proptab)){
    topdat = head(proptab, k)
    if (sum(topdat$prop_per)>levelSet){
      break
    }
  }
  topdat = topdat[-c(3)]
  topdat
}

# topdat = getFinalLvSet(tab, 20)

propLvSet <- function(dat, lvSet){
  topdat = getFinalLvSet(dat, levelSet = lvSet)
  for (i in max(dat$week):1){
    newdat = dat[which(dat$week <= i), ]
    newtopdat = getFinalLvSet(newdat,levelSet = lvSet)
    if (i == 1){
      return(i)
      break
    }
    if (identical(topdat,newtopdat) == FALSE){
      return(i)
      break
    }
  }
}

# resulted = propLvSet(dat,20)
# resulted

drawLevelSets <- function(dat, incre, anchorCount = TRUE, coloredLines = TRUE,
                          annotateLocation = FALSE){
  finalResults = c()
  bins = c()
  countAnchorLocations = c()
  for (j in seq(10, 100, incre)){
    results = propLvSet(dat, j)
    finalResults = append(finalResults, results)
    bins = append(bins, j)
    if(anchorCount == TRUE){
      locCounts = nrow(getFinalLvSet(dat, j))
      countAnchorLocations = append(countAnchorLocations, locCounts)
    }
  }
  if(dat$gen[1] == 1){
    cala = 16
  } else if(dat$gen[1] == 2){
    cala = 17
  } else{
    cala = 18
  }
  if(annotateLocation == TRUE){
    plot(finalResults~bins, bty = "l", type ="b", pch=cala, cex = 0.9, col="black",
         ylim=c(0,max(dat$week)*1.3), xlim=c(0,100), ylab="Estimator Results", xlab="LevelSets")
    text(bins, finalResults, countAnchorLocations,
         cex=0.9, pos=3,col="black")
    abline(h=max(dat$week), col="darkred", lwd = 2, lty = 2)
    text(x=88, y=max(dat$week + 3), labels="Max Length", col="blue")
  } else if(coloredLines == TRUE){
    lincol = c()
    for(m in 1:length(countAnchorLocations)){
      if(countAnchorLocations[m] == 1){
        lincol = append(lincol, "red")
      } else if(countAnchorLocations[m] == 2){
        lincol = append(lincol, "orange")
      } else if(countAnchorLocations[m] <= 4){
        lincol = append(lincol, "green")
      } else if(countAnchorLocations[m] > 4){
        lincol = append(lincol, "blue")
      }
    }
    plot(finalResults~bins, type ="b", pch=cala, cex = 0.6, col=lincol,
         ylim=c(0,max(dat$week)*1.3), xlim=c(0,100), ylab="Estimator Results (Week)", xlab="Level(%)")
    segments(bins[-length(bins)],finalResults[-length(finalResults)],bins[-1L],finalResults[-1L],
             col=lincol, lwd = 5.5)
    abline(h=max(dat$week), col="darkred", lwd = 2, lty = 2)
    text(x=88, y=max(dat$week + 3), labels="Max Length", col="blue")
  } else{
    aw <- data.frame(bins, finalResults)
    aw
  }
}

# Experimental Section
dat = person_6174
p2_4_7 = drawLevelSets(dat,10, anchorCount = FALSE, coloredLines = FALSE, annotateLocation = FALSE)

a1_total<- merge(a1_total,p1_3_9, by = "bins")

a2_total<- merge(a2_total,p2_5_11, by = "bins")

a3_total<- merge(a3_total,p2_, by = "bins")

summaryPlot <-function(aggData){
  plot(aggData$bins, aggData[,2], type ="l", cex = 0.3, col="grey",
       ylim=c(0,100), xlim=c(0,100), ylab="Estimator Results (Week)",
       lwd=0.3, xlab="Level(%)")
  for (j in 3:ncol(aggData)){
    lines(aggData$bins, aggData[,j], type ="l", cex = 0.3,
          lwd = 0.3, col="grey")
  }
  mean = c()
  sd = c()
  med = c()
  for (i in 1:10){
    meanr = sum(aggData[i,2:ncol(aggData)])/(ncol(aggData)-1)
    mean = append(mean, meanr)
    medr = median(as.numeric(aggData[i,2:ncol(aggData)]))
    med = append(med, medr)
    sdr = sd(aggData[i,2:ncol(aggData)])
    sd = append(sd, sdr)
  }
  ms = data.frame(bins, mean, sd, med)
  ms$bins = ms$bins*10
  lines(ms$bins, ms$mean, type="l", cex = 20, col="darkslategrey", lwd = 3)
  lines(ms$bins, ms$mean+ 1.96*ms$sd/sqrt(ncol(aggData)-1), type="l", lty = 2, cex = 24,
        lwd = 1.6, col="dodgerblue") #ci+95
  lines(ms$bins, ms$mean- 1.96*ms$sd/sqrt(ncol(aggData)-1), type="l", lty = 2,
        lwd = 1.6, cex = 24, col="dodgerblue") #ci-95
  #lines(ms$bins, ms$mean+ 1.96*ms$sd, type="l", lty = 2, cex = 24, col="brown") #pi+95
  #lines(ms$bins, ms$mean- 1.96*ms$sd, type="l", lty = 2, cex = 24, col="brown") #pi-95
  lines(ms$bins, ms$med, type="l", cex = 20, col="purple", lwd = 4)
  legend(0, 100, legend=c("Male Mean", "Male Median", "Male 95% C.I.", "Female Mean", "Female Median",
                         "AG2 95% C.I."),
         col=c("navyblue", "purple", "dodgerblue","orange", "pink", "red"
               ), lty=c(1,1,2,1,1, 2), cex=0.8, lwd = c(3,3,1.6,3,3,1.6))
}

summaryPlot2 <-function(aggData){
  mean = c()
  sd = c()
  med = c()
  for (i in 1:10){
    meanr = sum(aggData[i,2:ncol(aggData)])/(ncol(aggData)-1)
    mean = append(mean, meanr)
    medr = median(as.numeric(aggData[i,2:ncol(aggData)]))
    med = append(med, medr)
    sdr = sd(aggData[i,2:ncol(aggData)])
    sd = append(sd, sdr)
  }
  ms = data.frame(bins, mean, sd, med)
  ms$bins = ms$bins*10
  lines(ms$bins, ms$mean, type="l", cex = 20, col="orange", lwd = 3)
  lines(ms$bins, ms$mean+ 1.96*ms$sd/sqrt(ncol(aggData)-1), type="l", lty = 2, cex = 24,
        lwd = 1, col="red") #ci+95
  lines(ms$bins, ms$mean- 1.96*ms$sd/sqrt(ncol(aggData)-1), type="l", lty = 2,
        lwd = 1, cex = 24, col="red") #ci-95
  lines(ms$bins, ms$med, type="l", cex = 20, col="pink", lwd = 4)
}

summaryPlot(p1_total)
summaryPlot2(p2_total)

# Voila!


petsc_matmatmult <- function(A_dim, A_ldim, A_data, A_row_ptr, A_col_ind, B_dim, B_ldim, B_data, B_row_ptr, B_col_ind)
{
  storage.mode(A_dim) <- "integer"
  storage.mode(A_ldim) <- "integer"
  storage.mode(A_data) <- "double"
  storage.mode(A_row_ptr) <- "integer"
  storage.mode(A_col_ind) <- "integer"
  storage.mode(B_dim) <- "integer"
  storage.mode(B_ldim) <- "integer"
  storage.mode(B_data) <- "double"
  storage.mode(B_row_ptr) <- "integer"
  storage.mode(B_col_ind) <- "integer"

  out <- .Call(sbase_petsc_matmatmult,
            A_dim, A_ldim, A_data, A_row_ptr, A_col_ind,
            B_dim, B_ldim, B_data, B_row_ptr, B_col_ind)

  return( out )
}

petsc_mattranspose <- function(dim, ldim, Data, row_ptr, col_ind)
{
  storage.mode(dim) <- "integer"
  storage.mode(ldim) <- "integer"
  storage.mode(Data) <- "double"
  storage.mode(row_ptr) <- "integer"
  storage.mode(col_ind) <- "integer"

  out <- .Call(sbase_petsc_mattranspose, dim, ldim, Data, row_ptr, col_ind)

  return( out )
}


#!/usr/local/bin/env Rscript --vanilla
args <- commandArgs(TRUE)
if (!is.null(args[1])) {
    df <- read.csv(args[1], sep = "\t", row.names = 1, check.names=F)
    if (max(df, na.rm = TRUE) < 50) {
        df <- 2^df
    }
} else {
    stop("No expression matrix provided!")
}

library(matrixStats)
library(EPIC)

if (length(args) > 1) {
    ref <- read.csv(args[2], sep = "\t", row.names = 1, check.names=F) ### Need to change later
    cellTypeNames <- colnames(ref)
    markerGenes <- c()
    for (s in colnames(ref)) {
        tempMarkers <- rownames(ref[ref[,s] > quantile(ref[,s], prob=0.75),])
        markers <- c()
        if (length(args) > 2) {
            if (tolower(args[3]) == "pairwise") {
                tempTb = 2 * ref[,cellTypeNames[cellTypeNames != s]] - ref[,s]
                markers <- tempMarkers[tempMarkers %in% rownames(ref[rowSums(tempTb < 0) == (length(cellTypeNames) - 1), ])]
            } else {
                markers <- tempMarkers[tempMarkers %in% rownames(ref)[ref[,s] > 2 * rowMedians(as.matrix(ref))]]
            }
        } else {
            markers <- tempMarkers[tempMarkers %in% rownames(ref)[ref[,s] > 2 * rowMedians(as.matrix(ref))]]
        }
        if (length(markers) == 0) {
            # tempMarkers <- rownames(ref[ref[,s] > quantile(ref[,s], prob=0.5),])
            deOverMedians <- (ref[tempMarkers,s] - rowMedians(as.matrix(ref[tempMarkers,])))/rowMedians(as.matrix(ref[tempMarkers,]))
            markers <- tempMarkers[deOverMedians > quantile(deOverMedians, prob=0.9)]
        }
        genesTemp <- markers
        markerGenes <- c(markerGenes, genesTemp)
        # genesNull <- rownames(ref)[!(rownames(ref) %in% genesTemp)]
        # ref[genesNull, s] <- 0.0
    }
    markerGenes <- unique(markerGenes)

    epicRef <- list()
    epicRef$sigGenes <- markerGenes

    sigMatName <- tools::file_path_sans_ext(basename(args[2]))
    meanFile <- paste0("/data/", sigMatName, "_refMean.txt")
    stdFile <- paste0("/data/", sigMatName, "_refStd.txt")
    if (file.exists(meanFile)) {
        epicRef$refProfiles <- read.csv(meanFile, sep = "\t", row.names = 1, check.names=F)
    } else {
        epicRef$refProfiles <- ref
        warning("Reference profile was not provided or can not be found, using the signature matrix instead\n")
    }
    if (file.exists(stdFile)) {
        epicRef$refProfiles.var <- read.csv(stdFile, sep = "\t", row.names = 1, check.names=F)
    }
    tryCatch(
        expr = {
            xc <- EPIC(bulk = df, reference = epicRef)

            results <- xc$cellFractions
            results <- results[,colnames(results) != "otherCells"]
            write.table(t(results), file = "deconvoluted.tsv", quote = FALSE, col.names = NA, sep = "\t")
        },
        error = function(e){
            # (Optional)
            # Do this if an error is caught...
            print(e)
            Y <- read.csv(args[1], sep = "\t", check.names=F)
            X <- read.csv(args[2], sep = "\t", row.names = 1, check.names=F)
            results <- matrix(0, nrow = length(colnames(Y)) - 1, ncol = length(colnames(X)), dimnames = list(colnames(Y)[2:length(colnames(Y))], colnames(X)))
            write.table(t(results), file="deconvoluted.tsv", quote = FALSE, col.names = NA, sep = "\t")
        }
    )
} else {
    xc <- EPIC(bulk = df)

    results <- xc$cellFractions
    write.table(t(results), file = "deconvoluted.tsv", quote = FALSE, col.names = NA, sep = "\t")
}


#################################################################
#################################################################
#################################################################
#' @title Apply trinary map tools to a set of continuous range maps
#'
#' @description Run all aspects of modeling and plotting for a suite of previously run models
#'
#' @details
#' See Examples.
#'
#' @param dirs,
#' @param stats
#' @param env
#' @param pres.f
#' @param plot.f
#' @param map.dir
#' @param openFig=FALSE
#' @param expertRaster.f=NULL

# @keywords
#'
# @examples
#'
#'
#' @return a data.frame
#' @author Cory Merow <cory.merow@@gmail.com>
#' @note This is a convenience wrapper around each component of trinary map building.
# @seealso
# @references
# @aliases - a list of additional topic names that will be mapped to
# this documentation when the user looks them up from the command
# line.
# @family - a family name. All functions that have the same family tag will be linked in the documentation.
#' @export
trinaryMapWorkflow=function(pres,
													  background,
													  model,
													  modelNames,
													  species,
													  rasterOutputDir=NULL,
													  mapPlotDir=NULL,
													  ROCPlotDir=NULL,
													  #expertShpPath=NULL,
													  #expertRasterPath=NULL,
													  shapesToPlot=NULL,
													  openFig=T,
													  NATo0=TRUE){

	#  for testing
	#  background=bg1; doMapPlot=TRUE; doROCPlot=TRUE; ROCPlotDir=mapPlotDir

	p=raster::extract(model,pres)
	a=raster::extract(model,background)
	if(NATo0){
		a[is.na(a)]=0
		p[is.na(p)]=0
	}
	p <- na.omit(p)
	a <- na.omit(a)
	message(paste0(species,': ',length(p),' presences and ',length(a),' background points used for building trinary maps'))
	ins=rbind(data.frame(Y=1,X=p),data.frame(Y=0,X=a))

	# fit auc curves
	threshs=tryCatch(trinaryROCRoots(ins=ins),error=function(e) e)
	if(class(threshs)=='try-error'){
		message(paste0(species, ": couldn't find roots of the ROC curve; this often happens if you have  very few presence or background points. so you're not getting any trinary maps'"))
		return(list(threshs=NULL,trinary.rasters=NULL))
	}

	# make maps
	trinary.rasters=trinaryMap(model,
														 threshLo=threshs[[2]]$threshLo,
														 threshHi=threshs[[2]]$threshHi,
														 threshYouden=threshs[[2]]$threshYouden,
														 modelNames=modelNames,
														 species=species,
														 rasterOutputDir=rasterOutputDir,
														 overwrite=TRUE,format="GTiff",
												 		 datatype="INT1U", options=c("COMPRESS=DEFLATE"))

	# calculate stats
	range.size=trinaryRangeSize(trinary.rasters,model>threshs[[2]]$threshYouden)

	# plot
	plotFile=paste0(mapPlotDir,modelNames,'.pdf')

	if(!is.null(mapPlotDir)) trinaryMapPlot(trinary.rasters=trinary.rasters,plotFile=plotFile,pres=pres, species=species, shapesToPlot=shapesToPlot,openFig=openFig)

	ROCPlotFile=paste0(ROCPlotDir,species,'_',modelNames,'ROC.pdf')
	if(!is.null(ROCPlotDir)) trinaryROCPlot(ROCPlotFile,plotThings=threshs$plotThings,out1=threshs[[1]],openFig=openFig)

	return(list(threshs=threshs,trinary.rasters=trinary.rasters))

}


library(readr)

argv <- commandArgs(TRUE)
file <- "output200.csv"	#matrice d'occurences
valfile <- "./values200.csv"	#
r = 80			# nombre de concepts que l'on veut -> utilser la fonction clusters(n°du cluster) pour savoir ce qu'il y a dedans
pvalue <- 0.8		# seuil a depasser pour chacun des clusters
output = 'res200.txt'	# le fichier de resultat
headcsv = 'head.csv'	# le fichier contenant la 1ere ligne de file (TODO, ne plus utiliser ce fichier)
input <- readr::read_csv(file)

sink('log.txt')		# le fichier pour verifier combien de temps prend chaque etape

time1 = Sys.time()

### PARTIE LSA ###
print('begin lsa')

X <- svd(input)
S <- diag(X$d)
Sk = S[1:r,1:r]
Uk = X$u[,1:r]
M <- solve(Sk)%*%t(Uk)

print('end lsa')
time2 = Sys.time()
tdiff = difftime(time2,time1)
print(tdiff)

### AFFICHAGE ###
#attach(mtcars)/*{{{*/
#par(mfrow=c(r,1))
#barplot(M[1,])
#barplot(M[2,])
#barplot(M[3,])
#barplot(M[4,])
#barplot(M[5,])
#/*}}}*/

### TRI DES VALEURS DE M et PREMIERE RECHERCHE DE CLUSTER (sans pvalue) ###
print('begin tri')
am = abs(M)
f = t(apply(am, 1, function(line) {
sort(line, decreasing=TRUE, index.return=TRUE)$x
 } ) )

fid = t(apply(am, 1, function(line) {
sort(line, decreasing=TRUE, index.return=TRUE)$ix
 } ) )

labels = read.csv(file=valfile,head=FALSE,sep=",")	#les headers sont les labels donc laisser head=FALSE
qlabels = apply(fid, 1,function(x){ labels[1,x]})

print('end tri')
time3 = Sys.time()
tdiff = difftime(time3,time2)
print(tdiff)

#############################################

### AFFINAGE DES CLUSTERS ###
Sy <- rowSums(am)

nb_sign <- function (vecteur, pval)	#return int
{/*{{{*/
	s <- sum(vecteur)
	i <- 1
	while ((1/s*sum(vecteur[1:i])) < pval){
		i <- i + 1
	}
	i <- i +1 	#pour depasser la pvalue
	return (i)
}/*}}}*/

print('begin nb valeurs signicatives')

nb_val_sign = apply(f, 1, function(x){ nb_sign(x,pvalue) })

print('end nb valeurs signicatives')
time4 = Sys.time()
tdiff = difftime(time4,time3)
print(tdiff)

clusters <- function (concept_line)
{/*{{{*/
	gid <- fid[concept_line, (1:nb_val_sign[concept_line]) ]
	glabels = labels[1,gid]
	return (glabels)
}/*}}}*/

###########################################################

getfeat <- function (concept_line)	#return (feats$x,feats$ix), ie features + feature_id
{/*{{{*/
	gid <- fid[concept_line, (1:nb_val_sign[concept_line]) ]
	feats <- sort(colSums(input[gid[1:nb_val_sign[concept_line]],]),decreasing=TRUE,index.return=TRUE)
	return (feats)
}/*}}}*/

qualite <- function (concept_line)	#return int
{/*{{{*/
	gid <- fid[concept_line, (1:nb_val_sign[concept_line]) ] #pas encore assez optimise
	feats_freq <- getfeat(concept_line)
	feat_id <- feats_freq$ix[1] # la feature la plus representee dans le concept
	qual <- 1 - sum(input[gid,][feat_id] == 0)/nb_val_sign[concept_line]
	return (qual)
}/*}}}*/
head = read.csv(file=headcsv,head=FALSE,sep=",")

###########################################################

print ('begin ecrire les resultats')

for (i in 1:r){
	write.table(head[getfeat(i)$ix[1]],row.names=FALSE,col.names=FALSE,file=output,append=T)
	write.table(qualite(i),output,row.names=FALSE,col.names=FALSE,append=T)
	write.table(clusters(i),output,row.names=FALSE,col.names=FALSE,append=T)
	cat("\n",file=output,append=TRUE)
}

print('end ecrire les resultats')
time4 = Sys.time()
tdiff = difftime(time4,time3)
print(tdiff)
sink()

aff <- function (deb, fin, tsleep)	#return void
{ /*{{{*/
	for (i in deb:fin){
		Sys.sleep(tsleep)
		print(i)
		print(qualite(i))
		x = seq(from =1 , to=nb_val_sign[i]+10, by=1)
		tmp=f[i,1:(nb_val_sign[i]+10)]
		barplot(tmp,xlab=qualite(i) , ylab=i, col=ifelse(x<= nb_val_sign[i],'grey','blue'))
	}
}/*}}}*/


#!/usr/bin/env Rscript

suppressPackageStartupMessages({
  if(!require(optparse)){ install.packages("optparse"); library(optparse) }
  if(!require(data.table)){ install.packages("data.table"); library(data.table) }
  if(!require(dplyr)){ install.packages("dplyr"); library(dplyr) }
  if(!require(tidyr)){ install.packages("tidyr"); library(tidyr) }
  if(!require(ggplot2)){ install.packages("ggplot2"); library(ggplot2) }
  if(!require(scales)){ install.packages("scales"); library(scales) }
  if(!require(viridis)){ install.packages("viridis"); library(viridis) }
  if(!require(gridExtra)){ install.packages("gridExtra"); library(gridExtra) }
  if(!require(cowplot)){ install.packages("cowplot"); library(cowplot) }
})
#########################################
##
## Script used to run REGENIE/BOLT/fastGWA/SAIGE for GWAS
## for the analyses in the REGENIE 2020 paper
## For more details, visit: https://rgcgithub.github.io/regenie/
##
#########################################
option_list = list(
  make_option("--loadFuns", type="character", default="",
    help="script with functions for plotting"),
  make_option("--manColors", type="character", default="",
    help="file with colors for Manhattan plot"),
  make_option("--phenoNames", type="character", default="",
    help="phenotype file"),
  make_option("--figfolder", type="character", default="",
    help="prefix of output files"),
  make_option("--prefix", type="character", default="",
    help="prefix of output files")
  );
opt_parser = OptionParser(option_list=option_list);
opt = parse_args(opt_parser);

source(opt$loadFuns)
alpha <- .05
p.thr <- 324
plot.folder <- opt$figfolder
manP.colors <- fread( opt$manColors, header=FALSE) %>% pull
pheno.names <- fread(opt$phenoNames, header = FALSE) %>% pull
log10P.X <- expression(paste("SAIGE -log"[10],"P"))
log10P.Y <- expression(paste("-log"[10],"P"))
################################################

## Manhatthan plots
pdf(NULL, width = 36, height=24)
dev.control(displaylist="enable")

layout(matrix(1:16,4, byrow=T))
par(mar = c(6,5,8,0.1), oma = c(1.1, 1.1, 1.1, 0.1))
par(tcl = -0.25)
par(mgp = c(3, 0.6, 0))

dummy <- lapply(1:length(pheno.names), function(ipheno){
  panel.lab <- letters[ipheno]
  pn <- pheno.names[ipheno]

  ## SAIGE results
  fns <- system(paste0("ls ", opt$prefix, "_saige_phenoCol", ipheno,"_chr*\\.test" ), intern = TRUE)
  saige.df <- rbindlist(lapply(fns, function(fn){
      fread(fn, fill=TRUE, showProgress = FALSE) %>%
        mutate( pval.SAIGE = -log10(Tstat) ,
          SE.SAIGE = BETA, BETA.SAIGE = N,
          A1FREQ = as.numeric(AC_Allele2),
          MAF = pmin(A1FREQ, 1 - A1FREQ) ) %>%
      select( rsid, MAF, BETA.SAIGE, SE.SAIGE, pval.SAIGE )
}))

  ## BOLT results
  fns <- system(paste0("ls ", opt$prefix, "_bolt_phenoCol", ipheno,"\\.test" ), intern = TRUE)
  bolt.df <- rbindlist(lapply(fns, function(fn){
      fread(fn, colClasses="character", showProgress = FALSE) %>%
        mutate( pval.BOLT = -log10(as.numeric(P_BOLT_LMM)) ) %>%
        select( SNP, pval.BOLT )
}))

  # regenie-Firth
  fns <- system(paste0("ls ", opt$prefix, "_regenie_phenoCol", ipheno,"_Firth_chr*\\.regenie" ), intern = TRUE)
  regenie.firth.df <- rbindlist(lapply(fns, function(fn){
      m1 <-  fread(fn, showProgress = FALSE) %>%
        select( CHROM, GENPOS, ID, BETA, SE, CHISQ,LOG10P ) %>%
        rename_(.dots = setNames(c("BETA","SE","CHISQ","LOG10P"), c("beta.rg.firth","se.rg.firth","test.firth","pval.regenie.firth")) )
      m1 %>%
        select( CHROM, GENPOS, ID, beta.rg.firth, se.rg.firth, pval.regenie.firth)
}))

  # regenie-SPA
  fns <- system(paste0("ls ", opt$prefix, "_regenie_phenoCol", ipheno,"_SPA_chr*\\.regenie" ), intern = TRUE)
  regenie.spa.df <- rbindlist(lapply(fns, function(fn){
      m1 <-  fread(fn, showProgress = FALSE) %>%
        select( ID, BETA, SE, LOG10P ) %>%
        rename_(.dots = setNames(c("BETA","SE","LOG10P"), c("beta.rg.spa","se.rg.spa","pval.regenie.spa")) )
}))

  #     # combine info
  comb.df <- left_join(saige.df, regenie.firth.df, by = c( "rsid" = "ID") )  %>%
    left_join(., bolt.df, by = c( "rsid" = "SNP") )  %>%
    left_join(., regenie.spa.df, by = c( "rsid" = "ID") ) %>%
    as.data.table
  cat("# variants tested for", pn, "=", scales::number(comb.df %>% drop_na %>% nrow, big.mark=",") ,"\n")

  # convert infinite values to p.thr
  comb.df[ is.infinite(pval.BOLT), `:=`(pval.BOLT=p.thr)]
  comb.df[ is.infinite(pval.SAIGE), `:=`(pval.SAIGE=p.thr)]
  comb.df[ pval.regenie.firth > p.thr, `:=`(pval.regenie.firth = p.thr)]
  comb.df[ pval.regenie.spa > p.thr, `:=`(pval.regenie.spa = p.thr)]
  comb.df <- comb.df %>% drop_na

  man.data <- comb.df  %>%
    rename(POS = GENPOS, CHR = CHROM) %>%
    prep.data
  axisdf <- man.data %>% group_by(CHR) %>% summarize(center=( max(POScum) + min(POScum) ) / 2 )

  ## using base R
  all.methods <- c("BOLT-LMM","REGENIE-FIRTH","REGENIE-SPA","SAIGE")
  df0 <- man.data %>%
    gather(Method, pval, pval.BOLT, pval.SAIGE, pval.regenie.firth, pval.regenie.spa) %>%
    mutate( Method = factor(Method,
        levels = paste0("pval.",c("BOLT","regenie.firth", "regenie.spa","SAIGE")),
        labels= all.methods) ,
      pval = as.numeric(pval),
      pval = pval * (pval < p.thr) + p.thr * (pval >= p.thr ),
      color = manP.colors[CHR]
      ) %>%
    select(POScum,Method,pval,color)
  chrlims.min <- man.data %>% group_by(CHR) %>% summarize(bot= min(POScum))%>%pull
  chrlims.max <- man.data %>% group_by(CHR) %>% summarize(top= max(POScum))%>%pull
  chrlims <- rowMeans(cbind(head(chrlims.max,-1), chrlims.min[-1]))
  myaxis <- c(0, ifelse(max(df0$pval) < 70 , 70, 120))

  # compress plot for pvalues >= 20 by 8/100 => x/12.5 + 20 - 1.6
  df1 <- df0  %>%
    mutate(
      pval = pval * (pval <= 20) + (pval/12.5+20-1.6) * (pval > 20 )
    )
  tot.vals <- ifelse(max(df0$pval) < 70 , 7, 8)
  myaxis <- c(seq(0, 20, by = 4), seq(24, 60, by = 4))[1:tot.vals]
  myaxis.lab <- c(seq(0,20,by=4),70,120)[1:tot.vals]

  ddd <- lapply(all.methods, function(my.method){
    cat(my.method,"\n")

    df1.meth <- df1 %>% filter( Method == my.method ) %>%
      mutate(x=POScum,y=pval) %>%
      select(x,y,color)

    plot(df1.meth$x, df1.meth$y, axes = FALSE,
      col=df1.meth$color, cex=1, pch=16, cex.lab=1.7,
      xlab="", ylab="",
      main="", ylim=c(0,max(myaxis)))
    abline(h=-log10(5e-8),lty=2); abline(h=20,col="gray45",lty=2)

    # method name
    title(my.method, cex.main=3.5, font.main=1, line=0)

    if(my.method==all.methods[1]) {
      # add trait title, x/y-axis label
      fig_label(panel.lab, cex=4, font=2)
      title(main=pn, adj=0, line = 4.5, cex.main=3, font.main=1)
      # y-axis
      mtext(log10P.Y, side = 2, outer = F, cex = 1.6, line = 3.5)
      axis(2, at = myaxis, labels=myaxis.lab, las=1,, cex.axis=2)
      plotrix::axis.break(axis=2,breakpos=20,bgcol="white",breakcol="black", style="zigzag",brw=0.02)
    }

    # x-axis labels (jittered)
    abline(v=chrlims, lty=2, col="gray65") #chr breaks
    axis(1, at = c(0,chrlims, tail(chrlims.max,1)), labels=FALSE) #chr ticks
    # x-axes jittered
    x.c <- seq(1,22,2)
    text(x=axisdf$center[x.c],  par("usr")[3],
      labels = axisdf$CHR[x.c], pos = 1, xpd = TRUE, cex=2.1, offset=.7)
    text(x=axisdf$center[-x.c],  par("usr")[3],
      labels = axisdf$CHR[-x.c], pos = 1, xpd = TRUE, cex=2.1, offset=2.7)
    return(NULL)
    })
  return(NULL)
})

p1.base <- recordPlot()
invisible(dev.off())

## TIFF format
tiff(paste0(plot.folder, "Figure2.tiff"), width= 300*6*1.5*4 , height = 300*6*4, res = 300, compression ="lzw")
p1.base
dev.off()
## convert to pdf using imagemagick
system(paste0("convert -density 300 -units PixelsPerInch ", plot.folder, "Figure2.tiff ", plot.folder, "Figure2.pdf"))


library("testthat")
library("waddR")
library("SingleCellExperiment")

ts.names <- c(  "d.wass", "d.wass^2", "d.comp^2", "d.comp", "location",
                "size", "shape", "rho", "p.nonzero", "p.ad.gpd", "N.exc",
                "perc.loc", "perc.size", "perc.shape", "decomp.error",
                "p.zero", "p.combined", "p.adj.nonzero","p.adj.zero",
                "p.adj.combined")

os.names <-  c( "d.wass", "d.wass^2", "d.comp^2", "d.comp", "location",
                "size", "shape", "rho", "pval", "p.ad.gpd", "N.exc",
                "perc.loc", "perc.size", "perc.shape", "decomp.error",
                "pval.adj")

# input data: x, y represent the expression levels for 16 individuals in two
# conditions of one gene
x <- c( 0.0000000, 1.3389541, 0.9190596, 1.5532290, 0.0000000, 1.1014202,
        1.1996965, 1.6332709, 1.0306684, 0.8707476, 1.8230579, 0.0000000,
        1.7195922, 1.1879964, 1.2522275, 0.9516880, 1.0937032, 0.0000000,
        1.3789891, 1.6415560, 0.9871225, 0.0000000, 0.0000000, 0.0000000,
        1.2619454, 1.9289606, 1.7133513, 0.6006942, 0.0000000, 0.7728401,
        0.0000000, 1.2871939, 1.1820060, 1.2520491, 1.6542059, 1.7451944,
        1.3761065, 1.7967770, 0.0000000, 1.1881549, 0.0000000, 0.0000000,
        0.0000000, 2.5839898, 1.0142013, 0.0000000, 0.0000000, 0.0000000,
        1.0225624, 1.3279117, 2.0216712, 1.4012923, 1.7473436, 2.1056459,
        1.5037761, 0.7913705, 1.7510320, 1.0444939, 1.0175730, 1.4748563,
        1.8437568, 1.0464078, 1.6420597, 0.0000000, 1.2362845, 1.3117466,
        1.6209499, 2.8094403)

y <- c( 0.0000000, 2.0119764, 0.0000000, 1.9831597, 2.2152026, 0.0000000,
        0.0000000, 0.0000000, 1.1296844, 1.8360126, 0.0000000, 1.4104730,
        0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000, 0.0000000,
        1.1391284, 0.0000000, 2.1396251, 1.2516926, 0.0000000, 2.2685640,
        1.4992348, 0.0000000, 1.4731691, 0.0000000, 0.0000000, 1.0289898,
        2.3849062, 0.0000000, 0.9660725, 2.2188470, 1.2731303, 2.3051090,
        1.4992348, 0.0000000, 1.4731691, 0.0000000, 1.1537754, 2.0288085,
        2.1383034, 0.0000000, 1.5264164, 0.0000000 )

z <- rep(0, 40)
z2 <- rep(0,34)

# matrix input
dat <- matrix(c(x, y), nrow=1)
dat2 <- matrix(c(x, z), nrow=1)
dat3 <- matrix(c(z, z2), nrow=1)
condition1 <- c(rep(0, length(x)), rep(1, length(y)))
condition2 <- c(rep(0, length(x)), rep(1, length(z)))
condition3 <- c(rep(0, length(z)), rep(1, length(z2)))

# SingleCellExperiment input
sce.a <- SingleCellExperiment(assays=list(counts=matrix(x, nrow=1)))
sce.b <- SingleCellExperiment(assays=list(counts=matrix(y, nrow=1)))
sce.a2 <- SingleCellExperiment(assays=list(counts=matrix(z, nrow=1)))
sce.b2 <- SingleCellExperiment(assays=list(counts=matrix(z2, nrow=1)))

test_that("Correctness of wasserstein single cell output", {

    # these are the fields of the two stage output that don't depend on random
    # sampling (during permutation procedure), but purely on the input
    ts.stable.fields <- c(1, 2, 3, 4, 5, 6, 7, 8, 12, 13, 14, 15, 16, 19)

    # these fields of the two stage output involve random sampling and might
    # change slightly between runs
    ts.volatile.fields <- c(9, 10, 17, 18, 20)

    # these are the fields of the one stage output that depend on a random
    # sampling (during permutation procedure), but purely on the input
    os.stable.fields <- c(1, 2, 3, 4, 5, 6, 7, 8, 12, 13, 14, 15)

    # these fields of the two stage output involve random sampling and might
    # change slightly between runs
    os.volatile.fields <- c(9, 16)

    #########################################################################
    #           Test with two expressed genes in two conditions             #
    #########################################################################

    # TWO STAGE TEST
    res <- wasserstein.sc(dat, condition1)
    res.dup <- wasserstein.sc(sce.a, sce.b)

    # reference results
    res.values <- matrix(c(0.3237935, 0.1048422, 0.1060877, 0.3257111,
                           0.07527902, 0.0007892882, 0.0300194, 0.9262569,
                           0.01, NA, NA, 70.96, 0.74, 28.3,
                           0.01187943, 0.9434837, 0.05, 0.01,
                           0.9434837, 0.05), nrow=1)
    colnames(res.values) <- ts.names

    # compare
    expect_equal(res.dup[ts.stable.fields],
                 res.values[ts.stable.fields],
                 tolerance=0.00001)
    expect_equal(res.dup[ts.volatile.fields],
                 res.values[ts.volatile.fields],
                 tolerance=0.1)
    expect_equal(res[ts.stable.fields],
                 res.values[ts.stable.fields],
                 tolerance=0.00001)
    expect_equal(res[ts.volatile.fields],
                 res.values[ts.volatile.fields],
                 tolerance=0.1)

    # ONE STAGE TEST
    res.os1 <- wasserstein.sc(dat, condition1, "OS")
    res.os.dup1 <- wasserstein.sc(sce.a, sce.b, "OS")

    # reference results
    res.os.values1 <- matrix(c(0.4926601, 0.242714, 0.2474003, 0.4973935,
                               0.03169485, 0.03660025, 0.1791052, 0.8637736,
                               0.09090909, NA, NA, 12.81, 14.79, 72.39,
                               0.01930829, 0.09090909), nrow=1)
    colnames(res.os.values1) <- os.names

    expect_equal(res.os1[os.stable.fields],
                 res.os.values1[os.stable.fields],
                 tolerance=0.0000001)
    expect_equal(res.os1[os.volatile.fields],
                 res.os.values1[os.volatile.fields],
                 tolerance=0.1)
    expect_equal(res.os.dup1[os.stable.fields],
                 res.os.values1[os.stable.fields],
                 tolerance=0.0000001)
    expect_equal(res.os.dup1[os.volatile.fields],
                 res.os.values1[os.volatile.fields],
                 tolerance=0.1)

    #########################################################################
    #           Test with one expressed + one zero-epressed gene            #
    #########################################################################

    # TWO STAGE TEST
    res.ts.3 <- wasserstein.sc(dat2, condition2)
    res.ts.dup.3 <- wasserstein.sc(sce.a, sce.a2)

    # reference results
    res.ts.values.3 <- matrix(c( NA, NA, NA, NA, NA, NA, NA, NA, NA, NA, NA,
                                 NA, NA, NA, NA, 0.8148052, 0.8, NA,
                                 0.8148052, 0.8), nrow=1)

    # we only test the stable fields here, because the volatile fields are
    # unchanging
    expect_equal(res.ts.3[ts.stable.fields],
                 res.ts.values.3[ts.stable.fields],
                 tolerance=0.0000001)
    expect_equal(res.ts.dup.3[ts.stable.fields],
                 res.ts.values.3[ts.stable.fields],
                 tolerance=0.0000001)

    # ONE STAGE TEST
    res.os.3 <- wasserstein.sc(dat2, condition2, "OS")
    res.os.dup.3 <- wasserstein.sc(sce.a, sce.a2, "OS")

    # reference results
    res.os.values.3 <- matrix(c(1.27496, 1.625523, 1.633163, 1.277952,
                                1.113671, 0.5194913, 0, 0, 0, 0.3815411, 250,
                                68.19, 31.81, 0, 0.004699766, 0), nrow=1)

    expect_equal(res.os.3[os.stable.fields],
                 res.os.values.3[os.stable.fields],
                 tolerance=0.000001)
    expect_equal(res.os.3[os.volatile.fields],
                 res.os.values.3[os.volatile.fields],
                 tolerance=0.1)
    expect_equal(res.os.dup.3[os.stable.fields],
                 res.os.values.3[os.stable.fields],
                 tolerance=0.000001)
    expect_equal(res.os.dup.3[os.volatile.fields],
                 res.os.values.3[os.volatile.fields],
                 tolerance=0.1)

    #########################################################################
    #           Test with a zero-expressed gene in both conditions          #
    #########################################################################

    # TWO STAGE PROCEDURE
    res.ts <- wasserstein.sc(dat3, condition3, permnum=10)
    res.ts.dup <- wasserstein.sc(sce.a2, sce.b2, permnum=10)
    res.ts.values2 <- matrix(c(NA, NA, NA, NA, NA, NA, NA, NA, NA, NA, NA, NA,
                               NA, NA, NA, 1, 1, NA, 1, 1), nrow=1)
    colnames(res.ts.values2) <- ts.names
    expect_equal(res.ts, res.ts.values2)
    expect_equal(res.ts.dup, res.ts.values2)
    expect_equal(res.ts, res.ts.dup)

    # ONE STAGE PROCEDURE
    res.os2 <- wasserstein.sc(dat3, condition3, method="OS", permnum=10)
    res.os.dup2 <- wasserstein.sc(sce.a2, sce.b2, method="OS", permnum=10)
    res.os.values2 <- matrix(c(0, 0, 0, 0, 0, 0, 0, 0, 1, NA, NA, NaN, NaN,
                               NaN, 0, 1), nrow=1)
    colnames(res.os.values2) <- os.names
    expect_equal(res.os2, res.os.values2)
    expect_equal(res.os.dup2, res.os.values2)
    expect_equal(res.os2, res.os.dup2)
})

test_that("Example run wasserstein single cell", {

    expect_equal(   colnames(wasserstein.sc(dat, condition1, permnum=10)),
                    ts.names)
    expect_equal(   colnames(wasserstein.sc(dat2, condition2, "OS", 10)),
                    os.names)
    expect_equal(   colnames(wasserstein.sc(sce.a, sce.b, "OS", 10)),
                    os.names)
    expect_equal(   colnames(wasserstein.sc(sce.a, sce.b2, permnum=10)),
                    ts.names)
})


#################################################
#
# make_figure_4.r    16 FEB 2017
#
# Generates figure 4 in "Priority for the
# worse off and the Social Cost of Carbon". This
# code is meant to be sourced from make_figures.r
# and not necessarily run on its own.
#
#################################################

# Setup the plotting device
outfile <- paste(plotdir, "/figure4.pdf", sep="")
pdf(outfile, width=8, height=6)

# Make this a 4-panel plot
par(mfrow=c(2,2), mar=c(4.5,5,2.5,1)+0.1)

# y-axis limits
temp.ylims <- matrix(c(0,1100,0,5500,0,1100,0,1100), ncol=2, byrow=T)

# Loop over the regions
for(i in 1:length(norm.regions)) {

  # Define the current run
  this.run <- norm.regions[i]

  # Subset the data frame for this region
  util.subset <- my.data[which(with(my.data, focusregion == this.run &
                                      eta <= 3.0 & eta >= 0.0 &
                                      rho == 0.01 & gamma == 0 & c0 == 1)),]
  prior.subset <- my.data[which(with(my.data, focusregion == this.run &
                                       eta <= 3.0 & eta >= 0.0 &
                                       gamma == 1.0 & rho == 0 & c0 == 500)),]

  # Setup the plotting parameters
  my.ylim <- temp.ylims[i,]
  line.cols <- c("magenta", "cyan")
  region.label <- this.run
  #if(region.label == "World") {region.label <- "Global"}

  # Make the plot
  plot(util.subset$eta, util.subset$SCC, type="l", ylim=my.ylim, col=line.cols[1],
       xlab="eta [-]", ylab="", main=region.label, log="")
  points(util.subset$eta, util.subset$SCC, pch=15, col=line.cols[1])
  mtext("SCC", 2, line=3.7)
  mtext(expression(paste("[2015 USD per ton ", CO[2], "]", sep="")), 2, line=2.5)
  lines(prior.subset$eta, prior.subset$SCC, col=line.cols[2])
  points(prior.subset$eta, prior.subset$SCC, pch=16, col=line.cols[2])
  legend("topright", legend=c("Disc. Utilitarian", "Undisc. Prioritarian"),
         lty=1, pch=c(15,16), col=line.cols, bty="n")
  put.fig.letter(letters[i], font=2, cex=1.2, offset=c(0.01, -0.05))

}

# Turn off the plotting device if necessary
dev.off()


# By Yea-Hung Chen

# risk period 1
ii$y1start<-ii$start
ii$y1end<-ii$end-365.25*2

# risk period 2
ii$y2start<-ii$end-365.25*2
ii$y2end<-ii$end-365.25

# risk period 3
ii$y3start<-ii$end-365.25
ii$y3end<-ii$end

# start date for adult
ii$adult_start<-as.Date('1900-01-01')
is.na(ii$adult_start)<-TRUE
ii$adult_start<-if_else(ii$age_0<15,
                        ii$y1start+(15-ii$age_0)*365.25,
                        ii$adult_start)
is.na(ii$adult_start)<-(ii$adult_start<ii$y1start)
is.na(ii$adult_start)<-(ii$y3end<ii$adult_start)

# start date for entering risk
ii$enter_start<-pmax(ii$inmigrant_date_3,ii$adult_start,na.rm=TRUE)
ii$enter_start<-with(ii,if_else(is.na(enter_start),y1start-365.25*100,
                                enter_start))

# modify risk periods for period-1 entries
ii$en.iny1<-with(ii,(y1start<enter_start)&(enter_start<y1end))
ii$y1start[ii$en.iny1]<-ii$enter_start[ii$en.iny1]

# modify risk periods for period-2 entries
ii$en.iny2<-with(ii,(y2start<enter_start)&(enter_start<y2end))
ii$y2start[ii$en.iny2]<-ii$enter_start[ii$en.iny2]
is.na(ii$y1start)<-ii$en.iny2
is.na(ii$y1end)<-ii$en.iny2

# modify risk periods for period-3 entries
ii$en.iny3<-with(ii,(y3start<enter_start)&(enter_start<y3end))
ii$y3start[ii$en.iny3]<-ii$enter_start[ii$en.iny3]
is.na(ii$y1start)<-ii$en.iny3
is.na(ii$y2start)<-ii$en.iny3
is.na(ii$y1end)<-ii$en.iny3
is.na(ii$y2end)<-ii$en.iny3


library(rstan)
library(data.table)
library(lubridate)
library(gdata)
library(EnvStats)
library(tidyr)
library(stringr)
library(dplyr)
library(jsonlite)

# time.taken <- end.time - start.time
# time.taken

# start timer
start.time <- Sys.time()

# Rscript base.r us_base 150 4000 [--validate]
args = commandArgs(trailingOnly=TRUE)
StanModel = args[1]
minimumReportedDeaths = as.integer(args[2])
nStanIterations = as.integer(args[3])
validateFlag = args[4]
if (is.na(validateFlag)){validateFlag <- ""}

if (validateFlag == "--validate"){
  print("INFO: --validate flag passed - running validation routine")
}
print(sprintf("Running stan model %s",StanModel))
print(sprintf("Only running on counties with at least %d total reported deaths", minimumReportedDeaths))
print(sprintf("Running MCMC routine with %d iterations", nStanIterations))

## smoothing death data: https://github.com/ImperialCollegeLondon/covid19model/blob/v6.0/usa/code/utils/read-data-usa.r#L37-L40
## -> https://github.com/ImperialCollegeLondon/covid19model/blob/v6.0/usa/code/utils/read-data-usa.r#L24-L27
## -> something to try, for sure

library(zoo)
# smooth_fn = function(x, days=3) {
#   return(rollmean(x, days, align = "right"))
# }

# case-mortality table
d <- read.csv("../modelInput/ILCaseAndMortalityV1.csv", stringsAsFactors = FALSE)

# drop counties with fewer than cutoff cumulative deaths or cases
cumCaseAndDeath <- aggregate(cbind(d$deaths), by=list(Category=d$countryterritoryCode), FUN=sum)
dropCounties <- subset(cumCaseAndDeath, V1 < minimumReportedDeaths)$Category
d <- subset(d, !(countryterritoryCode %in% dropCounties))
# print(sprintf("nCounties with more than %d deaths before %s: %d", minimumReportedDeaths, dateCutoff, length(unique(d$countryterritoryCode))))
print(sprintf("nCounties with more than %d deaths: %d", minimumReportedDeaths, length(unique(d$countryterritoryCode))))

d$date = as.Date(d$dateRep,format='%m/%d/%y')

d$countryterritoryCode <- sapply(d$countryterritoryCode, as.character)
# trim US code prefix
d$countryterritoryCode <- sub("840", "", d$countryterritoryCode)

# county population
# pop = unique(d[c("countryterritoryCode", "popData2018")])

# IMPORTANT: this smoothing function is 100% broken
# SMOOTHING reported death and case counts
# try <steps> day moving average to smooth raw reported case and death counts
# "so the bars don't look so bad" - really to account for bias/periodic fluxuations in reporting
# steps = 7
# d$deaths = c(rep(0, steps-1), as.integer(smooth_fn(d$deaths, days=steps)))
# d$cases = c(rep(0, steps-1), as.integer(smooth_fn(d$cases, days=steps)))

codeToName <- unique(data.frame("countyCode" = d$countryterritoryCode, "countyName" = d$countriesAndTerritories))

# write list of counties used in this simulation
CountyCodeList <- unique(d$countryterritoryCode)
write.table(CountyCodeList, "../modelOutput/figures/CountyCodeList.txt", row.names=FALSE, col.names=FALSE)

countries <- unique(d$countryterritoryCode)

# weighted fatality table
cfr.by.country = read.csv("../modelInput/ILWeightedFatalityV1.csv")
cfr.by.country$country = as.character(cfr.by.country[,3])
cfr.by.country$country <-  sub("840", "", cfr.by.country$country) # cutoff US prefix code - note: maybe this should be in the python etl, not here

# serial interval discrete gamma distribution
serial.interval = read.csv("../modelInput/ILSerialIntervalV1.csv") # new table

# N2 is Number of time points with data plus forecast
N2 = 0

# >>>>>>>>>>>>>>> MOBILITY >>>>>>>>>>>>>>> #

# Read google mobility
source("./read-mobility.r")
mobility <- read_google_mobility(countries=countries, codeToName=codeToName)

# basic impute values for NA in google mobility
# see: https://github.com/ImperialCollegeLondon/covid19model/blob/v6.0/base-usa.r#L87-L88
for(i in 1:ncol(mobility)){
  if (is.numeric(mobility[,i])){
    mobility[is.na(mobility[,i]), i] <- mean(mobility[,i], na.rm = TRUE)
  }
}

# Read predicted mobility
google_pred <- read.csv('../modelInput/mobility/google-mobility-forecast.csv', stringsAsFactors = FALSE)
google_pred$date <- as.Date(google_pred$date, format = '%Y-%m-%d')

# replicate statewide prediction by county -> this can be MUCH more nuanced, but for now - just get something working
stateAndCounty <- codeToName
stateAndCounty$state <- "Illinois"
google_pred <- left_join(stateAndCounty, google_pred, "state" = "state")
colnames(google_pred)[colnames(google_pred) == 'state'] <- 'sub_region_1'

# Append predicted mobility
if (max(google_pred$date) > max(mobility$date)){
  google_pred <- google_pred[google_pred$date > max(mobility$date),]
  mobility <- rbind(as.data.frame(mobility),as.data.frame(google_pred[,colnames(mobility)]))
}

max_date <- max(mobility$date)

## validation cutoff
if (validateFlag == "--validate"){
  lastObs <- max_date - 7
  print(sprintf("validation: running model through date: %s", lastObs))
  print(sprintf("validation: validating on 7 days leading up to date of last observation: %s", max_date))
} else {
  lastObs <- max_date
  print(sprintf("MAX DATE : %s", lastObs))
}
d <- d[as.Date(d$dateRep, format = "%m/%d/%y") <= lastObs, ]

## need to take a close look @ this -> looks fine ##
# see: https://stackoverflow.com/questions/8055508/in-r-formulas-why-do-i-have-to-use-the-i-function-on-power-terms-like-y-i
formula_partial_county = as.formula('~ -1 + averageMobility + I(transit * transit_use) + residential')

# <<<<<<<<<<<<<<< MOBILITY <<<<<<<<<<<<<<<<<< #

dates = list()
reported_cases = list()
deaths_by_country = list()

stan_data = list(M=length(countries),
                N=NULL, # Number of time points with data
                y=NULL,
                deaths=NULL,
                f=NULL,
                N0=6, # Number of days in seeding
                cases=NULL,
                EpidemicStart = NULL # Date to start epidemic in each county
                )

# N2 is the length of time window to simulate
# adjust N2 before main procesesing routine - i.e., adjust N2 so that it's uniform across all counties
for(Country in countries) {

  tmp=d[d$countryterritoryCode==Country,]
  # tmp$date = as.Date(tmp$dateRep,format='%m/%d/%y')
  tmp$t = decimal_date(tmp$date)
  tmp=tmp[order(tmp$t),]

  index1 = which(cumsum(tmp$deaths)>=10)[1]
  index2 = index1-30

  tmp=tmp[index2:nrow(tmp),]

  N = length(tmp$cases)
  if(N2 - N < 0) {
    print(sprintf("raising N2 from %d to %d", N2, N))
    N2 = N
    # N2 = N + 7
  }
}

covariate_list_partial_county <- list()

# k is their counter
k <- 1
for(Country in countries) {

  CFR=cfr.by.country$weighted_fatality[cfr.by.country$country == Country]

  d1=d[d$countryterritoryCode==Country,]

  d1$t = decimal_date(d1$date)
  d1=d1[order(d1$t),]

  index = which(d1$cases>0)[1]
  index1 = which(cumsum(d1$deaths)>=10)[1]
  index2 = index1-30

  print(sprintf("First non-zero cases is on day %d, and 30 days before %d cumulative deaths is day %d",index,minimumReportedDeaths,index2))
  d1=d1[index2:nrow(d1),]
  stan_data$EpidemicStart = c(stan_data$EpidemicStart,index1+1-index2)

  # Selects population for each county
  # pop_county <- pop[pop$countryterritoryCode==Country,][[2]]

  # dates[[as.character(Country)]] = d1$date
  dates[[Country]] = d1$date

  # hazard estimation
  N = length(d1$cases)
  print(sprintf("%s has %d days of data",Country,N))

  short <- N2 - N

  # >>>>>>>>>>> mobility >>>>>>>>>>>>> #

  # Selects mobility data for each county
  covariates_county <- mobility[which(mobility$countyCode == Country),]

  # Find minimum date for the data
  min_date <- min(d1$date)
  padded_covariates <- nu_pad_mobility(min_date, covariates_county, Country, N2)

  # include transit
  transit_usage <- rep(1, (N + short))

  # creating features -> only want "partial_state"
  df_features <- create_features(padded_covariates, transit_usage)
  features_partial_county <- model.matrix(formula_partial_county, df_features)
  covariate_list_partial_county[[k]] <- features_partial_county

  # <<<<<<<<<<< mobility <<<<<<<<<<<<< #

  h = rep(0,short+N) # discrete hazard rate from time t = 1, ..., 100
  mean1 = 5.1; cv1 = 0.86; # infection to onset
  mean2 = 18.8; cv2 = 0.45 # onset to death

  # td: double check these comments
  # otherwise should be fine
  ## icl: assume that CFR is probability of dying given infection
  x1 = rgammaAlt(5e6,mean1,cv1) # icl: infection-to-onset ----> do all people who are infected get to onset?
  x2 = rgammaAlt(5e6,mean2,cv2) # icl: onset-to-death
  f = ecdf(x1+x2)

  convolution = function(u) (CFR * f(u))
  h[1] = (convolution(1.5) - convolution(0))
  for(i in 2:length(h)) {
    h[i] = (convolution(i+.5) - convolution(i-.5)) / (1-convolution(i-.5))
  }
  s = rep(0,N2)
  s[1] = 1
  for(i in 2:N2) {
    s[i] = s[i-1]*(1-h[i-1])
  }

  f = s * h

  y=c(as.vector(as.numeric(d1$cases)),rep(-1,short))
  reported_cases[[Country]] = as.vector(as.numeric(d1$cases))
  deaths=c(as.vector(as.numeric(d1$deaths)),rep(-1,short))
  cases=c(as.vector(as.numeric(d1$cases)),rep(-1,short))
  deaths_by_country[[Country]] = as.vector(as.numeric(d1$deaths))

  ## icl: append data
  stan_data$N = c(stan_data$N,N)
  stan_data$y = c(stan_data$y,y[1]) # icl: just the index case!

  stan_data$SI=serial.interval$fit[1:N2]

  stan_data$f = cbind(stan_data$f,f)

  stan_data$deaths = cbind(stan_data$deaths,deaths)

  stan_data$N2=N2
  if(length(stan_data$N) == 1) {
    stan_data$N = as.array(stan_data$N)
  }

  k <- k+1
}

# >>>>>>>>>>> mobility >>>>>>>>>>>>> #

# newSTAN - ok
stan_data$P_partial_county = dim(features_partial_county)[2]
# newSTAN - ok
stan_data$X_partial_county = array(NA, dim = c(stan_data$M , stan_data$N2 ,stan_data$P_partial_county))

# NOTE: mapped *_partial_state -> *_partial_county

for (i in 1:stan_data$M){
  stan_data$X_partial_county[i,,] = covariate_list_partial_county[[i]]
}

# newSTAN - ok
stan_data$W <- ceiling(stan_data$N2/7)
# newSTAN - ok
stan_data$week_index <- matrix(1,stan_data$M,stan_data$N2)
for(j in 1:stan_data$M) {
  stan_data$week_index[j,] <- rep(2:(stan_data$W+1),each=7)[1:stan_data$N2]
  last_ar_week = which(dates[[j]]==max(d$date) - 28)
  stan_data$week_index[j,last_ar_week:ncol(stan_data$week_index)] <-  stan_data$week_index[j,last_ar_week]
}

# <<<<<<<<<<< mobility <<<<<<<<<<<<< #

stan_data$y = t(stan_data$y)

options(mc.cores = parallel::detectCores())
rstan_options(auto_write = TRUE)
m = stan_model(paste0('../stan/',StanModel,'.stan'))

# it works!!
fit = sampling(m,data=stan_data,iter=nStanIterations,warmup=nStanIterations/2,chains=8,thin=4,control = list(adapt_delta = 0.90, max_treedepth = 10))

# stop timer
end.time <- Sys.time()

# calculate duration of routine
duration <- as.integer(difftime(end.time, start.time, units="secs"))

print(sprintf("DURATION: %d", duration))

out = rstan::extract(fit)
prediction = out$prediction
estimated.deaths = out$E_deaths

JOBID = Sys.getenv("PBS_JOBID")
if(JOBID == "")
  JOBID = as.character(abs(round(rnorm(1) * 1000000)))
print(sprintf("Jobid = %s",JOBID))
save.image(paste0('../modelOutput/results/',StanModel,'-',JOBID,'.Rdata'))
save(JOBID,nStanIterations,duration,minimumReportedDeaths,fit,prediction,dates,reported_cases,deaths_by_country,countries,estimated.deaths,out,lastObs,covariate_list_partial_county,file=paste0('../modelOutput/results/',StanModel,'-',JOBID,'-stanfit.Rdata'))

#### saving of simulation results is finished

## log ##
# create summary
summary <- list(
    jobid=JOBID,
    time=duration, # in seconds
    lastObs=lastObs,
    deathsCutoff=minimumReportedDeaths,
    nCounties=length(countries),
    nIter=nStanIterations
)
exportJSON <- toJSON(summary, pretty=TRUE, auto_unbox=TRUE)
# sent to stdout
print("--- summary ---")
print(exportJSON)
# write summary to log
write(exportJSON, "../modelOutput/log.json")

#### now -> visualize model results -> ####

# fixme: don't force other scripts to load R data -> unnecessary overhead
# import viz routine, call those functions here -> way, way better
# still save the R data and fit though, for backup, etc.

filename <- paste0(StanModel, '-', JOBID)
system(paste0("Rscript plot-trend.r ", filename,'.Rdata'))
# system(paste0("Rscript plot-forecast.r ", filename,'.Rdata')) ## icl: to run this code you will need to adjust manual values of forecast required

# suppressing for now
# system(paste0("Rscript plot-explore.r ", filename,'.Rdata'))

if (validateFlag == "--validate"){
  print("validating model output ...")
  system(paste0("Rscript validate.r ", filename,'.Rdata'))
}


trimfill.default <- function(yi, vi, weights = NULL, ni = NULL, maxiter = 100, estimator = "L0", verbose = F)
{
  idix <- sort(yi, index.return = TRUE)$ix
  yi <- yi[idix]
  vi <- vi[idix]
  weights <- weights[idix]
  ni <- ni[idix]
  k <- length(yi)
  k0.sav <- -1
  k0 <- 0
  side = "left"
  iter <- 0
  while (abs(k0 - k0.sav) > 0) {
    k0.sav <- k0
    iter <- iter + 1
    if (iter > maxiter)
      stop("Trim and fill algorithm did not converge.")
    yi.t <- yi[1:(k - k0)]
    vi.t <- vi[1:(k - k0)]
    weights.t <- weights[1:(k - k0)]
    res <- suppressWarnings(rma.uni(yi.t, vi.t, weights = weights.t))
    b <- c(res$b)
    yi.c <- yi - b
    yi.c.r <- rank(abs(yi.c), ties.method = "first")
    yi.c.r.s <- sign(yi.c) * yi.c.r
    if (estimator == "R0") {
      k0 <- (k - max(-1 * yi.c.r.s[yi.c.r.s < 0])) - 1
      se.k0 <- sqrt(2 * max(0, k0) + 2)
    }
    if (estimator == "L0") {
      Sr <- sum(yi.c.r.s[yi.c.r.s > 0])
      k0 <- (4 * Sr - k * (k + 1))/(2 * k - 1)
      varSr <- 1/24 * (k * (k + 1) * (2 * k + 1) + 10 *
                         k0^3 + 27 * k0^2 + 17 * k0 - 18 * k * k0^2 -
                         18 * k * k0 + 6 * k^2 * k0)
      se.k0 <- 4 * sqrt(varSr)/(2 * k - 1)
    }
    if (estimator == "Q0") {
      Sr <- sum(yi.c.r.s[yi.c.r.s > 0])
      k0 <- k - 1/2 - sqrt(2 * k^2 - 4 * Sr + 1/4)
      varSr <- 1/24 * (k * (k + 1) * (2 * k + 1) + 10 *
                         k0^3 + 27 * k0^2 + 17 * k0 - 18 * k * k0^2 -
                         18 * k * k0 + 6 * k^2 * k0)
      se.k0 <- 2 * sqrt(varSr)/sqrt((k - 1/2)^2 - k0 *
                                      (2 * k - k0 - 1))
    }
    k0 <- max(0, round(k0))
    se.k0 <- max(0, se.k0)
    if (verbose)
      cat("Iteration:", iter, "\tmissing =", k0, "\t  b =",
          ifelse(side == "right", -1 * b, b), "\n")
  }
  if (k0 > 0) {
    if (side == "right") {
      yi.c <- -1 * (yi.c - b)
    }
    else {
      yi.c <- yi.c - b
    }
    yi.fill <- c(yi, -1 * yi.c[(k - k0 + 1):k])
    vi.fill <- c(vi, vi[(k - k0 + 1):k])
    weights.fill <- c(weights, weights[(k - k0 + 1):k])
    ni.fill <- c(ni, ni[(k - k0 + 1):k])
    res <- suppressWarnings(rma.uni(yi.fill, vi.fill, weights = weights.fill,
                                    ni = ni.fill))
    res$fill <- c(rep(FALSE, k), rep(TRUE, k0))
    res$slab.null <- FALSE
  }
  else {
    res <- list()
    res$fill <- rep(FALSE, k)
  }
  res$k0 <- k0
  res$se.k0 <- se.k0
  res$side <- side
  res$k0.est <- estimator
  if (estimator == "R0") {
    m <- -1:(k0 - 1)
    res$p.k0 <- 1 - sum(choose(0 + m + 1, m + 1) * 0.5^(0 +
                                                          m + 2))
  }
  else {
    res$p.k0 <- NA
  }
  class(res) <- c("default.trimfill", class(res))
  return(res)
}


# Sample code to build models in the main analysis
# You can just run this script and it will call the python code directly
# try from the command line: R CMD BATCH sample_model_building.r
# or download Rstudio (www.rstudio.org) and run from within the IDE

# important libraries
require(data.table)
require(caret)
require(pROC)
require(Hmisc)
require(glmnet)
require(randomForest)
require(gbm)

# Let's build a sample data set
# The actual data used in this investigation can be accessed
# for free after signing a DUA at:
# https://mimic.physionet.org/

# number of observations
set.seed(20172017)
N <- 10000

# the outcome, death or ICU LOS >= 7 days
Y <- sample(0:1, size = N, replace = TRUE, prob = c(0.8,0.2))

# Let's make 3 variables that are somewhat correlated with the outcome
# and introduce some additional bias
age <- 50 + Y * 10 + rnorm(N, 3, sd = 4)
sbp <- 110 - Y * 10 + rnorm(N, 8, sd = 3)
any_vent <- ifelse(Y * runif(N, 0, 1) >= 0.5, 1, 0)

# and randomly corrupt a few of these to make model fitting a bit more challenging
age_crpt_indx <- sample(1:N, size = round(N * 0.3), replace = F)
age[age_crpt_indx] <- sqrt(age[age_crpt_indx]) * 5 + 10 * cos(age[age_crpt_indx])
sbp_crpt_indx <- sample(1:N, size = round(N * 0.4), replace = F)
sbp[sbp_crpt_indx] <- 20 + log(sbp[sbp_crpt_indx]) * 15 + 5 * sin(sbp[sbp_crpt_indx])

# Let's make 5 more variables that are essentially random
# i.e. uncorrelated with the outcome
gender <- sample(c('m','f'), size = N, replace = T)
pao2 <- rnorm(N, 120, 20)
temp <- rnorm(N, 37, 2)
any_cardiac_arrest <- sample(0:1, size = N, replace = T, prob = c(0.05,0.95))
creatinine <- rnorm(N, 1.5 , 0.2)

# Now let's make some sample unstructured data
note_template_terms <- ('This patient is a 53 year old female with a history
                        of hypertension, bipolar disorder, chronic kidney
                        disease and asthma who was admitted to the intensive
                        care unit for acute respiratory failure. Her course
                        has been complicated by sepsis and delirium and she
                        remains on broad-spectrum antibiotics. A plain chest
                        radiograph reveals an evolving infiltrate at the right
                        lung base, and her white blood cell count continues to
                        rise. Current drips include heparin, norepinephrine,
                        and vasopressin.')

# now let's scramble those words to make unique combinations
# the order, although non-sensical, won't matter for this sample analysis
note_tokens <- unlist(strsplit(note_template_terms,' '))
note_tokens <- gsub('\n', '', note_tokens) # remove newlines
note_tokens <- grep('^$', note_tokens, value = TRUE, invert = TRUE) # remove empty
sample_notes <- sapply(1:N, function(x) {
    wcount <- round(rnorm(1,150, 20))
    token_indices <- sample(1:length(note_tokens), size = wcount, replace = T)
    paste0(note_tokens[token_indices], collapse = ' ')
})

# now add in select text terms correlated with the outcome
sample_notes[Y == 1] <- paste(sample_notes[Y == 1] ,
                              sample(c('poor prognosis', 'very sick', ''),
                                     size=sum(Y),replace=TRUE,prob = c(0.8,0.15,0.05)))
sample_notes[Y == 0] <- paste(sample_notes[Y == 0],
                              sample(c('feeling well', ''),sum(Y==0),replace=TRUE))

# create a data frame
df <- data.table(pt_id = 1:N, age, sbp, any_vent, gender, pao2, temp,
                        any_cardiac_arrest, creatinine, Y)

# export the text data for analysis
dir.create('corpus_txt/', showWarnings = FALSE)
invisible(sapply(1:N, function(n) {
  writeLines(sample_notes[n], paste0('corpus_txt/',n,'.txt'))
}))

# Now run the Python code
# NB. this can take 10 minutes or more
system('python3 sample_text_analysis_main.py')

# now import the text output of the Python code
df_ngrams <- fread('ngrams_dtm.csv')
df_key_terms <- fread('text_results_keyterms.csv')

# merge text features together and ignore duplicates
key_terms_unique <- names(df_key_terms)[! names(df_key_terms) %in% names(df_ngrams)]
df_text_terms <- merge(df_ngrams,
                       df_key_terms[,c('pt_id',key_terms_unique),with=F],
                       by = 'pt_id')

# fix variable names
names(df_text_terms)[-1] <- paste0('term_',
                              gsub(' ','_',
                                   names(df_text_terms)[-1]))

# transform word counts per document to binary presence/absence
df_text_terms[,grep("term_",colnames(df_text_terms), value=TRUE) :=
                sapply(.SD, function(x) x > 0),
                                             .SDcols = grep('term_',names(df_text_terms))]

# transform word counts to yes/no presence for each document
df_text_terms_txf <- data.table(pt_id = df_text_terms$pt_id,
                                df_text_terms[,-1,with=F] > 0)

# set up training and testing data stratified on the outcome
training <- createDataPartition(df$Y, p=0.75, list=FALSE)[,1]
testing <- (1:nrow(df))[-training]

# do some variable selection on the text data
x <- df_text_terms_txf[pt_id %in% df[training]$pt_id] # just using TRAINING DATA!!!
x <- x[order(pt_id)]
y <- df[training]$Y
x_ptids <- x$pt_id
x <- x[, pt_id := NULL]
# also remove all columns already present in the data set
keep_cols <- which(! names(x) %in% names(df))
x <- x[, keep_cols, with = FALSE]

# now do variable selection
term_select <- cv.glmnet(as.matrix(x), y, family = 'binomial', parallel = FALSE)
# use cv.glmnet for CV in the future...
# then get the coefficients using:
terms_coeff <- coef(term_select, s = "lambda.min")[,1] # FOR CV ONLY
# keep up to the top 20 non-zero
terms_coeff_filt <- na.exclude(names(sort(abs(terms_coeff[terms_coeff != 0 & names(terms_coeff) != '(Intercept)']),
                               decreasing = T))[1:20])

# now get variable name list and merge with dt.train
add_terms <- df_text_terms_txf[, c('pt_id',terms_coeff_filt),with=F]
# merge text data into full data set
df <- merge(df, add_terms, by = 'pt_id')

# define limited variable set of structured data
var_names_struct <- grep('term_', names(df), value = TRUE, invert = TRUE)

# validation details
## 5-times repeated 10-fold repeated CV
fitControl <- trainControl(
  method = "repeatedcv",
  number = 10,
  repeats = 5,
  verboseIter = TRUE)

# set tuning grid for each model type for which it's needed
rfGrid <- data.frame(mtry = 5:15)
gbmGrid <- expand.grid(n.trees = c(100, 250, 500),
                       interaction.depth = c(5,10,15,20),
                       shrinkage = c(0.005, 0.01),
                       n.minobsinnode = c(5,10,15))
enetGrid <- expand.grid(alpha = seq(0.1,1,by=0.1),
                        lambda = c(0.001, 0.003, 0.005, 0.010, 0.015,
                                   0.020, 0.025, 0.030))

# Running the follow models may take a while (a few hours)
# To speed things up, uncomment the following lines to build
# multiple models in parallel
#require(doMC)
#registerDoMC(6)

# Now train a model using structured and structured + unstructured data for each type
# Logistic regression
lr.str <- train(as.numeric(Y) ~ .,
                method = 'glm',
                family = 'binomial',
                trControl = fitControl,
                data = df[training,var_names_struct,with=F])
lr.unstr <- train(as.numeric(Y) ~ .,
                method = 'glm',
                family = 'binomial',
                trControl = fitControl,
                data = df[training])

# Penalized logistic regression (elastic net)
en.str <- train(as.numeric(Y) ~ .,
                method = 'glmnet',
                tuneGrid = enetGrid,
                trControl = fitControl,
                data = df[training,var_names_struct,with=F])
en.unstr <- train(as.numeric(Y) ~ .,
                  method = 'glmnet',
                  tuneGrid = enetGrid,
                  trControl = fitControl,
                  data = df[training])

# Random forests
rf.str <- train(as.numeric(Y) ~ .,
                method = 'rf',
                tuneGrid = rfGrid,
                trControl = fitControl,
                data = df[training,var_names_struct,with=F])
rf.unstr <- train(as.numeric(Y) ~ .,
                  method = 'rf',
                  tuneGrid = rfGrid,
                  trControl = fitControl,
                  data = df[training])

# Gradient boosted machines
gbm.str <- train(as.numeric(Y) ~ .,
                method = 'gbm',
                tuneGrid = gbmGrid,
                trControl = fitControl,
                data = df[training,var_names_struct,with=F])
gbm.unstr <- train(as.numeric(Y) ~ .,
                  method = 'gbm',
                  tuneGrid = gbmGrid,
                  trControl = fitControl,
                  data = df[training])

# generate predictions on hold-out test set

# a helper function to make predicted probabilities in useful range
fix_range <- function(vals, low = 0, high = 1) {
  vals [vals < low] <- low
  vals [vals > high] <- high
  vals
}

## structured data models
lr.str.preds <- fix_range(predict(lr.str, df[testing]))
en.str.preds <- fix_range(predict(en.str, df[testing]))
rf.str.preds <- fix_range(predict(rf.str, df[testing]))
gbm.str.preds <- fix_range(predict(gbm.str, df[testing]))

## structured + unstructured data models
lr.unstr.preds <- fix_range(predict(lr.unstr, df[testing]))
en.unstr.preds <- fix_range(predict(en.unstr, df[testing]))
rf.unstr.preds <- fix_range(predict(rf.unstr, df[testing]))
gbm.unstr.preds <- fix_range(predict(gbm.unstr, df[testing]))

# generate ROC data for each model
## structured data models
lr.str.roc <- roc(predictor = lr.str.preds, response = df[testing]$Y, ci = TRUE)
en.str.roc <- roc(predictor = en.str.preds, response = df[testing]$Y, ci = TRUE)
rf.str.roc <- roc(predictor = rf.str.preds, response = df[testing]$Y, ci = TRUE)
gbm.str.roc <- roc(predictor = gbm.str.preds, response = df[testing]$Y, ci = TRUE)

## structured + unstructured data models
lr.unstr.roc <- roc(predictor = lr.unstr.preds, response = df[testing]$Y, ci = TRUE)
en.unstr.roc <- roc(predictor = en.unstr.preds, response = df[testing]$Y, ci = TRUE)
rf.unstr.roc <- roc(predictor = rf.unstr.preds, response = df[testing]$Y, ci = TRUE)
gbm.unstr.roc <- roc(predictor = gbm.unstr.preds, response = df[testing]$Y, ci = TRUE)

# make two-panel roc plot
par(mfrow = c(1,2))
plot_cols <- c('red', 'blue','forest green', 'dark orange')

# Plot structured group
plot(lr.str.roc, print.auc = TRUE, print.auc.cex = 0.8, col = plot_cols[1],
     main = 'Discrimination of models using structured data')
plot(en.str.roc, col = plot_cols[2], add = TRUE, print.auc = TRUE, print.auc.cex = 0.8, print.auc.y = 0.45)
plot(rf.str.roc, col = plot_cols[3], add = TRUE, print.auc = TRUE, print.auc.cex = 0.8, print.auc.y = 0.40)
plot(gbm.str.roc, col = plot_cols[4], add = TRUE, print.auc = TRUE, print.auc.cex = 0.8, print.auc.y = 0.35)
legend('bottomright', legend = c('LR', 'EN', 'RF', 'GBM'), bty = 'n',
       col = plot_cols, lwd = 2, cex = .8, y.intersp = 0.5, seg.len = 1)

# plot structured + unstructured group
plot(lr.unstr.roc, print.auc = TRUE, print.auc.cex = 0.8, col = plot_cols[1],
     main = 'Discrimination of models using structured and \nunstructured data')
plot(en.unstr.roc, col = plot_cols[2], add = TRUE, print.auc = TRUE, print.auc.cex = 0.8, print.auc.y = 0.45)
plot(rf.unstr.roc, col = plot_cols[3], add = TRUE, print.auc = TRUE, print.auc.cex = 0.8, print.auc.y = 0.40)
plot(gbm.unstr.roc, col = plot_cols[4], add = TRUE, print.auc = TRUE, print.auc.cex = 0.8, print.auc.y = 0.35)
legend('bottomright', legend = c('LR', 'EN', 'RF', 'GBM'), bty = 'n',
       col = plot_cols, lwd = 2, cex = .8, y.intersp = 0.5, seg.len = 1)


#==============================================================================================#

#intiial settings

#==============================================================================================#
# name pattern
run_id <- "E12_iD6"

# create directories
dir.create(paste0("./out/", run_id))
dir.create(paste0("./data/", run_id))
out_dir <- paste0("./out/", run_id, "/")
data_dir <- paste0("./data/", run_id, "/")

#load utilities
source("R/util.r")

#==============================================================================================#

# load data

#==============================================================================================#
data_paths <- c(
    E12 = "./data/mapped/E12",
    iD6 = "./data/mapped/induced_D6"
)

data_list <- NULL
for(i in seq(length(data_paths))){
    data_list <- c(
        data_list,
        list(scCountProc(data_path = data_paths[i], sample_name = names(data_paths)[i], normalization = FALSE))
    )
}
names(data_list) <- names(data_paths)

#==============================================================================================#

# QC

# using QC.r
#==============================================================================================#
dir.create(paste0(out_dir, "QC"))

qc_plot_list <- lapply(data_list, QC)

for(i in seq(length(data_paths))){
    data_list[[i]][["percent.mt"]] <- qc_plot_list[[i]]$percent.mt

    p <- plot_grid(qc_plot_list[[i]]$violin_plot, qc_plot_list[[i]]$scatter_plot, ncol = 1, nrow = 2)
    ggsave(p, file = paste0(out_dir, "QC/", names(data_list)[i], ".pdf"), width = 21, height = 14)
}

# cut-off is manually decided by violin plot &scatter plot
data_list[["E12"]] <- subset(data_list[["E12"]], subset = nFeature_RNA > 750 & nFeature_RNA < 6500 & nCount_RNA < 50000 & percent.mt < 25)
data_list[["iD6"]] <- subset(data_list[["iD6"]], subset = nFeature_RNA > 750 & nFeature_RNA < 5000 & nCount_RNA < 19000 & percent.mt < 25)

#==============================================================================================#

# Normalization

#==============================================================================================#
for(i in seq(length(data_paths))){
    data_list[[i]] <- NormalizeData(object = data_list[[i]], verbose = FALSE)
    data_list[[i]] <- FindVariableFeatures(object = data_list[[i]], selection.method = "vst", nfeatures = 1000, verbose = FALSE)
}

#==============================================================================================#

# Data Integration

#==============================================================================================#
data.anchors <- FindIntegrationAnchors(object.list = data_list, dims = 1:30)
all_genes <- data.anchors@object.list[[1]]@assays$RNA@counts@Dimnames[[1]]
data.integrated <- IntegrateData(anchorset = data.anchors, dims = 1:30, features.to.integrate = all_genes)

age <- data.integrated@meta.data$sample
age %>%
    replace(which(date == "E12"), 12) %>%
    replace(which(date == "iD6"), 12) -> age

vivo.vitro <- data.integrated@meta.data$sample
vivo.vitro %>%
    replace(which(vivo.vitro == "E12"), "vivo") %>%
    replace(which(vivo.vitro == "iD6"), "vitro") -> vivo.vitro

data.integrated@meta.data <- cbind(data.integrated@meta.data, age, vivo.vitro)
save(data.integrated, file = paste0(data_dir, "data.integrated.QC.RData"))

#==============================================================================================#

# Cell Cycle Regression

#==============================================================================================#
dir.create(paste0(out_dir, "ccReg_results"))

data.integrated <- ccReg(data = data.integrated)
save(data.integrated, file = paste0(out_dir, "ccReg_results/data.integrated.QC.ccreg.RData"))

#==============================================================================================#

# Clustering & UMAP (Supplementary Figure 6E)

#==============================================================================================#
#clustering
data.integrated <- RunPCA(object = data.integrated, npcs = 16, verbose = FALSE)
data.integrated <- RunUMAP(object = data.integrated, reduction = "pca", dims = 1:16)
data.integrated <- FindNeighbors(object = data.integrated, dims = 1:16)
data.integrated <- FindClusters(object = data.integrated, algorithm = 1, resolution = 0.2)

save(data.integrated, file = paste0(data_dir, "data.integrated.QCed.ccreg.clust.RData"))

# UMAP colored by cluster and splited by sample
pdf(paste0(out_dir, "by.sample_umap.pdf"), height=7, width=14)
DimPlot(object = data.integrated, reduction = "umap", split.by="sample", pt.size = 1.2)
dev.off()

#==============================================================================================#

#  Maeker gene expression (Supplementary Figure 6F)

#==============================================================================================#
dir.create(paste0(out_dir, "known_marker"))

#Germ cells
multiVlnPlot(object=data.integrated, features = MARKERS$germ, title="Germ Cell Markers", out_file =paste0(out_dir, "known_marker/germ_vln.pdf"))

#Endothelial
multiVlnPlot(object=data.integrated, features = MARKERS$endothelial, title="Endothelial Markers", out_file =paste0(out_dir, "known_marker/endothelial_Vln.pdf"))

#erythrovyte
multiVlnPlot(object=data.integrated, features =  MARKERS$erythro, title="Erythrocyte Markers",out_file =paste0(out_dir, "known_marker/ery_Vln.pdf"))

#megakaryocyte
multiVlnPlot(object=data.integrated, features = MARKERS$mk, title="MK Markers",out_file =paste0(out_dir, "known_marker/mk_vln.pdf"))

#==============================================================================================#

# Extraction of Gonadal Somatic Cells (Figure 2C, D)

#==============================================================================================#
data.integrated.sel <- subsetUmapClust(
    data =data.integrated,
    subset.name = "seurat_clusters",
    subset.value = 0:5,
    resolution=0.2,  # resolution shoud be optimized
    npcs = 20,
    dims = 1:20
)

save(data.integrated.sel, file = paste0(data_dir, "data.integrated.QCed.ccreg.clust.sel.RData"))

#UMAP colored by cluster and splited by sample
pdf(paste0(out_dir, "soma_by.sample_umap.pdf"), height=7, width=14)
DimPlot(object = data.integrated.sel, reduction = "umap", split.by="sample",pt.size = 1.2)
dev.off()

#UMAP colored by cell cycle phase and splited by sample
pdf(paste0(out_dir, "soma_cc.umap.pdf"), height=7, width=14)
DimPlot(object = data.integrated.sel, reduction = "umap", group.by="Phase", split.by="sample",pt.size = 1.2)
dev.off()

#==============================================================================================#

#Proportion of each cluster (Figure 2E)

#==============================================================================================#
#Reorder the clusters (1(EP), 5(EP), 0(STpro), 3(ST), 2(GR), 4(GR))
levels(data.integrated.sel@meta.data$seurat_cluster) <-  c(1, 5, 0, 3, 2, 4)
Idents(data.integrated.sel) <- data.integrated.sel@meta.data$seurat_cluster

#computtion of proportion
prop_cluster <- prop.table(x=table(data.integrated.sel@meta.data$seurat_cluster, data.integrated.sel@meta.data$sample), margin =2)

#color setting
ncluster <- length(levels(data.integrated.sel@meta.data$seurat_cluster))
clust_col = gg_color_hue(ncluster)
names(clust_col) = paste0("S", c(0:5))
clust_col <- clust_col[ c("S1", "S5", "S0", "S3", "S2", "S4")]

#plot
Cluster_name=factor(paste0("S", rownames(prop_cluster)), levels=c("S1", "S5", "S0", "S3", "S2", "S4"))    # To order, making a factor object for the cluster
df <- data.frame(Cluster=Cluster_name, D6=prop_cluster[,2], E12.5=prop_cluster[,1])
df2 <- tidyr::gather(df, key=Sample, value=Proportion, -Cluster, factor_key = TRUE)

##vidualization
theme <- theme(panel.background = element_blank(),    # initialization
            panel.border = element_blank(),
            panel.grid.major = element_blank(),
            panel.grid.minor = element_blank(),
            strip.background = element_blank(),
            plot.title = element_text(size = 20,hjust = 0.5),
            axis.text.x = element_text(size=30),
            axis.text.y = element_text(size=20),
            axis.title.x = element_blank(),
            axis.title.y = element_blank(),
            axis.ticks = element_line(size=0.5),
            axis.line = element_line(size=0.5),
            legend.text = element_text(size=20),
            legend.title =  element_text(size=20),
            plot.margin = unit(c(1,1,1,1),"line")
    )

g <- ggplot(df2, aes(x= Sample, y = Proportion, fill=Cluster))
        g <- g + geom_bar(stat = "identity")
        g <- g + scale_y_continuous(labels = percent)    # display as %
        g <- g + scale_fill_manual(values = clust_col, limits = c("S0", "S1", "S2","S3", "S4", "S5"))
        g <- g + theme

        g <- g + annotate("text",size = 10, x = 1, y = 0.07, label = "GR")
        g <- g + annotate("text",size = 10, x = 1, y = 0.19, label = "GR")
        g <- g + annotate("text",size = 10, x = 1, y = 0.34, label = "ST")
        g <- g + annotate("text",size = 10, x = 1, y = 0.56, label = "STpro")
        g <- g + annotate("text",size = 10, x = 1, y = 0.706, label = "EP")
        g <- g + annotate("text",size = 10, x = 1, y = 0.865, label = "EP")

        g <- g + annotate("text",size = 10, x = 2, y = 0.023, label = "GR")
        g <- g + annotate("text",size = 10, x = 2, y = 0.27, label = "GR")
        g <- g + annotate("text",size = 10, x = 2, y = 0.532, label = "ST")
        g <- g + annotate("text",size = 10, x = 2, y = 0.675, label = "STpro")
        g <- g + annotate("text",size = 10, x = 2, y = 0.815, label = "EP")
        g <- g + annotate("text",size = 10, x = 2, y = 0.935, label = "EP")

ggsave(filename = paste0(out_dir, "soma_proportion.pdf"), plot = g, width = 12, height = 14)

#==============================================================================================#

#Soma correlation matrix (Figure 2F)

#==============================================================================================#
#add a meta data
data.integrated.sel@meta.data$sample_clust <- paste0(data.integrated.sel@meta.data$sample, "-", data.integrated.sel@meta.data$seurat_clusters)
Idents(data.integrated.sel) <- data.integrated.sel@meta.data$sample_clust

#comutation of the average expression of a cluster of each sample
av.exp <- AverageExpression(data.integrated.sel)$integrated

#correlation coefficient
cor.exp <- cor(av.exp)

#hierarchical clustering
hc <- hclust(dist(cor.exp))

#color setting for R
breaks1 <- seq(0.96, 1, length=500)    #the range should be optimized
mycols1 <- colorRamp2(breaks1,colorRampPalette(brewer.pal(9, "GnBu"))(500)) # color for motif enrichment

#cluster color setting
data_set <- unique(data.integrated.sel@meta.data$sample_clust)

cluster <- sapply(strsplit(data_set, "-"), function(x) x[2])
names(cluster) <- data_set

ncluster=length(unique(cluster))
clust_col = gg_color_hue(ncluster)
names(clust_col) <- sort(unique(cluster))

#cluster annotation
ha_row = rowAnnotation(df = data.frame(cluster = cluster), col = list(cluster = clust_col), width = unit(1, "cm"), annotation_legend_param = list(title_gp = gpar(fontsize = 12), labels_gp = gpar(fontsize = 12)))
ha_col = HeatmapAnnotation(df = data.frame(cluster = cluster), col = list(cluster = clust_col), height = unit(1, "cm"), annotation_legend_param = list(title_gp = gpar(fontsize = 12), labels_gp = gpar(fontsize = 12)))

#correlation matrix
ht1 <- Heatmap(cor.exp,
                col = mycols1,
                name = "R", #title of legend,
                row_names_gp = gpar(fontsize = 12),    # Text size for row names
                column_names_gp = gpar(fontsize = 12),    # Text size for column names
                column_names_rot = 90,
                heatmap_legend_param =  list(title_gp = gpar(fontsize = 12), labels_gp = gpar(fontsize = 12)),
                row_dend_width = unit(2, "cm"),   #row dendrogram width
                column_dend_height = unit(2, "cm"),   #column dendrogram height
                row_dend_side = "left",
                column_dend_side = "top",
                cluster_rows = hc,
                cluster_columns = hc,
                top_annotation = ha_col,
                left_annotation = ha_row,
                width = unit(14, "cm"),
                height = unit(14, "cm"),
                cell_fun = function(j, i, x, y, width, height, fill){
                        grid.rect(x = x, y = y, width = width, height = height, gp = gpar(col = "black", fill = NA))
                }
            )

draw(ht1, merge_legend = TRUE)

pdf(paste0(out_dir, "soma_correlationplot.pdf"), height = 9, width = 9)
draw(ht1, merge_legend = TRUE)
dev.off()

#==============================================================================================#

# Soma Marker expression (Supplementary Figure 6G)

#==============================================================================================#
granulosa_markers <- c("Kitl", "Inha", "Foxl2", "Runx1")
stromal_markers <- c("Wnt5a", "Pdgfra", "Tcf21", "Acta2")
progenitor_marker <- c("Sox11",  "Ecm1", "Nr5a1")

pdf(paste0(out_dir, "known_marker/soma_granulosa.pdf"), width=28, height=7)
FeaturePlot(object = data.integrated.sel, features = granulosa_markers , min.cutoff = 0, max.cutoff = c(2.0, 2.5, 1.5, 2), cols=c("gray", "red"), ncol=4)
dev.off()

pdf(paste0(out_dir, "known_marker/soma_stromal.pdf"), width=28, height=7)
FeaturePlot(object = data.integrated.sel, features = stromal_markers, min.cutoff = 0, max.cutoff = c(3.0, 2.0, 3, 3.0), cols=c("gray", "red"), ncol=4)
dev.off()

pdf(paste0(out_dir, "known_marker/soma_progenitor.pdf"), width=21, height=7)
FeaturePlot(object = data.integrated.sel, features = progenitor_marker, min.cutoff = c(0,0,0),  max.cutoff = c(3.5, 3.0, 1.0), cols=c("gray", "red"), ncol=4)
dev.off()

#==============================================================================================#

#Soma Cluster specific gene expression (Supplementary Figure 6H)

#==============================================================================================#
#identification of cluster specific genes. We used MAST algorism.
library(dplyr)
integrated.sel.markers <- FindAllMarkers(
object = data.integrated.sel,
only.pos = FALSE,
min.pct = 0.25,
logfc.threshold = 0.25,
test.use = "MAST"
)

save(integrated.sel.markers, file=paste0(data_dir, "soma_markers.RData"))

#Heatmap of top5 genes
integrated.sel.markers %>%
filter(avg_logFC >= 0) %>%
group_by(cluster) %>%
top_n(5, avg_logFC) -> top5.posi

pdf(paste0(out_dir, "soma_top5.posi.heatmap.pdf"), , height=7, width=7)
DoHeatmap(object = data.integrated.sel, features=top5.posi$gene, size=3, label=FALSE)
dev.off()

#==============================================================================================#

# differentially expression analysis  (Supplementary Figure 6I)

#==============================================================================================#
#parameters
control <- "E12"
treatment <- "id6"
outname <-  "E12"

seurat_clusters <- levels(unlist(data.integrated.sel[["seurat_clusters"]]))

#DEG analysis (MAST)
diff_genes.list <- list()
for (i in seq(length(seurat_clusters))){
    data_subset <- eval(parse(text=paste0("subset(data.integrated.sel, subset = seurat_clusters == ", seurat_clusters[i], ")")))    #Extraction of the target cluster
    Idents(object =data_subset) <- data_subset@meta.data$sample

    diff_genes <- try(
            Deg(seurat_object=data_subset,
            control = control,
            treatment = treatment,
            outID = paste0(outname, "_C", seurat_clusters[i]),
            only.pos = FALSE,
            min.pct = 0.25,
            logfc.threshold = 0,
            test.use = "MAST",
            min.cells.feature = 1,
            min.cells.group = 1),
        silent = FALSE)    #DEG analysis
    if(class(diff_genes) == "try-error"){    # In case DEG analysis is error
        diff_genes <- NA
    }

    diff_genes.list[[as.numeric(i)]] <- diff_genes
    names(diff_genes.list)[as.numeric(i)] <- seurat_clusters[i]
}

#sig DEG list
sig_genes.list <- lapply(diff_genes.list, function(x){
    if(is.data.frame(x) && nrow(x) > 0){
        x %>%
            mutate(gene = rownames(x)) %>%
            dplyr::filter(avg_logFC >= 0.25 | avg_logFC <= -0.25) %>%
            dplyr::filter(p_val_adj < 10e-2)
    }
})

#MA plot
p1.list <- list()
for(i in seurat_clusters){
    sig_genes <- eval(parse(text=paste0("sig_genes.list$'", i, "'$gene")))
    if(length(sig_genes) == 0) sig_genes <- NULL

    title <- paste0("Cluster ", i)

    p1 <- averageDegMA(seurat_object=data.integrated.sel,
        cluster_key = "seurat_clusters",
        clust = i,
        comp_key = "sample",
        control = control,
        treatment =treatment,
        sig_genes = sig_genes,
        genes_to_label = NULL,
        title = title
    )
    listname <- paste0("Cluster_", i)
    p1.list <- eval(parse(text = paste0("c(p1.list, list(", listname,  " = p1))")))
}

#vidualization
p1.list2 <- list(p1.list$Cluster_0, p1.list$Cluster_1, p1.list$Cluster_5, p1.list$Cluster_3, p1.list$Cluster_2, p1.list$Cluster_4)
g1 <- plot_grid(plotlist = p1.list2, nrow = 1, align="h")
ggsave( paste0(out_dir, "soma_MA_all.pdf"), g1, height=7, width=42)


source("./metrics.r")

## Testing full results
#calculate <- dget("metrics.r")
test = read.csv("../sample-data/test_data.csv", header = TRUE)
#test['predictionscore'] <- order(test['D_PFS'])
pred = read.csv("../sample-data/predictions.csv", header = TRUE)
# rawscore, highrisk, PFStime, pfs_flag
#calculate.metrics(pred,  pred['highriskflag'], test['D_PFS'], test['D_PFS_FLAG'])

calculate.metrics_wrapper(pred, test['D_PFS'], test['D_PFS_FLAG'])

library(pec)
data(GBSG2,package="pec")
calculate.metrics_wrapper(pred, GBSG2['time'], GBSG2['cens'])

## Testing prodLim package kaplan-meier reverse

# T              : vector of observed failure times
# delta          : vector of indicator of status

#library(prodlim)
#library(rms)

#times = 30.5 * c(14, 16, 18, 20, 22)

#predicted = test['predictionscore']
#D_PFS = test['D_PFS']
#D_PFS_FLAG =  test['D_PFS_FLAG']

#T = test['D_PFS']
#colnames(T) <- c("T")

#delta = test['D_PFS_FLAG']
#colnames(delta) <- c("delta")

#T <- T[order(T$T),]

#order_T = order(T)
#print(order_T)
#T <- T[order_T,]
#delta <- delta[order_T,]

# reverse the censored events (making them non-censored)
#df = data.frame(failure_time=T,status=as.numeric(delta!=0))

# using the marginal Kaplan-Meier for the censoring times
#weights <- pec::ipcw(Surv(failure_time,status)~1,data=df,method="marginal",times=times,subjectTimes=T,subjectTimesLag=1)

#print(weights)

# timeROC library example
#suppressPackageStartupMessages(library("timeROC"))
#tempAUC <- timeROC(T = D_PFS, delta = D_PFS_FLAG, marker = predicted, cause = 1, times = times)

#print(tempAUC)


# script for ploting pretty heatmaps for contrast RNAseq data
# jmmunozp@usp.br
#install.packages("tidyverse")
#install.packages(c("pheatmap", "BiocManager"))
#BiocManager::install("DESeq2")

# variables
wd = "/home/j/BIOINFORMATICA/NRGSC_old"
mdata = "/metadata_complete.csv"
vst_path = "./DEA/complete/vst.Rdata"

# libraries
library(tidyverse)
library(DESeq2)
library(pheatmap)

# set working directory
setwd(wd)

sample_table <-read.table(mdata, sep = ",", header = T)
sample_table$Condition <- as.factor((sample_table$Condition))
sample_table$Genotype <- as.factor((sample_table$Genotype))
sample_table$DevStage <- as.factor((sample_table$DevStage))
sample_table$Individual <- as.factor((sample_table$Individual))
sample_table$Group <- as.factor(paste(sample_table$Genotype, '_', sample_table$Condition, '_', sample_table$DevStage, sep=''))

# read vst file
load(vst_path)
# function to graph heatmaps from contrast data

contrast_heatmaps <- function (vst, contrast_samples_index, title, contrast_table) {

dea_contrast <- read.table(contrast_table, header = T, sep = ",")

df <- as.data.frame(t(assay(vst)))
# NR_270_B (Sample_34, Sample_30, Sample_26)
# NR_10_B (Sample_46, Sample_42, Sample_38)
contrast <- df[contrast_samples_index,]

annotation_col <- sample_table[contrast_samples_index,]
anot <- select(annotation_col, "Group","DevStage","Genotype", "Condition")
dat <- as.data.frame(t(contrast[dea_contrast$names]))
rownames(anot) <- colnames(dat)
#heat map complete with all genes
png(paste("./DEA/hm_", title, ".png", sep = ""),  width = 15*1.3, height = 15, res = 320, units = "cm")
pheatmap(dat, main = title, annotation_col = anot, show_rownames = F , cellwidth = 20)
dev.off()

# upregulated genes
upregulated <- dea_contrast[tail(order(dea_contrast$log2FoldChange),20),]
dat <- as.data.frame(t(contrast[upregulated$names]))
png(paste("./DEA/upregulated_", title, ".png", sep = ""),  width = 15*1.3, height = 15, res = 320, units = "cm")
pheatmap(dat, main =  paste( title," upregulated", sep = ""), annotation_col = anot, show_rownames = T, fontsize_row = 8, cellwidth = 20)
dev.off()

# downregulated genes
downregulated <- dea_contrast[head(order(dea_contrast$log2FoldChange),20),]
dat <- as.data.frame(t(contrast[downregulated$names]))
png(paste("./DEA/downregulated_", title, ".png", sep = ""),  width = 15*1.3, height = 15, res = 320, units = "cm")
pheatmap(dat, main = paste( title," downregulated", sep = ""), annotation_col = anot, show_rownames = T, fontsize_row = 8, cellwidth = 20)
dev.off()
}

# call function with different contrast
contrast_heatmaps(vst,c(34,30,26,46,42,38), "NR_270_B vs NR_10_B", "/home/j/BIOINFORMATICA/NRGSC_old/DEA/complete/filter_2NR_270_B_vs_NR_10_B.csv")