## Usage: Rscript --vanilla analysis_bio3d.R cafemol_inputfile atom_type
## atom_type = choose one (sugar, phosphate, base)
## the output will be text files afmcc.txt, rmsf.txt and plot file bio3d_plots.pdf
## atom_type is used for calculating rmsf and residue-residue cross correlation


library(dplyr)
library(stringr)
library(bio3d)

#!/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)!=2) {
  stop("Please provide cafemol inputfile and atom_type", call.=FALSE)
}

## File arguments
# atom.type = "phosphate"
# cafe.input = "/Users/bhandariyr/Documents/projects/afm/cobalamin/miscen/scripts/stats/h20/input.inp"
cafe.input = args[1]
atom.type = args[2]
# 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(pdbfile); print(dcdfile); print(timestep.file)

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)
data.ts$frame = as.numeric(rownames(data.ts)) - 1
data.ts.sort = data.ts[,c("frame", "X2", "X15")]
colnames(data.ts.sort) = c("frame","time.step", "afmcc")

## Name the pdf file to plot
out.file = "bio3d_plots.pdf"
pdf(out.file) 
plot(data.ts.sort[2:nrow(data.ts.sort),]$frame, data.ts.sort[2:nrow(data.ts.sort),]$afmcc, xlab = "Frame No.", ylab = "correlation")
write.table(format(data.ts.sort), file = "afmcc.txt", sep = "\t", row.names = FALSE, quote = FALSE)

dcd <- read.dcd(dcdfile)
pdb <- read.pdb(pdbfile)

if (atom.type == "sugar") {
   ca.inds <- atom.select(pdb, elety = "S")
}   else if (atom.type == "phosphate") {
    ca.inds <- atom.select(pdb, elety = "P")
}  else if (atom.type == "base") {
    ca.inds <- atom.select(pdb, elety = c("Ab","Ub", "Gb","Cb"))
}  else if (atom.type == "all") {
    ca.inds <- atom.select(pdb)
}

ca.inds
# ca.inds <- atom.select(pdb)
# Print details of the first few selected atoms
# pdb$atom[ca.inds$atom, ]

xyz <- fit.xyz(fixed=pdb$xyz, mobile=dcd,
               fixed.inds=ca.inds$xyz,
               mobile.inds=ca.inds$xyz)
# dim(xyz) == dim(dcd)

## Root mean square deviation
#rd <- rmsd(xyz[1,ca.inds$xyz], xyz[,ca.inds$xyz])
#rmsd.frame = data.frame(frame = 1:length(rd), rmsd = rd)
#write.table(rmsd.frame, file = "rmsd.txt", sep = "\t", row.names = FALSE, quote = FALSE)
#plot(rd, typ="l", ylab="RMSD", xlab="Frame No.")
#points(lowess(rd), typ="l", col="red", lty=2, lwd=2)
#hist(rd, breaks=40, freq=FALSE, main="RMSD Histogram", xlab="RMSD", ylab = "population")
#lines(density(rd), col="gray", lwd=3)
## root mean square flucturation
rf <- rmsf(xyz[,ca.inds$xyz])
rmsf.frame = data.frame(residue = 1:length(rf), rmsf = rf)
write.table(rmsf.frame, file = "rmsf.txt", sep = "\t", row.names = FALSE, quote = FALSE)
plot(rf, ylab="RMSF", xlab="Residue Position", typ="l")
## PCA analysis
# pc <- pca.xyz(xyz[,ca.inds$xyz])
# plot(pc, col=bwr.colors(nrow(xyz)), use.svd = TRUE )
## Cross correlation analysis
cij<-dccm(xyz[,ca.inds$xyz])
plot(cij)

# Close the pdf file
dev.off() 
