## Usage: Rscript /data/PNAI/yuba/scripts/analysis2/write_dcd2pdb.R $cafe_infile $top_percent $top_pdb_path
## 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)!=3) {
  stop("Please provide the cafemol input file, and top_percent", call.=FALSE)
}
  
library(dplyr)
library(stringr)
library(bio3d)

cafe.input = args[1]
top.percent = as.numeric(args[2])
top.pdb.path = args[3]
# setwd("/Users/bhandariyr/Documents/projects/afm/cobalamin/miscen/scripts/analysis")

## Read the cafemol input file to parse paths for pdb, dcd and timestep file
flag.pdb = 0 ## Flag to read the pdb file name
for (lines in readLines(cafe.input)){
  if (startsWith(lines, "path ")){
    result.path = trimws(unlist(strsplit(lines,"="))[2])
  }
  if (startsWith(lines, "filename ")){
    result.prefix = trimws(unlist(strsplit(lines,"="))[2])
  }
  if (startsWith(lines, "path_pdb")){
    pdb.path = trimws(unlist(strsplit(lines,"="))[2])
  }
  
  if (startsWith(lines, "<<<< unit_and_state")){
    flag.pdb = 1
  }
  if (flag.pdb == 1 & (startsWith(lines, "1 ") | startsWith(lines, "1-2 "))){
    pdb.name = trimws(unlist(strsplit(lines,"rna"))[2])
    flag.pdb = 0
  }
}
dcdfile = paste(result.path, paste0(result.prefix, ".dcd"), sep = "/")
timestep.file = paste(result.path, paste0(result.prefix, ".ts"), sep = "/")
pdbfile = paste(pdb.path, pdb.name, sep = "/")
print(paste("PDB DCD and Timestep files read are: ", pdbfile, dcdfile, timestep.file, sep=", "))

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

afmcc.file = paste("top_", top.percent, "percent_cc.txt", sep = "")
afmcc.file.path = paste(top.pdb.path, afmcc.file, sep = "/")
top.frames = round(top.percent*nrow(data.ts)/100) 
data.ts.sort = data.ts[order(data.ts$X15, decreasing = TRUE), ]
data.ts.sort$frame = as.numeric(rownames(data.ts.sort)) - 1
data.ts.sort = data.ts.sort[1:top.frames,c("frame", "X2", "X15")]
colnames(data.ts.sort) = c("frame","time.step", "afmcc")
rownames(data.ts.sort) = 1:nrow(data.ts.sort)
write.table(format(data.ts.sort), file = afmcc.file.path, 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 = "")
  pdb.name.path = paste(top.pdb.path, pdb.name, sep = "/")
  write.pdb(pdb, dcd[val, ], file=pdb.name.path)
  }
}

