#!/usr/bin/python

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

tps = tpsinput.Tpsinput()
shootfactor = tps.SHOOTFACTOR
timefactor = tps.TIMEFACTOR

cfile = 'trajectory.dat'
ufile = open(cfile, 'r')
xdat = cPickle.load(ufile)
ufile.close()

nsw = len(xdat)
inpt = takeinp.TakeInput('POSCAR')
lattmat = inpt.lattmat
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)

rnd = random.random()
f = open('report', 'a')
if rnd > timefactor:
    x_f, v_f, cros = abstracts.newShootingPath(
        nsw, xdat, lattmat, allmasses, shootfactor)
    x_f = x_f % 1
    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-cros)
    abstracts.write_incar('INCAR.b', cros)
    print 1
else:
    x, v, cros, test = abstracts.newTimeshiftPath(inpt, xdat, nsw)
    x = x % 1
    if test == 1:
        abstracts.write_poscar('POSCAR.f', x, v, lattmat)
        abstracts.write_incar('INCAR.f', nsw-cros)
        abstracts.write_tmpXdat(xdat, cros, test)
    elif test == 2:
        abstracts.write_poscar('POSCAR.b', x, v, lattmat)
        abstracts.write_incar('INCAR.b', nsw-cros)
        abstracts.write_tmpXdat(xdat, cros, test)
f.write('The crossing point:'+' ' +`cros`+ '\n')
f.close()
