## Calculate the statistics for 3D plot of Kappa, afmcc and Energy
## Usage:
## Rscript /data/PNAI/yuba/scripts/analysis/get_Rgstats_3D_short.R kappa_file top_percent1 %frames_cutoff

args = commandArgs(trailingOnly=TRUE)
if (length(args) != 3) {
  cat("Calculate aggregate statistics of Kappa, afmcc, Rg, ccNorm, ccNormxEtot for top X% afmcc structures \n")
  cat("Usage:\n")
  cat("       Rscript get_Rgstats_3D_short.R kappa_file top_percent %frames_cutoff\n")
  cat("       kappa_file        file with list of kappa values\n")
  cat("       top_percent       top X percent structures with lowest afmCC\n")
  cat("       %frames_cutoff    percent of early frames to remove\n")
  cat("       Output files:     3Dstats_filterEnRg_topXpercent.txt\n") 
  quit()
  }

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


kfile = args[1]
conn = file(kfile,open="r")
kvalues = readLines(conn)
input.folders = paste0("k", kvalues)
ind = kvalues
#input.folders = c(unlist(strsplit(str_sub(args[1], 2, -2),",")))
#ind <- c(unlist(strsplit(str_sub(args[2], 2, -2),",")))

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

## initialize data frames for each energy term
df_corr_tot = data.frame(kappa = numeric(), correlation = numeric()) ## correlation between afmcc and nromCCxEtot
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())
df_CE_mean_tot = data.frame(cc=numeric(), id=integer(), sd=numeric())
df_radg_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")
path = "/data/PNAI/yxwang/CafeMol/cafe_runs/trRNaseP/experimental/f5511199d-10pc-1m-1A"
for (i in 1:(length(input.folders))){
  
ext1 = paste(path, 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])

## calculate correlation between afmcc and nromCCxEtot
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))
ccNormxEtot = cc_norm*toten
data.ts.sort$ccNorm = cc_norm
data.ts.sort$ccNormxEtot = ccNormxEtot

## remove the top frames
frames.cutoff.percent = as.numeric(args[3])
frames.cutoff = round(frames.cutoff.percent*nrow(data.ts.sort)/100)

#print(dim(data.ts.sort))
data.ts.sort = data.ts.sort[-(1:frames.cutoff), , drop=FALSE]
print(paste("current kappa is: ", i, sep = ""))
#print(frames.cutoff)
#print(dim(data.ts.sort))

topenergy.file = paste("en_top" , top_percent1, "percent.txt", sep = "")
top.frames = round(top_percent1*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))


corr = cor(afmcc, ccNormxEtot, method = "pearson")
df_corr = data.frame(kappa = as.numeric(ind[i]), correlation = corr) 
df_corr_tot = rbind(df_corr_tot, df_corr)

## 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)
## ccNormxEtot
df_CE_mean = data.frame(cc=mean(as.numeric(data.ts.rank$ccNormxEtot)), id = i, sd = sd(as.numeric(data.ts.rank$ccNormxEtot)))
df_CE_mean_tot = rbind(df_CE_mean_tot, df_CE_mean)
df_radg_mean = data.frame(cc=mean(as.numeric(data.ts.rank$radg)), id = i, sd = sd(as.numeric(data.ts.rank$radg)))
df_radg_mean_tot = rbind(df_radg_mean_tot, df_radg_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_mean = (df_afmcc_mean_tot$cc-min(df_afmcc_mean_tot$cc))/(max(df_afmcc_mean_tot$cc)-min(df_afmcc_mean_tot$cc))
ccNormxEtot.mean = cc_norm_mean*df_toten_mean_tot$cc 
matrix.3d = data.frame(kappa = ind, afmcc = df_afmcc_mean_tot$cc, afmcc.sd = df_afmcc_mean_tot$sd,
                       etot = df_toten_mean_tot$cc, etot.sd = df_toten_mean_tot$sd,
                       radg.mean = df_radg_mean_tot$cc, radg.sd = df_radg_mean_tot$sd,
                       ccNorm = cc_norm_mean, ccNormxEtot = ccNormxEtot.mean
                       )


file.3d = paste("3Dstats_filterEnRg_top", top_percent1, "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 = ".")
