from visual import *
from random import random
from time import clock

scene.width = 800
scene.height = 400
scene.autoscale = 0
scene.up = (0,0,1)
scene.forward = (0,1,-.5)
scene.lights = [vector(.5,-.5,-.5), vector(-.5,.5,-.5)]

spins = []

def lattice(Ntotx=8, Ntoty=8, deltax=1.0, deltay=1.0):
    xmin = -Ntotx*deltax/2.
    ymin = -Ntoty*deltay/2.
    nz = 0
    
    for nx in range(Ntotx):
          x = xmin + nx*deltax
          for ny in range(Ntoty):
              y = ymin + ny*deltay

              spins.append(frame())
              spins[-1].pos = vector(x,y,0)
              spins[-1].spin = vector()
              spins[-1].spinvec = arrow(pos=(x,y,0), axis=(0,0,1), color=(.9,.9,.9), shaftwidth=0.2)
          
              spins[-1].torque = vector()
#              spins[-1].torquevec = arrow(pos=(x,y,0), axis=(0,0,1), color=(.3,.3,1), shaftwidth=0.05)
          
              spins[-1].nearx = range(2)
              spins[-1].neary = range(2)
              spins[-1].indices = (nx,ny,nz)
          
    for s in spins:
          nx, ny, nz = s.indices
          if nx == 0:                       # leftmost spin in a row
              nspinl = Ntotx*ny + Ntotx-1       # wrap around to spin on right side
              s.nearx[0] = nspinl               # reference this by its list element number, given by the order in which the spins are added to the list...
          else:
              nspinl = Ntotx*ny + nx-1
              s.nearx[0] = nspinl
          if nx == Ntotx-1:                 # rightmost spin in a row
              nspinr = Ntotx*ny
              s.nearx[1] = nspinr
          else:
              nspinr = Ntotx*ny + nx+1 
              s.nearx[1] = nspinr

          if ny == 0:                       # bottom spin in a column
              nspind = Ntotx*(Ntoty-1) + nx 
              s.neary[0] = nspind
          else:
              nspind = Ntotx*(ny-1) + nx
              s.neary[0] = nspind
          if ny == Ntoty-1:                 # top spin in a column
              nspinu = nx
              s.neary[1] = nspinu
          else:
              nspinu = Ntotx*(ny+1) + nx
              s.neary[1] = nspinu
    return spins


Ntotx = 40
Ntoty = 40
Ntotz = 1
Ntotal = Ntotx*Ntoty*Ntotz

dx = 2
dy = 2

Lx = dx*Ntotx
Ly = dy*Ntoty


spins = lattice(Ntotx, Ntoty, deltax=dx, deltay=dy)

dt = .05
J = -1
k = 2*pi*vector(7./Lx, 3./Ly, 0)

spinvec_length = 2.2
spinvec_sigma = 1
spinvecz_mag = (spinvec_length**2-spinvec_sigma**2)**(.5)

spin_length = 1.
spin_sigma = .01
spinz_mag = (spin_length**2-spin_sigma**2)**(.5)


for s in spins:
    spinx = spin_sigma*cos(dot(k, s.pos))
    spiny = spin_sigma*sin(dot(k, s.pos))
    spinz = 1*(J/abs(J))**(s.indices[0] + s.indices[1])
    
    s.spin = vector(spinx,spiny,spinz)

#    s.pos.z = s.pos.z + (-1)**(s.indices[0] + s.indices[1])
    
    s.spinvec.axis = s.spin*spinvec_length
    s.spinvec.pos = s.pos - s.spinvec.axis/2.


tt = clock()
Nsteps = 0
while 1:
    for s in spins:
        s.torque = vector(0,0,0)

        for nnx in range(2):
            nspin = s.nearx[nnx]
            s.torque = s.torque+cross(s.spin,spins[nspin].spin)
        for nny in range(2):
            npsin = s.neary[nny]
            s.torque = s.torque+cross(s.spin,spins[nspin].spin)
            
        s.torque.z = 0
        
    for s in spins:
        spinz_dir = s.spin.z/abs(s.spin.z)                      # preserve spin z direction
        s.spin.z = 0                                            # set z component to zero
        s.spin = s.spin+J*s.torque*dt                           # add torque (torque z component zero)
        spinxy = s.spin/mag(s.spin)                             # normed spin xy plane projection vector for use with next statement and for spinvec

        s.spin = spin_sigma*spinxy                              # make spin xy plane projection magnitude sigma
        s.spin.z = spinz_mag*spinz_dir                          # reset z component of spin

        s.spin = s.spin/mag(s.spin)                             # ensure spin remains magnitude 1
        
        s.spinvec.axis = spinxy*spinvec_sigma + vector(0,0,spinvecz_mag*spinz_dir)
        s.spinvec.pos = s.pos - s.spinvec.axis/2.

#        s.torquevec.axis = .5*J*s.torque/(mag(s.torque)+.001)
#        s.torquevec.pos = s.pos + s.spinvec.axis/2.



    if Nsteps == 1000:
        tt = clock()-tt
        print '%0.1f' % tt, 'seconds for', Nsteps, 'steps with', Ntotal, 'spins'
    Nsteps = Nsteps+1

