from visual import *




class Tracer:
    def __init__(self, length):
        self.length = length
        self.curve = curve(color=color.black)
       
    def step(self, updatePosition):
        if len(self.curve.pos) >= 332:
            self.curve.pos = self.curve.pos[32:]

        self.curve.append(pos=updatePosition)

    def delTracer(self):
        points = [[0.0,0.0,0.0]]*self.length
        self.curve.pos = points

    def getTracer(self):
        return self.points



class Spin:   
    def __init__(self, tracerLength):
        self.tracerLength = tracerLength

        self.tracer = Tracer(tracerLength)
        self.fieldTracer = Tracer(tracerLength)
        self.point = sphere(visible=0)
        self.vector = arrow(visible=0)



 
class VisualSphere:
    def __init__(self, nSpinsX, nSpinsY, tracerLength, nLattice):
        self.window = display(title='Spin Sphere', width=600, height=600, x=15, y=15, center=(0,0,0), background=(.0,.0,.0))
        self.window.fov = pi/900.0
        self.window.lights = [vector(.5,.5,.5), vector(-.5,-.5,-.5)]
        self.window.up = (0.0,0.001,1.0)
        self.window.forward = (-0.5,1.0,-0.8)
        self.window.range = 2.0
        self.window.autoscale = 0
        self.window.visible = 1

        self.window.select()


        self.nSpinsX = nSpinsX
        self.nSpinsY = nSpinsY
        
        self.tracerLength = tracerLength
        
        self.nLattice = nLattice

        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.tracerLength)


        self.tvTracers = 1
        self.tvPoints = 0
        self.tvVectors = 0
        self.tvVectorsPhaseColor = 0
        self.tvSWAxis = 0

        self.swAxis = []  
        self.swAxisLine = curve(color=(.8,.8,.95), visible=self.tvSWAxis)
        self.equatorLine = curve(x=cos(arange(0,2.1*pi,.1565)), y=sin(arange(0,2.1*pi,.1565)), color=(.5,.5,.5))

        
        self.xAxisLine = curve(pos=[(-1.4,-1.3,0.0),(-0.8,-1.3,0.0)], color=(.9,.9,.9))
        self.yAxisLine = curve(pos=[(-1.3,-1.4,0.0),(-1.3,-0.8,0.0)], color=(.9,.9,.9))
        self.zAxisLine = curve(pos=[(-1.3,-1.3,-0.1),(-1.3,-1.3,0.5)], color=(.9,.9,.9))

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

        self.addPoints()
        self.addTracers()
        self.addVectors()




    ## TRACER METHODS

    def addTracers(self):
        spinArray = self.nLattice.returnState()
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                z = spinArray[x][y][0][2]
                self.visualArray[x][y].tracer.curve = curve(color=((1+z)/2,0.0,(1-z)/2))
                self.visualArray[x][y].fieldTracer.curve = curve(color=(0.9,0.9,0.9))

    def updateTracers(self):
        spinArray = self.nLattice.returnState()
        fieldArray = self.nLattice.meanFieldArray
##        for x in range(self.nSpinsX):
##            for y in range(self.nSpinsY):
##                self.visualArray[x][y].tracer.step(spinArray[x][y][0])
##                self.visualArray[x][y].fieldTracer.step(fieldArray[x][y][0])

        self.visualArray[1][0].tracer.step(spinArray[1][0][0])
        self.visualArray[0][0].tracer.step(spinArray[1][0][0]+spinArray[0][0][0])
        self.visualArray[2][0].tracer.step(spinArray[1][0][0]+spinArray[0][0][0]+spinArray[2][0][0])
 
 
                
    def resetTracers(self):
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.visualArray[x][y].tracer.delTracer()
                self.visualArray[x][y].fieldTracer.delTracer()
                
    def toggleTracers(self):
        self.tvTracers = (self.tvTracers+1)%2
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.visualArray[x][y].tracer.curve.visible = self.tvTracers
                self.visualArray[x][y].fieldTracer.curve.visible = self.tvTracers



    ## POINT METHODS

    def addPoints(self):
        spinArray = self.nLattice.returnState()
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                z = spinArray[x][y][0][2]
                self.visualArray[x][y].point = sphere(pos=(0.0,0.0,0.0), color=((1+z)/2,0.0,(1-z)/2), radius=0.03, visible=self.tvPoints)

    def updatePoints(self):
        spinArray = self.nLattice.returnState()
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.visualArray[x][y].point.pos = spinArray[x][y][0]

    def togglePoints(self):
        self.tvPoints = (self.tvPoints+1)%2
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.visualArray[x][y].point.visible = self.tvPoints



    ## VECTOR METHODS

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

    def updateVectors(self):
        spinArray = self.nLattice.returnState()
##        for x in range(self.nSpinsX):
##            for y in range(self.nSpinsY):
##                self.visualArray[x][y].vector.axis = spinArray[x][y][0]
##                
##                if self.tvVectorsPhaseColor == 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)

        self.visualArray[1][0].vector.axis = spinArray[1][0][0]

        self.visualArray[0][0].vector.pos = spinArray[1][0][0]
        self.visualArray[0][0].vector.axis = spinArray[0][0][0]

        self.visualArray[2][0].vector.pos = spinArray[1][0][0]+spinArray[0][0][0]
        self.visualArray[2][0].vector.axis = spinArray[2][0][0]
        
       
        if self.tvVectorsPhaseColor == 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
                


    ## AXIS METHODS
 
    def updateAxis(self):
        axis = self.nLattice.returnCommonAxis()
        self.swAxis = [-1.5*axis,1.5*axis]
        self.swAxisLine.pos = self.swAxis

    def toggleAxis(self):
        self.tvSWAxis = (self.tvSWAxis+1)%2
        self.swAxisLine.visible = self.tvSWAxis


    def updateViewAxis(self):
        axis = self.nLattice.returnCommonAxis()
        self.window.forward = norm(axis+vector(0.0,0.0,0.000001))

    def returnArray(self):
        return self.visualArray











##    def changeTracerLength(self, tracerLength):
##        for x in range(self.nSpinsX):
##            for y in range(self.nSpinsY):
##                self.visualArray[x][y].tracer.


##    def resetVisuals(self):
##        for x in range(self.nSpinsX):
##            for y in range(self.nSpinsY):
##                self.visualArray[x][y].tracer.visible = 0
##                self.visualArray[x][y].point.visible = 0
##                self.visualArray[x][y].vector.visible = 0















