library("mosaic")
library("dplyr")
library("foreach")
library("doParallel")

#setup parallel backend to use many processors
cores=detectCores()
cl <- makeCluster(cores[1]-1) #not to overload your computer
registerDoParallel(cl)

args = commandArgs(trailingOnly=TRUE)

if (length(args)==0) {
  runtype="timedump_253048_1873f6_all/timedump"
  target="waters_int"
  outputpath="~/code/FRET/LibAFL/fuzzers/FRET/benchmark/"
  #MY_SELECTION <- c('state', 'afl', 'graph', 'random')
  SAVE_FILE=TRUE
} else {
  runtype=args[1]
  target=args[2]
  outputpath=args[3]
  MY_SELECTION <- args[4:length(args)]
  SAVE_FILE=TRUE
}
worst_cases <- list(waters=0, waters_int=0, tmr=405669, micro_longint=0)
worst_case <- worst_cases[[target]]
if (is.null(worst_case)) {
  worst_case = 0
}

#MY_COLORS=c("green","blue","red", "orange", "pink", "black")
MY_COLORS <- c("green", "blue", "red", "magenta", "orange", "cyan", "pink", "gray", "orange", "black", "yellow","brown")
BENCHDIR=sprintf("~/code/FRET/LibAFL/fuzzers/FRET/benchmark/%s",runtype)
BASENAMES=Filter(function(x) x!="" && substr(x,1,1)!='.',list.dirs(BENCHDIR,full.names=FALSE))
PATTERNS="%s.[0-9]*$"
#RIBBON='sd'
#RIBBON='span'
RIBBON='both'
DRAW_WC = worst_case > 0
LEGEND_POS="topright"
#LEGEND_POS="bottomright"
CONTINUE_LINE_TO_END=FALSE

# https://www.r-bloggers.com/2013/04/how-to-change-the-alpha-value-of-colours-in-r/
alpha <- function(col, alpha=1){
  if(missing(col))
    stop("Please provide a vector of colours.")
  apply(sapply(col, col2rgb)/255, 2, 
        function(x) 
          rgb(x[1], x[2], x[3], alpha=alpha))  
}

# Trimm a list of data frames to common length
trim_data <- function(input,len=NULL) {
  if (is.null(len)) {
    len <- min(sapply(input, function(v) dim(v)[1]))
  }
  return(lapply(input, function(d) slice_head(d,n=len)))
}

length_of_data <- function(input) {
  min(sapply(input, function(v) dim(v)[1]))
}

# Takes a flat list
trace2maxline <- function(tr) {
  maxline = tr
  for (var in seq_len(length(maxline))[2:length(maxline)]) {
    #if (maxline[var]>1000000000) {
    #  maxline[var]=maxline[var-1]
    #} else {
      maxline[var] = max(maxline[var],maxline[var-1]) 
    #}
  }
  #plot(seq_len(length(maxline)),maxline,"l",xlab="Index",ylab="WOET")
  return(maxline)
}

# Take a list of data frames, output same form but maxlines
data2maxlines <- function(tr) {
  min_length <- min(sapply(tr, function(v) dim(v)[1]))
  maxline <- tr
  for (var in seq_len(length(tr))) {
    maxline[[var]][[1]]=trace2maxline(tr[[var]][[1]])
  }
  return(maxline)
}
# Take a multi-column data frame, output same form but maxlines
frame2maxlines <- function(tr) {
  for (var in seq_len(length(tr))) {
    tr[[var]]=trace2maxline(tr[[var]])
  }
  return(tr)
}

trace2maxpoints <- function(tr) {
  minval = tr[1,1]
  collect = tr[1,]
  for (i in seq_len(dim(tr)[1])) {
    if (minval < tr[i,1]) {
      collect = rbind(collect,tr[i,])
      minval = tr[i,1]
    }
  }
  tmp = tr[dim(tr)[1],]
  tmp[1] = minval[1]
  collect = rbind(collect,tmp)
  return(collect)
}

sample_maxpoints <- function(tr,po) {
  index = 1
  collect=NULL
  endpoint = dim(tr)[1]
  for (p in po) {
    if (p<=tr[1,2]) {
      tmp = tr[index,]
      tmp[2] = p
      collect = rbind(collect, tmp)
    } else if (p>=tr[endpoint,2]) {
      tmp = tr[endpoint,]
      tmp[2] = p
      collect = rbind(collect, tmp)
    } else {
      for (i in seq(index,endpoint)-1) {
        if (p >= tr[i,2] && p<tr[i+1,2]) {
          tmp = tr[i,]
          tmp[2] = p
          collect = rbind(collect, tmp)
          index = i
          break
        }
      } 
    }
  }
  return(collect)
}

#https://www.r-bloggers.com/2012/01/parallel-r-loops-for-windows-and-linux/
all_runtypetables <- foreach (bn=BASENAMES) %do% {
  runtypefiles <- list.files(file.path(BENCHDIR,bn),pattern=sprintf(PATTERNS,target),full.names = TRUE)
  if (length(runtypefiles) > 0) {
    runtypetables_reduced <- foreach(i=seq_len(length(runtypefiles))) %dopar% {
      rtable = read.csv(runtypefiles[[i]], col.names=c(sprintf("%s%d",bn,i),sprintf("times%d",i)))
      trace2maxpoints(rtable)
    }
    #runtypetables <- lapply(seq_len(length(runtypefiles)),
    #    function(i)read.csv(runtypefiles[[i]], col.names=c(sprintf("%s%d",bn,i),sprintf("times%d",i))))
    #runtypetables_reduced <- lapply(runtypetables, trace2maxpoints)
    runtypetables_reduced
    #all_runtypetables = c(all_runtypetables, list(runtypetables_reduced))
  }
}
all_runtypetables = all_runtypetables[lapply(all_runtypetables, length) > 0]
all_min_points = foreach(rtt=all_runtypetables,.combine = cbind) %do% {
  bn = substr(names(rtt[[1]])[1],1,nchar(names(rtt[[1]])[1])-1)
  ret = data.frame(min(unlist(lapply(rtt, function(v) v[dim(v)[1],2]))))
  names(ret)[1] = bn
  ret/(3600 * 1000)
}
all_max_points = foreach(rtt=all_runtypetables,.combine = cbind) %do% {
  bn = substr(names(rtt[[1]])[1],1,nchar(names(rtt[[1]])[1])-1)
  ret = data.frame(max(unlist(lapply(rtt, function(v) v[dim(v)[1],2]))))
  names(ret)[1] = bn
  ret/(3600 * 1000)
}
all_points = sort(unique(Reduce(c, lapply(all_runtypetables, function(v) Reduce(c, lapply(v, function(w) w[[2]]))))))
all_maxlines <- foreach (rtt=all_runtypetables) %do% {
  bn = substr(names(rtt[[1]])[1],1,nchar(names(rtt[[1]])[1])-1)
  runtypetables_sampled = foreach(v=rtt) %dopar% {
    sample_maxpoints(v, all_points)[1]
  }
  #runtypetables_sampled = lapply(rtt, function(v) sample_maxpoints(v, all_points)[1])
  tmp_frame <- Reduce(cbind, runtypetables_sampled)
  statframe <- data.frame(rowMeans(tmp_frame),apply(tmp_frame, 1, sd),apply(tmp_frame, 1, min),apply(tmp_frame, 1, max), apply(tmp_frame, 1, median))
  names(statframe) <- c(bn, sprintf("%s_sd",bn), sprintf("%s_min",bn), sprintf("%s_max",bn), sprintf("%s_med",bn))
  #statframe[sprintf("%s_times",bn)] = all_points
  round(statframe)
  #all_maxlines = c(all_maxlines, list(round(statframe))) 
}
one_frame<-data.frame(all_maxlines)
one_frame[length(one_frame)+1] <- all_points/(3600 * 1000)
names(one_frame)[length(one_frame)] <- 'time'

typenames = names(one_frame)[which(names(one_frame) != 'time')]
typenames = typenames[which(!endsWith(typenames, "_sd"))]
typenames = typenames[which(!endsWith(typenames, "_med"))]
ylow=min(one_frame[typenames])
yhigh=max(one_frame[typenames],worst_case)
typenames = typenames[which(!endsWith(typenames, "_min"))]
typenames = typenames[which(!endsWith(typenames, "_max"))]

ml2lines <- function(ml,lim) {
  lines = NULL
  last = 0
  for (i in seq_len(dim(ml)[1])) {
    if (!CONTINUE_LINE_TO_END && lim<ml[i,2]) {
      break
    }
    lines = rbind(lines, cbind(X=last, Y=ml[i,1]))
    lines = rbind(lines, cbind(X=ml[i,2], Y=ml[i,1]))
    last = ml[i,2]
  }
  return(lines)
}

plotting <- function(selection, filename, MY_COLORS_) {
# filter out names of iters and sd cols
typenames = names(one_frame)[which(names(one_frame) != 'times')]
typenames = typenames[which(!endsWith(typenames, "_sd"))]
typenames = typenames[which(!endsWith(typenames, "_med"))]
typenames = typenames[which(!endsWith(typenames, "_min"))]
typenames = typenames[which(!endsWith(typenames, "_max"))]
typenames = selection[which(selection %in% typenames)]
if (length(typenames) == 0) {return()}

h_ = 500
w_ = h_*4/3

if (SAVE_FILE) {png(file=sprintf("%s%s_%s.png",outputpath,target,filename), width=w_, height=h_)}
par(mar=c(4,4,1,1))
par(oma=c(0,0,0,0))

plot(c(1,max(one_frame['time'])),c(ylow,yhigh), col='white', xlab="Time [h]", ylab="WORT [insn]", pch='.')

for (t in seq_len(length(typenames))) {
  #proj = one_frame[seq(1, dim(one_frame)[1], by=max(1, length(one_frame[[1]])/(10*w_))),]
  #points(proj[c('iters',typenames[t])], col=MY_COLORS_[t], pch='.')
  avglines = ml2lines(one_frame[c(typenames[t],'time')],all_max_points[typenames[t]])
  #lines(avglines, col=MY_COLORS_[t])
  medlines = ml2lines(one_frame[c(sprintf("%s_med",typenames[t]),'time')],all_max_points[typenames[t]])
  lines(medlines, col=MY_COLORS_[t], lty='solid')
  milines = NULL
  malines = NULL
  milines = ml2lines(one_frame[c(sprintf("%s_min",typenames[t]),'time')],all_max_points[typenames[t]])
  malines = ml2lines(one_frame[c(sprintf("%s_max",typenames[t]),'time')],all_max_points[typenames[t]])
  if (exists("RIBBON") && ( RIBBON=='max' )) {
    #lines(milines, col=MY_COLORS_[t], lty='dashed')
    lines(malines, col=MY_COLORS_[t], lty='dashed')
    #points(proj[c('iters',sprintf("%s_min",typenames[t]))], col=MY_COLORS_[t], pch='.')
    #points(proj[c('iters',sprintf("%s_max",typenames[t]))], col=MY_COLORS_[t], pch='.')
  }
  if (exists("RIBBON") && RIBBON != '') {
  for (i in seq_len(dim(avglines)[1]-1)) {
    if (RIBBON=='both') {
      # draw boxes
      x_l <- milines[i,][['X']]
      x_r <- milines[i+1,][['X']]
      y_l <- milines[i,][['Y']]
      y_h <- malines[i,][['Y']]
      rect(x_l, y_l, x_r, y_h, col=alpha(MY_COLORS_[t], alpha=0.1), lwd=0)
    }
    if (FALSE && RIBBON=='span') {
      # draw boxes
      x_l <- milines[i,][['X']]
      x_r <- milines[i+1,][['X']]
      y_l <- milines[i,][['Y']]
      y_h <- malines[i,][['Y']]
      rect(x_l, y_l, x_r, y_h, col=alpha(MY_COLORS_[t], alpha=0.1), lwd=0)
    }
    #if (FALSE && RIBBON=='both' || RIBBON=='sd') {
    #  # draw sd
    #  x_l <- avglines[i,][['X']]
    #  x_r <- avglines[i+1,][['X']]
    #  y_l <- avglines[i,][['Y']]-one_frame[ceiling(i/2),][[sprintf("%s_sd",typenames[t])]]
    #  y_h <- avglines[i,][['Y']]+one_frame[ceiling(i/2),][[sprintf("%s_sd",typenames[t])]]
    #  if (x_r != x_l) {
    #    rect(x_l, y_l, x_r, y_h, col=alpha(MY_COLORS_[t], alpha=0.1), lwd=0) 
    #  }
    #}
    #sd_ <- row[sprintf("%s_sd",typenames[t])][[1]]
    #min_ <- row[sprintf("%s_min",typenames[t])][[1]]
    #max_ <- row[sprintf("%s_max",typenames[t])][[1]]
    #if (exists("RIBBON")) {
    #  switch (RIBBON,
    #    'sd' = arrows(x_, y_-sd_, x_, y_+sd_, length=0, angle=90, code=3, col=alpha(MY_COLORS_[t], alpha=0.03)),
    #    'both' = arrows(x_, y_-sd_, x_, y_+sd_, length=0, angle=90, code=3, col=alpha(MY_COLORS_[t], alpha=0.05)),
    #    'span' = #arrows(x_, min_, x_, max_, length=0, angle=90, code=3, col=alpha(MY_COLORS_[t], alpha=0.03))
    #  )
    #}
    ##arrows(x_, y_-sd_, x_, y_+sd_, length=0.05, angle=90, code=3, col=alpha(MY_COLORS[t], alpha=0.1))
  }
  }
}
leglines=typenames
if (DRAW_WC) {
  lines(c(0,length(one_frame[[1]])),y=c(worst_case,worst_case), lty='dotted')
  leglines=c(typenames, 'worst observed') 
}
legend(LEGEND_POS, legend=leglines,#"topleft"
       col=c(MY_COLORS_[1:length(typenames)],"black"),
       lty=c(rep("solid",length(typenames)),"dotted"))

if (SAVE_FILE) {dev.off()}
}
stopCluster(cl)

par(mar=c(3.8,3.8,0,0))
par(oma=c(0,0,0,0))

#RIBBON='both'
#MY_SELECTION = c('state_int','generation100_int')
#MY_SELECTION = c('state_int','frafl_int')

if (exists("MY_SELECTION")) {
  plotting(MY_SELECTION, 'custom', MY_COLORS[c(1,2)])
} else {
  # MY_SELECTION=c('state', 'afl', 'random', 'feedlongest', 'feedgeneration', 'feedgeneration10')
  #MY_SELECTION=c('state_int', 'afl_int', 'random_int', 'feedlongest_int', 'feedgeneration_int', 'feedgeneration10_int')
  #MY_SELECTION=c('state', 'frAFL', 'statenohash', 'feedgeneration10')
  #MY_SELECTION=c('state_int', 'frAFL_int', 'statenohash_int', 'feedgeneration10_int')
  MY_SELECTION=typenames
  RIBBON='both'
  for (i in seq_len(length(MY_SELECTION))) {
    n <- MY_SELECTION[i]
    plotting(c(n), n, c(MY_COLORS[i]))
  }
  RIBBON='max'
  plotting(MY_SELECTION,'all', MY_COLORS)
}

for (t in seq_len(length(typenames))) {
  li = one_frame[dim(one_frame)[1],]
  pear = (li[[typenames[[t]]]]-li[[sprintf("%s_med",typenames[[t]])]])/li[[sprintf("%s_sd",typenames[[t]])]]
  print(sprintf("%s pearson: %g",typenames[[t]],pear))
}