## Calculate the statistics given various kappa values
## Usage:
## Rscript  get_stats.R output1_0 output1_1 output1_2 output1_3 output1_4 output1_5 output1_6 output1_7 output1_8 output1_9 output1_10 [0.5,2,3,4,5,6,7,8,9,10,11] 10

args = commandArgs(trailingOnly=TRUE)
if (length(args) == 0) {
  cat("Plot mean and standard deviation of various energy terms of the top X percent afmcc strutures for each kappa values\n")
  cat("Usage:\n")
  cat("       Rscript get_stats.R output1 output2 output3  [k1,k2,k3,..] top_percent\n\n")
  cat("       output1 output2 output3 etc. are the folder names corresonding to kappa values\n")
  cat("       [k1,k2,k3,..]    k1, k2, k3 etc are the kappa values\n")
  cat("       top_percent      integer number for the top percent of afmcc structures\n\n")
  cat("       The number of folders and kappa values should match\n")
  cat("       Each folder contains the timestep file and another folder named top_X%_pdbs (contains the rmsd file)\n")
  cat("       Please see example at /data/PNAI/yuba/cafe_runs/stats/contact10_5\n")
  cat("       The final output is a pdf file with summary statistics figures for energy terms such as correlation,\n")
  cat("       total energy, local energy, etc against kappa\n")
  quit()
}

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

#print(args[1])
#print(length(args))

ind <- c(unlist(strsplit(str_sub(args[length(args)-1], 2, -2),",")))
#print(length(ind))
#print(length(args)-2)
if (length(ind) != (length(args)-2)) {
  stop("Number of input folders must match the number of kappa values", call.=FALSE)
}
#setwd("/Users/bhandariyr/Documents/projects/afm/cobalamin/miscen/scripts/stats/contact10_5")
top_percent = as.numeric(args[length(args)])
# print(paste("top percent is", top_percent,sep = ": "))

## initialize data frames for each energy term
df_afmcc_tot = data.frame(cc=numeric(), id=integer())
df_afmcc_mean_tot = data.frame(cc=numeric(), id=integer(), sd=numeric())
df_toten_mean_tot = data.frame(cc=numeric(), id=integer(), sd=numeric())
df_local_mean_tot = data.frame(cc=numeric(), id=integer(), sd=numeric())
df_go_mean_tot = data.frame(cc=numeric(), id=integer(), sd=numeric())
df_rmsd_mean_tot = data.frame(cc=numeric(), id=integer(), sd=numeric())
df_repul_mean_tot = data.frame(cc=numeric(), id=integer(), sd=numeric())
df_stack_mean_tot = data.frame(cc=numeric(), id=integer(), sd=numeric())
df_hbond_mean_tot = data.frame(cc=numeric(), id=integer(), sd=numeric())
df_elect_mean_tot = data.frame(cc=numeric(), id=integer(), sd=numeric())
df_afmfit_mean_tot = data.frame(cc=numeric(), id=integer(), sd=numeric())
df_stage_mean_tot = data.frame(cc=numeric(), id=integer(), sd=numeric())

## Loop through the given folders
# folders = list("output1_0","output1_1", "output1_2","output1_3","output1_4","output1_5",
#                "output1_6","output1_7","output1_8","output1_9","output1_10")
for (i in 1:(length(args)-2)){
  
ext1 = paste(args[i],"*.ts", sep = "/")
file = Sys.glob(ext1)
ext2 = paste("top_",top_percent, "%_pdbs", sep = "")
rmsd_path = paste(args[i],ext2,"rmsd_all.txt", sep = "/")
#print(rmsd_path)
df.rmsd = read.table(rmsd_path, header = TRUE)

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())
## 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
# View(data.ts.sort)
## print the energy corresponding to top X percent afmcc sturctures
#top_percent = as.numeric(args[2])

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)
#print(head(data.ts.rank))

## afm correlation
df_afmcc = data.frame(cc=as.numeric(data.ts.rank$afmcc), id = rep.int(i, nrow(data.ts.rank)))
df_afmcc_tot = rbind(df_afmcc_tot, df_afmcc)
df_afmcc_mean = data.frame(cc=mean(as.numeric(data.ts.rank$afmcc)), id = i, sd = sd(as.numeric(data.ts.rank$afmcc)))
df_afmcc_mean_tot = rbind(df_afmcc_mean_tot, df_afmcc_mean)
## total energy
df_toten_mean = data.frame(cc=mean(as.numeric(data.ts.rank$etot)), id = i, sd = sd(as.numeric(data.ts.rank$etot)))
df_toten_mean_tot = rbind(df_toten_mean_tot, df_toten_mean)
## local energy
df_local_mean = data.frame(cc=mean(as.numeric(data.ts.rank$local)), id = i, sd = sd(as.numeric(data.ts.rank$local)))
df_local_mean_tot = rbind(df_local_mean_tot, df_local_mean)
## go energy term
df_go_mean = data.frame(cc=mean(as.numeric(data.ts.rank$go)), id = i, sd = sd(as.numeric(data.ts.rank$go)))
df_go_mean_tot = rbind(df_go_mean_tot, df_go_mean)
## rmsd term calculated by external script
df_rmsd_mean = data.frame(cc=mean(as.numeric(data.ts.rank$calc_rmsd)), id = i, sd = sd(as.numeric(data.ts.rank$calc_rmsd)))
df_rmsd_mean_tot = rbind(df_rmsd_mean_tot, df_rmsd_mean)
## repulsion energy
df_repul_mean = data.frame(cc=mean(as.numeric(data.ts.rank$repul)), id = i, sd = sd(as.numeric(data.ts.rank$repul)))
df_repul_mean_tot = rbind(df_repul_mean_tot, df_repul_mean)
## stack energy
df_stack_mean = data.frame(cc=mean(as.numeric(data.ts.rank$stack)), id = i, sd = sd(as.numeric(data.ts.rank$stack)))
df_stack_mean_tot = rbind(df_stack_mean_tot, df_stack_mean)
## hbond energy
df_hbond_mean = data.frame(cc=mean(as.numeric(data.ts.rank$hbond)), id = i, sd = sd(as.numeric(data.ts.rank$hbond)))
df_hbond_mean_tot = rbind(df_hbond_mean_tot, df_hbond_mean)
## electrostatic energy
df_elect_mean = data.frame(cc=mean(as.numeric(data.ts.rank$elect)), id = i, sd = sd(as.numeric(data.ts.rank$elect)))
df_elect_mean_tot = rbind(df_elect_mean_tot, df_elect_mean)
## afmfit energy
df_afmfit_mean = data.frame(cc=mean(as.numeric(data.ts.rank$afmfit)), id = i, sd = sd(as.numeric(data.ts.rank$afmfit)))
df_afmfit_mean_tot = rbind(df_afmfit_mean_tot, df_afmfit_mean)
## afm stage energy
df_stage_mean = data.frame(cc=mean(as.numeric(data.ts.rank$stage)), id = i, sd = sd(as.numeric(data.ts.rank$stage)))
df_stage_mean_tot = rbind(df_stage_mean_tot, df_stage_mean)

}
cc_norm = (df_afmcc_mean_tot$cc-min(df_afmcc_mean_tot$cc))/(max(df_afmcc_mean_tot$cc)-min(df_afmcc_mean_tot$cc))
toten_norm = (df_toten_mean_tot$cc-max(df_toten_mean_tot$cc))/(max(df_toten_mean_tot$cc)-min(df_toten_mean_tot$cc))
df_toten_mean_tot$cc_toten = cc_norm*toten_norm

#print(df_toten_mean_tot)
#print(df_afmcc_mean_tot)
pdf(paste("top",top_percent,"percent_en_plots.pdf",sep = ""))
library(ggplot2)
p1 = ggplot(df_afmcc_mean_tot, aes(x=id, y=cc)) +
       geom_pointrange(aes(ymin=cc-sd, ymax=cc+sd)) +
       theme_classic(base_size = 10)+ ylab("correlation") + 
       scale_x_discrete(name = "kappa", limits = ind)

p2 = ggplot(df_toten_mean_tot, aes(x=id, y=cc )) +
       geom_pointrange(aes(ymin=cc-sd, ymax=cc+sd)) +
       theme_classic(base_size = 10)+ ylab("total energy") + 
       scale_x_discrete(name = "kappa", limits = ind)

p3 = ggplot(df_toten_mean_tot, aes(x=id, y=cc_toten )) +
  geom_point() +
  theme_classic(base_size = 10)+ ylab("score") + 
  scale_x_discrete(name = "kappa", limits = ind)

p4 = ggplot(df_rmsd_mean_tot, aes(x=id, y=cc )) +
       geom_pointrange(aes(ymin=cc-sd, ymax=cc+sd)) +
       theme_classic(base_size = 10)+ ylab("rmsd") + 
       scale_x_discrete(name = "kappa", limits = ind)

p5 = ggplot(df_local_mean_tot, aes(x=id, y=cc )) +
       geom_pointrange(aes(ymin=cc-sd, ymax=cc+sd)) +
       theme_classic(base_size = 10)+ ylab("local energy") + 
       scale_x_discrete(name = "kappa", limits = ind)

p6 = ggplot(df_go_mean_tot, aes(x=id, y=cc )) +
  geom_pointrange(aes(ymin=cc-sd, ymax=cc+sd)) +
  theme_classic(base_size = 10)+ ylab("go energy") + 
  scale_x_discrete(name = "kappa", limits = ind)

p7 = ggplot(df_repul_mean_tot, aes(x=id, y=cc )) +
  geom_pointrange(aes(ymin=cc-sd, ymax=cc+sd)) +
  theme_classic(base_size = 10)+ ylab("repulsive energy") + 
  scale_x_discrete(name = "kappa", limits = ind)

p8 = ggplot(df_stack_mean_tot, aes(x=id, y=cc )) +
  geom_pointrange(aes(ymin=cc-sd, ymax=cc+sd)) +
  theme_classic(base_size = 10)+ ylab("stack energy") + 
  scale_x_discrete(name = "kappa", limits = ind)

p9 = ggplot(df_hbond_mean_tot, aes(x=id, y=cc )) +
  geom_pointrange(aes(ymin=cc-sd, ymax=cc+sd)) +
  theme_classic(base_size = 10)+ ylab("hbond energy") + 
  scale_x_discrete(name = "kappa", limits = ind)

p10 = ggplot(df_elect_mean_tot, aes(x=id, y=cc )) +
  geom_pointrange(aes(ymin=cc-sd, ymax=cc+sd)) +
  theme_classic(base_size = 10)+ ylab("electrostatic energy") + 
  scale_x_discrete(name = "kappa", limits = ind)

p11 = ggplot(df_afmfit_mean_tot, aes(x=id, y=cc )) +
  geom_pointrange(aes(ymin=cc-sd, ymax=cc+sd)) +
  theme_classic(base_size = 10)+ ylab("afmfit energy") + 
  scale_x_discrete(name = "kappa", limits = ind)

p12 = ggplot(df_stage_mean_tot, aes(x=id, y=cc )) +
  geom_pointrange(aes(ymin=cc-sd, ymax=cc+sd)) +
  theme_classic(base_size = 10)+ ylab("stage energy") + 
  scale_x_discrete(name = "kappa", limits = ind)

grid.arrange(p1,p2,p3,p4,p5,p6,p7,p8,p9,p10,p11,p12, nrow=4, ncol=3)
p1+theme_classic(base_size = 15)
p2+theme_classic(base_size = 15)
p3+theme_classic(base_size = 15)
p4+theme_classic(base_size = 15)
p5+theme_classic(base_size = 15)
p6+theme_classic(base_size = 15)
p7+theme_classic(base_size = 15)
p8+theme_classic(base_size = 15)
p9+theme_classic(base_size = 15)
p10+theme_classic(base_size = 15)
p11+theme_classic(base_size = 15)
p12+theme_classic(base_size = 15)

dev.off()




