from visual import *
import time

from mpmath import *

scene = display(title='Spin Simulation', width=600, height=400, center=(0,0,0), background=(0.0,0.0,0.0), autoscale=0, forward=vector(0.001,1.0,.001))

scene.range = 32
scene.up = vector(0.0,0.0,1.0)

nSpinsX = 4
nSpinsY = 24

scale = 1.0

vectorarray = [ None ] * nSpinsX
for x in xrange(nSpinsX):
    vectorarray[x] = [ None ] * nSpinsY
    for y in xrange(nSpinsY):
        vectorarray[x][y] = arrow(axis=(0.0,0.0,1.0), pos=(x*scale-nSpinsX/2.0,0.0,y*scale-nSpinsY/2.0),color=(0.85,0.85,0.85), width = .125)


class SpinLattice:
    def __init__(self, nSpinsX, nSpinsY):
        self.nSpinsX = int(nSpinsX)
        self.nSpinsY = int(nSpinsY)

        self.Ja = -1.0
        self.Jbx = -1.0
        self.Jby = -1.0
        self.Jad = -1.0
     
        self.initSpins()


    def initSpins(self):
        t = 0.0
        
        kx = .02
        ky = .6
        omega = .5

        alpha = .12
        beta = .16
        gamma = sqrt(1-(alpha)**2-(beta)**2)

        for x in xrange(nSpinsX):
            for y in xrange(nSpinsY):

                spinX = alpha*sin(kx*x+ky*y-omega*t)
                spinY = beta*cos(kx*x+ky*y-omega*t)
                spinZ = gamma

                direction = vector(spinX,spinY,spinZ)

                vectorarray[x][y].axis = direction

    def animateSpins(self):
        t = 0.0
        deltat = 0.05
        tend = 10.0

        kx = .02
        ky = .6
        omega = .5

        alpha = .12
        beta = .16
        gamma = sqrt(1-(alpha)**2-(beta)**2)

        while t <= tend:
            t += deltat
            for x in xrange(nSpinsX):
                for y in xrange(nSpinsY):

                    spinX = alpha*sin(kx*x+ky*y-omega*t)
                    spinY = beta*cos(kx*x+ky*y-omega*t)
                    spinZ = gamma
            
                    direction = vector(spinX,spinY,spinZ)

                    vectorarray[x][y].axis = direction


    def setCouplings(self, Ja, Jbx, Jby, Jad):          # sets the couplings of neighboring spins and creates two coupling arrays.
        self.Ja = Ja                                    # Unit Cell interior region coupling. Unit Cell Object specific. Currently a constant.
        self.Jbx = Jbx                                  # Unit Cell boundary coupling., between unit cells, in the x direction.
        self.Jby = Jby                                  # in the y direction...
        self.Jad = Jad                                  # separation of couplings between middle spins.

        self.couplingXArray = zeros((self.nSpinsX,self.nSpinsX))        # resets to all zeros
        self.couplingYArray = zeros((self.nSpinsY,self.nSpinsY))

        for i in range(self.nSpinsX):
            for j in range(self.nSpinsX):
                if      i - j == 1:
                    if      i%self.stripeSpacingX == 0 and j%self.stripeSpacingX == (self.stripeSpacingX-1):    self.couplingXArray[i][j] += self.Jbx
                    else:                                                                                       self.couplingXArray[i][j] += self.Ja 
                if      j - i == 1:
                    if      j%self.stripeSpacingX == 0 and i%self.stripeSpacingX == (self.stripeSpacingX-1):    self.couplingXArray[i][j] += self.Jbx
                    else:                                                                                       self.couplingXArray[i][j] += self.Ja
                if      i - j == self.nSpinsX - 1 or j - i == self.nSpinsX - 1:                                 self.couplingXArray[i][j] += self.Jbx

        i=0
        j=0

        for i in range(self.nSpinsY):
            for j in range(self.nSpinsY):
                if      i - j == 1:
                    if      i%self.stripeSpacingY == 0 and j%self.stripeSpacingY == (self.stripeSpacingY-1):    self.couplingYArray[i][j] += self.Jby
                    else:                                                                                       self.couplingYArray[i][j] += self.Ja
                if      j - i == 1:
                    if      j%self.stripeSpacingY == 0 and i%self.stripeSpacingY == (self.stripeSpacingY-1):    self.couplingYArray[i][j] += self.Jby
                    else:                                                                                       self.couplingYArray[i][j] += self.Ja
                if      i - j == self.nSpinsY - 1 or j - i == self.nSpinsY - 1:                                 self.couplingYArray[i][j] += self.Jby


spinLattice = SpinLattice(nSpinsX,nSpinsY)
spinLattice.setCouplings(1.0,1.0,1.0,1.0)
spinLattice.animateSpins()





