from Spin import *

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

        self.addVectors()
        self.addTorqueVectors()
        self.addPoints()
    




    ## 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 = (z,0.0,1.0-z)
                width = 0.1
                self.spinLattice[x][y].setArrow(arrowPos, color, width)

    def updateVectors(self, nLattice):
        currentState = nLattice.returnState()
        phiScale = nLattice.returnPhiScale()+.001
        phiArray = nLattice.returnPhiArray()
        phaseArray = nLattice.returnPhaseArray()
        zArray = nLattice.returnzArray()

        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                xModXSpacing = x%nLattice.stripeSpacingX
                yModYSpacing = y%nLattice.stripeSpacingY

                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

                tAngle = phaseArray[xModXSpacing][yModYSpacing]
                zFlip = zArray[xModXSpacing][yModYSpacing]
                
                new = currentState[x][y]
                k1 = norm(self.k)
                z = zFlip*dot(new, k1)[0]/(phiScale*phiArray[xModXSpacing][yModYSpacing])
                color = ((1.0+z)/2.0,1.0-(z)**2,(1.0-z)/2.0)

                width = 0.1
                
                self.spinLattice[x][y].setArrow(arrowPos, color, width)
                self.spinLattice[x][y].updateVectorAxis(currentState[x][y][0])
             
    def toggleVectors(self):
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.spinLattice[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.spinLattice[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.spinLattice[x][y].updateTorqueVectorPos(self.spinLattice[x][y].getVectorPos()+ self.spinLattice[x][y].getVectorAxis())
                self.spinLattice[x][y].updateTorqueVectorAxis(currentState[x][y][0])

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



    ## POINT METHODS

    def addPoints(self):
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                z = 0.7
                color = (z,0.3,1.0-z)
                radius = 0.09
                self.spinLattice[x][y].setPoint(color, radius)

    def updatePoints(self, nLattice):
        phiScale = nLattice.returnPhiScale()+.001
        
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                location = self.spinLattice[x][y].getVectorPos() + self.spinLattice[x][y].getVectorAxis()
                self.spinLattice[x][y].updatePointPos(location)
                
                k1 = (1.0,0.0,0.0)
                z = dot(self.spinLattice[x][y].getVectorAxis(), k1)/phiScale
                color = (z,0.0,1.0-z)
                
                radius = 0.09
                self.spinLattice[x][y].setPoint(color, radius)

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






    def returnSpinLattice(self):
        return self.spinLattice







    
