# Copyright (c) 2018: Elodie Laine
# Copyright (c) 2022: Mustafa Tekpinar
# This code is part of the gemme 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("GEMME_PATH"),"/pred.R",sep=""))

# 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("GEMME_PATH"),"/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]

# 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)
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]])
wt = wt[[1]][1:n]
#indicesQ = getIndicesAli(ali$seq,1)

if(simple==FALSE){
  print("reading selected mutations...")
  res=readMut(fname)
  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...")
#normWeightMode="trace+pc+cv"
#normWeightMode="trace+pc"
#normWeightMode="trace+cv"
#normWeightMode="trace"
#print(normWeightMode)
#In future, you may need to comment line 44 to use this functionality if you do
#weighting not just in normalization.
#This part of the code obtains max values of Trace, PC, or CV for weighting the 
#normalized results. 
trace = c()
if((normWeightMode=="max-trace-pc") | (normWeightMode=="max-pc-trace")){
  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"]))
    }else{
        trace<-append(trace, max(jet[row, "trace"], jet[row, "pc"]))
    }
  }
} else if ((normWeightMode=="max-trace-cv") | (normWeightMode=="max-cv-trace")){
  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, "cv"]))
    }else{
        trace<-append(trace, max(jet[row, "trace"], jet[row, "cv"]))
    } 
  }
} else if ((normWeightMode=="max-trace-pc-cv")|(normWeightMode=="max-trace-cv-pc")){
  print(paste("Using ", normWeightMode))
  for (row in 1:nrow(jet)) {
    if(sum(colnames(jet)=="traceMax")==1){
      trace<-append(trace, max(jet[row, "traceMax"], max(jet[row, "pc"], jet[row, "cv"])))
    }else{
      trace<-append(trace, max(jet[row, "trace"], max(jet[row, "pc"], jet[row, "cv"])))
    }
  }
} else if ((normWeightMode=="half-cv+pc") | (normWeightMode=="half-pc+cv")){
  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, "cv"]))
    }else{
        trace<-append(trace, (jet[row, "pc"]+jet[row, "cv"])/2.0)
    } 
  }  
} else if (normWeightMode=="trace"){
  print("Using only 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{
  print("ERROR: Unknown --normWeightMode selected!")
  print("It can only be 'trace', 'max-trace-pc', 'max-trace-cv', 'max-trace-pc-cv' or 'half-cv+pc'!")
}
print(trace)

if(simple){
  #Independent model normalization
  normPredInd = normalizePredWithNbSeqsZero(predInd,trace,wt)
  rownames(normPredInd)=aa
  
  # Epistatic model normalization
  normPred=normalizePred(pred, trace, wt)
  rownames(normPred)=aa

}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"))

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

print("running combined model...")
alpha = 0.6
alphabet = "lz-bl.7"
if(simple){
  normPredCombi = normalizePredWithNbSeqsPC(pred,trace,wt,alpha,nbSeqs,alphabet)
  rownames(normPredCombi)=aa
}else{
  normPredCombi = normalizePredWithNbSeqsPCSelMult(pred,trace,wt,list(pos,subsaa),alpha,nbSeqs[[1]],alphabet)
  names(normPredCombi)=rawMut
}
# output the predicted mutational effects based on sequence counts (conservation at the bottom)
write.table(normPredCombi,paste0(prot,"_normPred_evolCombi.txt"))
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=""))
# }


	
