from visual import *

from Tracer import *
from Spin import *


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

  
class VisualLattice:
    def __init__(self, nSpinsX, nSpinsY, tracerLength):
        self.window = display(title='Spin Lattice', width=480, height=420, x=60, y=435, 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.forward = (-0.5,0.5,-0.3)
        self.window.range = 2.0
        self.window.autoscale = 0
        self.window.visible = 0
        self.window.select()

        self.nSpinsX = nSpinsX
        self.nSpinsY = nSpinsY

        self.tvWindow = 1
        self.tvBonds = 1

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

        self.tracerArray = [ None ] * nSpinsX
        for x in range(nSpinsX):
            self.tracerArray[x] = [0] * nSpinsY
            for y in range(nSpinsY):
                self.tracerArray[x][y] = Tracer(tracerLength)



        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.visualBondXRArray = [ None ] * nSpinsX                               # a vector of null elements, nSpinsX long
        for x in range(nSpinsX):
            self.visualBondXRArray[x] = 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 = (-self.nSpinsX/2.0,-self.nSpinsY/2.0,0.0)



        self.axisOrigin = vector(-self.nSpinsX-.3,-self.nSpinsY-.3,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.addTorqueVectors()

        self.addBonds()

        self.toggleVectors()









    ## BOND METHODS

    def addBonds(self):       
        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+1.0,y-self.nSpinsY+1.0/2.0,0.0)
                    self.visualBondXArray[x][y].box = box(pos=boxPos, length=0.5, width=0.05, height=0.05, 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+1.0/2.0,y-self.nSpinsY+1.0,0.0)
                    self.visualBondYArray[x][y].box = box(pos=boxPos, length=0.05, width=0.05, height=0.5, color=(.9,.9,.9), visible=self.tvBonds)

        for y in range(self.nSpinsY):
            boxPos = (-self.nSpinsX,y-self.nSpinsY+1.0/2.0,0.0)
            self.visualBondXLArray[y].box = box(pos=boxPos, length=0.5, width=0.05, height=0.05, color=(.9,.9,.9), visible=self.tvBonds)

        for x in range(self.nSpinsX):
            boxPos = (x-self.nSpinsX+1.0/2.0,-self.nSpinsY,0.0)
            self.visualBondXRArray[x].box = box(pos=boxPos, length=0.05, width=0.05, height=0.5, color=(.9,.9,.9), visible=self.tvBonds)


            
    def updateBonds(self, nLattice):
        currentState = nLattice.returnState()
        
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
##                coupling = couplingX[x][(x+1)%self.nSpinsX]
                z = dot(currentState[x][y][0],currentState[(x+1)%self.nSpinsX][y][0])#*self.bondEnergyScale
                color = ((1+z)/2,1.0-(z-1.0/2.0)**2,(1-z)/2)
                self.visualBondXArray[x][y].box.color = color
                
##                coupling = couplingY[y][(y+1)%self.nSpinsY]
                z = dot(currentState[x][y][0],currentState[x][(y+1)%self.nSpinsY][0])#*self.bondEnergyScale
                color = ((1+z)/2,1.0-(z-1.0/2.0)**2,(1-z)/2)
                self.visualBondYArray[x][y].box.color = color
                
        for y in range(self.nSpinsY):
##            coupling = couplingX[0][(-1)%self.nSpinsX]
            z = dot(currentState[0][y][0],currentState[(-1)%self.nSpinsX][y][0])#*self.bondEnergyScale
            color = ((1+z)/2,1.0-(z-1.0/2.0)**2,(1-z)/2)
            self.visualBondXLArray[y].box.color = color

        for x in range(self.nSpinsX):
##            coupling = couplingX[0][(-1)%self.nSpinsX]
            z = dot(currentState[x][0][0],currentState[x][(-1)%self.nSpinsY][0])#*self.bondEnergyScale
            color = ((1+z)/2,1.0-(z-1.0/2.0)**2,(1-z)/2)
            self.visualBondXRArray[x].box.color = color


    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

        for x in range(self.nSpinsX):
            self.visualBondXRArray[x].box.visible = self.tvBonds







    ## VECTOR METHODS


    def addVectors(self):
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                arrowPos = ((x-self.nSpinsX+1.0/2.0),(y-self.nSpinsY+1.0/2.0),0.0)
                z=.7
                color = ((1+z)/2,0.0,(1-z)/2)
                width = 0.1
                self.spinArray[x][y].setArrow(arrowPos, color, width)

    def updateVectors(self, nLattice):
        currentState = nLattice.returnState()
        energyArray = nLattice.returnEnergies()
        
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                offsetVector = vector((-self.nSpinsX+1.0/2.0),(-self.nSpinsY+1.0/2.0),0.0) + vector(x,y,0.0)

                arrowPos = offsetVector - currentState[x][y][0]/2.0

                z = energyArray[x][y]/12.0
                color = ((1+z)/2,1.0-(z-1.0/2.0)**2,(1-z)/2)
                
                width = 0.1
                
                self.spinArray[x][y].setArrow(arrowPos, color, width)

                self.spinArray[x][y].updateVectorAxis(currentState[x][y][0])
                                
             
    def toggleVectors(self):
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.spinArray[x][y].toggleVectorVisibility()




    ## TORQUE VECTOR METHODS

    def addTorqueVectors(self):
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                arrowPos = ((x-self.nSpinsX+1),(y-self.nSpinsY+1),0.0)
                color=(.5,.5,1)
                width=0.02
                self.spinArray[x][y].setTorqueArrow(arrowPos, color, width)


    def updateTorqueVectors(self, nLattice):
        currentState = nLattice.returnTorqueArray()
        
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.spinArray[x][y].updateTorqueVectorPos(self.spinArray[x][y].getVectorPos()+ self.spinArray[x][y].getVectorAxis())
                self.spinArray[x][y].updateTorqueVectorAxis(currentState[x][y][0])


    def toggleTorqueVectors(self):
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.spinArray[x][y].toggleTorqueVectorVisibility()



    ## TRACER METHODS

          
    def updateTracers(self, nLattice):
        energyArray = nLattice.returnEnergies()
        
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):               
                location = self.spinArray[x][y].getVectorPos() + self.spinArray[x][y].getVectorAxis() 
                z = energyArray[x][y]/12.0
                color = ((1+z)/2,1.0-(z-1.0/2.0)**2,(1-z)/2)

                self.tracerArray[x][y].step(location, color)
              
    def resetTracers(self):
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.tracerArray[x][y].delTracer()

                
    def toggleTracers(self):
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.tracerArray[x][y].toggleVisibility()






    ## POINT METHODS

    def addPoints(self):
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
##                z = spinArray[x][y][0][2]
                z=.7
                color = ((1+z)/2,1.0-(z-1.0/2.0)**2,(1-z)/2)
                radius = 0.03
                self.spinArray[x][y].setPoint(color, radius)

    def updatePoints(self, nLattice):
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.spinArray[x][y].updatePointPos(nLattice.returnState()[x][y][0])
                z=.7
                color = ((1+z)/2,1.0-(z-1.0/2.0)**2,(1-z)/2)
                radius = 0.03
                self.spinArray[x][y].setPoint(color, radius)

    def togglePoints(self):
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.spinArray[x][y].togglePointVisibility()





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

##    def updateState(self):
##        self.spinArray = self.numerics.returnState()
##        self.meanFieldArray = self.numerics.returnMeanFieldArray()
##        self.energies = self.numerics.returnEnergies()
##        self.bondEnergies = self.numerics.returnBondEnergies()
        














