#!/usr/bin/Rscripy

# install.packages(c("ggplot2", "data.table", "plyr", "reshape2", "grid", "scales"))
library("ggplot2")
library("data.table")
library("plyr")
library("reshape2") # for melt
library("grid")     # for unit
library("scales")   # for rescale

# TODO: speed up data loading using RSQLite?

seconds.per.tick <- 40/1024

csv2vector <- function (line) {
  return( sapply(strsplit(line, ",", fixed=TRUE), as.numeric) )
}

mean.na.rm <- function (v) {
  mean(v, na.rm=TRUE)
}

parse.movement.data <- function (lines) {
  stats.list <- list()
  airplane.number <- 1
  for (line in lines) {
      first.tick <- line[1]
      positions <- line[-1]
      displacements <- positions[-1] - positions[-length(positions)]
      speeds <- c(abs(displacements) / seconds.per.tick, NA)
      N <- length(positions)
      new.stats <- data.table(airplane=rep(airplane.number, N),
                              tick=seq(from=first.tick, length.out=N),
                              position=positions, speed=speeds)
      stats.list <- c(stats.list, list(new.stats))
      airplane.number <- airplane.number + 1
  }
  stats <- rbindlist(stats.list)
  stats$airplane <- factor(stats$airplane)
  return(stats)
}

extract.mean.speeds <- function (N, movements) {
  movements <- subset(movements, select=c("tick", "speed"))
  movements <- rbind(movements, data.table(tick=seq(from=0, length.out=N),
                                           speed=rep(NA, N)))
  setkey(movements, "tick")
  return(movements[,mean(speed, na.rm=TRUE),by="tick"]$V1)
}

load.stats.from.file <- function (filepath) {
  lines <- sapply(readLines(filepath), csv2vector)

  free.counts <- lines[[1]]
  total.counts <- lines[[2]]
  free.areas <- lines[[3]]
  summed.areas <- lines[[4]]
  N <- length(free.counts)

  movement.stats <- parse.movement.data(lines[-(1:4)])
  mean.speeds <- extract.mean.speeds(N, movement.stats)

  summary.stats <- data.table(tick=seq(from=0, length.out=N),
                              mean.speed=mean.speeds,
                              free=free.counts, total=total.counts,
                              nonfree=total.counts - free.counts,
                              free.area=free.areas, summed.area=summed.areas)
  setnames(summary.stats, colnames(summary.stats),  # data.table insists on appending
           sub(".V1", "", colnames(summary.stats))) # ".V1" to column names :S
                                                    # so change it back
  averages <- data.table(speed=mean.na.rm(mean.speeds),
                         free.count.fraction=mean.na.rm(free.counts/total.counts),
                         free.area.fraction=mean.na.rm(free.areas/summed.areas))

  return(list(averages, summary.stats, movement.stats))
}

load.multi.run.averages <- function () {
  averages.list <- list()
  directories <- c("untrimmed-stats/", "trimmed-stats/")
  for (i in c(1, 2)) {
    directory <- directories[i]
    trimmed <- c(FALSE, TRUE)[i]
    for (filename in list.files(directory, pattern="stats-[.0-9]+-[0-9]+\\.csv")) {
      filepath <- paste(directory, filename, sep="")
      matches <- as.numeric(regmatches(filename, gregexpr("[0-9.]*[0-9]", filename))[[1]])
      timestep <- matches[1]
      run <- matches[2]
      l <- load.stats.from.file(filepath)
      averages <- l[[1]]
      averages$timestep <- rep(timestep, nrow(averages))
      averages$run <- rep(run, nrow(averages))
      averages$trimmed <- rep(trimmed, nrow(averages))
      averages.list <- c(averages.list, list(averages))
    }
  }
  averages <- rbindlist(averages.list)
  #averages$timestep <- factor(averages$timestep)
  averages$run <- factor(averages$run)
  return(averages)
}

load.run.details <- function (run) {
  summary.list <- list()
  movements.list <- list()
  directories <- c("untrimmed-stats/", "trimmed-stats/")
  for (i in c(1, 2)) {
    directory <- directories[i]
    trimmed <- c(FALSE, TRUE)[i]
    for (filename in list.files(directory, pattern=paste("stats-[.0-9]+-", run, "\\.csv", sep=""))) {
      filepath <- paste(directory, filename, sep="")
      matches <- as.numeric(regmatches(filename, gregexpr("[0-9.]*[0-9]", filename))[[1]])
      timestep <- matches[1]
      run <- matches[2]
      l <- load.stats.from.file(filepath)
      summary <- l[[2]]
      movements <- l[[3]]
      summary$timestep <- rep(timestep, nrow(summary))
      movements$timestep <- rep(timestep, nrow(movements))
      summary$run <- rep(run, nrow(summary))
      movements$run <- rep(run, nrow(movements))
      summary$trimmed <- rep(trimmed, nrow(summary))
      movements$trimmed <- rep(trimmed, nrow(movements))
      summary.list <- c(summary.list, list(summary))
      movements.list <- c(movements.list, list(movements))
    }
  }
  summary <- rbindlist(summary.list)
  movements <- rbindlist(movements.list)
  #summary$timestep <- factor(summary$timestep)
  #movements$timestep <- factor(movements$timestep)
  summary$run <- factor(summary$run)
  movements$run <- factor(movements$run)
  return(list(summary, movements))
}

########################################################################
# LOAD THE DATA
########################################################################

indexes <- c("timestep", "run", "trimmed")
averages <- load.multi.run.averages()
averages.long <- data.table( melt(averages, indexes) )
averages.long.speed <- data.table( melt(averages, c(indexes, "speed")) )
averages.long.quantiles <- data.table( ddply(averages.long, c("timestep", "trimmed", "variable"),
    function (x) { c(mean=mean(x$value, na.rm=TRUE),
                     min=min(x$value, na.rm=TRUE), max=max(x$value, na.rm=TRUE),
                     q=unname(quantile(x$value, c(0.25,0.75), na.rm=TRUE))) }) )


#    timestep run            variable     value
# 1:       10  34 free.count.fraction 0.8171971  # max
# 2:       10  46 free.count.fraction 0.4984742  # min

l <- load.run.details(46)
run46_summary <- l[[1]]
run46_movements <- l[[2]]

########################################################################
# GRAPHING HELPERS
########################################################################

timestep.as.expression <- function (timestep) {
  if (timestep == 0) {
    return(expression("" %~~% 0 * "    "))
  } else if (0 < timestep && timestep < 1) {
    return(paste("1/", 1/timestep, sep=""))
    #return(as.expression(call("/", 1, 1/timestep)))
  } else if (timestep > 60) {
    return(">60")
    #return(expression(infinity))
  } else {
    return(timestep)
    #return(as.expression(timestep))
  }
}
label_timestep <- function (ignored, timesteps) {
  return(sapply(timesteps, timestep.as.expression))  
}

prettier.variable <- function (variable) {
  if (variable == "speed") {
    return("Average label speed (px/s)")
  } else if (variable == "free.count.fraction") {
    return("Average free label fraction")
  } else if (variable == "free.area.fraction") {
    return("Average free label area fraction")
  }
}
label_variable <- function (ignored, variables) {
  return(sapply(variables, prettier.variable))
}

text.width <- 345/72
text.height <- 550/72
pt <- 25.4/72 # a Postscript point is 1/72 inch, an inch is 25.4 mm

theme_set(theme_bw(base_size=10))
theme_update(axis.text = element_text(size=6), plot.margin = unit(c(0,0,0,0),"lines"))

untrimmed_color = "red"
untrimmed_linewidth = 1*pt
trimmed_color = "black"
trimmed_linewidth = 1*pt

timesteps <- sort(unique(averages$timestep))
timestep_colors <- paste("grey", seq(20, 80, along.with=timesteps), sep="")
selected_timesteps <- timesteps == 0 | timesteps >= 1
timestep_labels <- ifelse(selected_timesteps, label_timestep("", timesteps), "")

time_x_scale <- scale_x_continuous(name="Time (s)", breaks=seq(0, 1536, 128), labels=seq(0, 60, 5))
timestep_x_scale <- scale_x_continuous(name="\nTimestep (s)", breaks=timesteps, labels=timestep_labels)
low_speeds_x_scale <- scale_x_continuous(name="Average label speed (px/s)\n", breaks=seq(0, 150, 10))
log10_speed_ticks <- apply(expand.grid(1:9, 10^seq(-5, 3)), 1, prod)
log10_speed_labels <- ifelse(seq(from=0, along.with=log10_speed_ticks) %% 9 == 0, 
                             gsub("\\.?0*$", "", sprintf("%f", log10_speed_ticks)), "")
high_speeds_x_scale <- scale_x_log10(name="Label speed (px/s)", breaks=log10_speed_ticks, labels=log10_speed_labels)

#high_speeds_y_scale <- scale_y_continuous(breaks=seq(0, 2000, 500), labels=c(0,"",1000,"",2000),
#                                          name="Label speed (px/s)\n") + coord_cartesian(ylim=c(0,2000))
high_speeds_y_scale <- scale_y_sqrt(name="Label speed (px/s)\n", breaks=c(0, 10, seq(10, 50, by=10)^2))
free_count_y_scale <- scale_y_continuous(name = "Free label fraction", breaks=seq(0,1,by=0.25), labels=c(0.0,"",0.5,"",1.0))
free_area_y_scale <- scale_y_continuous(name = "Free label area fraction", breaks=seq(0.5,1,by=0.125), labels=c(0.5,"",0.75,"",1.0))

timestep_color_scale <- scale_color_gradientn(name="Timestep (s)", breaks=timesteps,
		                              colours=timestep_colors, values=rescale(timesteps))
trimming_color_scale <- scale_color_manual(values=c(untrimmed_color, trimmed_color))

facet_by_timestep <- facet_grid(timestep ~ ., labeller=label_timestep)
facet_by_variable <- facet_grid(variable ~ ., scales="free_y", labeller=label_variable) 

no_y_axis_title <- theme(axis.title.y = element_blank())
no_color_legend <- guides(color = FALSE)
no_y_axis_text <- theme(axis.title.y = element_blank(), axis.text.y = element_blank(),
                        axis.ticks.y = element_blank(), axis.line.y = element_blank(),
			plot.margin = unit(c(0,0,0,-1), "lines"))

########################################################################
# DO THE GRAPHING
########################################################################

# Plot average label speeds + freeness for min/25%/mean/75%/max runs
dataset <- averages.long.quantiles
p <- ggplot(dataset, aes(color=trimmed, group=trimmed))
p <- p + facet_by_variable + timestep_x_scale + no_y_axis_title
p <- p + trimming_color_scale + no_color_legend
p <- p + geom_line(aes(timestep, mean))
p <- p + geom_line(aes(timestep, min), linetype="dotted")
p <- p + geom_line(aes(timestep, max), linetype="dotted")
p <- p + geom_line(aes(timestep, q1), linetype="dashed")
p <- p + geom_line(aes(timestep, q2), linetype="dashed")
p
ggsave(file = "graph-quantile-runs.pdf", width=text.width, height=0.8*text.height)

# Plot average label speeds + freeness for all runs (UNTRIMMED)
dataset <- averages.long[trimmed==FALSE]
p <- ggplot(dataset, aes(timestep, value, group=run)) + geom_line(color=alpha("darkgray", 1/4))
#p <- ggplot(dataset, aes(timestep, value, group=run)) + geom_boxplot()
p <- p + facet_by_variable + timestep_x_scale + no_y_axis_title
p <- p + geom_line(aes(timestep, value), data=dataset[run==46], size=0.5*pt, color="black")
p
ggsave(file = "graph-all-runs-untrimmed.pdf", width=text.width, height=0.8*text.height)

# Plot average label speeds + freeness for all runs (TRIMMED)
dataset <- averages.long[trimmed==TRUE]
p <- ggplot(dataset, aes(timestep, value, group=run)) + geom_line(color=alpha("darkgray", 1/4))
#p <- ggplot(dataset, aes(timestep, value, group=run)) + geom_boxplot()
p <- p + facet_by_variable + timestep_x_scale + no_y_axis_title
p <- p + geom_line(aes(timestep, value), data=dataset[run==46], size=0.5*pt, color="black")
p
ggsave(file = "graph-all-runs-trimmed.pdf", width=text.width, height=0.8*text.height)

# Plot all label speeds for one run
dataset <- run46_movements
p <- ggplot(dataset, aes(speed, group=trimmed, color=trimmed)) + geom_density(size=0.5*pt)
p <- p + facet_by_timestep + high_speeds_x_scale
p <- p + trimming_color_scale + no_color_legend + no_y_axis_text
p
ggsave(file = "graph-run46-speed-density.pdf", width=text.width, height=0.8*text.height)

# Plot all label speeds for one run (another way)
dataset <- run46_movements
p <- ggplot(dataset, aes(tick, speed, group=trimmed, color=trimmed)) + geom_line(size=0.5*pt)
p <- p + facet_by_timestep + time_x_scale + high_speeds_y_scale
p <- p + trimming_color_scale + no_color_legend
p
ggsave(file = "graph-run46-speeds.pdf", width=text.width, height=0.8*text.height)

# Plot all label free counts for one run
dataset <- run46_summary
p <- ggplot(dataset, aes(tick, free/total, group=trimmed, color=trimmed)) + geom_line(size=0.5*pt)
p <- p + facet_by_timestep + time_x_scale + free_count_y_scale
p <- p + trimming_color_scale + no_color_legend
p
ggsave(file = "graph-run46-counts.pdf", width=text.width, height=0.8*text.height)

# Plot all label free areas for one run
dataset <- run46_summary
p <- ggplot(dataset, aes(tick, free.area/summed.area, group=trimmed, color=trimmed)) + geom_line(size=0.5*pt)
p <- p + facet_by_timestep + time_x_scale + free_area_y_scale
p <- p + trimming_color_scale + no_color_legend
p
ggsave(file = "graph-run46-areas.pdf", width=text.width, height=0.8*text.height)

# Plot the trade-off between label speed and freeness (TRIMMED)
dataset <- averages.long.speed[trimmed==TRUE]
dataset <- dataset[order(dataset$timestep),]
p <- ggplot(dataset, aes(speed, value, group=run, color=timestep)) + geom_path(size=0.5*pt)
p <- p + facet_by_variable + no_y_axis_title + low_speeds_x_scale 
p <- p + timestep_color_scale + no_color_legend
p
ggsave(file = "graph-speed-freeness-tradeoff-trimmed.pdf", width=text.width, height=0.8*text.height)
