#file name   :  glide_output_StructureWithGoodScore.py
#Function    : output structures and properties, e. g. docking properties, ligand effiecncy, ... with the docking score better than a input value.
#file needed : *mmgbsa-out.mae or *_pv.maegz
#
#Usage       : run glide_output_StructureWithGoodScore.py -g no/yes -t out.txt -i 3G51_MODEL_1_RSK2_mmgbsa-out.mae -o aa.mae/aa.sdf
#
#Author      : zjxu@simm.ac.cn
#Date        : 3/18/2012
#Update      : 2/5/2018 change r_i_glide_gscore to r_i_docking_score: %s/r_i_glide_gscore/r_i_docking_score/g
#Update      : 6/15/2020 add -s to output the structure with the docking score better than the input score.


import sys
def main(argv):
  from schrodinger import structure
  #from schrodinger.maestro import maestro
  #from schrodinger import project
  import getopt, sys
  import os
  
  if (len(sys.argv) != 11):
    print 'Usage: run glide_output_StructureWithGoodScore.py -s -6.0 -g no -t out.txt -i 3G51_MODEL_1_RSK2_mmgbsa-out.mae -o aa.sdf'
    sys.exit()

  try:
    opts, args = getopt.getopt(sys.argv[1:], "s:g:t:i:o:")
  except getopt.GetoptError:
    print 'Usage: run glide_output_StructureWithGoodScore.py -s -6.0 -g no -t out.txt -i 3G51_MODEL_1_RSK2_mmgbsa-out.mae -o aa.sdf'
   
    print 'run glide_output_StructureWithGoodScore.py -s -6.0 -g no -t out.txt -i 3G51_MODEL_1_RSK2_mmgbsa-out.mae -o aa.sdf'
    sys.exit()

  for opt, value in opts:
    if opt == "-s":
      score = float(value)
    if opt == "-g":
      gbsa = value
    elif opt == "-t":
      out_txt = value
    elif opt == "-i":
      in_file = value
    elif opt == "-o":
      out_file = value
  print "score: ", score
  #print "rank: ", rank
  #print "in_file: ", in_file
  #print "out_file: ", out_file
  if os.path.exists(out_file):
    os.remove(out_file)
  #pt = maestro.project_table_get()
  # Loop over the selected entries
  #for irow in xrange(1,len(pt)+1):
   # print irow.split('\t')[0]
    # Do something for all entries using pt[irow]
  myfile = open(out_txt, 'w')
  for (ix, st) in enumerate(structure.StructureReader(in_file)):
    #if ix == 0:
      #st.append(out_file)
#    if (ix > 0):
    if (ix > 0) and st.property['r_i_docking_score'] < score:
      st.append(out_file)
      #print "title: ", st.property['s_m_title'], "\t", "glide_gscore: ", st.property['r_i_docking_score'], "\t", "glide_ligand_efficiency: ", st.property['r_i_glide_ligand_efficiency'], "\t", "MMGBSA_DG_bind: ", st.property['r_psp_Prime_MMGBSA_DG_bind'] 
      if gbsa == "yes":
        #print st.property['s_m_title'], "\t", st.property['r_i_docking_score'], "\t", st.property['r_i_glide_ligand_efficiency'], "\t", st.property['r_psp_Prime_MMGBSA_DG_bind'] 
        myfile.write( str(st.property['s_m_title']) + "\t" + str(st.property['r_i_docking_score']) + "\t" + str(st.property['r_i_glide_ligand_efficiency']) + "\t" + str(st.property['r_psp_Prime_MMGBSA_DG_bind']) + "\n" )
      elif gbsa == "no":
        #print st.property['s_m_title'], "\t", st.property['r_i_docking_score'], "\t", st.property['r_i_glide_ligand_efficiency']
        myfile.write( str(st.property['s_m_title']) + "\t" + str(st.property['r_i_docking_score']) + "\t" + str(st.property['r_i_glide_ligand_efficiency']) + "\n" )
      else:
        print "Warning: you should set -g yes or no. Here, the option no was used"
        myfile.write( str(st.property['s_m_title']) + "\t" + str(st.property['r_i_docking_score']) + "\t" + str(st.property['r_i_glide_ligand_efficiency']) + "\n" )


    #elif (ix > 0):
     # if st.property['r_i_docking_score'] < gscore:
      #if st.property['s_m_title'] == "4720_01.mol2":
       # st.append(out_file)
  myfile.close()

if __name__ == "__main__":
  main(sys.argv[1:])