## Find the top Correlation-Energy (ce) Score structures for the top kappa cohort
## Usage: Rscript input_folder top_percent pdbfile rmsd_file 

args = commandArgs(trailingOnly=TRUE)
if (length(args) != 4) {
  cat("Extract PDBs and plot various energy type based on correlation-energy (ce) score\n")
  cat("ce score is the product of normalized afmcc and normalized total Energy\n")
  cat("Usage:\n")
  cat("       Rscript get_stats_CE.R input_folder top_percent pdbfile rmsd_file\n\n")
  cat("       input_folder    this folder should contain the timestep and dcd file\n")
  cat("       top_percent     integer representing the top percent of PDBs based on ce score\n")
  cat("       pdbfile         pdbfile is a template PDB file for extracting coordinates from dcd\n")
  cat("       rmsd_file       path to rmsd_file\n\n")
  cat("       \n")
  cat("       Folder top_Xpercent_ce_pdbs and plot top_Xpercent_ce_score_plots.pdf are generated\n")
  cat("       ce_score and afmcc against RMSDs are plotted\n")
  quit()
}


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

#cohort = "/Users/bhandariyr/Documents/projects/afm/cobalamin/miscen/scripts/stats/contact10_5/output1_0"
cohort = args[1]
top_percent = as.numeric(args[2])
pdbfile = args[3]
rmsd_path = args[4]

ext1 = paste(cohort,"*.ts", sep = "/")
file.ts = Sys.glob(ext1)
ext2 = paste("top_",top_percent, "%_pdbs", sep = "")
#rmsd_path = paste(cohort,ext2,"rmsd.txt", sep = "/")
#print(rmsd_path)
df.rmsd = read.table(rmsd_path, header = TRUE)
#head(df.rmsd)
## Read and clean timestep file
data.ts.lines = readLines(file.ts)[grep("#all", readLines(file.ts))]
data.ts =  data.ts.lines %>%
  strsplit("\\s+") %>%
  do.call(rbind, .) %>%
  data.frame()
dim(data.ts)
colnames(data.ts) = unlist(strsplit(readLines(file.ts)[grep("#unit", readLines(file.ts))], "\\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())
## subset to remove frame zero
data.ts.sort = data.ts.sort[2:nrow(data.ts.sort),]
## add the rmsd column
data.ts.sort$calc_rmsd = df.rmsd$rmsd

## 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))
toten_norm = (toten - max(toten))/(max(toten) - min(toten))
data.ts.sort$ce_score = cc_norm*toten_norm
#View(data.ts.sort)

## Remove the first frame as it is almost an outlier
data.ts.sort = data.ts.sort[-1,]

## Get the top X percent structures based on ce score
top.frames = round(top_percent*nrow(data.ts.sort)/100)
data.ts.rank = data.ts.sort[order(data.ts.sort$ce_score, decreasing = FALSE), ]
data.ts.top = data.ts.rank[1:top.frames, ]
#View(data.ts.top)

## Do the plotting
dir_name = paste("top_", top_percent, "percent_ce_pdbs", sep = "")
dir.create(file.path(dir_name), recursive = FALSE, showWarnings = TRUE)

plot.name = paste("top",top_percent,"percent_ce_score_plots.pdf",sep = "")
plot.path = paste(dir_name, plot.name, sep = "/")
pdf(plot.path)
library(ggplot2)
p1 = ggplot(data.ts.sort, aes(x=calc_rmsd, y=ce_score)) +
  geom_point() +
  theme_classic(base_size = 10)+ ylab("ce_score") + 
  ggtitle("all frames") +
  xlab("rmsd")

p2 = ggplot(data.ts.top, aes(x=calc_rmsd, y=ce_score)) +
  geom_point() +
  theme_classic(base_size = 10)+ ylab("ce_score") + 
  ggtitle(paste("top_", top_percent, "_percent_frames", sep = "")) +
  xlab("rmsd")

data.ts.sort$afmcc = as.numeric(data.ts.sort$afmcc)
data.ts.top$afmcc = as.numeric(data.ts.top$afmcc)

p3 = ggplot(data.ts.sort, aes(x=calc_rmsd, y=afmcc)) +
  geom_point() +
  theme_classic(base_size = 10)+ ylab("afmcc") +
  ggtitle("all frames") +
  xlab("rmsd")

data.ts.sort.afmcc = data.ts.sort[order(data.ts.sort$afmcc, decreasing = TRUE), ]
data.ts.top.afmcc = data.ts.sort.afmcc[1:60,]
p4 = ggplot(data.ts.top.afmcc, aes(x=calc_rmsd, y=afmcc)) +
  geom_point() +
  theme_classic(base_size = 10)+ ylab("afmcc") +
  ggtitle(paste("top_", top_percent, "_percent_frames", sep = "")) +
  xlab("rmsd")

grid.arrange(p1,p2,p3,p4, nrow=2, ncol=2)

dev.off()


#cor.test(data.ts.sort$calc_rmsd, data.ts.sort$ce_score, method = "pearson")

## Output the pdb files corresponding to top X percent ce scores.
ext.dcd = paste(cohort,"*.dcd", sep = "/")

dcdfile = Sys.glob(ext.dcd)
# pdbfile = "/Users/bhandariyr/Documents/projects/afm/cobalamin/miscen/scripts/stats/contact10_5/pdb/cgRNaseP_FL_ori_center.pdb"
# pdbfile = args[1]
# dcdfile = args[2]
dcd <- read.dcd(dcdfile)
pdb <- read.pdb(pdbfile)
# ca.inds <- atom.select(pdb)
#setwd = "/Users/bhandariyr/Documents/projects/afm/cobalamin/miscen/scripts/stats/contact10_5"

for (val in data.ts.top$frame){
  if (val != 0) {
    pdb.name = paste("best_cc_frame_",val, ".pdb", sep = "")
    pdb.path = paste(dir_name, pdb.name, sep = "/")
    write.pdb(pdb, dcd[val, ], file=pdb.path)
  }
}

file1 = paste(dir_name, "energy_all.txt", sep = "/")
file2 = paste(dir_name, "/energy_top_", top_percent, "percent.txt", sep = "")
write.table(format(data.ts.sort), file = file1, sep = "\t", row.names = FALSE, quote = FALSE)
write.table(format(data.ts.top), file = file2, sep = "\t", row.names = FALSE, quote = FALSE)

print(paste("Correlation coef. of rmsd vs ce_score for all frames ",
      cor(data.ts.sort$calc_rmsd, data.ts.sort$ce_score, method = "pearson"), sep = ""))
print(paste("Correlation coef. of rmsd vs ce_score for top_", top_percent, "_percent_frames: ",
      cor(data.ts.top$calc_rmsd, data.ts.top$ce_score, method = "pearson"), sep = ""))




