from visual import *
from random import random

import time,sys,thread


spinUnitCellW = display(title='Spin Unit Cell', width=680, height=400, x=20, y=620, center=(0,0,0), background=(.0,.0,.0))
spinUnitCellW.fov = pi/90.0
spinUnitCellW.lights = [vector(.5,.5,.5), vector(-.5,-.5,-.5)]

spinUnitCellW.up = (0,0,1)
spinUnitCellW.visible = 1
spinUnitCellW.select()






        

class Spin:
    vector = arrow(visible=0)
    
    def __init__(self):
        self.vector = arrow(visible=0)


 
class VisualUnitCell:
    visualArray = []

    toggleVectorsVar = 0
    toggleVectorsPhaseColor = 1


    def __init__(self, nSpinsX, nSpinsY):
        self.nSpinsX = nSpinsX
        self.nSpinsY = nSpinsY

        visualArray = [ None ] * nSpinsX                               # a vector of null elements, nSpinsX long
        for x in range(nSpinsX):
            visualArray[x] = [0] * nSpinsY                             # a vector of zeros nSpinsY long
            for y in range(nSpinsY):
                visualArray[x][y] = Spin()      
               
        self.visualArray = visualArray
        
        spinUnitCellW.range = .75*(self.nSpinsX*self.nSpinsY)**0.66
        spinUnitCellW.forward = norm(vector(-1.0/self.nSpinsY-.5,1.0/self.nSpinsX,-(1.0/self.nSpinsX + 1.0/self.nSpinsY)**1.7))


    def addAxis(self):
        axisOrigin = vector(-self.nSpinsX/2.0,-self.nSpinsY/2.0,0.0)
        
        xAxis = []
        for x in arange(-0.1,1.01,1.1):
            xAxis.append(vector(x,0,0)+axisOrigin)
        xAxisLine = curve(pos=xAxis, color=(.9,.9,.9))

        yAxis = []
        for y in arange(-0.1,1.01,1.1):
            yAxis.append(vector(0,y,0)+axisOrigin)
        yAxisLine = curve(pos=yAxis, color=(.9,.9,.9))

        zAxis = []
        for z in arange(-0.5,0.51,1.0):
            zAxis.append(vector(0,0,z)+axisOrigin)
        zAxisLine = curve(pos=zAxis, color=(.9,.9,.9))

        xAxislabel = label(pos=vector(1.0,0,0) + axisOrigin, text='x Axis', xoffset=0, yoffset=-.001, space=.1, height=10, border=4)
        yAxislabel = label(pos=vector(0,1.0,0) + axisOrigin, text='y Axis', xoffset=-.001, yoffset=.001, space=.1, height=10, border=4)
        zAxislabel = label(pos=vector(0,0,0.5) + axisOrigin, text='z Axis', xoffset=.001, yoffset=.001, space=.1, height=10, border=4)



    def addVectors(self,spinArray):
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                arrowPos = (x-self.nSpinsX/2.0+1/2.0-spinArray[x][y][0][0]/2.0,y-self.nSpinsY/2.0+1/2.0-spinArray[x][y][0][1]/2.0,0.0)
                z = spinArray[x][y][0][2]
                self.visualArray[x][y].vector = arrow(pos=arrowPos, axis=spinArray[x][y][0], color=((1+z)/2,0.0,(1-z)/2), shaftwidth=0.125, visible=self.toggleVectorsVar)

        
    def updateVectors(self,spinArray):
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.visualArray[x][y].vector.pos = vector(x-self.nSpinsX/2.0+1/2.0,y-self.nSpinsY/2.0+1/2.0,0.0) - spinArray[x][y][0]/2.0
                self.visualArray[x][y].vector.axis = spinArray[x][y][0]
                colorArg = spinArray[x][y][0][2]*norm(vector(spinArray[x][y][0][0],spinArray[x][y][0][1],0)).x
                if self.toggleVectorsPhaseColor == 1:   self.visualArray[x][y].vector.color = ((colorArg+1.0)/2.0,0.0,1.0-(colorArg+1.0)/2.0)


    def toggleVectors(self):
        self.toggleVectorsVar = (self.toggleVectorsVar+1)%2
        
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.visualArray[x][y].vector.visible = self.toggleVectorsVar

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
















