from scipy import *
from numpy import *
from visual import *
from random import random
import time

class ElectronWave:
    def __init__(self, sizeOfLattice, center, direction, mode):
        self.sizeOfLattice = sizeOfLattice
        self.center = center
        self.direction = direction
        self.mode = mode
        self.tonorm = 1.0
        
    def spinor(self,x,y):
        if self.direction == 0:
            p = vector(self.mode*2*pi/self.sizeOfLattice,0.0,0.0)
        if self.direction == 1:
            p = vector(0.0,self.mode*2*pi/self.sizeOfLattice,0.0)
        
        s1upWf = self.tonorm/(8*sqrt(self.sizeOfLattice))*e**(-1j*dot(p,vector(x,y,0)))
        s1downWf = 0
        s2upWf = self.tonorm/(8*sqrt(self.sizeOfLattice))*e**(-1j*dot(p,vector(x,y,0)))
        s2downWf = 0
                
        return [s1upWf,s1downWf,s2upWf,s2downWf]

    def getDensity(self,x,y):
        return dot(conjugate(self.spinor(x,y)),self.spinor(x,y)).real

    def getPhase(self,x,y):
        return atan2(self.spinor[0](x,y).imag,self.spinor[0](x,y).real)

    def setNorm(self,tonorm):
        self.tonorm = tonorm


class Basis:
    def __init__(self, sizeOfLattice, sizeOfBasis, center, Z):
        self.sizeOfLattice = sizeOfLattice
        self.sizeOfBasis = sizeOfBasis
        
        self.center = center
        self.Z = Z

        self.maxIterations = 24
        self.tolerance = 1e-3

        self.basis = [None]*2        
        for direction in xrange(2):
            self.basis[direction] = [0]*self.sizeOfBasis
            for mode in xrange(self.sizeOfBasis):
                self.basis[direction][mode] = ElectronWave(sizeOfLattice, center, direction, mode)

        self.coulombicRepulsion = [None]*2
        for direction1 in xrange(2):
            self.coulombicRepulsion[direction1] = [0]*2
            for direction2 in xrange(2):
                self.coulombicRepulsion[direction1][direction2] = [0]*self.sizeOfBasis
                for mode1 in xrange(self.sizeOfBasis):
                    self.coulombicRepulsion[direction1][direction2][mode1] = [0]*self.sizeOfBasis
                    for mode2 in xrange(self.sizeOfBasis):
                        self.coulombicRepulsion[direction1][direction2][mode1][mode2] = 0.0

    def normalizeSpinors(self):
        norms = [0]*2
        for direction in xrange(2):
            norms[direction] = [0]*self.sizeOfBasis
            for mode in xrange(self.sizeOfBasis):
                norms[direction][mode] = 1/sqrt(integrate.dblquad(lambda x, y: self.basis[direction][mode].getDensity(x,y), 0, self.sizeOfLattice, lambda x: 0, lambda x: self.sizeOfLattice)[0])
                self.basis[direction][mode].setNorm(1/sqrt(norms[direction][mode]))
                print "Normalization constant was:", 1/sqrt(norms[direction][mode])


    def integrateCoulombicRepulsionE2(self,direction1,direction2,mode1,mode2,x1,y1):
        return integrate.dblquad(lambda x2, y2: self.basis[direction1][mode1].getDensity(x1,y1)*self.basis[direction2][mode2].getDensity(x2,y2)/(mag(vector(x1,y1,0)-vector(x2,y2,0))+.0000001), 0, self.sizeOfLattice, lambda x2: 0, lambda x2: self.sizeOfLattice)[0]

    def getCoulombicRepulsions(self):
        for direction1 in xrange(2):
            for direction2 in xrange(2):
                for mode1 in xrange(self.sizeOfBasis):
                    for mode2 in xrange(self.sizeOfBasis):
                        self.coulombicRepulsion[direction1][direction2][mode1][mode2] = integrate.dblquad(lambda x1, y1: self.integrateCoulombicRepulsionE2(direction1,direction2,mode1,mode2,x1,y1), 0, self.sizeOfLattice, lambda x1: 0, lambda x1: self.sizeOfLattice)[0]
        
        dump = open("dumpCoulombicRepulsion", mode='w')
        dump.write(str(self.coulombicRepulsion))
        dump.close()



class DensityPoint:
    def __init__(self):
        self.visibility = 1
        self.point = sphere(visible=self.visibility, radius=.001)

    def setAttributes(self,position,size,color):
        self.point.pos = position
        self.point.radius = size
        self.point.color = color

    def toggleVisibility(self):
        self.visibility = (1+self.visibility)%2
        

class Lattice:
    def __init__(self,basis,sizeOfLattice,sizeOfBasis,center):
        self.sizeOfLattice = sizeOfLattice
        self.sizeOfBasis = sizeOfBasis
        self.center = center

        scale = sizeOfLattice*10

        self.electronLattice=[None]*sizeOfLattice
        for x in xrange(sizeOfLattice):
            self.electronLattice[x]=[None]*sizeOfLattice
            for y in xrange(sizeOfLattice):
                position = vector(x,y,0.0)
                density = 0.0
                count = 0
                for direction in xrange(2):
                    for mode in xrange(self.sizeOfBasis):
                        density += basis.basis[direction][mode].getDensity(x,y)
                        count += 1
                        density = density/count
                        
                size = scale*density*count
                color = (1.0,1.0,1.0)
                self.electronLattice[x][y] = DensityPoint()
                self.electronLattice[x][y].setAttributes(position,size,color)


##    def getDifferentialXup(self,elec,x,y):
##        if x+1 == sizeOfLattice:
##            return (self.electronLattice[-1][y][elec].spinor[0]-self.electronLattice[x-1][y][elec].spinor[0])/4.0
##        else:
##            return (self.electronLattice[x+1][y][elec].spinor[0]-self.electronLattice[x-1][y][elec].spinor[0])/4.0
##
##    def getDifferentialXdown(self,elec,x,y):
##        if x+1 == sizeOfLattice:
##            return (self.electronLattice[-1][y][elec].spinor[1]-self.electronLattice[x-1][y][elec].spinor[1])/4.0
##        else:
##            return (self.electronLattice[x+1][y][elec].spinor[1]-self.electronLattice[x-1][y][elec].spinor[1])/4.0
##
##    def getDifferentialYup(self,elec,x,y):
##        if y+1 == sizeOfLattice:
##            return (self.electronLattice[x][-1][elec].spinor[2]-self.electronLattice[x][y-1][elec].spinor[2])/4.0
##        else:
##            return (self.electronLattice[x][y+1][elec].spinor[2]-self.electronLattice[x][y-1][elec].spinor[2])/4.0
##
##    def getDifferentialYdown(self,elec,x,y):
##        if y+1 == sizeOfLattice:
##            return (self.electronLattice[x][-1][elec].spinor[3]-self.electronLattice[x][y-1][elec].spinor[3])/4.0
##        else:
##            return (self.electronLattice[x][y+1][elec].spinor[3]-self.electronLattice[x][y-1][elec].spinor[3])/4.0





sizeOfLattice = 32
sizeOfBasis = 8

m = 1.0
c = 1.0

center = vector((sizeOfLattice/2.0),(sizeOfLattice/2.0),0)
Z = 2.0
sigma = sizeOfLattice/6.0


scene = display(title='Density Plot', x=0, y=0, width=600, height=600, center=center, background=(0,0,0))

basis = Basis(sizeOfLattice, sizeOfBasis, center, Z)
##basis.normalizeSpinors()
##basis.getCoulombicRepulsions()

lattice = Lattice(basis, sizeOfLattice, sizeOfBasis, center)





