#!/usr/bin/python

from numpy import *
import sys
import math
import tpsinput
import takeinp
import abstracts

T = float(sys.argv[1])
nsw = int(sys.argv[2])

inpt = takeinp.TakeInput('POSCAR')
lattmat = inpt.lattmat
x_f = array(inpt.coords_d)
atoms = inpt.atoms

masses = abstracts.readMasses('MASSES')
allmasses = []
for i in range(len(atoms)):
    tmp = atoms[i]*[masses[i]]
    allmasses = allmasses+tmp
allmasses = array(allmasses)
allmasses = allmasses/6.02214e23/1000

x_f = x_f % 1
v_f = abstracts.brandNewVelocity(T, lattmat, allmasses)
x_b = x_f-v_f
x_b = x_b % 1
v_b = -v_f


abstracts.write_poscar('POSCAR.f', x_f, v_f, lattmat)
abstracts.write_poscar('POSCAR.b', x_b, v_b, lattmat)
abstracts.write_incar('INCAR.f', nsw)
abstracts.write_incar('INCAR.b', nsw-1)
