# Copyright (c) 2018: Elodie Laine
# Copyright (c) 2022-2023: Mustafa Tekpinar
# This code is part of the prescott package and governed by its license.
# Please see the LICENSE.txt file included as part of this package.

# Usage: Rscript --save computePred.R XXXX

library("seqinr")
#source(paste(Sys.getenv("PRESCOTT_PATH"),"/pred.R",sep=""))
source("pred.R")
# amino acid full alphabet
aa = c("a","c","d","e","f","g","h","i","k","l","m","n","p","q","r","s","t","v","w","y")

# # BLOSUM62  matrix
# blosum62p = as.matrix(read.table(paste0(Sys.getenv("PRESCOTT_PATH"),"/data/blosum62p.txt"),row.names=1))
# bgp = blosum62p[order(rownames(blosum62p)),order(colnames(blosum62p))]
# bg = apply(bgp,1,sum)

# get arguemnts 
args = commandArgs(trailingOnly = TRUE)
prot = args[1]
aliFile = args[2]
simple = args[3]
fname = args[4]
normWeightMode = args[5]
selectedAlphabet = args[6]

# read alignment and convert to matrix
ali=as.matrix.alignment(read.alignment(aliFile,format="fasta"))
ali[!ali%in%aa] = "-"
npos = dim(ali)[[2]]
# convert to binary matrix with respect to the query (0: identical, 1: different)
binAli = t(apply(ali,1,f<-function(x){return(x!=ali[1,])}))
pId =apply(binAli,1,f<-function(x){return(1-sum(x)/npos)})
aliCons = ali[pId>0.6,]
aliVeryCons = ali[pId>0.8,]
# number of sequences
N = c(dim(ali)[[1]],dim(aliCons)[[1]],dim(aliVeryCons)[[1]])
# #resAliCons = computePSSM(aliCons)
# res = list(computePSSM(ali,N[1],npos),computePSSM(aliCons,N[2],npos),computePSSM(aliVeryCons,N[3],npos))
# # write.table(res[[1]][[3]],paste0(prot,"_pssm.txt"))
# # write.table(res[[2]][[3]],paste0(prot,"_pssm60.txt"))
# # write.table(res[[3]][[3]],paste0(prot,"_pssm80.txt"))

# read evolutionary traces computed by JET
jet=read.table(paste(prot,"_jet.res",sep=""),head=TRUE)
############################################Added by MT to test effect of PC, CV, DFI, Bfactor etc.
trace = c()
if ((normWeightMode=="max")){
  print(paste("Using ", normWeightMode))
  for (row in 1:nrow(jet)) {
    if(sum(colnames(jet)=="traceMax")==1){
      trace<-append(trace, max((jet[row, "traceMax"]+jet[row, "pc"])/2.0, max((jet[row, "traceMax"]+jet[row, "cv"])/2.0, (jet[row, "pc"]+jet[row, "cv"])/2.0 )))
    }else{
      trace<-append(trace, max((jet[row, "trace"]+jet[row, "pc"])/2.0, max((jet[row, "trace"]+jet[row, "cv"])/2.0, (jet[row, "pc"]+jet[row, "cv"])/2.0 )))
    }
  }
}else if (normWeightMode=="tjet"){
  print("Using only T_JET2 traces")
  for (row in 1:nrow(jet)) {
    if(sum(colnames(jet)=="traceMax")==1){
      trace<-append(trace, jet[row, "traceMax"])
    }else{
      trace<-append(trace, jet[row, "trace"])
    } 
  }
}else if (normWeightMode=="cv"){
  print("Using only CV traces")
  for (row in 1:nrow(jet)) {
    trace<-append(trace, jet[row, "cv"]) 
  }
}else if (normWeightMode=="pc"){
  print("Using only PC traces")
  for (row in 1:nrow(jet)) {
    trace<-append(trace, jet[row, "pc"]) 
  }
}else if (normWeightMode=="sstjetormax"){ 
  print("Using only sstjetormax")
  for (row in 1:nrow(jet)) {
    if(sum(colnames(jet)=="sstjetormax")==1){
      trace<-append(trace, jet[row, "sstjetormax"])
    }else{
      print("No field called sstjetormax in the JET output!")
      quit(status=-1)
    } 
  }
}else if (normWeightMode=="sstjetormaxtwocomponent"){ 
  print("Using only sstjetormaxtwocomponent")
  for (row in 1:nrow(jet)) {
    if(sum(colnames(jet)=="sstjetormaxtwocomponent")==1){
      trace<-append(trace, jet[row, "sstjetormaxtwocomponent"])
    }else{
      print("No field called sstjetormaxtwocomponent in the JET output!")
      quit(status=-1)
    } 
  }
}else if (normWeightMode=="sstjetormaxthirdchanged"){ 
  print("Using only sstjetormaxthirdchanged")
  for (row in 1:nrow(jet)) {
    if(sum(colnames(jet)=="sstjetormaxthirdchanged")==1){
      trace<-append(trace, jet[row, "sstjetormaxthirdchanged"])
    }else{
      print("No field called sstjetormaxthirdchanged in the JET output!")
      quit(status=-1)
    } 
  }
}else if (normWeightMode=="sstjetormaxthirdpcstar"){ 
  print("Using only sstjetormaxthirdpcstar")
  for (row in 1:nrow(jet)) {
    if(sum(colnames(jet)=="sstjetormaxthirdpcstar")==1){
      trace<-append(trace, jet[row, "sstjetormaxthirdpcstar"])
    }else{
      print("No field called sstjetormaxthirdpcstar in the JET output!")
      quit(status=-1)
    } 
  }
}else if (normWeightMode=="tjetormax"){ 
  print("Using tjetormax with inverse")
  for (row in 1:nrow(jet)) {
    if(sum(colnames(jet)=="traceMax")==1){
	    trace<-append(trace, max((jet[row, "traceMax"]+jet[row, "pc"])/2.0, max((jet[row, "traceMax"]+1.0-jet[row, "cv"])/2.0, (jet[row, "pc"]+1.0-jet[row, "cv"])/2.0 )))
    }else{
	    trace<-append(trace, max((jet[row, "trace"]+jet[row, "pc"])/2.0, max((jet[row, "trace"]+1.0-jet[row, "cv"])/2.0, (jet[row, "pc"]+1.0-jet[row, "cv"])/2.0 ))) 
    }
  }
}else{
  print("ERROR: Unknown --normWeightMode selected!")
  print("It can only be 'tjet', 'pc', 'cv',")
  print("'max', 'tjetormax' or 'sstjetormax'!")
}
print(trace)
#####################################################################################################################



#if(sum(colnames(jet)=="traceMax")==1){trace=jet[,"traceMax"]}else{trace=jet[,"trace"]}

#traceAli = sweep(binAli, MARGIN=2, trace, `*`)
# compute evolutionary distances of all sequences with respect to the query
distTrace = binAli[2:N[1],] %*% trace^2

wt=read.fasta(paste(prot,".fasta",sep=""))

n = length(wt[[1]])
# # print(wt)
# for ( residue in wt){ 
#   print(toupper(residue))
# }
# print(length(wt[[1]]))
myWTSequence <- wt[[1]]
#create empty list of length  of the wt sequence
myColumnNames <- vector(mode='list', length=length(wt[[1]]))
# myColumnNames <- list()
for ( k in 1:length(wt[[1]])) {
  # print(toupper())
  myColumnNames[[k]] <-paste(toupper(myWTSequence[k]), k, sep = "")
}
# print(myColumnNames)

wt = wt[[1]][1:n]
#indicesQ = getIndicesAli(ali$seq,1)

if(simple==FALSE){
  print("reading selected mutations...")
  res=readMut(fname)
  #print(res)
  pos = res[[1]]
  subsaa = res[[2]]
  rawMut = res[[3]]
  print("done")
}

print("running independent model...")
# compute sequence counts for each position and each aa
nbSeqs = computeNbSeqs(ali)
# compute gap counts for each position
nbGaps = N[1] - apply(nbSeqs,2,sum)
# output the conservation values
# dat=rbind(trace,KL=res[[1]][[2]],SE=res[[1]][[1]],gap=nbGaps/N[1],KL60=res[[2]][[2]],SE60=res[[2]][[1]],KL80=res[[3]][[2]],SE80=res[[3]][[1]])
# write.table(dat,paste0(prot,"_conservation.txt"))
# compute log-odd ratios between mutated and wt sequence counts
predInd = computePredNbSeqs(wt,nbSeqs)
rownames(predInd)=aa

# output the sequence counts log-odd ratios
# write.table(predInd,paste0(prot,"_pred_evolInd.txt"))
print("done")

print("running global epistatic model...")
pred=computePredSimple(ali,distTrace,wt,5)
rownames(pred)=aa

# output the evolutionary distances between the query and the closest variants
evolDist = pred/sum(trace^2)
evolDist[is.na(evolDist)] = 1
# write.table(evolDist,paste0(prot,"_pred_evolEpi.txt"))
print("done")

print("running normalization...")

if(simple){
  #Independent model normalization
  normPredInd = normalizePredWithNbSeqsZero(predInd,trace,wt)
  rownames(normPredInd)=aa
  
  # Epistatic model normalization
  normPred=normalizePred(pred, trace, wt)
  rownames(normPred)=aa
  # output the normalized predicted mutational effects based on sequence counts (conservation at the bottom)
  write.table(normPredInd,paste0(prot,"_normPred_evolInd.txt"), col.names=myColumnNames)

  # output the predicted normalized mutational effects based on evolutionary distance (conservation at the bottom)
  write.table(normPred,paste0(prot,"_normPred_evolEpi.txt"), col.names=myColumnNames)
}else{
  #Independent model normalization
  normPredInd = normalizePredWithNbSeqsZeroSelMult(predInd, trace, wt, list(pos,subsaa))
  names(normPredInd)=rawMut

  # Epistatic model normalization
  normPred=normalizePredSelMult(pred, trace, wt, list(pos,subsaa))
  names(normPred)=rawMut

  # output the normalized predicted mutational effects based on sequence counts (conservation at the bottom)
  write.table(normPredInd,paste0(prot,"_normPred_evolInd.txt"), quote = FALSE, col.names=FALSE)

  # output the predicted normalized mutational effects based on evolutionary distance (conservation at the bottom)
  write.table(normPred,paste0(prot,"_normPred_evolEpi.txt"), quote = FALSE, col.names=FALSE) 
}
print("done")


print("running combined model...")
# #Read frequencies from the jet output file. 
# frequencies = c()
# print(paste("Reading frequencies:"))
# for (row in 1:nrow(jet)) {
#   frequencies<-append(frequencies, jet[row, "freq"] )
# }
# #print(frequencies)
# freq_mean = mean(frequencies)

alpha = 0.6
# alpha = c()

# #If we want to assign each aa the same weight of epistatic and independent 
# #contribution, that's the way to do it. 
# for (row in 1:nrow(jet)) {
#   alpha<-append(alpha, 0.6)
# }
# alpha_val = 0.6
# for (row in 1:nrow(jet)) {
#   if(frequencies[row]>freq_mean){
#     alpha<-append(alpha, alpha_val)
#   } else{
#     alpha<-append(alpha, (1.0 - alpha_val))
#   }  
# }
#print(alpha)

#alphabet = "lz-bl.7"
alphabet = selectedAlphabet
if(simple){
  normPredCombi = normalizePredWithNbSeqsPC(pred,trace,wt,alpha,nbSeqs,alphabet)
  rownames(normPredCombi)=aa
  # output the predicted mutational effects based on sequence counts (conservation at the bottom)
  write.table(normPredCombi,paste0(prot,"_normPred_evolCombi.txt"), col.names=myColumnNames)
}else{
  #normPredCombi = normalizePredWithNbSeqsPCSelMult(pred,trace,wt,list(pos,subsaa),alpha,nbSeqs[[1]],alphabet)
  normPredCombi = normalizePredWithNbSeqsPCSelMult(pred,trace,wt,list(pos,subsaa),alpha,nbSeqs,alphabet)
  # names(normPredCombi)=rawMut
  names(normPredCombi)=rawMut

  # output the predicted mutational effects based on sequence counts (conservation at the bottom)
  write.table(normPredCombi,paste0(prot,"_normPred_evolCombi.txt"), quote = FALSE, col.names=FALSE)
}

print("done")

# Since I am doing the plots with Python, I am commenting this part
# In addition, I want to get rid of R dependency completely!
# if(simple){
#   plotMatBlues(paste(prot,"_normPred_evolEpi",sep=""))
#   plotMatGreens(paste(prot,"_normPred_evolInd",sep=""))
#   plotMatOranges(paste(prot,"_normPred_evolCombi",sep=""))
# }


	
