from visual import *
from random import random
from Numeric import *
import time,sys,thread

class Spin:                                                         # A single spin object for numeric use only.  It is really just a specific set of data structures
    indices = [0,0,0]
    couplings = [0,0,0,0]
    neighbors = array(((0,0,0,0),(0,0,0,0),(0,0,0,0),(0,0,0,0)))    # first two entries are the unit cell indices m and n, to be filled in when composing the lattice
                                                                    # second two entries are the spin indices x and y, within the unit cell.
    spin = vector()
    torque = vector()


class NumericUnitCell:                          # Numeric unit cells are purely devoted to the numeric aspects of the simulation.
    nSpinsX = 2
    nSpinsY = 3

    dx = 1.0                                    # needed for the scale of k.  But, we assume natural units, so this shouldnt need to be changed.
    dy = 1.0

    Ja = -1.0                                   # Unit Cell interior region coupling. Unit Cell Object specific. Currently a constant.
    Jbx = -1.0                                  # Unit Cell boundary coupling., between unit cells, in the x direction.
    Jby = 0.8                                   # in the y direction...

    numericArray = []                           # A list (will really be 2d) storing all relevant numeric information about the unit cell

    def __init__(self, nSpinsX, nSpinsY):
        self.nSpinsX = nSpinsX
        self.nSpinsY = nSpinsY
        
        numericArray = [ None ] * nSpinsX                               # a vector of null elements, nSpinsX long
        for x in range(nSpinsX):
            numericArray[x] = [0] * nSpinsY                             # a vector of zeros nSpinsY long
            for y in range(nSpinsY):
                numericArray[x][y] = Spin()                             # Individual elements are individual spin objects

                numericArray[x][y].indices = [x,y,1-2*((x+y)%2)]        # x and y, then +-1 for spin up and down
                for nn in range(4):
                    if nn == 0: numericArray[x][y].neighbors[nn] = [0,0,(x-1)%nSpinsX,y]        # establishes periodic boundary conditions in order up down left right...
                    if nn == 1: numericArray[x][y].neighbors[nn] = [0,0,(x+1)%nSpinsX,y]        # nn=0, up; nn=1, down; nn=2, left; nn=3, right;, nn%2 -> 0101; more elegant way?
                    if nn == 2: numericArray[x][y].neighbors[nn] = [0,0,x,(y-1)%nSpinsY]        # note this makes no modification to the upper level m or n indices
                    if nn == 3: numericArray[x][y].neighbors[nn] = [0,0,x,(y+1)%nSpinsY]

        self.numericArray = numericArray

    def setCouplings(self, Ja, Jbx, Jby):           # sets the couplings and also all indices for referencing of neighboring spins.
        self.Ja = Ja
        self.Jbx = Jbx
        self.Jby = Jby

        for row in self.numericArray:
            for spin in row:
                for nn in range(4):
                    spin.couplings[nn] = self.Ja                                        # default to interior coupling
                    if nn == 0 and spin.indices[0]%self.nSpinsX == 0:                   spin.couplings[nn] = self.Jbx
                    if nn == 1 and spin.indices[0]%self.nSpinsX == (self.nSpinsX-1):    spin.couplings[nn] = self.Jbx
                    if nn == 2 and spin.indices[1]%self.nSpinsY == 0:                   spin.couplings[nn] = self.Jby
                    if nn == 3 and spin.indices[1]%self.nSpinsY == (self.nSpinsY-1):    spin.couplings[nn] = self.Jby

##                    print 'spin (',spin.indices[0],spin.indices[1],') has coupling', spin.couplings[nn], 'to spin with index', spin.neighbors[nn]

        self.resetCommonAxis()


    def setState(self, k, baseSigma):           # sets up the initial state of the spins for a particular k.
        if k.x != 0:
            kAngle = arctan(k.y/k.x)
        else:
            kAngle = pi/2.       

        chi = abs((cos(k.x*self.dx)+cos(k.y*self.dy))/2)
        sigmaRatio = .5*(1+1/chi-((1/chi-1)*(1/chi+3))**(.5))       # calculated from k,... could be an array in general

        if abs(sigmaRatio) > 1:
            sigmaRatio = 1/sigmaRatio
         
        for row in self.numericArray:
            for spin in row:
                tempSigma = (1-2.0*((spin.indices[0]+spin.indices[1])%2))*baseSigma + ((spin.indices[0]+spin.indices[1])%2)*baseSigma*sigmaRatio
                tempAngle = dot(k,vector(spin.indices[0]*self.dx, spin.indices[1]*self.dy, 0.0)) + kAngle + spin.indices[2]*pi/2.0
                tempSpinZ = spin.indices[2]*(1.0-tempSigma**2)**(1/2.0)
                spin.spin = vector(tempSigma*cos(tempAngle),tempSigma*sin(tempAngle),tempSpinZ)

        self.resetCommonAxis()



    def rotateSpins(self, thetaZ, thetaX):
        rotationZArray = array(((cos(-thetaZ),sin(-thetaZ),0),(-sin(-thetaZ),cos(-thetaZ),0),(0,0,1)))  # CCW rotation about Z
        rotationXArray = array(((1,0,0),(0,cos(-thetaX),sin(-thetaX)),(0,-sin(-thetaX),cos(-thetaX))))  # CCW rotation about X
##        rotationYArray = array(((cos(-thetaY),0,-sin(-thetaY)),(0,1,0),(sin(-thetaY),0,cos(-thetaY))))

        for row in self.numericArray:
            for spin in row:
                spin.spin = matrixmultiply(rotationZArray,spin.spin)
                spin.spin = matrixmultiply(rotationXArray,spin.spin)

                

    def getCommonAxis(self):
        self.getTorques()                           # it is required that the torque matrix be populated with valid data for the current state of the system, especially if torques have yet to be computed
        commonAxis = vector(0.0,0.0,1.0)
        
        if self.nSpinsX > 1:    commonAxis = cross(self.numericArray[0][0].torque,self.numericArray[1][0].torque)       ## ??? only utilizes two spins ???
        elif self.nSpinsY > 1:  commonAxis = cross(self.numericArray[0][0].torque,self.numericArray[0][1].torque)
        else:                   commonAxis = self.numericArray[0][0].spin

        if commonAxis[2] <= 0.0: commonAxis = -commonAxis
        if commonAxis[2] == 0.0: commonAxis = (0.0,0.0,1.0)
        
        return norm(commonAxis)

    

    def resetCommonAxis(self):
        currentAxis = self.getCommonAxis()          # more direct measures could be used...

##        print 'currentAxis is:', currentAxis

        if currentAxis[1] != 0.0:  thetaZ = arctan(currentAxis[0]/currentAxis[1])
        else: thetaZ = pi/2.0
        
        rotationZArray = array(((cos(-thetaZ),sin(-thetaZ),0),(-sin(-thetaZ),cos(-thetaZ),0),(0,0,1)))
        currentAxisTwo = matrixmultiply(rotationZArray,currentAxis)

##        print 'currentAxisTwo is:', currentAxisTwo

        if currentAxisTwo[2] != 0.0: thetaX = arctan(currentAxisTwo[1]/currentAxisTwo[2])
        else: thetaX = pi/2.0

##        rotationXArray = array(((1,0,0),(0,cos(-thetaX),sin(-thetaX)),(0,-sin(-thetaX),cos(-thetaX))))
##        finalCommonAxis = matrixmultiply(rotationXArray,currentAxisTwo)
##        print 'finalCommonAxis is:', finalCommonAxis

        self.rotateSpins(thetaZ,thetaX)

##        for row in self.numericArray:
##            for spin in row:
##                spin.spin = matrixmultiply(rotationZArray,spin.spin)
##                spin.spin = matrixmultiply(rotationXArray,spin.spin)



    def getTorques(self):
        for row in self.numericArray:
            for spin in row:
                spin.torque = vector(0,0,0)
                for nn in range(4):
                    nnx = spin.neighbors[nn][2]
                    nny = spin.neighbors[nn][3]
                    spin.torque = spin.torque + spin.couplings[nn]*cross(spin.spin,self.numericArray[nnx][nny].spin)


##                spin.torque = vector(0.0,0.0,0.0)
##                neighbor_total = vector(0.0,0.0,0.0)
##                nx = spin.indices[0]
##                ny = spin.indices[1]
##            
##                left_neighbor = spin.couplings[0]*self.numericArray[(nx-1)%self.nSpinsX][ny].spin
##                right_neighbor = spin.couplings[1]*self.numericArray[(nx+1)%self.nSpinsX][ny].spin
##                up_neighbor = spin.couplings[2]*self.numericArray[nx][(ny-1)%self.nSpinsY].spin
##                down_neighbor = spin.couplings[3]*self.numericArray[nx][(ny+1)%self.nSpinsY].spin
##
##                if ny == 0: up_neighbor = vector(up_neighbor.x*cos(k.y*self.nSpinsY)+up_neighbor.y*sin(k.y*self.nSpinsY),up_neighbor.y*cos(k.y*self.nSpinsY)-up_neighbor.x*sin(k.y*self.nSpinsY),up_neighbor.z)
##                if ny == self.nSpinsY-1: down_neighbor = vector(down_neighbor.x*cos(-k.y*self.nSpinsY)+down_neighbor.y*sin(-k.y*self.nSpinsY),down_neighbor.y*cos(-k.y*self.nSpinsY)-down_neighbor.x*sin(-k.y*self.nSpinsY),down_neighbor.z)
##                if nx == 0: left_neighbor = vector(left_neighbor.x*cos(k.x*self.nSpinsX)+left_neighbor.y*sin(k.x*self.nSpinsX),left_neighbor.y*cos(k.x*self.nSpinsX)-left_neighbor.x*sin(k.x*self.nSpinsX),left_neighbor.z)
##                if nx == self.nSpinsX-1: right_neighbor = vector(right_neighbor.x*cos(-k.x*self.nSpinsX)+right_neighbor.y*sin(-k.x*self.nSpinsX),right_neighbor.y*cos(-k.x*self.nSpinsX)-right_neighbor.x*sin(-k.x*self.nSpinsX),right_neighbor.z)
##                        
##                neighbor_total = left_neighbor + right_neighbor + up_neighbor + down_neighbor
##                
##                spin.torque = spin.torque + cross(spin.spin, neighbor_total)


    def thermalize(self, temp):                             # adds a random torque <--> kick to each spin, with a scale set by temp.
        if temp != 0:
            for row in self.numericArray:
                for spin in row:                                # perturbed by unit ball
                    randomU = random()
                    randomV = random()                          # we produce a uniform distribution on the unit ball
                    theta = 2.0*3.1416*randomU
                    phi = arccos(2.0*randomV-1.0)
                    spin.kick = temp*norm(vector(sin(theta)*cos(phi),cos(theta)*sin(phi),cos(phi)))

                    spin.torque = spin.torque + spin.kick       # this kick is constant in magnitude but random in direction


    def addExternalField(self, MFd, MFu):
        if MFd != 0 or MFu != 0:
            for row in self.numericArray:
                for spin in row:
                    nz = spin.indices[2]
                    externalField = vector(0.0,0.0,(1.0-nz)*MFd + (1.0+nz)*MFu)           # applied mean field, vertical only
                    
                    spin.torque = spin.torque + cross(spin.spin,externalField)



    def timeEvolveState(self, dt):                          # This takes us one step forward in time by adding the torque to our spins and renormalizing.
        for row in self.numericArray:                       # A Linear Approximation, with interval spacing in t of dt.
            for spin in row:
                spin.spin = spin.spin + spin.torque*dt                 # (Should we use RUNGA KUTTA?)
                spin.spin = norm(spin.spin)


                                 
    def randomizeState(self):
        for row in self.numericArray:
            for spin in row:
                    randomU = random()
                    randomV = random()
                    theta = 2.0*3.1416*randomU
                    phi = arccos(2.0*randomV-1.0)
                    
                    spin.spin = norm(vector(sin(theta)*cos(phi),cos(theta)*sin(phi),cos(phi)))
        self.resetCommonAxis()


    def getSpinSum(self):
        sum = vector(0.0,0.0,0.0)
        for row in self.numericArray:
            for spin in row:
                sum = sum + spin.spin
        return sum



        

    def getState(self):                         # handy, since we dont want to have to know the name of the object's actual variable!
        return self.numericArray



###############################################################
##      TESTING ROUTINES

##
##dt = .001
##
##
##Nx = 800
##Ny = 400
##uc = NumericUnitCell(Nx,Ny)
##
##Ja = -1.0
##Jbx = 1.8
##Jby = -1.0
##uc.setCouplings(Ja,Jbx,Jby)
##
##ks = vector(1.0,.5,0.0)
##baseSigma = 0.001
##uc.setState(ks,baseSigma)
##
##
##
##
##
##
##
##print '\n-----------------------------'
##print "Calculation Instance Started"
##print time.ctime()
##
##
##
##k = 0
##loops = 10
##while k <= loops:
##    k+=1
##
##    uc.getTorques()
##    uc.timeEvolveState(dt)
##
##
##print "\nCalculation Instance Ended For",loops,"Loops"
##print time.ctime()
##print '-----------------------------'




##  63 seconds for 100,000 iterations with a 3 by 2 unit cell = 6 spins
##   6 seconds for 10,000 iterations with a 3 by 2 unit cell = 6 spins

##   4 seconds for 1,000 iterations with an 8 by 4 unit cell = 32 spins
##  33 seconds for 100 iterations with an 80 by 40 unit cell = 3,200 spins

## 301 seconds for 10 iterations with an 800 by 400 unit cell = 320,000 spins


# effective rate is 10,000 ops/spin/second

