#!/usr/bin/env python
# syntax remcmd1.py (optional -q to queue new runs)


import sys, os, glob, re, random, time, numpy; from math import exp, sqrt


## routine to get final F (free energy including kinetic and electronic entopy)
def getF(file):
    with open(file, 'r') as f:
        for line in f:
            if "F=" in line: work=line.split("F=")[1]
        return float(work.split()[0])

## routine to alter INCAR and POSCAR files
def fix_dir(dir,Told,Tnew):

# new T in INCAR
    with open(dir+"/Trun",'w+') as fh: fh.write(str(int(Tnew))+"\n"); fh.close()
    with open(dir+"/INCAR",'r') as fh: incar = fh.readlines(); fh.close()
    for line in incar:
        if "TEBEG" in line: incar[incar.index(line)] = "TEBEG = "+str(Tnew)+"\n"
        if "TEEND" in line: incar[incar.index(line)] = "TEEND = "+str(Tnew)+"\n"
        if "SIGMA" in line: incar[incar.index(line)] = "SIGMA = "+str(kB*Tnew)+"\n"
    with open(dir+"/INCAR",'w') as fh: fh.writelines(incar); fh.close()

# scale velocities in POSCAR
    with open(dir+"/POSCAR",'r') as fh: poscar = fh.readlines(); fh.close()
    natoms = sum(map(int, poscar[6].split()))
    velo = sqrt(Tnew/Told)*numpy.loadtxt(poscar[8+natoms+1:8+natoms+1+natoms])
    v_list = []
    for v in velo: v_list.append(' '.join(map(str,v))+'\n')
    poscar[8+natoms+1:8+natoms+1+natoms] = v_list
    with open(dir+"/POSCAR",'w+') as fh: fh.writelines(poscar); fh.close()

# preliminaries

kB = 1./11604.


# check that all runs are completed, construct T hash and get parity
runs = glob.glob('run*'); nruns = len(runs); Thash = {}; Tlist = []
if os.path.isfile(runs[0]+"/Nrun"):
    with open(runs[0]+"/Nrun") as fh: Nrun = int(fh.readline())
else:
    Nrun = 0

# if more than two runs (different T's) then apply even/odd swaps
parity = 0 if nruns == 2 else Nrun%2

for run in runs:
    with open(run+"/RunAt") as fh:
        for line in fh: pass; last = line
        if not "Done" in last: print run+" not done. Exiting"; sys.exit()

        with open(run+"/Trun") as fh:
            Trun = float(fh.readline())
            Thash[Trun] = run
            Tlist.append(Trun)
            fh.close()

Tlist.sort()

#
for a in xrange(parity, nruns-1, 2):
    Ta = Tlist[a]; Tb = Tlist[a+1]
    adir = Thash[Ta]; bdir = Thash[Tb]
    Fa = getF(adir+"/output")
    Fb = getF(bdir+"/output")
    dF = Fa-Fb; dbeta=11604.0/Ta-11604/Tb
    prob = 1.000 if (dbeta*dF) > 0 else exp(dbeta*dF)
    ran = random.random()
    print "Cycle %d swap of Ta=%d with Tb=%d dbeta*dF=%.4f prob=%.4f ran=%.4f" % (Nrun,Ta,Tb,dbeta*dF,prob,ran),

    if ran < prob:
        print "accept"
        fix_dir(adir,Ta,Tb); fix_dir(bdir,Tb,Ta)
    else:
        print "reject"

if "-q" in sys.argv:
    for run in runs:
        os.chdir(run)
        os.system("qsub ../pbs_remcmd1 >> qsub.out")
        time.sleep(1) # pause 1 second for queue
        os.chdir("..")