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

class NumericUnitCell:
    nSpinsX = 2
    nSpinsY = 3

    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...

    spinArray = []          # spin vectors in a matrix
    meanFieldArray = []     # torque vectors in a matrix
    torqueArray = []

    sigmaArray = []         # an array of sigma values
    
    couplingXArray = []     # coupling array for X direction (Up and Down)
    couplingYArray = []     # coupling array for Y direction (Left and Right)

    
    
    def __init__(self, nSpinsX, nSpinsY):
        self.nSpinsX = nSpinsX
        self.nSpinsY = nSpinsY
        
        self.spinArray = zeros((nSpinsX,nSpinsY,1,3),Float32)           # vectors
        self.meanFieldArray = zeros((nSpinsX,nSpinsY,1,3),Float32)      # vectors
        self.torqueArray = zeros((nSpinsX,nSpinsY,1,3),Float32)         # vectors

        self.sigmaUPArray = zeros((nSpinsX,nSpinsY),Float32)            # scalars
        self.sigmaDOWNArray = zeros((nSpinsX,nSpinsY),Float32)          # scalars
        self.sigmaArray = zeros((nSpinsX,nSpinsY),Float32)              # scalars
        
        self.couplingXArray = zeros((nSpinsX,nSpinsX),Float32)          # scalars
        self.couplingYArray = zeros((nSpinsY,nSpinsY),Float32)          # scalars



    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 i in range(self.nSpinsX):
            for j in range(self.nSpinsX):
                self.couplingXArray[i][j] = 0.0
                if i - j == 1 or j - i == 1:                                    self.couplingXArray[i][j] += self.Ja
                if i - j == self.nSpinsX - 1 or j - i == self.nSpinsX - 1:      self.couplingXArray[i][j] += self.Jbx

        for i in range(self.nSpinsY):
            for j in range(self.nSpinsY):
                self.couplingYArray[i][j] = 0.0
                if i - j == 1 or j - i == 1:                                    self.couplingYArray[i][j] += self.Ja
                if i - j == self.nSpinsY - 1 or j - i == self.nSpinsY - 1:      self.couplingYArray[i][j] += self.Jby


                
    def setState(self, k, baseSigma):
        chi = cos(k.x)+cos(k.y)
        if chi != 0.0:  sigmaRatio = abs((2.0 + (4.0-chi**2.0)**.5)/chi)
        else: sigmaRatio = 0.0

        if sigmaRatio > 1.0: sigmaRatio = 1.0/sigmaRatio

        sigmaRatio = 1.0

        self.sigmaArray = [[1.0,sigmaRatio],[0,0],[sigmaRatio,1.0]]


##        self.phaseArray = [[0,0],[0,0],[0,0]]
    
##        for x in range(self.nSpinsX):
##            for y in range(self.nSpinsY):
##                self.sigmaUPArray[x][y] = ((x+y+1)%2)
##                self.sigmaDOWNArray[x][y] = sigmaRatio*((x+y)%2)
##                self.sigmaArray[x][y] = baseSigma*(self.sigmaUPArray[x][y] + self.sigmaDOWNArray[x][y])
##
      
        if k.x != 0:    kAngle = arctan(k.y/k.x)
        else:           kAngle = pi/2. 

        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                z = (-1.0)**(x+y)
                tempSigma = baseSigma*self.sigmaArray[x][y]
                tempAngle = dot(k,vector(float(x)/float(self.nSpinsX-1), float(y)/float(self.nSpinsY-1), 0.0)) + kAngle + (z+1)*pi/2.0
                tempSpinZ = z*(1.0-tempSigma**2)**.5
                
                self.spinArray[x][y] = (tempSigma*cos(tempAngle),tempSigma*sin(tempAngle),tempSpinZ)


    def normalizeState(self):
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.spinArray[x][y][0] = norm(self.spinArray[x][y][0])



    def getMeanFields(self):
        flippedLatticeArray = transpose(self.spinArray)
        twiceflippedLatticeArray = transpose(transpose(self.spinArray,(1,0,2,3)))

        meanFieldXArray = matrixmultiply(flippedLatticeArray,self.couplingXArray)
        meanFieldYArray = matrixmultiply(twiceflippedLatticeArray,self.couplingYArray)
        
        meanFieldXArray = transpose(meanFieldXArray)
        meanFieldYArray = transpose(meanFieldYArray, (2,3,1,0))
        
        self.meanFieldArray = meanFieldXArray + meanFieldYArray




        
    def addExternalField(self, MFu, MFd):           ########## to move to array language, make array, then operate with array
        if MFd != 0 or MFu != 0:
            for x in range(self.nSpinsX):
                for y in range(self.nSpinsY):
                    nz = -(-1)**((1+x+y)%2)
                    externalField = vector(0.0,0.0,0.5*(1.0-nz)*MFd + 0.5*(1.0+nz)*MFu)           # applied mean field, vertical only
                    self.meanFieldArray[x][y] = self.meanFieldArray[x][y] + externalField

    def thermalize(self, temp):                             # adds a random torque <--> kick to each spin, with a scale set by temp.
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):                             # 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)
                kick = temp*norm(vector(sin(theta)*cos(phi),cos(theta)*sin(phi),cos(phi)))

                self.meanFieldArray[x][y] = self.meanFieldArray[x][y] + kick       # this kick is constant in magnitude but random in direction



 

    def getTorques(self):
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.torqueArray[x][y] = cross(self.spinArray[x][y][0],self.meanFieldArray[x][y][0])


    def measurek(self):
        commonAxis = self.getCommonAxis()

        projectedUL = norm(self.spinArray[0][0][0] - dot(self.spinArray[0][0][0],commonAxis)*commonAxis)
        projectedUR = norm(self.spinArray[self.nSpinsX-1][0][0] - dot(self.spinArray[self.nSpinsX-1][0][0],commonAxis)*commonAxis)
        projectedBL = norm(self.spinArray[0][self.nSpinsY-1][0] - dot(self.spinArray[0][self.nSpinsY-1][0],commonAxis)*commonAxis)
        
        dotx = dot(projectedUL,((-1.0)**(self.nSpinsX-1))*projectedUR)
        doty = dot(projectedUL,((1.0)**(self.nSpinsY-1))*projectedBL)
       
        if self.nSpinsX > 1:
            if abs(dotx) < 1.0:     kx = arccos(dotx)
            else:                   kx = pi
        else:                       kx = 0.0

        if self.nSpinsY > 1:
            if abs(doty) < 1.0:     ky = arccos(doty)
            else:                   ky = pi
        else:                       ky = 0.0

        return vector((kx%(2*pi)),(ky%(2*pi)),0.0)



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


    def randomizeState(self):
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                randomU = random()
                randomV = random()
                theta = 2.0*3.1416*randomU
                phi = arccos(2.0*randomV-1.0)
                
                self.spinArray[x][y][0] = norm(vector(sin(theta)*cos(phi),cos(theta)*sin(phi),cos(phi)))



    def getSpinSum(self):
        sum = vector(0.0,0.0,0.0)
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                sum = sum + self.spinArray[x][y][0]
        return sum


    def getCommonAxis(self):
        commonAxis = vector(0.0,0.0,0.0)

        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                commonAxis = commonAxis + (-1)**(x+y)*cross(self.torqueArray[x][y][0], self.torqueArray[(x+1)%self.nSpinsX][y][0])
                commonAxis = commonAxis + (-1)**(x+y)*cross(self.torqueArray[x][y][0], self.torqueArray[x][(y+1)%self.nSpinsY][0])
        
        return norm(commonAxis)




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





##############################################################
##  TESTING CODE:                                           ##

##dt = .001
##
##
##Nx = 3
##Ny = 2
##uc = NumericUnitCell(Nx,Ny)
##
##Ja = -1.0
##Jbx = 1.8
##Jby = -1.0
##uc.setCouplings(Ja,Jbx,Jby)
##
##k = vector(1.0,.5,0.0)
##baseSigma = 0.001
##uc.setState(k,baseSigma)


##
##print '\nspinArray:'
##print uc.getState()
##
##uc.getMeanFields()
##print '\nmeanFieldArray:'
##print uc.meanFieldArray
##
##uc.getTorques()
##print '\ntorqueArray:'
##print uc.torqueArray
##
##uc.timeEvolve(dt)
##print '\ntimeEvolved State:'
##print uc.getState()
##


##print '\n-----------------------------'
##print "Calculation Instance Started"
##print time.ctime()
##
##
##
##k = 0
##loops = 100000
##while k <= loops:
##    k+=1
##
##    uc.getMeanFields()
##    uc.getTorques()
##    uc.timeEvolve(dt)
##
##
##print "\nCalculation Instance Ended For",loops,"Loops"
##print time.ctime()
##print '-----------------------------'
##
##print uc.getState()


#  NumericUnitCell:
##  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



#   newUnitCell:
##  43 seconds for 100,000 iterations with a 3 by 2 unit cell = 6 spins
##   4 seconds for 10,000 iterations with a 3 by 2 unit cell = 6 spins

##   2 seconds for 1,000 iterations with an 8 by 4 unit cell = 32 spins
##  18 seconds for 100 iterations with an 80 by 40 unit cell = 3,200 spins
## 343 seconds for 10 iterations with an 800 by 400 unit cell = 320,000 spins


## Notes:
## I would expect the new method (numeric matrices) will be faster any time we can actually utilize a matrix to do an operation.
## This makes sense and seems obvious, but since this is the difining quality, we must now ask what places we can use the matrix method.
## Matrix method works for:

## Coupling:  Since the coupling can be generated as a single matrix and does not need to be 'recreated' with each iteration.
## Torques:  This should be more straightforward than how it is currently implemented, where I iterate through the spin entries
## 




