from Numeric import *
from scipy import *                                                                                                     # module which contains some useful eigensystem operations


class Matrix:                                                                                                           # my main Matrix class
    basisDim = 8
    array = []

    def __init__(self, basisDim):
        self.basisDim = basisDim
        array = zeros((basisDim,basisDim))                                                                              # array of Float objects 
        self.array = array

    def ensureSymmetry(self):                                                                                           # averages the matrix with its transpose
        for p in range(self.basisDim):                                                                                  # producing a net symmetric matrix
            for q in range(self.basisDim):
                self.array[p][q] = self.array[p][q] - (self.array[p][q] - self.array[q][p])/2.0

    def fillOverlapMatrix(self, alphas, centers):                                                                       # Fills in the overlap matrix
        for p in range(self.basisDim):
            for q in range(self.basisDim):



        self.ensureSymmetry()

    def fillKineticMatrix(self, alphas, centers, overlapMatrix):                                                        # Fills in the kinetic matrix
        for p in range(self.basisDim):
            for q in range(self.basisDim):



        self.ensureSymmetry()

    def fillNuclearMatrix(self, alphas, centers):                                                                       # Fills in the nuclear potential matrix
        for p in range(self.basisDim):
            for q in range(self.basisDim):
                                                                                                                        # nuclear potential
        self.ensureSymmetry()


    def makehMatrix(self, kineticMatrix, nuclearMatrix):                                                                # this adds corresponding entries of the kinetic and nuclear matrices
        for p in range(self.basisDim):                                                                                  # to produce the one electron matrix
            for q in range(self.basisDim):
                self.array[p][q] = kineticMatrix.getArray()[p][q] + nuclearMatrix.getArray()[p][q]
        self.ensureSymmetry()

    def makeqMatrix(self, QMatrix, overlapMatrix, initialCr, initialCs):                                                # two electron contribution, generated from the Q and overlap matrices, and the current expansion for the electron
        initialCr = initialCr/(matrixmultiply(initialCr,matrixmultiply(overlapMatrix.getArray(),initialCr)))**0.5       # ensures normality
        initialCs = initialCs/(matrixmultiply(initialCs,matrixmultiply(overlapMatrix.getArray(),initialCs)))**0.5       # ensures normality
       
        for p in range(self.basisDim):
            for q in range(self.basisDim):
                self.array[p][q] = matrixmultiply(matrixmultiply(QMatrix.getArray()[p][q], initialCr), initialCs)
        self.ensureSymmetry()

    def makefockMatrix(self, hMatrix, qMatrix):                                                                         # this adds the corresponding entries of the single and double electron matrices together
        for p in range(self.basisDim):
            for q in range(self.basisDim):
                self.array[p][q] = hMatrix.getArray()[p][q] + qMatrix.getArray()[p][q]
        self.ensureSymmetry()
    


    def makeVMatrix(self, aMatrix):                                                                                     # creates the transformation matrix for solving the generalized eigenvalue problem
        U = aMatrix.returnEigenVectors()
        b = aMatrix.returnSEigenValueArray()
        self.array = matrixmultiply(U,b)

        
    def makeTransformedMatrix(self, originalMatrix, transformationMatrix):                                              # transforms a given matrix with another via Vdagger*H*V
        transformationMatrixdagger = transpose(conjugate(transformationMatrix.getArray()))
        self.array = matrixmultiply(matrixmultiply(transformationMatrixdagger,originalMatrix.getArray()),transformationMatrix.getArray())


    def returnInverse(self):
        return linalg.inv(self.array)

    def returnEigenValues(self):
        return linalg.eig(self.array)[0]

    def returnEigenVectors(self):
        return linalg.eig(self.array)[1]

    def returnSEigenValueArray(self):                                                                                   # for use in the generation of our transformation matrix
        eigenValues = self.returnEigenValues()
        newarray = zeros((self.basisDim,self.basisDim))
        
        for p in range(self.basisDim):
            for q in range(self.basisDim):
                if p == q: newarray[p][q] = abs(eigenValues[p])**(-0.5)
                else: newarray[p][q] = 0.0
        return newarray
    

    def getArray(self):                                                                                                 # returns the array object
        return self.array

    def printMatrix(self):                                                                                              # prints the matrix to output
        for p in range(self.basisDim):
            print self.array[p].real




class QMatrix:                                                                                                          # two electron matrix, of size NxNxNxN where N is the dimension of our basis
    basisDim = 8
    array = []

    def __init__(self, basisDim):
        self.basisDim = basisDim

        array = zeros((basisDim,basisDim,basisDim,basisDim))      
        self.array = array


    def fillQMatrix(self, alphas, centers):                                                                             # fill the two electron matrix via the equation on page 81 of Thijssen
        for p in range(self.basisDim):
            for q in range(self.basisDim):
                for r in range(self.basisDim):
                    for s in range(self.basisDim):
                        self.array[p][q][r][s] = (2.0*pi**(2.5))/((alphas[p]+alphas[q])*(alphas[r]+alphas[s])*(alphas[p]+alphas[q]+alphas[r]+alphas[s])**0.5)


    def getArray(self):                                                                                                 # returns the array object
        return self.array

    def printMatrix(self):                                                                                              # prints the matrix to output
        for p in range(self.basisDim):
            print self.array[p].real



