#!@PYTHON@
#
# converts mpqc output files to pcp.energy.all, pcp.forces.all and pcp.speed for joiner/analyzer to work on

import sys, string, re, math
wrerr = sys.stderr.write

if len(sys.argv) < 3:
  wrerr("Usage: "+sys.argv[0]+" <configbase> <prefix>\n")
  sys.exit(1)

output = open(sys.argv[1]+".conf.out", "r")
config = open(sys.argv[1]+".conf", "r")
energy = open(sys.argv[2]+".energy.all", "w")
forces = open(sys.argv[2]+".forces.all", "w")
speed = open(sys.argv[2]+".speed", "w")

# prepare the headers
print >>energy,"Time\t\tTotal"
print >>forces,"Type\tNo\t\tPos0\t\tPos1\t\tPos2\t\tTotal0\t\tTotal1\t\tTotal2"
print >>speed,"\nMaxG[] = nA nA nA -> N[] = nA nA nA"
print >>speed,"LevNo   MaxG    RMaxG   MaxN    N[0]    N[1]    N[2]    RLevSStep"
print >>speed,"0       5147691 10295381        25165824        384     256    256     0"
print >>speed,"1       643665  1287329 3145728 192     128     128     267"
print >>speed,"2       80442   160883  393216  96      64      64      267"
print >>speed,"3       10074   20147   49152   48      32      32      333"
print >>speed,"4       1264    2527    6144    24      16      16      862"
print >>speed,"procs   proc[0] proc[1] PsiMax  PsiDoub PsiUp   PsiDown NRStep  TotalIons"
print >>speed,"1       1       1       $Psi    $Psi+0    0+0     0+0     1       $nr"

# parse the config file
nr=0
limit=0
for line in config:
    nr=nr+1
    if "# ====== MD step" in line:
        limit=nr
config.seek(0)
#print "Beginning with line %d ..." % (limit)
n=0
nr=0
atompos = []
for line in config:
    nr=nr+1
    #print "%d > %d ?" % (limit, nr)
    if limit > nr:
        continue
    #print "%d: %s" % (nr,line)
    if re.compile("Ion_Type[0-9]+[_]+").match(line):
        words = line.split("Type")
        words2 = (words[0]+"Type_"+words[1]).split()
        #print words2
        #print words
        atompos.append(words2[0]+"_"+words2[1]+"_"+words2[2]+"_"+words2[3])
        n=n+1
print "I found ",n," atoms."

if n == 0:
    wrerr("Quitting.\n")
    sys.exit(1)

# parse the output file for ...
line = output.readline()
EnergyWritten = 0
while (line != ''):
    # energy
    if "total scf energy" in line:
        if EnergyWritten == 0:
            words = line.split()
            #print words
            print >>energy,"0.000000e+00\t",words[4]		
            EnergyWritten = 1
    # forces
    if "Total Gradient:" in line:
        #print "found gradient"
        mean=0.
        for i in range(n):
             line = output.readline()
             words = line.split()
             atome = atompos[i].split("_")
             #print atome
             #print words
             print >>forces,"%d\t%d\t%s\t%s\t%s\t%s\t%s\t%s" % ( (int(atome[2])-1),(int(atome[3])-1),atome[4],atome[5],atome[6],words[2],words[3],words[4] )
             mean += math.sqrt(math.pow(float(words[2]),2)+math.pow(float(words[3]),2)+math.pow(float(words[4]),2))
        print >>forces,"MeanForce:\t",mean/n
    # speed
    if "mpqc:" in line:
        types=["Types"]
        cpu=["CPU"]
        wall=["Wall"]
        for i in range(25):
             words = line.split(":")
             times = words[1].split()
             types.append(words[0])
             cpu.append(times[0])
             wall.append(times[1])
             line = output.readline()
        for j in range(25):
              types[j] = types[j].replace('_', '')	# TeX does not like underscores in non-math env
              speed.write(types[j]+"\t")
        speed.write("\n")
        for j in range(25):
              speed.write(cpu[j]+"\t")
        speed.write("\n")
        for j in range(25):
              speed.write(wall[j]+"\t")
        speed.write("\n")
    line = output.readline()