from Numeric import *
from ArrayPrinter import *

class LuDecomposition:
        def __init__(self, aSquareArray):
                assert len(aSquareArray) == len(aSquareArray[0])
                self.array = aSquareArray
                self.n = len(aSquareArray)

                self.l = self.makeEmptySquareMatrix(self.n)
                self.u = self.makeEmptySquareMatrix(self.n)

                for i in range(self.n):
                        self.u[i][i] = 1.0

        def makeEmptySquareMatrix(self, n):
                matrix = []
                for i in range(0,n):
                        row = []
                        for j in range(0,n):
                                row.append(0.0)
                        matrix.append(row)
                return matrix

        def perform(self):
                for i in range(self.n):
                        self._writeL(i)
                        self._writeU(i)
                return self.l, self.u

        def _writeL(self, aCol):
                for eachRow in range(aCol, self.n):
                        self.l[eachRow][aCol] = self.array[eachRow][aCol] - self._minusForWriteL(eachRow, aCol)

        def _minusForWriteL(self, aRow, aCol):
                totalMinus = 0
                for n in range(0, aCol):
                        totalMinus += self.l[aRow][n] * self.u[n][aCol]
                return totalMinus

        def _writeU(self, aRow):
                for eachCol in range(aRow+1, self.n):
                        self.u[aRow][eachCol] = float(self.array[aRow][eachCol] - self._minusForWriteU(eachCol, aRow)) / float(self.l[aRow][aRow])

        def _minusForWriteU(self, aCol, aRow):
                totalMinus = 0
                for n in range(0, aRow):
                        totalMinus += self.l[aRow][n] * self.u[n][aCol]
                return totalMinus

        def _getL(self):
                return self.l

        def _getU(self):
                return self.u

if __name__=="__main__":
                a = [[1.0,2.0,3.0],[4.0,5.0,6.0],[7.0,8.0,9.0]]
                l, u = LuDecomposition(a).perform()
                print "array : "
                print array2string(array(a))
                print "Lu Decomposition - L"
                print array2string(array(l))
                print "Lu Decomposition - U"
                print array2string(array(u))