## Calculate the statistics for 3D plot of Kappa, afmcc and Energy
## Usage:
## Rscript /data/PNAI/yuba/scripts/analysis/get_stats_3D.R [k1,k10,k17] [1,10,17] 5

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

args = commandArgs(trailingOnly=TRUE)
#print(args[1])
#print(length(args))

input.folders = c(unlist(strsplit(str_sub(args[1], 2, -2),",")))
ind <- c(unlist(strsplit(str_sub(args[2], 2, -2),",")))

#print(length(ind))
#print(length(args)-2)
if (length(args) != 3) {
  stop("Please provide input folders, kappa values and top_percent", call.=FALSE)
}
if (length(input.folders) != length(ind)) {
  stop("total number of input folders and kappa values are different, they should match", call.=FALSE)
}

#setwd("/Users/bhandariyr/Documents/projects/afm/cobalamin/miscen/scripts/stats/contact10_5")
top_percent = as.numeric(args[3])
# print(paste("top percent is", top_percent,sep = ": "))

## initialize data frames for each energy term
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_go_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_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(input.folders))){
  
ext1 = paste(input.folders[i],"output", "*.ts", sep = "/")
file = Sys.glob(ext1)
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),]
# 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_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+go energy
local_go_en = as.numeric(data.ts.rank$local) + as.numeric(data.ts.rank$go)
df_local_go_mean = data.frame(cc=mean(local_go_en), id = i, sd = sd(local_go_en))
df_local_go_mean_tot = rbind(df_local_go_mean_tot, df_local_go_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)
## 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))

matrix.3d = data.frame(kappa = ind, afmcc = df_afmcc_mean_tot$cc, afmcc.sd = df_afmcc_mean_tot$sd,
                       energy.tot = df_toten_mean_tot$cc, energy.tot.sd = df_toten_mean_tot$sd,
                       energy.local=df_local_mean_tot$cc, energy.local.sd = df_local_mean_tot$sd,
                       energy.go=df_go_mean_tot$cc, energy.go.sd = df_go_mean_tot$sd,
                       energy.hbond=df_hbond_mean_tot$cc, energy.hbond.sd = df_hbond_mean_tot$sd,
                       energy.stack=df_stack_mean_tot$cc, energy.stack.sd = df_stack_mean_tot$sd,
                       energy.afmfit=df_afmfit_mean_tot$cc, energy.afmfit.sd = df_afmfit_mean_tot$sd,
                       energy.elect=df_elect_mean_tot$cc, energy.elect.sd = df_elect_mean_tot$sd,
                       cc.norm = cc_norm)


file.3d = paste("3Dstats_top", top_percent, "percent.txt", sep = "")
write.table(format(matrix.3d), file = file.3d, sep = "\t", row.names = FALSE, quote = FALSE)

## Reading data table
## data = read.delim(file, header = TRUE, sep = "\t", dec = ".")
