from visual import *
from random import random

import time,sys,thread




class Bond:
    def __init__(self):
        self.box = box(visible=0)

class Spin:
    def __init__(self):
        self.vector = arrow(visible=0)
 
class VisualLattice:
    def __init__(self, nSpinsX, nSpinsY, stripeSpacingX, stripeSpacingY, couplingX, couplingY, nLattice):
        self.window = display(title='Spin Lattice', width=480, height=320, x=5, y=480, center=(0,0,0), background=(.0,.0,.0))
        self.window.fov = pi/90.0
        self.window.lights = [vector(.5,.5,.5), vector(-.5,-.5,-.5)]
        self.window.up = (0,0.001,1)
        self.window.visible = 0
        self.window.select()


        self.nSpinsX = nSpinsX
        self.nSpinsY = nSpinsY
        self.stripeSpacingX = stripeSpacingX
        self.stripeSpacingY = stripeSpacingY
        self.couplingX = couplingX
        self.couplingY = couplingY

        self.nLattice = nLattice

        self.tvWindow = 1
        self.tvVectors = 1
        self.tvBonds = 0
        self.tvVectorPhaseColor = 0

    
        self.visualArray = [ None ] * nSpinsX                               # a vector of null elements, nSpinsX long
        for x in range(nSpinsX):
            self.visualArray[x] = [0] * nSpinsY                             # a vector of zeros nSpinsY long
            for y in range(nSpinsY):
                self.visualArray[x][y] = Spin()

        self.visualBondXArray = [ None ] * nSpinsX                               # a vector of null elements, nSpinsX long
        for x in range(nSpinsX):
            self.visualBondXArray[x] = [0] * nSpinsY                             # a vector of zeros nSpinsY long
            for y in range(nSpinsY):
                self.visualBondXArray[x][y] = Bond()

        self.visualBondYArray = [ None ] * nSpinsX                               # a vector of null elements, nSpinsX long
        for x in range(nSpinsX):
            self.visualBondYArray[x] = [0] * nSpinsY                             # a vector of zeros nSpinsY long
            for y in range(nSpinsY):
                self.visualBondYArray[x][y] = Bond()

        self.visualBondXLArray = [ None ] * nSpinsY                               # a vector of null elements, nSpinsX long
        for y in range(nSpinsY):
            self.visualBondXLArray[y] = Bond()

               
        self.window.range = 1.2*(self.nSpinsX*self.nSpinsY)**0.5
        self.window.forward = norm(vector(-1.0/self.nSpinsY,1.5/self.nSpinsX+1.0,-(1.0/self.nSpinsX + 1.0/self.nSpinsY)**1.7))
        self.window.visible = self.tvWindow
        self.window.center = (-0.25,0.0,-0.30)


        self.axisOrigin = vector(-self.nSpinsX/2.0-1.0/2.0,-self.nSpinsY/2.0-1.0/2.0,0.0)
        
        self.xAxis = []
        self.yAxis = []
        self.zAxis = []
        
        for x in arange(-0.1,1.01,1.1):     self.xAxis.append(vector(x,0,0) + self.axisOrigin)
        for y in arange(-0.1,1.01,1.1):     self.yAxis.append(vector(0,y,0) + self.axisOrigin)
        for z in arange(-0.5,0.51,1.0):     self.zAxis.append(vector(0,0,z) + self.axisOrigin)
        
        self.xAxisLine = curve(pos=self.xAxis, color=(.9,.9,.9), visible=self.tvWindow)
        self.yAxisLine = curve(pos=self.yAxis, color=(.9,.9,.9), visible=self.tvWindow)
        self.zAxisLine = curve(pos=self.zAxis, color=(.9,.9,.9), visible=self.tvWindow)

        self.xAxisLabel = label(pos=vector(1.0,0,0) + self.axisOrigin, text='x', xoffset=0, yoffset=-.001, space=.1, height=10, box=0, opacity=0, line=0, border=4, color=(.9,.9,.9), visible=self.tvWindow)
        self.yAxisLabel = label(pos=vector(0,1.0,0) + self.axisOrigin, text='y', xoffset=-.001, yoffset=.001, space=.1, height=10, box=0, opacity=0, line=0, border=4, color=(.9,.9,.9), visible=self.tvWindow)
        self.zAxisLabel = label(pos=vector(0,0,0.5) + self.axisOrigin, text='z', xoffset=.001, yoffset=.001, space=.1, height=10, box=0, opacity=0, line=0, border=4, color=(.9,.9,.9), visible=self.tvWindow)

        self.addVectors()
        self.addBonds()



    def addBonds(self):
        spinArray = self.nLattice.returnState()
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):              
                if self.nSpinsX < 3 and x == self.nSpinsX - 1:
                    self.visualBondXArray[x][y].box = box(length=0.0, width=0.0, height=0.0, visible=0)
                else:
                    boxPos = (x-self.nSpinsX/2.0+1/2.0+1/2.0,y-self.nSpinsY/2.0+1/2.0,-1.5)
                    self.visualBondXArray[x][y].box = box(pos=boxPos, length=0.5, width=0.1, height=0.1, color=(.9,.9,.9), visible=self.tvBonds)


                if self.nSpinsY < 3 and y == self.nSpinsY - 1:
                    self.visualBondYArray[x][y].box = box(length=0.0, width=0.0, height=0.0, visible=0)
                else:
                    boxPos = (x-self.nSpinsX/2.0+1/2.0,y-self.nSpinsY/2.0+1/2.0 + 1/2.0,-1.5)
                    self.visualBondYArray[x][y].box = box(pos=boxPos, length=0.1, width=0.1, height=0.5, color=(.9,.9,.9), visible=self.tvBonds)

        for y in range(self.nSpinsY):
            boxPos = (-self.nSpinsX/2.0+1/2.0-1/2.0,y-self.nSpinsY/2.0+1/2.0,-1.5)
            self.visualBondXLArray[y].box = box(pos=boxPos, length=0.5, width=0.1, height=0.1, color=(.9,.9,.9), visible=self.tvBonds)
            
    def updateBonds(self):
        spinArray = self.nLattice.returnState()
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                coupling = self.couplingX[x][(x+1)%self.nSpinsX]
                colorArg = coupling*dot(spinArray[x][y][0],spinArray[(x+1)%self.nSpinsX][y][0])/2.0
                self.visualBondXArray[x][y].box.color = ((colorArg+1.0)/2.0,0.0,1.0-(colorArg+1.0)/2.0)
                
                coupling = self.couplingY[y][(y+1)%self.nSpinsY]
                colorArg = coupling*dot(spinArray[x][y][0],spinArray[x][(y+1)%self.nSpinsY][0])/2.0
                self.visualBondYArray[x][y].box.color = ((colorArg+1.0)/2.0,0.0,1.0-(colorArg+1.0)/2.0)
                
        for y in range(self.nSpinsY):
            coupling = self.couplingX[0][(-1)%self.nSpinsX]
            colorArg = coupling*dot(spinArray[0][y][0],spinArray[(-1)%self.nSpinsX][y][0])/2.0
            self.visualBondXLArray[y].box.color = ((colorArg+1.0)/2.0,0.0,1.0-(colorArg+1.0)/2.0)

    def toggleBonds(self):
        self.tvBonds = (self.tvBonds+1)%2

        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.visualBondXArray[x][y].box.visible = self.tvBonds
                self.visualBondYArray[x][y].box.visible = self.tvBonds
                
        for y in range(self.nSpinsY):
            self.visualBondXLArray[y].box.visible = self.tvBonds





    def addVectors(self):
        spinArray = self.nLattice.returnState()
        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.tvVectors)

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

    def toggleVectors(self):
        self.tvVectors = (self.tvVectors+1)%2

        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.visualArray[x][y].vector.visible = self.tvVectors



    def toggleWindow(self):
        self.tvWindow = (self.tvWindow+1)%2

        if self.tvWindow == 1:
            self.window.visible = self.tvWindow
            
            self.xAxisLine.visible = self.tvWindow
            self.yAxisLine.visible = self.tvWindow
            self.zAxisLine.visible = self.tvWindow

            self.xAxisLabel.visible = self.tvWindow
            self.yAxisLabel.visible = self.tvWindow
            self.zAxisLabel.visible = self.tvWindow
        else:
            self.xAxisLine.visible = self.tvWindow
            self.yAxisLine.visible = self.tvWindow
            self.zAxisLine.visible = self.tvWindow

            self.xAxisLabel.visible = self.tvWindow
            self.yAxisLabel.visible = self.tvWindow
            self.zAxisLabel.visible = self.tvWindow
            
            self.window.visible = self.tvWindow

           
    def returnArray(self):
        return self.visualArray
















