import sys

def getenergies(outputfilename):
   import sys
   marker = 0  # a marker used to keep track of where we are in the mdout file

   # Parsing the mdout file. We are parsing potential terms that look like:
   #
   #   NSTEP       ENERGY          RMS            GMAX         NAME    NUMBER
   #      1      -5.0085E+03     1.7869E+01     1.0861E+02     C        3395
   #
   # BOND    =      747.3975  ANGLE   =     2159.7999  DIHED      =     2674.1632
   # VDWAALS =    -2025.6491  EEL     =   -17003.5331  EGB        =    -3086.6385
   # 1-4 VDW =      905.5446  1-4 EEL =    10529.1160  RESTRAINT  =        0.0000
   # ESURF   =       91.3192
   bond=[]
   angle=[]
   dihed=[]
   vdwaals=[]
   eel=[]
   egb=[]
   vdwof=[]
   eelof=[]
   esurf=[]

   e_int=[]
   g_solv=[]
   total=[]

   outputfile = open(outputfilename, 'r')

   for outputline in outputfile: # outputfile was passed as an open file object. Loop through the lines

      if 'error' in outputline.lower(): # look for errors in the sander mdout
         print >> sys.stderr, 'Error: Sander error ({0})'.format(outputline.strip())
         return -1
      if 'BOND' in outputline:  # if the BOND potential term label is in this line:
         words = outputline.split() # split along whitespace
         marker = 1  # we are in the first line we are parsing values from
         try:
            bond.append(float(words[2]))  # bond value is 3rd in array after BOND and =
            angle.append(float(words[5])) # angle is 6th
            dihed.append(float(words[8])) # dihedral is 9th
         except ValueError: # this will occur if the values are ********** (incompatible prmtop/mdcrds)
            print >> sys.stderr, 'Error: Sander output is missing values!' # print error message
            print >> sys.stderr, outputline.strip() # print offending line and return error code
            return -1
      elif marker == 1:   # this occurs the cycle after BOND was found. It is the next line
         words = outputline.split() # split along whitespace
         marker += 1 # move the marker along (should be 2 now)
         try:
            vdwaals.append(float(words[2])) # vdw is 3rd word after VDWAALS and =
            eel.append(float(words[5]))   # eel is 6th
            egb.append(float(words[8]))   # egb is 9th
         except ValueError:  # find ********* in mdout
            print >> sys.stderr, 'Error: Sander output is missing values!' # print error message
            print >> sys.stderr, outputline.strip()  # and the offending line and return error code
            return -1
      elif marker == 2: # this occurs 2 lines after BOND was found.
         words = outputline.split() # split along whitespace
         marker = marker + 1 # move marker along
         try:
            vdwof.append(float(words[3])) # 1-4vdw is the 4th word (after 1-4, VDW, and =)
            eelof.append(float(words[7])) # 1-4 eel is the 8th word
         except ValueError: # find ********** in mdout
            print >> sys.stderr, 'Error: Sander output is missing values!' # print error message
            print >> sys.stderr, outputline.strip() # print offending line and return error code
            return -1
      elif marker == 3: # this occurs 3 lines after BOND was found
         words = outputline.split() # split along whitespace
         marker = marker + 1 # move marker along
         try:
            esurf.append(float(words[2])) # esurf is the 3rd word after ESURF and =
         except ValueError: # find ******* in mdout
            print >> sys.stderr, 'Error: Sander output is missing values!' # print error message
            print >> sys.stderr, outputline.strip() # print offending line and return error code
            return -1

   if len(bond) == 0: # if we went through whole mdout and didn't find BOND once...
      print >> sys.stderr, 'Error: No potential terms in sander output! Check output files.'
      return -1 # return error code after printing error message
   

   for x in range(len(bond)): # all arrays will be same length if script made it this far
      try: # add all potential terms
          sum = bond[x] + angle[x] + dihed[x] + vdwaals[x] + eel[x] + egb[x] + vdwof[x] + eelof[x] + esurf[x]
          e_int_sum = bond[x] + angle[x] + dihed[x] + vdwaals[x] + eel[x] +  vdwof[x] + eelof[x]
          g_solv_sum = egb[x] + esurf[x]
      except IndexError: # catch this very strange error if it occurs... it never should
         print >> sys.stderr, 'Error: Inconsistent number of potential terms in mdouts. This is a strange error'
         return -1
      
      e_int.append(e_int_sum)
      g_solv.append(g_solv_sum)
      total.append(sum) # append sum to the total array


   file = open(outputfilename+"_all_energies.dat", 'w')

   h1=""
   for i in range(1, len(bond)+1):
       h1 = h1 + "frame" + str(i)+ ", "
   h1=h1.strip(", ")+"\n"
   
   file.write(h1)
   
   file.write("bond, " +str(bond).strip("[").strip("]")+'\n')
   file.write("angle, " + str(angle).strip("[").strip("]")+'\n')
   file.write("dihed, " + str(dihed).strip("[").strip("]")+'\n')
   file.write("vdwaals, " + str(vdwaals).strip("[").strip("]")+'\n')
   file.write("eel, " + str(eel).strip("[").strip("]")+'\n')
   file.write("egb, " + str(egb).strip("[").strip("]")+'\n')
   file.write("vdwof, " + str(vdwof).strip("[").strip("]")+'\n')
   file.write("eelof, " + str(eelof).strip("[").strip("]")+'\n')
   file.write("esurf, " + str(esurf).strip("[").strip("]")+'\n')

   file.write("e_int, " + str(e_int).strip("[").strip("]")+'\n')
   file.write("g_solv, " + str(g_solv).strip("[").strip("]")+'\n')
   file.write("total, " + str(total).strip("[").strip("]")+'\n')
   file.close()


   return 1





if __name__ == '__main__':
    
    getenergies(sys.argv[1])