from visual import *
from random import random
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 = [[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 = 1.0                                   # 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]


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

        sigmaRatio = .3                       # calculated from k,... is an array in general
        
        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)) + k_angle + 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)

    def permuteState(self, dt):                     # This takes us one step forward in time.  A Linear Approximation (NEED TO CONVERT TO RUNGA KUTTA), with interval spacing in t of dt.
        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)

        for row in self.numericArray:
            for spin in row:
                spin.spin = spin.spin + spin.torque*dt
                spin.spin = norm(spin.spin)                 # PLEASE NOTE THE NUMERIC METHOD BEING USED.  SHOULD WE USE RUNGA KUTTA?

    def perturbState(self, dt, temp, Ja):
        for row in self.numericArray:
            for spin in row:                                # perturbed by unit ball
                randomU = random()
                randomV = random()
                theta = 2*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.spin = spin.spin + spin.kick*dt         # what kind of kick is this?
                spin.spin = norm(spin.spin)
                                 
    def randomizeState(self):
        for row in self.numericArray:
            for spin in row:
                spin.spin = norm(vector(random()-.5,random()-.5,spin.indices[2]*random()))

    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 getCommonAxis(self):
        commonAxis = vector(0.0,0.0,0.0)
        commonAxis = cross(self.numericArray[0][0].torque,self.numericArray[0][1].torque)       ## ??? not any more complicated..... ???
        return commonAxis

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



##    def saveState(self):
##        copy
##
##    def isCloseToStart(self):
##        
##        if
##            return 0
##
##            return 1





















