## Usage: Rscript pdbfile dcdfile time_stepfile top_percent
## the output will be a text file of timestep and cc, as well as 
## top percent of the pdbs written out

#!/usr/bin/env Rscript
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)!=4) {
  stop("Please provide the pdb, dcd, timestep file, and top_percent", call.=FALSE)
}
  
library(dplyr)
library(stringr)
library(bio3d)

#file = "/Users/bhandariyr/Documents/projects/afm/cafemol/dimer/twochain/run2_rescaled/output1/cbl_afm.ts"
file = args[3]
data.ts.lines = readLines(file)[grep("#all", readLines(file))]
data.ts =  data.ts.lines %>%
              strsplit("\\s+") %>%
              do.call(rbind, .) %>%
              data.frame()
dim(data.ts)
#View(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


top.percent = as.numeric(args[4])
afmcc_top.file = paste("top_", top.percent, "percent_normCCxEtot.txt", sep = "")
afmcc_all.file = "energy_all.txt"

top.frames = round(top.percent*nrow(data.ts)/100) 
## Calculate the correlation energy (ce) score, and add to the data frame
afmcc = as.numeric(data.ts$afmcc) #Afmcc score
toten = as.numeric(data.ts$etot)  #total energy
cc_norm = (afmcc - min(afmcc))/(max(afmcc) - min(afmcc))
data.ts$ce_score = cc_norm*toten

data.ts.sort = data.ts[order(data.ts$ce_score, decreasing = FALSE), ]
#data.ts.sort = data.ts[order(data.ts$X15, decreasing = TRUE), ]
data.ts.sort = data.ts.sort[1:top.frames,c("frame", "step", "afmcc","etot", "ce_score")]
colnames(data.ts.sort) = c("frame","time.step", "afmcc","etot", "normCC*Etot")
rownames(data.ts.sort) = 1:nrow(data.ts.sort)
write.table(format(data.ts), file = afmcc_all.file, sep = "\t", row.names = FALSE, quote = FALSE)
write.table(format(data.ts.sort), file = afmcc_top.file, sep = "\t", row.names = FALSE, quote = FALSE)

# dcdfile = "~/Documents/projects/afm/cafemol/dimer/twochain/run2_rescaled/output1/cbl_afm.dcd"
# pdbfile = "~/Documents/projects/afm/cafemol/dimer/twochain/run1/pdb/cbl_dimer_offset.pdb"
pdbfile = args[1]
dcdfile = args[2]
dcd <- read.dcd(dcdfile)
pdb <- read.pdb(pdbfile)
ca.inds <- atom.select(pdb)
xyz <- fit.xyz(fixed=pdb$xyz, mobile=dcd,
               fixed.inds=ca.inds$xyz,
               mobile.inds=ca.inds$xyz)
dim(xyz) == dim(dcd)
for (val in data.ts.sort$frame){
  if (val != 0) {
  pdb.name = paste("best_cc_frame_",val, ".pdb", sep = "")
  write.pdb(pdb, dcd[val, ], file=pdb.name)
  }
}

