## Usage: Rscript --vanilla plot_energy.R time_stepfile top_percent frames_cutoff
## the output will be text files energy_all.txt energy_topX.txt

#!/usr/bin/env Rscript

args = commandArgs(trailingOnly=TRUE)
if (length(args) !=3) {
  cat("Plot Energy and output various energies as text files\n")
  cat("Usage:\n")
  cat("       Rscript plot_energy.R timestep_file top_percent frames_cutoff\n\n")
  cat("       timestep_file    timestep file from cafemol output\n")
  cat("       top_percent      integer number for the top percent of afmcc structures\n")
  cat("       frames_cutoff    number of first frames to remove for plotting\n\n")
  cat("       Out files are en_all.txt (all frames), en_topNpercent.txt (top frames) and en_plots.pdf\n")
  cat("       total energy vs afmcc plot: red dots showing top X percent pdbs with lowest afmcc\n")
  cat("       normCC x Etot vs afmcc plot: red dots showing top 20 pdbs with lowest  normCC x Etot\n")
  quit()
}


#install.packages("gdata")
library(gdata)
library(dplyr)
library(stringr)
library(bio3d)

args = commandArgs(trailingOnly=TRUE)
#print(args[1])
#length(args)
## test if there is at least one argument: if not, return an error
if (length(args)!=3) {
  stop("Please provide timestep file and top_percent and frames_cutoff", call.=FALSE)
}


#setwd("/Users/bhandariyr/Documents/projects/afm/cafemol/miscen/scripts/energy")
# file = "/Users/bhandariyr/Documents/projects/afm/cafemol/monomer/ring/output/cbl_afm.ts"
file = args[1]
data.ts.lines = readLines(file)[grep("#all", readLines(file))]
data.ts =  data.ts.lines %>%
  strsplit("\\s+") %>%
  do.call(rbind, .) %>%
  data.frame()
dim(data.ts)
colnames(data.ts) = unlist(strsplit(readLines(file)[grep("#unit", readLines(file))], "\\s+"))
#head(data.ts)
colnames(data.ts)[12] = "stack"
#View(data.ts)
data.ts$frame = as.numeric(rownames(data.ts)) - 1

data.ts.sort <- data.ts %>% select("frame", everything())
data.ts.sort = subset(data.ts.sort, select = -c(velet, qscore, rmsd, tempk))
data.ts.sort = data.ts.sort[, -2] ## remove the column named #unit

## Calculate the correlation energy (ce) score, and add to the data frame
afmcc = as.numeric(data.ts.sort$afmcc) #Afmcc score
toten = as.numeric(data.ts.sort$etot)  #total energy

cc_norm = (afmcc - min(afmcc))/(max(afmcc) - min(afmcc))
data.ts.sort$ccNorm = cc_norm
data.ts.sort$ccNormxEtot = cc_norm*toten

## write the data table
write.table(format(data.ts.sort), file = "en_all.txt", sep = "\t", row.names = FALSE, quote = FALSE)
## remove the top frames
frames.to.remove = as.numeric(args[3])
data.ts.sort = data.ts.sort[-(1:frames.to.remove), , drop=FALSE]

## print the energy corresponding to top X percent afmcc sturctures
top.percent = as.numeric(args[2])
# top.percent = 10
topenergy.file = paste("en_top" , top.percent, "percent.txt", sep = "")
top.frames = round(top.percent*nrow(data.ts.sort)/100) 
data.ts.rank = data.ts.sort[order(data.ts.sort$afmcc, decreasing = TRUE), ]
data.ts.rank = data.ts.rank[1:top.frames, ]
write.table(format(data.ts.rank), file = topenergy.file, sep = "\t", row.names = FALSE, quote = FALSE)

## Name the pdf file to plot
pdf("en_plots.pdf") 
plot(data.ts.sort$afmcc, data.ts.sort$etot, xlab = "afmcc", ylab = "total energy", lty = 1)
points(data.ts.rank$afmcc, data.ts.rank$etot, col="red")

data.ts.sort.ce = data.ts.sort[order(data.ts.sort$ccNormxEtot, decreasing = FALSE), ]
data.ts.sort.ce.top20 = data.ts.sort.ce[1:20, ]
plot(data.ts.sort.ce$afmcc, data.ts.sort.ce$ccNormxEtot, xlab = "afmcc", ylab = "ccNorm x Etot", lty = 1)
points(data.ts.sort.ce.top20$afmcc, data.ts.sort.ce.top20$ccNormxEtot, col="red")

dev.off()

