#!/usr/bin/env python
import random, os, re, shutil, subprocess, math

## 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 scale velocities by power of mass ratio
def vScale(vel,ratio):
# power = 0, 0.5, 1 for constant velocity, kinetic energy or momentum
    power = 0.5
    vtmp = [float(vel[i])*ratio**power for i in range(3)]
    return " ".join(format(v, ".6e") for v in vtmp)+'\n'


## preliminary setup
random.seed()
kB = 1./11604.
with open('../vasp.cmd') as f: vcmd = f.readline().split()

with open('./POSCAR', 'r') as f:
    poscar = f.readlines(); f.close()

with open('./INCAR', 'r') as f:
    incar = f.readlines(); f.close()

types = poscar[5].split()
ntypes = len(types)
numbs = map(int, poscar[6].split())
natoms = sum(numbs)
start = [0]
for t in xrange(1,ntypes): start.append(start[t-1]+numbs[t-1])

# get masses for velocity rescale
mass = []
with open('./POTCAR', 'r') as f:
    potcar = f.readlines(); f.close()
    for line in potcar:
        if "POMASS" in line: mass.append (float((line.split()[2]).replace(';', '')))

# prepare subdirectory
if not os.path.exists("mcdir"): os.mkdir("mcdir")
shutil.copy("KPOINTS","mcdir"); shutil.copy("POTCAR","mcdir")

for line in incar:
    if "NSW" in line: incar[incar.index(line)] = "NSW = 1\n"
    if "TEBEG" in line: Temp = float(line.split()[-1])
    
with open('./mcdir/INCAR', 'w') as f:
    f.writelines(incar); f.close()

# create list of swaps
swaps = []
for t1 in xrange(0,ntypes):
    for t2 in xrange(t1+1,ntypes):
        swaps.append([t1, t2])
nswaps = len(swaps)


## Create static run of swapped structure
# choose swap
swap = swaps[random.randint(0,nswaps-1)]
swap0 = random.randint(0,numbs[swap[0]]-1); swap1 = random.randint(0,numbs[swap[1]]-1)
atom0 = start[swap[0]]+swap0; atom1 = start[swap[1]]+swap1
print("attempt t0=%d a0=%d with t1=%d a1=%d" % (swap[0],swap0,swap[1],swap1)),

# POSCAR offset for first atom is 8 lines
atom0 = atom0+8; atom1=atom1+8
# swap coordinates
poscar[atom0], poscar[atom1] = poscar[atom1], poscar[atom0]
## Swapping velocities is allowed by unnecessary
## Uncomment following lines to swap and scale velocities if desired
#v0_new = vScale(poscar[atom1+natoms+1].split(),mass[swap[1]]/mass[swap[0]])
#v1_new = vScale(poscar[atom0+natoms+1].split(),mass[swap[0]]/mass[swap[1]])
#poscar[atom0] = v0_new
#poscar[atom1] = v1_new

# omit predictors and correctors since they are incorrect
with open('./mcdir/POSCAR', 'w') as f:
    f.writelines(poscar[0:8+1+2*natoms]); f.close()


## Run VASP
#vcmd=["mpirun","-n","4","--bind-to","none","pgvasp5.3.5"]
with open('./mcdir/output', 'w') as f:
    vpid = subprocess.Popen(vcmd, stdout=f, cwd="./mcdir")
    vpid.wait()
    f.close()

## find dF and accept or reject
Fi = getF("./output"); Ff = getF("./mcdir/output"); dF=Ff-Fi
print("Ff=%.3f Fi=%.3f dF=%.3f" % (Ff, Fi, dF)),
if dF < 0:
    prob = 1.000; r=1.000; accept = 1
else:
    prob = math.exp(-dF/(kB*Temp))
    r = random.random()
    if prob < r:
        accept = 0
        print("prob=%.3f rand=%.3f reject" % (prob,r))
        os.remove("mcdir/WAVECAR")
    else:
        accept = 1

if accept == 1:
    print("prob=%.3f rand=%.3f accept" % (prob,r))
    shutil.copy("mcdir/POSCAR",".")
    shutil.copy("mcdir/output",".")
    shutil.copy("mcdir/WAVECAR",".")