def usage():
 print("Available command-line options:\n  -A protein sequence\n  -p Resid of the mutation\n  -w One letter of wild amino acid\n  -m One letter of mutant amino acid\n  -f Input file\n  -o Output file\n  -h Display this command-line summary")
def net_volume(wild,mutation):
 vol = {'A':'88.6','R':'173.4','N':'114.1','D':'111.1','C':'108.5','E':'138.4','Q':'143.8','G':'60.1','H':'153.2','I':'166.7','L':'166.7','K':'168.6','M':'162.9','F':'189.9','P':'112.7','S':'89.0','T':'116.1','W':'227.8','Y':'193.6','V':'140.0'}
 return ('{:.1f}'.format(float(vol.get(mutation,'0'))-float(vol.get(wild,'0'))))
def net_hydrophobicity(wild,mutation):
 hyd = {'A':'0','R':'3.71','N':'3.47','D':'2.95','C':'0.49','E':'1.64','Q':'3.01','G':'1.72','H':'4.76','I':'-1.56','L':'-1.81','K':'5.39','M':'-0.76','F':'-2.2','P':'-1.52','S':'1.83','T':'1.78','W':'-0.38','Y':'-1.09','V':'-0.78'}
 return ('{:.1f}'.format(float(hyd.get(mutation,'0'))-float(hyd.get(wild,'0'))))
def flexibility(wild,mutation):
 flex = {'A':'1','R':'81','N':'36','D':'18','C':'3','E':'54','Q':'108','G':'1','H':'36','I':'9','L':'9','K':'81','M':'27','F':'18','P':'2','S':'3','T':'3','W':'36','Y':'18','V':'3'}
 return (int(flex.get(mutation,'0'))-int(flex.get(wild,'0')))
def mutation_hydrophobicity(wild,mutation):
 if wild in ('I','V','L','F','C','M','A','W'):
  if mutation in ('G','T','S','Y','P','H'):
   return 0
  if mutation in ('N','D','Q','E','K','R'):
   return 1
  if mutation in ('I','V','L','F','C','M','A','W'):
   return 2
 if wild in ('G','T','S','Y','P','H'):
  if mutation in ('G','T','S','Y','P','H'):
   return 3
  if mutation in ('N','D','Q','E','K','R'):
   return 4
  if mutation in ('I','V','L','F','C','M','A','W'):
   return 5
 if wild in ('N','D','Q','E','K','R'):
  if mutation in ('G','T','S','Y','P','H'):
   return 6
  if mutation in ('N','D','Q','E','K','R'):
   return 7
  if mutation in ('I','V','L','F','C','M','A','W'):
   return 8
def mutation_polarity(wild,mutation):
 if wild in ('R','H','K'):
  if mutation in ('A','C','G','I','L','M','F','P','W','V'):
   return 0
  if mutation in ('R','H','K'):
   return 1
  if mutation in ('N','Q','S','T','Y'):
   return 2
  if mutation in ('D','E'):
   return 3
 if wild in ('A','C','G','I','L','M','F','P','W','V'):
  if mutation in ('A','C','G','I','L','M','F','P','W','V'):
   return 4
  if mutation in ('R','H','K'):
   return 5
  if mutation in ('N','Q','S','T','Y'):
   return 6
  if mutation in ('D','E'):
   return 7
 if wild in ('N','Q','S','T','Y'):
  if mutation in ('A','C','G','I','L','M','F','P','W','V'):
   return 8
  if mutation in ('R','H','K'):
   return 9
  if mutation in ('N','Q','S','T','Y'):
   return 10
  if mutation in ('D','E'):
   return 11
 if wild in ('D','E'):
  if mutation in ('A','C','G','I','L','M','F','P','W','V'):
   return 12
  if mutation in ('R','H','K'):
   return 13
  if mutation in ('N','Q','S','T','Y'):
   return 14
  if mutation in ('D','E'):
   return 15
def mutation_size(wild,mutation):
 if wild in ('G','A','S'):
  if mutation in ('C','D','P','N','T'):
   return 0
  if mutation in ('E','V','Q','H'):
   return 1
  if mutation in ('M','I','L','K','R'):
   return 2
  if mutation in ('F','Y','W'):
   return 3
  if mutation in ('G','A','S'):
   return 4
 if wild in ('C','D','P','N','T'):
  if mutation in ('C','D','P','N','T'):
   return 5
  if mutation in ('E','V','Q','H'):
   return 6
  if mutation in ('M','I','L','K','R'):
   return 7
  if mutation in ('F','Y','W'):
   return 8
  if mutation in ('G','A','S'):
   return 9
 if wild in ('E','V','Q','H'):
  if mutation in ('C','D','P','N','T'):
   return 10
  if mutation in ('E','V','Q','H'):
   return 11
  if mutation in ('M','I','L','K','R'):
   return 12
  if mutation in ('F','Y','W'):
   return 13
  if mutation in ('G','A','S'):
   return 14
 if wild in ('M','I','L','K','R'):
  if mutation in ('C','D','P','N','T'):
   return 15
  if mutation in ('E','V','Q','H'):
   return 16
  if mutation in ('M','I','L','K','R'):
   return 17
  if mutation in ('F','Y','W'):
   return 18
  if mutation in ('G','A','S'):
   return 19
 if wild in ('F','Y','W'):
  if mutation in ('C','D','P','N','T'):
   return 20
  if mutation in ('E','V','Q','H'):
   return 21
  if mutation in ('M','I','L','K','R'):
   return 22
  if mutation in ('F','Y','W'):
   return 23
  if mutation in ('G','A','S'):
   return 24
def mutation_hbonds(wild,mutation):
 if wild in ('R','W','K'):
  if mutation in ('A','C','G','I','L','M','F','P','V'):
   return 0 
  if mutation in ('R','W','K'):
   return 1   
  if mutation in ('N','Q','S','T','H','Y'):
   return 2   
  if mutation in ('D','E'):
   return 3   
 if wild in ('A','C','G','I','L','M','F','P','V'):
  if mutation in ('A','C','G','I','L','M','F','P','V'):
   return 4  
  if mutation in ('R','W','K'):
   return 5   
  if mutation in ('N','Q','S','T','Y','H'):
   return 6   
  if mutation in ('D','E'):
   return 7   
 if wild in ('N','Q','S','T','Y','H'):
  if mutation in ('A','C','G','I','L','M','F','P','V'):
   return 8   
  if mutation in ('R','W','K'):
   return 9   
  if mutation in ('N','Q','S','T','Y','H'):
   return 10   
  if mutation in ('D','E'):
   return 11  
 if wild in ('D','E'):
  if mutation in ('A','C','G','I','L','M','F','P','V'):
   return 12   
  if mutation in ('R','W','K'):
   return 13   
  if mutation in ('N','Q','S','T','Y','H'):
   return 14   
  if mutation in ('D','E'):
   return 15
def mutation_chemical(wild,mutation):
 if wild in ('A','G','I','L','P','V'):
  if mutation in ('R','H','K'):
   return 0   
  if mutation in ('N','Q'):
   return 1   
  if mutation in ('D','E'):
   return 2   
  if mutation in ('C','M'):
   return 3   
  if mutation in ('S','T'):
   return 4   
  if mutation in ('F','W','Y'):
   return 5   
  if mutation in ('A','G','I','L','P','V'):
   return 6
 if wild in ('R','H','K'):
  if mutation in ('R','H','K'):
   return 7   
  if mutation in ('N','Q'):
   return 8   
  if mutation in ('D','E'):
   return 9  
  if mutation in ('C','M'):
   return 10   
  if mutation in ('S','T'):
   return 11   
  if mutation in ('F','W','Y'):
   return 12   
  if mutation in ('A','G','I','L','P','V'):
   return 13
 if wild in ('N','Q'):
  if mutation in ('R','H','K'):
   return 14   
  if mutation in ('N','Q'):
   return 15   
  if mutation in ('D','E'):
   return 16   
  if mutation in ('C','M'):
   return 17   
  if mutation in ('S','T'):
   return 18   
  if mutation in ('F','W','Y'):
   return 19   
  if mutation in ('A','G','I','L','P','V'):
   return 20
 if wild in ('D','E'):
  if mutation in ('R','H','K'):
   return 21   
  if mutation in ('N','Q'):
   return 22   
  if mutation in ('D','E'):
   return 23   
  if mutation in ('C','M'):
   return 24   
  if mutation in ('S','T'):
   return 25   
  if mutation in ('F','W','Y'):
   return 26   
  if mutation in ('A','G','I','L','P','V'):
   return 27   
 if wild in ('C','M'):
  if mutation in ('R','H','K'):
   return 28   
  if mutation in ('N','Q'):
   return 29   
  if mutation in ('D','E'):
   return 30   
  if mutation in ('C','M'):
   return 31   
  if mutation in ('S','T'):
   return 32   
  if mutation in ('F','W','Y'):
   return 33   
  if mutation in ('A','G','I','L','P','V'):
   return 34      
 if wild in ('S','T'):
  if mutation in ('R','H','K'):
   return 35   
  if mutation in ('N','Q'):
   return 36   
  if mutation in ('D','E'):
   return 37   
  if mutation in ('C','M'):
   return 38   
  if mutation in ('S','T'):
   return 39   
  if mutation in ('F','W','Y'):
   return 40   
  if mutation in ('A','G','I','L','P','V'):
   return 41   
 if wild in ('F','W','Y'):
  if mutation in ('R','H','K'):
   return 42   
  if mutation in ('N','Q'):
   return 43   
  if mutation in ('D','E'):
   return 44   
  if mutation in ('C','M'):
   return 45  
  if mutation in ('S','T'):
   return 46   
  if mutation in ('F','W','Y'):
   return 47   
  if mutation in ('A','G','I','L','P','V'):
   return 48
def sequences_features(chainA,mutation_resid,wild_aa):
 mutation_resid=int(mutation_resid)-1
 features=[]
 target_sequence=""
 for line in open(chainA):
  str=line.strip()
  if str[0]!=">":
   target_sequence=target_sequence+str
 if len(target_sequence)<20:
  print("The fasta sequence length is at least 20!")
 for i in target_sequence:
  if i.upper() not in "ARNDCQEGHILKMFPSTWYV":
   print("Please check the input sequence!")
   sys.exit()
 if wild_aa.upper() != target_sequence[mutation_resid]:
  print("Wild type is not same as input sequence!")
  sys.exit()
 for i in range(int(mutation_resid)-5,int(mutation_resid)+6):
  if i < 0:
   features.append("0")
   continue
  try:
   target_sequence[i]
  except IndexError:
   features.append("0")
  else:
   features.append(mutation_aa_label(target_sequence[i]))
 return(features)
def translate_aa(three_letter):
 trans = {'ALA':'A','ARG':'R','ASN':'N','ASP':'D','CYS':'C','GLU':'E','GLN':'Q','GLY':'G','HIS':'H','ILE':'I','LEU':'L','LYS':'K','MET':'M','PHE':'F','PRO':'P','SER':'S','THR':'T','TRP':'W','TYR':'Y','VAL':'V'}
 return (trans[three_letter])
def mutation_aa_label(three_letter):
 aa = {'A':'1','R':'2','N':'3','D':'4','C':'5','E':'6','Q':'7','G':'8','H':'9','I':'10','L':'11','K':'12','M':'13','F':'14','P':'15','S':'16','T':'17','W':'18','Y':'19','V':'20'}
 return (aa.get(three_letter,0))
def mutation_type(wild,mutation):
 wild_lists = ['A','F','C','D','N','E','Q','G','H','L','I','K','M','P','R','S','T','V','W','Y']
 mutation_lists = ['A','F','C','D','N','E','Q','G','H','L','I','K','M','P','R','S','T','V','W','Y']
 label_1=0
 label_2=0
 for i in wild_lists:
  for j in mutation_lists:
   if i != j:
    label_1 += 1
    if wild == i:
     if mutation == j and i != j:
      label_2 = label_1
      break
 return label_2
def job_submit(file_name):
 print("Build "+file_name+" PSSM...")
 global PSIBLASTDIR
 global BLASTDATABASE
 global BLAST_NUM_THREADS
 os.system(PSIBLASTDIR+"psiblast -query "+file_name+" -num_threads "+str(BLAST_NUM_THREADS)+" -db "+BLASTDATABASE+" -num_iterations 3 -out "+file_name+".out -out_ascii_pssm "+file_name+".pssm 2>/dev/null")
def pssm_check(file_name):
 if os.path.exists(file_name+".pssm"):
  print("Build "+file_name+" PSSM done...")
 else:
  print("Can't get the PSSM, please check the sequence")
  sys.exit()
def psepssm(file):
 import math
 lamda=7
 pssm=[]
 features=[]
 for line1 in open(file+".pssm"):
  info1=line1.strip().split()
  if len(info1)>43:
   if info1[1].isupper():
    length=int(info1[0])
    for i in range(2,22):
     pssm.append(float(1)/(1+math.e**(-int(info1[i]))))
 for i in range(20):
  sum_pssm=0
  for j in range(int(length)):
   sum_pssm=sum_pssm+float(pssm[i+20*j])
  features.append("%.2f"%(float(sum_pssm)/length))
 for i in range(1,lamda+1):
   for j in range(20):
    s_pssm=0
    for k in range(int(length)-i):
     s_pssm=(float(pssm[20*k+j])-float(pssm[20*(k+i)+j]))**2+s_pssm
    features.append("%.2f"%(float(s_pssm)/(int(length)-i)))
 return(features)
def mpssmscores(file_name,position):
 windows=7
 pssm=[]
 line1=[]
 for line2 in open(file_name+".pssm"):
  line1.append(line2.strip())
 for j in range(int(position)-(windows//2),int(position)+windows//2+1):
  if j<0:
   for i in range(2,22):
    pssm.append(0)
   continue
  index=0
  for info1 in line1:
   info2=info1.split()
   if len(info2)>43:
    if int(info2[0])==j:
     for i in range(2,22):
      pssm.append(info2[i])
     index=1
    if info1==line1[-7] and index==0:
     for i in range(2,22):
      pssm.append(0)
 return(pssm)
def file_loop(chainA,mutation_resid,wild_aa,mutation_aa,mutation_list):
 if len(wild_aa) == 3:
  wild_aa=translate_aa(wild_aa)
 if len(mutation_aa) ==3:
  mutation_aa=translate_aa(mutation_aa)
 label = []
 label.append(net_volume(wild_aa,mutation_aa))
 label.append(net_hydrophobicity(wild_aa,mutation_aa))
 label.append(flexibility(wild_aa,mutation_aa))
 label.append(mutation_hydrophobicity(wild_aa,mutation_aa))
 label.append(mutation_polarity(wild_aa,mutation_aa))
 label.append(mutation_type(wild_aa,mutation_aa))
 label.append(mutation_size(wild_aa,mutation_aa))
 label.append(mutation_hbonds(wild_aa,mutation_aa))
 label.append(mutation_chemical(wild_aa,mutation_aa))
 label=label+sequences_features(chainA,mutation_resid,wild_aa)
 if mutation_list==0:
  job_submit(chainA)
  pssm_check(chainA)
 label=label+mpssmscores(chainA,mutation_resid)
 label=label+psepssm(chainA)
 if wild_aa == mutation_aa:
  return("0 Neutral")
 #for i in resid_label_aa:
  #label.append(mutation_aa_label(i))
 return(pred_feature(label))
def pred_feature(label):
 model=xgb.Booster(model_file='regression.model')
 x=np.array(label)
 #x=pd.DataFrame(x)
 #x=x.iloc[:,[-1]].values
 x=x.reshape((1,len(label)))
 x=xgb.DMatrix(x)
 y_pred = model.predict(x)
 if y_pred[0]>0:
  return ("%.2f"%y_pred[0]+" Stabilizing")
 else:
  return ("%.2f"%y_pred[0]+" Destabilizing")
 #return (y_pred[0])
def job_done(i):
 if i==0:
  print("##########################")
  print("#       Job done !       #")
  print("##########################")
 else:
  print("##########################")
  print("#         Result         #")
  print("##########################") 
def main(a,b,c):
 global PSIBLASTDIR
 global BLASTDATABASE
 global BLAST_NUM_THREADS
 PSIBLASTDIR=a
 BLASTDATABASE=b
 BLAST_NUM_THREADS=c
 #import pandas as pd
 opts, args = getopt.getopt(sys.argv[1:], "hA:p:w:m:f:o:", ["help", "inputA=","position=","wild=", "mutation", "file", "output="])
 for op, value in opts:
  if op == "-A":
   seq1_file = value
  elif op == "-p":
   position = value
  elif op == "-w":
   wild_aa= value
  elif op == "-m":
   mutation_aa = value
  elif op == "-o":
   output_name = value
  elif op == "-f":
   file_name = value
  elif op == "-h":
   usage()
   sys.exit()
 try:
  file_name
 except NameError:
  try:
   output_name
  except NameError:
   pred=file_loop(seq1_file,position,wild_aa,mutation_aa,0)
   job_done(1)
   print(pred)
   job_done(0)
  else:
   f = open(output_name, 'w')
   print("Position Wild Mutant ddG Type", file=f)
   pred=file_loop(seq1_file,position,wild_aa,mutation_aa,0)
   job_done(0)
   print(position,wild_aa,mutation_aa,pred, file=f)
   f.close()
 else:
  try:
   output_name
  except NameError:
   list_index=0
   for line in open(file_name):
    info=line.strip().split(' ')
    if list_index==0:
     pred=file_loop(seq1_file,info[1],info[0],info[2],0)
     job_done(1)
     print(pred)
     list_index=1
    else:
     pred=file_loop(seq1_file,info[1],info[0],info[2],1)
     print(pred)
   job_done(0)
   sys.exit()
  else:
   list_index=0
   f = open(output_name, 'w')
   print("Position Wild Mutant ddG Type", file=f)
   for line in open(file_name):
    info=line.strip().split(' ')
    if list_index==0:
     pred=file_loop(seq1_file,info[1],info[0],info[2],0)
     print(info[1],info[0],info[2],pred, file=f)
     list_index=1
    else:
     pred=file_loop(seq1_file,info[1],info[0],info[2],1)
     print(info[1],info[0],info[2],pred, file=f)
   f.close()
   job_done(0)
 #pred=file_loop(pdb_id,mutation_chain,mutation_resid,wild_aa,mutation_aa)
import sys, getopt
import os
import numpy as np
import xgboost as xgb
