#!/usr/bin/env python

import string,os,sys,profile,cPickle,math,socket,time,shutil,fnmatch,glob

# SCRIPT ERROR
# ============
# DISPLAY ERROR MESSAGE AND STOP SCRIPT
def error(err):
  text="ERROR - %s\nScript has been stopped." % err
  print text
  #log_message(text);
  sys.exit(1)

# FUNCTION
# ============
# READ XPLOR LOG FILE FOR RESTRAINTS AND WRITE CONCOORD FORMAT
# INFORMAT SHOULD BE XPLOR OR IUPAC
def xplor2concoord(infile,informat,outfile):
  lines=open(infile,'r').readlines()
  outfile_disre=open(outfile,'w')
  
  # CHECK WHICH INFORMAT WE HAVE, IN THE PSEUDOATOMS_DICT XPLOR IS AT POSITION 0
  # IUPAC AT POSITION 1, CORRECTION AT POSITION 2
  if string.upper(informat)=='XPLOR':
    pos=0
  elif string.upper(informat)=='IUPAC':
    pos=1  
  # READ THE PSEUDOATOMS DICTIONARY
  pseudoatoms_dict=pseudocorrect()
  
  # WRITE THE HEADER OF THE OUTPUT FILE
  outfile_disre.write('[ distance_restraints ]\n')
  outfile_disre.write('[ resid1 resname1 atomname1 segid1 ')
  outfile_disre.write('resid2 resname2 atomname2 segid2 ')
  outfile_disre.write('lowbound uppbound upppscor index restrnum ]\n')
  
  i = 0
  while i < len(lines):
    # GO THROUGH THE DISTANCES IN THE RESTRAINT LIST
    try:  
      if string.split(lines[i])[0]=='==========' and string.split(lines[i])[3]=='restraint':
        restr_dict = {}
        restr_dict['i']= {}
        restr_dict['j']= {}
        # X IS USED TO CREATE AN INDEX IN EACH RESTRAINT FOR THE DIFFERENT CONTRIBUTIONS
        # CNTI AND CNTJ COUNT THE I- AND J CONTRIBUTIONS RESPECTIVELY
        x = 0
        cnti = 0
        cntj = 0
        restr_num=int(string.split(lines[i])[4])
        while not string.split(lines[i+x])[0]=='R<average>=':
          x = x + 1
          # CREATE THE DICTIONARIES FOR THE DIFFERENT CONTRIBUTIONS IN EACH RESTRAINT
          if string.split(lines[i+x])[0]=='set-i-atoms':
            x = x + 1
            while not string.split(lines[i+x])[0]=='set-j-atoms':
              line=string.split(lines[i+x])
              atomname=line[-1]
              resname=line[-2]
              resnum=int(line[-3])
              try:
                segid=line[-4]
              except IndexError:
                segid='.'
              restr_dict['i'][x]= {}
              restr_dict['i'][x]['atomname']=atomname
              restr_dict['i'][x]['resname']=resname
              restr_dict['i'][x]['resnum']=resnum
              restr_dict['i'][x]['segid']=segid
              x = x + 1
              cnti = cnti + 1
          if string.split(lines[i+x])[0]=='set-j-atoms':
            x = x + 1
            while not string.split(lines[i+x])[0]=='R<average>=':
              line=string.split(lines[i+x])
              atomname=line[-1]
              resname=line[-2]
              resnum=int(line[-3])
              try:
                segid=line[-4]
              except IndexError:
                segid='.'
              restr_dict['j'][x]= {}
              restr_dict['j'][x]['atomname']=atomname
              restr_dict['j'][x]['resname']=resname
              restr_dict['j'][x]['resnum']=resnum
              restr_dict['j'][x]['segid']=segid
              x = x + 1
              cntj = cntj + 1
          # CALCULATE THE UPPER AND LOWER BOUNDS
          if string.split(lines[i+x])[0]=='R<average>=' and string.split(lines[i+x])[2]=='NOE=':
            line=string.split(lines[i+x])
            restr_arg1=float(line[3])
            restr_arg2=float(string.split(line[5],'/')[0])
            restr_arg3=float(string.split(line[6],')')[0])
            low=restr_arg1 - restr_arg2
            upl=restr_arg1 + restr_arg3
            
            index = 0
            # CHECK WHETHER WE NEED UPPER BOUND CORRECTIONS FOR THE BOUND SMOOTHING IN CONCOORD
            # IF CHECKI=1, CORRECTION IS APPLIED TO SET OF I ATOMS
            # IF CHECKJ=1, CORRECTION IS APPLIED TO SET OF J ATOMS
            # 
            # IN CASE OF ONLY ONE CONTRIBUTION, NO CORRECTION IS APPLIED, CHECKI,J=0
            if cnti==1:
              ipass='true'
              pseudoatomcorri=0.0
            if cntj==1:
              jpass='true'
              pseudoatomcorrj=0.0
            # IN CASE OF MULTIPLE CONTRIBUTIONS:
            if cnti>1:
              # GET THE INFORMATION FOR THE FIRST CONTRIBUTION IN THE RESTRAINT
              contrib1_index=restr_dict['i'].keys()[0]
              atomname1=restr_dict['i'][contrib1_index]['atomname']
              resname1=restr_dict['i'][contrib1_index]['resname']
              resnum1=restr_dict['i'][contrib1_index]['resnum']
              segid1=restr_dict['i'][contrib1_index]['segid']
              pseudoatomi=None
              checki=0
              # check whether all i contributions are in the same residue
              # if not do not use pseudoatom correction 
              for sel in restr_dict['i'].keys():
                atomnamei=restr_dict['i'][sel]['atomname']
                resnamei=restr_dict['i'][sel]['resname']
                resnumi=restr_dict['i'][sel]['resnum']
                segidi=restr_dict['i'][sel]['segid']
                if not resnamei==resname1 or not resnumi==resnum1 or not segidi==segid1:
                  checki=0
                  print 'i check1  ',resnamei,resname1,resnumi,resnum1,segidi,segid1,restr_num
                  break
                else:
                  checki=1
              # check whether a pseudoatom exists for the first atom in i-contributions;
              # if not, we have more than one different groups in the
              # i contributions: don't correct
              # check the number of contributions to identify which of the pseudoatoms
              # are involved
              if checki==1:
                for key in pseudoatoms_dict[resname1].keys():
                  if atomname1 in pseudoatoms_dict[resname1][key][pos] and cnti==len(pseudoatoms_dict[resname1][key][pos])-1:
                    pseudoatomi=key
                if not pseudoatomi:
                  checki=0
                  print 'i check2  ',resname1,atomname1,pseudoatomi,cnti,restr_num
              # Now check all contributions 
              if checki==1:
                for sel in restr_dict['i'].keys():
                  atomnamei=restr_dict['i'][sel]['atomname']
                  resnamei=restr_dict['i'][sel]['resname']
                  resnumi=restr_dict['i'][sel]['resnum']
                  segidi=restr_dict['i'][sel]['segid']
                  if not atomnamei in pseudoatoms_dict[resname1][pseudoatomi][pos]:
                    checki=0
                    print 'i check3  ',resnamei,atomnamei,pseudoatomi,restr_num
              if checki==1:
                ipass='true'
                print 'i pass  ',resnamei,atomnamei,pseudoatomi,restr_num
                pseudoatomcorri=pseudoatoms_dict[resname1][pseudoatomi][2]
              else:
                ipass='false'
                pseudoatomcorri=0.0
            # do the same for the j-contributions
            if cntj>1:
              contrib1_index=restr_dict['j'].keys()[0]
              atomname1=restr_dict['j'][contrib1_index]['atomname']
              resname1=restr_dict['j'][contrib1_index]['resname']
              resnum1=restr_dict['j'][contrib1_index]['resnum']
              segid1=restr_dict['j'][contrib1_index]['segid']
              pseudoatomj=None
              checkj=0
              for sel in restr_dict['j'].keys():
                atomnamej=restr_dict['j'][sel]['atomname']
                resnamej=restr_dict['j'][sel]['resname']
                resnumj=restr_dict['j'][sel]['resnum']
                segidj=restr_dict['j'][sel]['segid']
                if not resnamej==resname1 or not resnumj==resnum1 or not segidj==segid1:
                  checkj=0
                  print 'j check1  ',resnamej,resname1,resnumj,resnum1,segidj,segid1,restr_num
                  break
                else:
                  checkj=1
              if checkj==1:
                for key in pseudoatoms_dict[resname1].keys():
                  if atomname1 in pseudoatoms_dict[resname1][key][pos] and cntj==len(pseudoatoms_dict[resname1][key][pos])-1:
                    pseudoatomj=key
                if not pseudoatomj:
                  print 'j check2  ',resname1,atomname1,pseudoatomj,cntj,restr_num
                  checkj=0
              if checkj==1:
                for sel in restr_dict['j'].keys():
                  atomnamej=restr_dict['j'][sel]['atomname']
                  resnamej=restr_dict['j'][sel]['resname']
                  resnumj=restr_dict['j'][sel]['resnum']
                  segidj=restr_dict['j'][sel]['segid']
                  if not atomnamej in pseudoatoms_dict[resname1][pseudoatomj][pos]:
                    checkj=0
                    print 'j check3  ',resnamej,atomnamej,pseudoatomj,restr_num
              if checkj==1:
                jpass='true'
                print 'j pass  ',resnamej,atomnamej,pseudoatomj,restr_num
                pseudoatomcorrj=pseudoatoms_dict[resname1][pseudoatomj][2]
              else:
                jpass='false'
                pseudoatomcorrj=0.0
            
            # calculate correction
            if ipass=='true' and jpass=='true':
              pseudoatomcorrection = pseudoatomcorri + pseudoatomcorrj
            else:
              pseudoatomcorrection = 999.0 - upl
            # WRITE ALL THE INFORMATION TO THE FILE IN SORTED FASHION
            ilist=restr_dict['i'].keys()
            ilist.sort()
            jlist=restr_dict['j'].keys()
            jlist.sort()
            for sel1 in ilist:
              for sel2 in jlist:
                if cnti+cntj==2:
                  index=0
                elif cnti+cntj>2:
                  index=index+1
                outfile_disre.write('%8i'%(restr_dict['i'][sel1]['resnum']))
                outfile_disre.write('%9s'%(restr_dict['i'][sel1]['resname']))
                outfile_disre.write('%10s'%(restr_dict['i'][sel1]['atomname']))
                outfile_disre.write('%7s'%(restr_dict['i'][sel1]['segid']))
                outfile_disre.write('%7i'%(restr_dict['j'][sel2]['resnum']))
                outfile_disre.write('%9s'%(restr_dict['j'][sel2]['resname']))
                outfile_disre.write('%10s'%(restr_dict['j'][sel2]['atomname']))
                outfile_disre.write('%7s'%(restr_dict['j'][sel2]['segid']))
                outfile_disre.write('%9.3f'%(low))
                outfile_disre.write('%9.3f'%(upl))
                outfile_disre.write('%9.3f'%(upl+pseudoatomcorrection))
                outfile_disre.write('%6i'%(index))
                outfile_disre.write('%9i\n'%(restr_num))
    except IndexError:
      pass
    i = i + 1
  outfile_disre.close()

  
def pseudocorrect():  
  # CREATE A DICTIONARY PER RESIDUE THAT CONTAINS ALL POSSIBLE PSEUDOATOM DEFINITIONS
  # FOR EACH PSEUDOATOM A LIST OF LISTS IS CREATED:
  # LIST[0] CONTAINS LIST OF ATOMS IN XPLOR NOMENCLATURE
  # LIST[1] CONTAINS LIST OF ATOMS IN IUPAC NOMENCLATURE
  # LIST[2] CONTAINS THE ACTUAL PSEUDOATOM CORRECTION IN ANGSTROM
  # THE CORRECTIONS ARE NOT THE SAME AS THE TRADITIONAL CORRECTIONS
  # THE CORRECTIONS HERE ARE THE MAXIMUM DISTANCE BETWEEN TWO EQUIVALENT PROTONS IN 
  # A GROUP
  pseudoatoms_dict = {}
  # ALA
  pseudoatoms_dict['ALA']={}
  pseudoatoms_dict['ALA']['HB*']=[['HB1','HB2','HB3','HB*'],['HB1','HB2','HB3','MB'],1.8]
  # ARG
  pseudoatoms_dict['ARG']={}
  pseudoatoms_dict['ARG']['HB*']=[['HB1','HB2','HB*'],['HB3','HB2','QB'],1.8]
  pseudoatoms_dict['ARG']['HG*']=[['HG1','HG2','HG*'],['HG3','HG2','QG'],1.8]
  pseudoatoms_dict['ARG']['HD*']=[['HD1','HD2','HD*'],['HD3','HD2','QD'],1.8]
  pseudoatoms_dict['ARG']['HH1*']=[['HH11','HH12','HH1*'],['HH11','HH12','QH1'],1.8]
  pseudoatoms_dict['ARG']['HH2*']=[['HH21','HH22','HH2*'],['HH21','HH22','QH2'],1.8]
  pseudoatoms_dict['ARG']['HH*']=[['HH11','HH12','HH21','HH22','HH*'],['HH11','HH12','HH21','HH22','QH'],4.0]
  # ASN
  pseudoatoms_dict['ASN']={}
  pseudoatoms_dict['ASN']['HB*']=[['HB1','HB2','HB*'],['HB3','HB2','QB'],1.8]
  pseudoatoms_dict['ASN']['HD*']=[['HD21','HD22','HD*'],['HD21','HD22','QD'],1.8]
  # ASP
  pseudoatoms_dict['ASP']={}
  pseudoatoms_dict['ASP']['HB*']=[['HB1','HB2','HB*'],['HB3','HB2','QB'],1.8]
  # CYS
  pseudoatoms_dict['CYS']={}
  pseudoatoms_dict['CYS']['HB*']=[['HB1','HB2','HB*'],['HB3','HB2','QB'],1.8]
  # GLN
  pseudoatoms_dict['GLN']={}
  pseudoatoms_dict['GLN']['HB*']=[['HB1','HB2','HB*'],['HB3','HB2','QB'],1.8]
  pseudoatoms_dict['GLN']['HG*']=[['HG1','HG2','HG*'],['HG3','HG2','QG'],1.8]
  pseudoatoms_dict['GLN']['HE*']=[['HE21','HE22','HE*'],['HE21','HE22','QE'],1.8]
  # GLU
  pseudoatoms_dict['GLU']={}
  pseudoatoms_dict['GLU']['HB*']=[['HB1','HB2','HB*'],['HB3','HB2','QB'],1.8]
  pseudoatoms_dict['GLU']['HG*']=[['HG1','HG2','HG*'],['HG3','HG2','QG'],1.8]
  # GLY
  pseudoatoms_dict['GLY']={}
  pseudoatoms_dict['GLY']['HA*']=[['HA1','HA2','HA*'],['HA3','HA2','QA'],1.8]
  # HIS
  pseudoatoms_dict['HIS']={}
  pseudoatoms_dict['HIS']['HB*']=[['HB1','HB2','HB*'],['HB3','HB2','QB'],1.8]
  # ILE
  pseudoatoms_dict['ILE']={}
  pseudoatoms_dict['ILE']['HG1*']=[['HG11','HG12','HG1*'],['HG13','HG12','QG'],1.8]
  pseudoatoms_dict['ILE']['HG2*']=[['HG21','HG22','HG23','HG2*'],['HG21','HG22','HG23','MG'],1.8]
  pseudoatoms_dict['ILE']['HD1*']=[['HD11','HD12','HD13','HD1*'],['HD11','HD12','HD13','MD'],1.8]
  # LEU
  pseudoatoms_dict['LEU']={}
  pseudoatoms_dict['LEU']['HB*']=[['HB1','HB2','HB*'],['HB3','HB2','QB'],1.8]
  pseudoatoms_dict['LEU']['HD1*']=[['HD11','HD12','HD13','HD1*'],['HD11','HD12','HD13','MD1'],1.8]
  pseudoatoms_dict['LEU']['HD2*']=[['HD21','HD22','HD23','HD2*'],['HD21','HD22','HD23','MD2'],1.8]
  pseudoatoms_dict['LEU']['HD*']=[['HD11','HD12','HD13','HD21','HD22','HD23','HD*'],['HD11','HD12','HD13','HD21','HD22','HD23','QD'],4.4]
  # LYS
  pseudoatoms_dict['LYS']={}
  pseudoatoms_dict['LYS']['HB*']=[['HB1','HB2','HB*'],['HB3','HB2','QB'],1.8]
  pseudoatoms_dict['LYS']['HG*']=[['HG1','HG2','HG*'],['HG3','HG2','QG'],1.8]
  pseudoatoms_dict['LYS']['HD*']=[['HD1','HD2','HD*'],['HD3','HD2','QD'],1.8]
  pseudoatoms_dict['LYS']['HE*']=[['HE1','HE2','HE*'],['HE3','HE2','QE'],1.8]
  pseudoatoms_dict['LYS']['HZ*']=[['HZ1','HZ2','HZ3','HZ*'],['HZ1','HZ2','HZ3','QZ'],1.8]
  # MET
  pseudoatoms_dict['MET']={}
  pseudoatoms_dict['MET']['HB*']=[['HB1','HB2','HB*'],['HB3','HB2','QB'],1.8]
  pseudoatoms_dict['MET']['HG*']=[['HG1','HG2','HG*'],['HG3','HG2','QG'],1.8]
  pseudoatoms_dict['MET']['HE*']=[['HE1','HE2','HE3','HE*'],['HE1','HE2','HE3','ME'],1.8]
  # PHE
  pseudoatoms_dict['PHE']={}
  pseudoatoms_dict['PHE']['HB*']=[['HB1','HB2','HB*'],['HB3','HB2','QB'],1.8]
  pseudoatoms_dict['PHE']['HD*']=[['HD1','HD2','HD*'],['HD1','HD2','QD'],4.3]
  pseudoatoms_dict['PHE']['HE*']=[['HE1','HE2','HE*'],['HE1','HE2','QE'],4.3]
  pseudoatoms_dict['PHE']['QR']=[['HD1','HD2','HE1','HE2','HZ','name HD* OR name HE* OR name HZ'],['HD1','HD2','HE1','HE2','HZ','QR'],5.0]
  # PRO
  pseudoatoms_dict['PRO']={}
  pseudoatoms_dict['PRO']['HB*']=[['HB1','HB2','HB*'],['HB3','HB2','QB'],1.8]
  pseudoatoms_dict['PRO']['HG*']=[['HG1','HG2','HG*'],['HG3','HG2','QG'],1.8]
  pseudoatoms_dict['PRO']['HD*']=[['HD1','HD2','HD*'],['HD3','HD2','QD'],1.8]
  # SER
  pseudoatoms_dict['SER']={}
  pseudoatoms_dict['SER']['HB*']=[['HB1','HB2','HB*'],['HB3','HB2','QB'],1.8]
  # THR
  pseudoatoms_dict['THR']={}
  pseudoatoms_dict['THR']['HG2*']=[['HG21','HG22','HG23','HG2*'],['HG21','HG22','HG23','MG'],1.8]
  # TRP
  pseudoatoms_dict['TRP']={}
  pseudoatoms_dict['TRP']['HB*']=[['HB1','HB2','HB*'],['HB3','HB2','QB'],1.8]
  # TYR
  pseudoatoms_dict['TYR']={}
  pseudoatoms_dict['TYR']['HB*']=[['HB1','HB2','HB*'],['HB3','HB2','QB'],1.8]
  pseudoatoms_dict['TYR']['HD*']=[['HD1','HD2','HD*'],['HD1','HD2','QD'],4.3]
  pseudoatoms_dict['TYR']['HE*']=[['HE1','HE2','HE*'],['HE1','HE2','QE'],4.3]
  pseudoatoms_dict['TYR']['QR']=[['HD1','HD2','HE1','HE2','name HD* OR name HE*'],['HD1','HD2','HE1','HE2','QR'],5.0]
  # VAL
  pseudoatoms_dict['VAL']={}
  pseudoatoms_dict['VAL']['HG1*']=[['HG11','HG12','HG13','HG1*'],['HG11','HG12','HG13','MG1'],1.8]
  pseudoatoms_dict['VAL']['HG2*']=[['HG21','HG22','HG23','HG2*'],['HG21','HG22','HG23','MG2'],1.8]
  pseudoatoms_dict['VAL']['HG*']=[['HG11','HG12','HG13','HG21','HG22','HG23','HG*'],['HG11','HG12','HG13','HG21','HG22','HG23','QG'],4.4]
  
  return pseudoatoms_dict
  
# MAIN
# ==========
#

l=len(sys.argv)
if (l!=4):
  error("Give the xplor restraint input file, input format, concoord restraint output file.\
  \ne.g. xplor2concoord restraints.log xplor concoord_restr.out");

infile=sys.argv[1]
informat=sys.argv[2]
outfile=sys.argv[3]
xplor2concoord(infile,informat,outfile)

raise SystemExit