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

class waveFunction:
    def __init__(self, Anaught, n, m):
        self.Anaught = Anaught
        self.n = n
        self.m = m

    def fullWaveFunction(self, x, y):
        r = sqrt((x-center[0])**2+(y-center[1])**2)
        self.radialWavefunction = special.genlaguerre(2*self.n,2*self.n-1)(2.0*r/(self.Anaught*self.n))*exp(-r/(self.Anaught*self.n))*(2.0*r/(self.Anaught*self.n))**(self.n-1)
        self.thetaWavefunction = special.sph_harm(self.m,2,atan2(y-center[1],x-center[0]),pi/2)
        self.fullWavefunction = self.radialWavefunction*self.thetaWavefunction
        return self.fullWavefunction

class basis:
    def __init__(self, Anaught, maxn, coeffs):
        self.maxn = maxn
        self.coeffs = coeffs

        count = 0
        for n in xrange(maxn):
            for m in xrange(2*n+1):
                print "Radial QN", n+1
                print "Angular QN", m-n
                count += 1

        self.maxNbasis = count

        i = 0
        self.basis = [0]*self.maxNbasis
        for n in xrange(maxn):
            for m in xrange(2*n+1):
                self.basis[i] = waveFunction(Anaught,n+1,m-n)
                i += 1

    def getFullWavefunction(self, x, y):
        i = 0
        self.fullWavefunction = 0.0
        for n in xrange(maxn):
            for m in xrange(2*n+1):
                self.fullWavefunction += self.coeffs[i]*self.basis[i].fullWaveFunction(x,y)
                i += 1
        return self.fullWavefunction


class electron:
    def __init__(self, basis, maxn):
        self.basis = basis
        self.maxn = maxn

    def getDensity(self,x,y):
        return (conjugate(self.basis.getFullWavefunction(x,y))*self.basis.getFullWavefunction(x,y)).real

    def getPhase(self,x,y):
        self.phase = atan2(self.basis.getFullWavefunction(x,y).imag,self.basis.getFullWavefunction(x,y).real)
        return self.phase



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 visualLattice:
    def __init__(self, electronOne, sizeOfLattice, center):
        self.electronOne = electronOne
        self.sizeOfLattice = sizeOfLattice
        self.center = center
        self.scale = sizeOfLattice*.0001

        self.pointLattice=[None]*sizeOfLattice
        for x in xrange(sizeOfLattice):
            self.pointLattice[x]=[None]*sizeOfLattice
            for y in xrange(sizeOfLattice):
                position = vector(x,y,0.0)
                density = self.electronOne.getDensity(x,y)
                size = self.scale*density
                phase = self.electronOne.getPhase(x,y)
                colorArg = phase/pi
                color = ((1+abs(colorArg))/2.0,0.0,(1-abs(colorArg))/2.0)
                
                self.pointLattice[x][y] = densityPoint()
                self.pointLattice[x][y].setAttributes(position,size,color)

    def updateLattice(self,electronOne):
        for x in xrange(sizeOfLattice):
            for y in xrange(sizeOfLattice):
                position = vector(x,y,0.0)
                density = self.electronOne.getDensity(x,y)
                size = self.scale*density
                phase = self.electronOne.getPhase(x,y)
                colorArg = phase/pi
                color = ((1+abs(colorArg))/2.0,0.0,(1-abs(colorArg))/2.0)
                
                self.pointLattice[x][y].setAttributes(position,size,color)







maxn = 3
Anaught = 1.0

count = 0
for n in xrange(maxn):
    for m in xrange(2*n+1):
        print "Radial QN", n+1
        print "Angular QN", m-n
        count += 1

i = 0
coeffs = [None]*count
for n in xrange(maxn):
    for m in xrange(2*n+1):
        coeffs[i] = float(m-n)/((float(n+1))**2+1.0)
        i += 1

print coeffs

Basis = basis(Anaught, maxn, coeffs)
electronOne = electron(Basis, maxn)

sizeOfLattice = 32

center = (sizeOfLattice/2.0,sizeOfLattice/2.0,0.0)

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

visualLattice = visualLattice(electronOne, sizeOfLattice, center)

visualLattice.updateLattice(electronOne)








