import difast
import numpy as np

muB = 9.27400915e-24
xy = 0.96929e-9

r = [[0,0,0]]
mu = [[0,muB,0]]


n_x = 50
n_y = 30

min_x = -2.5*0.96929e-9 #CuPF6 lattice spacing
max_x = 2.5*0.96929e-9 #CuPF6 lattice spacing
min_y = -1.5*0.96929e-9 #CuPF6 lattice spacing
max_y = 1.5*0.96929e-9 #CuPF6 lattice spacing

rfactor = 2.7e9

f = open('spingrid-single.gle','w')

xsize = 13.3
ysize = 8

f.write("""
xsize = """+str(xsize)+"""
ysize = """+str(ysize)+"""
size xsize ysize

pt=0.0352777778 ! 1 postscript point

set lwidth 0.5*pt
set arrowangle 18
set arrowsize 3*pt

begin translate xsize/2 ysize/2
""")

r_fast = difast.reshape_array(r)
mu_fast = difast.reshape_array(mu)

for j in range(n_x+1):
	r_test_x = min_x + (max_x-min_x)*np.float(j)/np.float(n_x)
	for k in range(n_y+1):
		r_test_y = min_y + (max_y-min_y)*np.float(k)/np.float(n_y)
		if np.abs(r_test_x*rfactor) < xsize/2 and np.abs(r_test_y*rfactor) < ysize/2:
			r_test = np.array([r_test_x,r_test_y,0])
			#B = difn.calculate_dipole(r_test, r, mu) #the old function
			r_test_fast = difast.reshape_vector(r_test)
			B = difast.calculate_dipole(r_test_fast, r_fast, mu_fast)
			
			B_unit = [B[0],B[1]]/np.sqrt(np.dot(B,B))
			B_unit *= 0.2
			if str(B_unit[0]) != 'nan':
				if abs(r_test_x*rfactor) < xsize and abs(r_test_y*rfactor) < ysize:
					col = 1-(np.log(np.dot(B,B))/-12)
					if col > 0:
						print col
						if col <= 0.25:
							colour = str(1-col*4)+',1,1'
						elif col > 0.25 and col <= 0.75:
							colour = '0,'+str(1-(col-0.25)*2)+',1'
						else:
							colour = '0,0,'+str(1-(col-0.75)*4)
						f.write('set color rgb('+colour+')\n')
						f.write('amove '+str(r_test_x*rfactor-B_unit[0]/2)+' '+str(r_test_y*rfactor-B_unit[1]/2)+'\n')
						f.write('rline '+str(B_unit[0])+' '+str(B_unit[1])+' arrow end\n')

f.write("""
set lwidth 0.2
set arrowangle 25
set arrowsize 0.5
""")

for i in range(len(r)):
	if np.abs(r[i][0]*rfactor) < xsize/2 and np.abs(r[i][1]*rfactor) < ysize/2:
		spin_unit = [mu[i][0],mu[i][1]]/np.sqrt(np.dot(mu[i],mu[i]))
		spin_unit *= 1.5
		if mu[i][1] > 0:
			color = 'rgb255(255,209,35)'
		else:
			color = 'purple'
		f.write('set color '+color+'\n')
		f.write('amove '+str(r[i][0]*rfactor-spin_unit[0]/2)+' '+str(r[i][1]*rfactor-spin_unit[1]/2)+'\n')
		f.write('rline '+str(spin_unit[0])+' '+str(spin_unit[1])+' arrow end\n')

f.write("""
end translate""")

f.close()