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

args = commandArgs(trailingOnly=TRUE)
if (length(args) == 0) {
  cat("Calculate afmCC and Energy statistics including normalization \n")
  cat("for the top X percent structures with lowest afmCC\n")
  cat("Usage:\n")
  cat("       Rscript get_stats_3D_gbfilter.R kappa_file top_percent path_to_kappa_folders(optional)\n")
  cat("       kappa_file    file with list of kappa values\n")
  cat("       top_percent   top X percent structures with lowest afmCC\n\n")
  cat("       generates a plot and table of correlation between afmCC and (normalzed afmCC x total Energy) for each kappa\n")
  cat("       first finds the top X percent frames with lowest afmCC, and calculates the mean afmCC and mean Etot from them \n")
  cat("       secondly find the top 15 kappa cohorts based on minimum norm(mean afmCC)x(mean Etot) with normalization done across all kappas \n") 
  cat("       thirdly for above top 15 kappa cohorts, calculates nomalization of afmCC inside each kappa cohort and globally across top 15 cohorts \n")
  cat("       Output files: 3Dstats_topXpercent.txt, correlation.txt, correlation.pdf, En_topN.txt, normEn_topN.txt,gnormEn_topN.txt, rg_topN.txt\n")
  cat("       normEn_topN.txt (energy normalized within each kappa), gnormEn_topN.txt (energy normalized across top 15 kappa cohorts \n")
  cat("       If path_to_kappa_folders is given, only the kappa folders in this path will be analyzed, default path is current directory\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_percent = 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=character(), 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")

#path = "/data/PNAI/yxwang/CafeMol/cafe_runs/trRNaseP/experimental/f5511199d-10pc-1m-1A/"
## Define path if it exists
if (length(args) == 3) {
  path = args[3]
}

for (i in 1:(length(input.folders))){

## check if path exists
if (exists("path")){
ext1 = paste(path, input.folders[i],"output", "*.ts", sep = "/")
} else {
ext1 = paste(input.folders[i],"output", "*.ts", sep = "/")
}
print(paste0("Now reading ", ext1))
file = Sys.glob(ext1)
data.ts.lines = readLines(file)[grep("#all", readLines(file))]
#print("Hi I am here")
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)

## Check for the existence of unstable calculation by looking at last Etot line of timestep file
#print(data.ts.sort[nrow(data.ts.sort),c("etot","afmcc", "afmfit") ])
last.Enstatus = any(is.na(as.numeric(data.ts.sort$etot))) 
if (last.Enstatus) {
  print(paste0("found NA in kappa: ", input.folders[i]))
  print(paste("skipping", input.folders[i], "because of unstable cafemol calculation", sep = " "))
  next #go to next iteration
}
## 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))

## 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
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 = input.folders[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)

#print(paste0("selected kappa is: ", input.folders[i]))

}

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))
normEtot = cc_norm*df_toten_mean_tot$cc

matrix.3d = data.frame(kappa = df_afmcc_mean_tot$id, 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, ccNormxEtot = normEtot)
## remove NAs
#print(paste0("dimension of matrix.3d is: ", dim(matrix.3d)))
matrix.3d = matrix.3d[complete.cases(matrix.3d), ]
print(paste("dimension of matrix.3d is: row by col: ", nrow(matrix.3d), "by", ncol(matrix.3d), sep = " "))
## sort based on ccNormxEtot
matrix.3d = matrix.3d[order(matrix.3d$ccNormxEtot, decreasing = FALSE), ]
print(df_afmcc_mean_tot$id)
print(matrix.3d$kappa)

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

## remove NAs
df_corr_tot = df_corr_tot[complete.cases(df_corr_tot), ]
## Plot correlation
pdf("correlation.pdf")
plot(df_corr_tot$kappa, df_corr_tot$correlation, xlab = "kappa", ylab = "correlation", lty = 1)
dev.off()

## select top 15 structures based on ccNormxEtot
top_cohorts = 15
rg.list = list()
En.list = list()
normEn.list = list()
gbnormEn.list = list()
count = 1 ## counter for  list
print(paste("######## Now reading the top", top_cohorts, "cohorts ##########", sep = " "))

## Find global minimum and maximum of afmcc among top 15 structures
cc.gmin = 2
cc.gmax = -2
for (i in 1:top_cohorts){
  kappa = matrix.3d$kappa[i]
  print(kappa)
  ## check if path exists
  if (exists("path")){
  ext2 = paste(path, kappa,"output/en_all.txt", sep = "/")
  } else {
  ext2 = paste(kappa,"/output/en_all.txt", sep = "")
  }
  data.en = read.delim(ext2, header = TRUE, sep = "\t", dec = ".")
  cc.min = min(as.numeric(data.en$afmcc))
  cc.max = max(as.numeric(data.en$afmcc))
  #print(paste0("Now reading ", ext2))
  #print(paste(cc.min, cc.max, sep=" "))
  if(cc.min < cc.gmin){
    cc.gmin = cc.min
  }
  if (cc.max > cc.gmax){
    cc.gmax = cc.max
  }
}
print(paste("minimum AFMCC among top", top_cohorts, "cohorts is:", cc.gmin, sep=" "))
print(paste("maximum AFMCC among top", top_cohorts, "cohorts is:", cc.gmax, sep=" "))
######

for (i in 1:top_cohorts){
  #print(matrix.3d$kappa[i])
  kappa = matrix.3d$kappa[i]
  ## check if path exists
  if (exists("path")){
  ext2 = paste(path, kappa,"output/en_all.txt", sep = "/")
  } else {
  ext2 = paste(kappa,"/output/en_all.txt", sep = "")
  }
  print(paste0("Now reading ", ext2))
  data.en = read.delim(ext2, header = TRUE, sep = "\t", dec = ".")
 
  df.rg = data.frame(rg = data.en$radg)
  col1.rg = paste0("rg_",kappa)
  colnames(df.rg) = c(col1.rg)
  rg.list[[count]] = df.rg
  #print(dim(df.rg))  
 
  df.En = data.frame(cc = data.en$afmcc, etot = data.en$etot)
  col1.En = paste0("CC_",kappa)
  col2.En = paste0("Etot_",kappa)
  colnames(df.En) = c(col1.En, col2.En)
  En.list[[count]] = df.En

  df.normEn = data.frame(cc = data.en$ccNorm, etot = data.en$ccNormxEtot)
  col1.nEn = paste0("normCC_",kappa)
  col2.nEn = paste0("normEtot_",kappa)
  colnames(df.normEn) = c(col1.nEn, col2.nEn)
  normEn.list[[count]] = df.normEn

  ## use global normalization for AFMCC for top cohorts
  gbnorm.afmcc = (data.en$afmcc - cc.gmin)/(cc.gmax - cc.gmin)
  #print(gbnorm.afmcc)
  #print(data.en$etot)
  gbnorm.En = gbnorm.afmcc*data.en$etot
  df.gbnormEn = data.frame(cc = gbnorm.afmcc, etot = gbnorm.En)
  col1.gnEn = paste0("gnormCC_",kappa)
  col2.gnEn = paste0("gnormEtot_",kappa)
  colnames(df.gbnormEn) = c(col1.gnEn, col2.gnEn)
  gbnormEn.list[[count]] = df.gbnormEn
  
  
  count = count + 1
  #print(head(data.en$afmcc))  
}

df.rg = do.call(cbind, rg.list)
df.rg$frame = as.numeric(rownames(df.rg))-1
df.rg = df.rg %>%  select("frame", everything())
#print(df.rg)
df.En = do.call(cbind, En.list)
df.En$frame = as.numeric(rownames(df.rg))-1
df.En = df.En %>%  select("frame", everything())

df.normEn = do.call(cbind, normEn.list)
df.normEn$frame = as.numeric(rownames(df.rg))-1
df.normEn = df.normEn %>%  select("frame", everything())

df.gnormEn = do.call(cbind, gbnormEn.list)
df.gnormEn$frame = as.numeric(rownames(df.rg))-1
df.gnormEn = df.gnormEn %>%  select("frame", everything())

file.rg = paste("rg_top",top_cohorts,".txt",sep="")
file.En = paste("En_top",top_cohorts,".txt",sep="")
file.normEn = paste("normEn_top",top_cohorts,".txt",sep="")
file.gnormEn = paste("gnormEn_top",top_cohorts,".txt",sep="")
write.table(format(df.rg), file = file.rg, sep = "\t", row.names = FALSE, quote = FALSE)
write.table(format(df.En), file = file.En, sep = "\t", row.names = FALSE, quote = FALSE)
write.table(format(df.normEn), file = file.normEn, sep = "\t", row.names = FALSE, quote = FALSE)
write.table(format(df.gnormEn), file = file.gnormEn, sep = "\t", row.names = FALSE, quote = FALSE)


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