import pyMoleCuilder as mol
import sys, os, math
try:
	import numpy
	numpy_present = True
except ImportError:
	numpy_present = False

if len(sys.argv) < 5:
	print 'Usage: '+sys.argv[0]+' <input> <path> <steps> <no_atoms> <use bondgraph>'
	sys.exit(1)

steps=int(sys.argv[3])
equilibrium_distance=1.6
no_atoms=int(sys.argv[4])
inputfile=sys.argv[1]
forcespath=sys.argv[2]
forcesfile="ising.forces"
use_bondgraph=sys.argv[5]

# creating input file
atomstart=7.6-1.6*math.floor(no_atoms/2)
print "Creating "+inputfile
with open(inputfile, 'w') as f:
	f.write("# ATOMDATA\ttype\tId\tx=3\tu=3\tF=3\tneighbors=4\n")
	f.write("# Box\t20\t0\t0\t0\t20\t0\t0\t0\t20\n")
	for i in range(1, no_atoms+1):
		atompos=atomstart+1.6*float(i)
		if i==math.floor(no_atoms/2+1):
			atompos=atompos-.5
		if i==1:
			f.write("C\t%d\t%lg\t10\t10\t0\t0\t0\t0\t0\t0\t%d\t0\t0\t0\n" % (i, atompos, i+1));
		elif i==no_atoms:
			f.write("C\t%d\t%lg\t10\t10\t0\t0\t0\t0\t0\t0\t%d\t0\t0\t0\n" % (i, atompos, i-1));
		else:
			f.write("C\t%d\t%lg\t10\t10\t0\t0\t0\t0\t0\t0\t%d\t%d\t0\t0\n" % (i, atompos, i-1, i+1));

print "Parsing from "+inputfile
mol.WorldInput(inputfile)
mol.SelectionAllAtoms()
mol.CommandVerbose("4")

# calculate damping factor from finite geometric series
# s_n/a = \sum^{n-1}_{k=0} r^k = (1-r^n)/(1-r) -> s_(n+1)/a -1 = \sum^{n}_{k=1} r^k = (1-r^(n+1))/(1-r) - 1
# \sum^{n}_{k=1} r^k := 1 and 1 = (1-r^(n+1))/(1-r) - 1 -> 2*(1-r) = 1 - r^(n+1) -> 1 - 2*r + r^(n+1) = 0
# find root: p[0] is coefficient of monomial with highest power 
if numpy_present:
	p=[0.] * (no_atoms+1)
	p[0]=1.
	p[no_atoms-1]=-2.
	p[no_atoms]=1.
	zeros=numpy.roots(p)
	print("Roots of p "+str(p)+" are "+str(zeros))
	damping=numpy.real(zeros[-1])
	print "Using damping factor of "+str(damping)
else:
	if no_atoms == 2:
		damping=0.5
	elif no_atoms == 5:
		damping=0.5436890126920764

for i in range(0, steps):
	# TODO: Python interface should have something to iterate over selected atoms
	# and molecules and get information on their internal status
	
	# read current atomic positions
	outputfile=forcespath+'/'+forcesfile+'.xyz'
	try:
		os.remove(outputfile)
	except: OSError
		#
	mol.WorldOutputAs(outputfile)
	mol.wait()
	distances=[]
	coords=[0.,0.,0.]
	try:
		skiplines=2+i*(1+1+no_atoms+1) # no_atoms, comment, no_atoms atoms, empty line
		with open(outputfile) as f:
			for line in f:
				if skiplines != 0:
					skiplines=skiplines-1
					continue
				line=line.replace('\t',' ')
				print "LINE: "+line
				[elementtype, X, Y, Z] = line.split(' ', 4)
				if coords!=[0.,0.,0.]:
					distances.append(math.sqrt((coords[0]-float(X))**2+(coords[1]-float(Y))**2+(coords[2]-float(Z))**2))
				coords=[float(X),float(Y),float(Z)]
	except IOError:
		print 'Warning: '+outputfile+' not readable.'
		sys.exit(1)
	
	assert(len(distances)==no_atoms-1)
	
	# 
	# generate Ising model forces and store in file
	#
	# i.e. we have spring forces between neighboring atoms depending on their distance
	forces=[]
	for d in distances:
		forces.append( d - equilibrium_distance );
		
	# generate new forces file
	
	with open(forcespath+'/'+forcesfile, 'w') as f:
		f.write('# atom\tf_x\tf_y\tf_z\n')
		for i in range(len(distances)+1):
			force=0
			if i!=0:
				force=force-forces[i-1]
			if (i != len(distances)):
				force=force+forces[i]
			f.write("%d\t%f\t0.\t0.\n" % (i+1, force))

	mol.WorldStepWorldTime("1")
	mol.MoleculeForceAnnealing(forcespath+'/'+forcesfile, ".1", "%d" % (steps), "%d" % (no_atoms-1), "%lg" % (damping), use_bondgraph)
	mol.wait()
	
sys.exit(0)