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


class basisWavefunction:
    def __init__(self, center, Anaught, coeffs, n, m):
        self.center = center
        self.Anaught = Anaught
        self.n = n
        self.m = m

    def returnFullWavefunction(self,x,y):
        r = sqrt((x-center[0])**2+(y-center[1])**2)
        self.radialWavefunction = special.genlaguerre(self.n-self.m-1,2*self.m+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)
        return coeffs[i]*self.radialWavefunction*self.thetaWavefunction

    def setAnaught(self,Anaught):
        self.Anaught = Anaught


class electronPiece:
    def __init__(self, probabilityAmplitude):
        self.tonorm = 1.0
        self.probabilityAmplitude = self.tonorm*probabilityAmplitude
        self.spinor = [probabilityAmplitude,0,probabilityAmplitude,0]

    def getProbabilityAmplitude(self):
        return self.probabilityAmplitude
        
    def getDensity(self):
        self.probability = (dot(conjugate(self.spinor),self.spinor)).real
        return self.probability

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

    def normalize(self,tonorm):
        self.tonorm = tonorm
        self.probabilityAmplitude = self.tonorm*self.probabilityAmplitude
        self.spinor = [self.probabilityAmplitude,0,self.probabilityAmplitude,0]


class basis:
    def __init__(self, maxN, center, coeffs):

        for n in xrange(maxN):
            for m in xrange(2*n+1):
                
            


class numericLattice:
    def __init__(self, sizeOfLattice, initialWavefunction, center):
        self.sizeOfLattice = sizeOfLattice
        self.initialWavefunction = initialWavefunction
        self.center = center

        self.numericLattice = [0]*sizeOfLattice
        self.probabilityAmplitudes = [0]*sizeOfLattice
        self.differential = [0]*sizeOfLattice
        self.laplacian = [0]*sizeOfLattice
        for x in xrange(sizeOfLattice):
            self.numericLattice[x] = [0]*sizeOfLattice
            self.probabilityAmplitudes[x] = [0]*sizeOfLattice
            self.differential[x] = [0]*sizeOfLattice
            self.laplacian[x] = [0]*sizeOfLattice
            for y in xrange(sizeOfLattice):
                self.probabilityAmplitudes[x][y] = initialWavefunction.returnFullWavefunction(x,y)
                self.numericLattice[x][y] = electronPiece(self.probabilityAmplitudes[x][y])
                self.differential[x][y] = vector(0.0,0.0,0.0)
                self.laplacian[x][y] = 0.0

    def setWavefunction(self):
        for x in xrange(sizeOfLattice):
            for y in xrange(sizeOfLattice):
                self.probabilityAmplitudes[x][y] = self.initialWavefunction.returnFullWavefunction(x,y)
                self.numericLattice[x][y] = electronPiece(self.probabilityAmplitudes[x][y])
                self.differential[x][y] = vector(0.0,0.0,0.0)
                self.laplacian[x][y] = 0.0

    def normalizeDensity(self):
        density = 0.0
        for x in xrange(self.sizeOfLattice):
            for y in xrange(self.sizeOfLattice):
                density += self.numericLattice[x][y].getDensity()
        for x in xrange(self.sizeOfLattice):
            for y in xrange(self.sizeOfLattice):
                self.numericLattice[x][y].normalize(1/sqrt(density))

    def getIonicPotential(self,Z):
        self.ionicPotential = 0.0
        for x in xrange(self.sizeOfLattice):
            for y in xrange(self.sizeOfLattice):
                if x != self.center[0] or y != self.center[1]:
                    self.ionicPotential += -Z*self.numericLattice[x][y].getDensity()/(mag(vector(x,y,0)-self.center))
        return self.ionicPotential

    def getDifferential(self,lattice):
        for x in xrange(self.sizeOfLattice):
            for y in xrange(self.sizeOfLattice):
                differentialX = 0.0
                differentialY = 0.0
                if x+1 == self.sizeOfLattice:
                    differentialX += (lattice[-1][y]-lattice[x-1][y])/2.0
                else:
                    differentialX += (lattice[x+1][y]-lattice[x-1][y])/2.0
                if y+1 == self.sizeOfLattice:
                    differentialY += (lattice[x][-1]-lattice[x][y-1])/2.0
                else:
                    differentialY += (lattice[x][y+1]-lattice[x][y-1])/2.0
                self.differential[x][y] = (differentialX, differentialY, 0)
        return self.differential

    def getLaplacian(self,lattice):
        for x in xrange(self.sizeOfLattice):
            for y in xrange(self.sizeOfLattice):
                differentialX = 0.0
                differentialY = 0.0
                if x+1 == self.sizeOfLattice:
                    differentialX += (lattice[-1][y][0]-lattice[x-1][y][0])/2.0
                else:
                    differentialX += (lattice[x+1][y][0]-lattice[x-1][y][0])/2.0
                if y+1 == self.sizeOfLattice:
                    differentialY += (lattice[x][-1][1]-lattice[x][y-1][1])/2.0
                else:
                    differentialY += (lattice[x][y+1][1]-lattice[x][y-1][1])/2.0
                self.laplacian[x][y] = (differentialX + differentialY)
        return self.laplacian

    def getKineticEnergy(self):
        self.kineticEnergy = 0.0
        differential = self.getDifferential(self.probabilityAmplitudes)
        laplacian = self.getLaplacian(differential)
        for x in xrange(self.sizeOfLattice):
            for y in xrange(self.sizeOfLattice):
                self.kineticEnergy += -(conjugate(self.probabilityAmplitudes[x][y])*laplacian[x][y]).real/2.0
        return self.kineticEnergy




class interactions:
    def __init__(self, sizeOfLattice, numericLatticeOne, numericLatticeTwo):
        self.sizeOfLattice = sizeOfLattice
        self.numericLatticeOne = numericLatticeOne
        self.numericLatticeTwo = numericLatticeTwo

    def getCoulombicRepulsion(self):
        self.coulombicRepulsionPotential = 0.0
        for x1 in xrange(self.sizeOfLattice):
            for y1 in xrange(self.sizeOfLattice):
                for x2 in xrange(self.sizeOfLattice):
                    for y2 in xrange(self.sizeOfLattice):
                        if x1 != x2 or y1 != y2:
                            self.coulombicRepulsionPotential += self.numericLatticeOne.numericLattice[x1][y1].getDensity()*self.numericLatticeTwo.numericLattice[x2][y2].getDensity()/(mag(vector(x1,y1,0)-vector(x2,y2,0)))        
        return self.coulombicRepulsionPotential


    def getCoulombicExchange(self):
        self.coulombicExchangePotential = 0.0
        for x1 in xrange(self.sizeOfLattice):
            for y1 in xrange(self.sizeOfLattice):
                for x2 in xrange(self.sizeOfLattice):
                    for y2 in xrange(self.sizeOfLattice):
                        if x1 != x2 or y1 != y2:
                            self.coulombicExchangePotential += conjugate(self.numericLatticeOne.numericLattice[x1][y1].getProbabilityAmplitude()*self.numericLatticeTwo.numericLattice[x2][y2].getProbabilityAmplitude())*self.numericLatticeTwo.numericLattice[x1][y1].getProbabilityAmplitude()*self.numericLatticeOne.numericLattice[x2][y2].getProbabilityAmplitude()/(mag(vector(x1,y1,0)-vector(x2,y2,0)))        
        return self.coulombicExchangePotential.real        












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, numericLattice, sizeOfLattice, center):
        self.numericLattice = numericLattice
        self.sizeOfLattice = sizeOfLattice
        self.center = center
        self.scale = sizeOfLattice*.1

        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.numericLattice.numericLattice[x][y].getDensity()
                size = self.scale*density
                phase = self.numericLattice.numericLattice[x][y].getPhase()
                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,numericLattice):
        for x in xrange(sizeOfLattice):
            for y in xrange(sizeOfLattice):
                position = vector(x,y,0.0)
                density = self.numericLattice.numericLattice[x][y].getDensity()
                size = self.scale*density
                phase = self.numericLattice.numericLattice[x][y].getPhase()
                colorArg = phase/pi
                color = ((1+abs(colorArg))/2.0,0.0,(1-abs(colorArg))/2.0)
                
                self.pointLattice[x][y].setAttributes(position,size,color)








            
print '\n-----------------------------'
print "Calculation Instance Started"
print time.ctime()

sizeOfLattice = 16

center = vector(sizeOfLattice/2.0,sizeOfLattice/2.0,0.0)
sigma = 4.0
mOne = 2.0
mTwo = -2.0

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

initialWavefunctionOne = initialWavefunction(center, sigma, mOne)
initialWavefunctionTwo = initialWavefunction(center, sigma, mTwo)

numericLatticeOne = numericLattice(sizeOfLattice, initialWavefunctionOne, center)
numericLatticeTwo = numericLattice(sizeOfLattice, initialWavefunctionTwo, center)
numericLatticeOne.normalizeDensity()
numericLatticeTwo.normalizeDensity()

interactions = interactions(sizeOfLattice, numericLatticeOne, numericLatticeTwo)

##visualLattice = visualLattice(numericLattice, sizeOfLattice, center)
print "Initialization finalized at:", time.ctime()
tolerance = .0001 
print "Tolerance:", tolerance

energies = list()

for M in xrange(0,16,1):
    Z = 3.0+M/5.0
    if Z <= 5.0:
        sigma = 8.0
        maxIter = int(sigma*8)
        step = sigma/maxIter
    else:
        sigma = 4.0
        maxIter = int(sigma*8)
        step = sigma/maxIter
        
    i = 0
    oldEnergy = -.5
    newEnergy = 0.0
    lowestEnergy = 0.0
        
    while tolerance < abs(newEnergy-oldEnergy) and i < maxIter-1:
        oldEnergy = newEnergy

        initialWavefunctionOne.setSigma(sigma)
        numericLatticeOne.setWavefunction()
        numericLatticeOne.normalizeDensity()
        initialWavefunctionTwo.setSigma(sigma)
        numericLatticeTwo.setWavefunction()
        numericLatticeTwo.normalizeDensity()
    ##    visualLattice.updateLattice(numericLattice)

        exchangeEnergy = interactions.getCoulombicExchange()
        newEnergy = numericLatticeOne.getIonicPotential(Z) + numericLatticeOne.getKineticEnergy() + numericLatticeTwo.getIonicPotential(Z) + numericLatticeTwo.getKineticEnergy() - 2.0*exchangeEnergy + interactions.getCoulombicRepulsion()

        if newEnergy < lowestEnergy:
            lowestEnergy = newEnergy
            lowestExchangeEnergy = exchangeEnergy
            lowestSigma = sigma

        if newEnergy >= oldEnergy:
            i = maxIter

        print "Energy:", newEnergy, "for ion", Z, "and trial number:", i, "with sigma:", sigma
        print "Energy difference between present and last:", abs(newEnergy-oldEnergy)
        sigma -= step
        i += 1
        
    energies.append([Z,lowestEnergy,lowestExchangeEnergy,lowestSigma])
    print energies

dump = open("dumpAnalysis", mode='w')
dump.write(str(energies))
dump.close()

print "\nCalculation Instance Ended"
print time.ctime()
print '-----------------------------'

##
