from Numeric import *
from Matrix import *
from LinearAlgebra import *

from visual import *


class Hamiltonian:
    nSpinsX = 4
    nSpinsY = 4
    nSpinsTot = nSpinsX*nSpinsY

    Ja = -1.0
    Jbx = -1.0
    Jby = -1.0


    zArray = []
    
    couplingXArray = []
    couplingYArray = []

    couplingPhaseXArray = []
    couplingPhaseYArray = []

    meanFieldFromZArray = []

    HArray = []                             # hurray!!!!


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

        self.zArray = zeros((nSpinsX,nSpinsY),Float32)

        
        self.couplingXArray = zeros((nSpinsX,nSpinsX),Complex32)
        self.couplingYArray = zeros((nSpinsY,nSpinsY),Complex32)

        self.couplingPhaseXArray = zeros((nSpinsX,nSpinsX),Complex32)
        self.couplingPhaseYArray = zeros((nSpinsY,nSpinsY),Complex32)

        self.meanFieldFromZArray = zeros((nSpinsX,nSpinsY),Float32)

        self.HArray = zeros((self.nSpinsTot*2,self.nSpinsTot*2),Complex32)



    def setzArray(self):
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.zArray[x][y] = (-1)**(x+y)     # your typical interior antiferromagnetically textured unit cell
                

    def setCouplings(self, Ja, Jbx, Jby):
        self.Ja = Ja
        self.Jbx = Jbx
        self.Jby = Jby
              
        for i in range(self.nSpinsX):
            for j in range(self.nSpinsX):
                self.couplingXArray[i][j] = 0.0
                if i - j == 1:                      self.couplingXArray[i][j] += self.Ja
                if j - i == 1:                      self.couplingXArray[i][j] += self.Ja
                if i - j == self.nSpinsX - 1:       self.couplingXArray[i][j] += self.Jbx
                if j - i == self.nSpinsX - 1:       self.couplingXArray[i][j] += self.Jbx

        for i in range(self.nSpinsY):
            for j in range(self.nSpinsY):
                self.couplingYArray[i][j] = 0.0
                if i - j == 1:                      self.couplingYArray[i][j] += self.Ja
                if j - i == 1:                      self.couplingYArray[i][j] += self.Ja
                if i - j == self.nSpinsY - 1:       self.couplingYArray[i][j] += self.Jby
                if j - i == self.nSpinsY - 1:       self.couplingYArray[i][j] += self.Jby



    def getMeanFieldFromZ(self):
        latticeArray = self.zArray
        flippedLatticeArray = transpose(self.zArray)

        meanFieldXArray = matrixmultiply(flippedLatticeArray,self.couplingXArray)
        meanFieldYArray = matrixmultiply(latticeArray,self.couplingYArray)
        
        meanFieldXArray = transpose(meanFieldXArray)
        meanFieldYArray = meanFieldYArray
        
        self.meanFieldFromZArray = meanFieldXArray + meanFieldYArray





    def makeHamiltonian(self, k):
        for i in range(self.nSpinsTot*2):                           # i is target spin and component
            for j in range(self.nSpinsTot*2):                       # j is source spin and component
                self.HArray[i][j] = 0.0 + 1j*0.0                    # i = j implies xTarget == xSource and yTarget == ySource (target is source)
                
                if i < self.nSpinsTot and j >= self.nSpinsTot:      # x components in terms of y components
                    isub = i%self.nSpinsTot
                    jsub = j%self.nSpinsTot

                    xTarget = (isub-isub%self.nSpinsY)/self.nSpinsY
                    yTarget = isub%self.nSpinsY
                    targetZ = self.zArray[xTarget][yTarget]

                    xSource = (jsub-jsub%self.nSpinsY)/self.nSpinsY
                    ySource = jsub%self.nSpinsY
                    sourceZ = self.zArray[xSource][ySource]

                    if isub == jsub:                            self.HArray[i][j] = self.meanFieldFromZArray[xTarget,yTarget]

                    if xTarget - xSource == 1:                  self.HArray[i][j] += -self.Ja*e**(-1j*k.x)*targetZ
                    if xSource - xTarget == 1:                  self.HArray[i][j] += -self.Ja*e**(1j*k.x)*targetZ
                    if xTarget - xSource == self.nSpinsX - 1:   self.HArray[i][j] += -self.Jbx*e**(1j*k.x)*targetZ
                    if xSource - xTarget == self.nSpinsX - 1:   self.HArray[i][j] += -self.Jbx*e**(-1j*k.x)*targetZ
                    
                    if yTarget - ySource == 1:                  self.HArray[i][j] += -self.Ja*e**(-1j*k.y)*targetZ
                    if ySource - yTarget == 1:                  self.HArray[i][j] += -self.Ja*e**(1j*k.y)*targetZ
                    if yTarget - ySource == self.nSpinsY - 1:   self.HArray[i][j] += -self.Jby*e**(1j*k.y)*targetZ
                    if ySource - yTarget == self.nSpinsY - 1:   self.HArray[i][j] += -self.Jby*e**(-1j*k.y)*targetZ

                if j < self.nSpinsTot and i >= self.nSpinsTot:                                  # y components in terms of x components
                    self.HArray[i][j] = -self.HArray[i-self.nSpinsTot][j+self.nSpinsTot]        # sign change due to cross product equation


        for i in range(self.nSpinsTot*2):                           # i is target spin and component
            for j in range(self.nSpinsTot*2):                       # j is source spin and component               
                self.HArray[i][j] = 1j*self.HArray[i][j]            # takes care of division by i and multiplication by -1 (iw -> w and sign in torque equation)
                
##        print '\nHermitian if identically Zero: \n', self.HArray - transpose(conjugate(self.HArray))


    def getHArray(self):
        return self.HArray

        
    def getEigenValues(self):
        evalues, evectors = eigenvectors(self.HArray)
        evaluesReal = evalues.real

        return evaluesReal

    
    def getEigenVectors(self):
        evalues, evectors = eigenvectors(self.HArray)
        evectorsReal = evectors.real

        return evectorsReal


    def getGroundStateEnergy(self):
        evalues, evectors = eigenvectors(self.HArray)
        evectorsReal = evectors.real
        evaluesReal = evalues.real
       
        energylist = evaluesReal.tolist()
        energylist.sort()
        indexOfLosestEigenState = evaluesReal.tolist().index(energylist[0])     # finds lowest energy state

        lowestEval = evaluesReal[indexOfLosestEigenState]

        return lowestEval


    def getGroundState(self):
        evalues, evectors = eigenvectors(self.HArray)
        evectorsReal = evectors.real
        evaluesReal = evalues.real

        energylist = evaluesReal.tolist()
        energylist.sort()
        indexOfLosestEigenState = evaluesReal.tolist().index(energylist[0])     # finds lowest energy state

        lowestEvec = evectorsReal[indexOfLosestEigenState]

        lowestEvecXcomp = lowestEvec[:self.nSpinsTot]
        lowestEvecYcomp = lowestEvec[self.nSpinsTot:]

        xComponentUCArray = reshape(lowestEvecXcomp, (self.nSpinsX,self.nSpinsY))
        yComponentUCArray = reshape(lowestEvecYcomp, (self.nSpinsX,self.nSpinsY))
        
        return xComponentUCArray + 1j*yComponentUCArray




    def getState(self, n):
        evalues, evectors = eigenvectors(self.HArray)
        evectorsReal = evectors.real

        Evec = evectorsReal[n]

        EvecXcomp = Evec[:self.nSpinsTot]
        EvecYcomp = Evec[self.nSpinsTot:]

        xComponentUCArray = reshape(EvecXcomp, (self.nSpinsX,self.nSpinsY))
        yComponentUCArray = reshape(EvecYcomp, (self.nSpinsX,self.nSpinsY))
        
        return xComponentUCArray + 1j*yComponentUCArray









##Debug code
##        
##        
##nSpinsX = 2
##nSpinsY = 2
##nSpinsTot = nSpinsX*nSpinsY
##
##Ja = -1.0
##Jbx = -1.0
##Jby = -1.0
##
##k = pi*vector(0.01,0.0,0.0)
##
##H = Hamiltonian(nSpinsX, nSpinsY)
##H.setzArray()
##H.setCouplings(Ja, Jbx, Jby)
##H.getMeanFieldFromZ()
##
##H.makeHamiltonian(k)
##
##
##
##print '\n--------------------'
##
##print '\nzArray:\n',H.zArray
##print '\ncouplingXArray: \n', H.couplingXArray
##print '\ncouplingYArray: \n', H.couplingYArray
##print '\nmeanFieldFromZArray: \n', H.meanFieldFromZArray
##print '\nHArray: \n', H.HArray
##print '\nHermitian iff identically Zero: \n', H.HArray - transpose(conjugate(H.HArray))
##
##
##
##print '\n--------------------'
##
##print '\nFor the Unit Cell with Pattern:\n',H.zArray
##print '\nWe obtain for the Lowest Energy State:',H.getGroundStateEnergy()
##
##print '\nWith x Components:\n',H.getGroundState().real
##print '\nAnd y Components:\n',H.getGroundState().imaginary
##print '\n\n\n'




