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 = []
    
    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.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):
                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):
                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):
        if k.x != 0:
            kAngle = arctan(k.y/k.x)
        else:
            kAngle = pi/2.       

        chi = abs((cos(k.x)+cos(k.y))/2)                                # Antiferromagnetic values
        sigmaRatio = 0.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 x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                z = 1.0 - 2.0*((x+y)%2)
                
                tempSigma = (1.0-2.0*((x+y)%2))*baseSigma + ((x+y)%2)*baseSigma*sigmaRatio
                tempAngle = dot(k,vector(x, y, 0.0)) + kAngle + z*pi/2.0
                tempSpinZ = z*(1.0-tempSigma**2.0)**(1.0/2.0)
                self.spinArray[x][y] = (tempSigma*cos(tempAngle),tempSigma*sin(tempAngle),tempSpinZ)




    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 getTorques(self):
        for i in range(self.nSpinsX):
            for j in range(self.nSpinsY):
                self.torqueArray[i][j] = cross(self.spinArray[i][j][0],self.meanFieldArray[i][j][0])








##    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 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?)

##        self.spinArray = norm(self.spinArray)

##    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)))
##
##
##    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 getAxis(self):
##        commonAxis = vector(0.0,0.0,0.0)
##        if self.nSpinsX > 1:
##            commonAxis = cross(self.numericArray[0][0].torque,self.numericArray[1][0].torque)       ## ??? only utilizes two spins ???
##        if self.nSpinsY > 1:
##            commonAxis = cross(self.numericArray[0][0].torque,self.numericArray[0][1].torque)
##        else:
##            commonAxis = self.numericArray[0][0].spin
##        return 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
## 




