import numpy as np
import laff
import flame

# Later this week, you will find out that applying Gaussian Elimination to a matrix
# is equivalent to computing a unit lower triangular matrix, L, and upper 
# triangular matrix, U, such that A = L U.  Here we use that fact to create a 
# matrix with which to perform Gaussian elimination that doesn't have nasty fractions.

L = np.matrix( ' 1, 0, 0. 0;\
                -2, 1, 0, 0;\
                 1,-3, 1, 0;\
                 2, 3,-1, 1' )

U = np.matrix( ' 2,-1, 3,-2;\
                 0,-2, 1,-1;\
                 0, 0, 1, 2;\
                 0, 0, 0, 3' )

A = L * U

print( 'A = ' )
print( A )

# create a solution vector x

x = np.matrix( '-1;\
                 2;\
                 1;\
                -2' )

# store the original value of x

xold = np.matrix( np.copy( x ) )

# create a solution vector b so that A x = b

b = A * x
print( 'b = ' )
print( b )

L0 = np.matrix( ' 1, 0, 0, 0;\
                  2, 1, 0, 0;\
                 -1, 0, 1, 0;\
                 -2, 0, 0, 1' );
A = L0 * A
print( A )

L1 = np.matrix( ' 1, 0, 0, 0;\
                  0, 1, 0, 0;\
                  0, 3, 1, 0;\
                  0,-3, 0, 1' );
A = L1 * A
print( A )

L2 = np.matrix( ' 1, 0, 0, 0;\
                  0, 1, 0, 0;\
                  0, 0, 1, 0;\
                  0, 0, 1, 1' );
A = L2 * A
print( A )

b = L0 * b
b = L1 * b
b = L2 * b

print( 'b after forward substitution (application of Gauss transforms):' )
print( b )

x[ 3 ] = b[ 3 ] / A[ 3,3 ]
x[ 2 ] = ( b[ 2 ] - A[ 2,3 ] * x[ 3 ] ) / A[ 2,2 ]
x[ 1 ] = ( b[ 1 ] - A[ 1,2 ] * x[ 2 ] - A[ 1,3 ] * x[ 3 ] ) / A[ 1,1 ]
x[ 0 ] = ( b[ 0 ] - A[ 0,1 ] * x[ 1 ] - A[ 0,2 ] * x[ 2 ]- A[ 0,3 ] * x[ 3 ] ) / A[ 0,0 ]

print( 'x = ' )
print( x )

print( 'x - xold' )
print( x - xold )

import flame
import laff

def GaussianElimination_unb(A):

    ATL, ATR, \
    ABL, ABR  = flame.part_2x2(A, \
                               0, 0, 'TL')

    while ATL.shape[0] < A.shape[0]:

        A00,  a01,     A02,  \
        a10t, alpha11, a12t, \
        A20,  a21,     A22   = flame.repart_2x2_to_3x3(ATL, ATR, \
                                                       ABL, ABR, \
                                                       1, 1, 'BR')

        #------------------------------------------------------------#

        laff.invscal( alpha11, a21 )        #  a21 := a21 / alpha11
        laff.ger( -1.0, a21, a12t, A22 )    #  A22 := A22 - a21 * a12t

        #------------------------------------------------------------#

        ATL, ATR, \
        ABL, ABR  = flame.cont_with_3x3_to_2x2(A00,  a01,     A02,  \
                                               a10t, alpha11, a12t, \
                                               A20,  a21,     A22,  \
                                               'TL')

    flame.merge_2x2(ATL, ATR, \
                    ABL, ABR, A)



# recreate matrix A
A = L * U

# recreate the right-hand side
b = A * xold

# apply Gaussian elimination to matrix A
GaussianElimination_unb( A )


print( A )
print( 'Original L' )
print( L )
print( 'Original U' )
print( U )

import flame

def ForwardSubstitution_unb(A, b):

    ATL, ATR, \
    ABL, ABR  = flame.part_2x2(A, \
                               0, 0, 'TL')

    bT, \
    bB  = flame.part_2x1(b, \
                         0, 'TOP')

    while ATL.shape[0] < A.shape[0]:

        A00,  a01,     A02,  \
        a10t, alpha11, a12t, \
        A20,  a21,     A22   = flame.repart_2x2_to_3x3(ATL, ATR, \
                                                       ABL, ABR, \
                                                       1, 1, 'BR')

        b0,    \
        beta1, \
        b2     = flame.repart_2x1_to_3x1(bT, \
                                         bB, \
                                         1, 'BOTTOM')

        #------------------------------------------------------------#

        laff.axpy( -beta1, a21, b2 )

        #------------------------------------------------------------#

        ATL, ATR, \
        ABL, ABR  = flame.cont_with_3x3_to_2x2(A00,  a01,     A02,  \
                                               a10t, alpha11, a12t, \
                                               A20,  a21,     A22,  \
                                               'TL')

        bT, \
        bB  = flame.cont_with_3x1_to_2x1(b0,    \
                                         beta1, \
                                         b2,    \
                                         'TOP')

    flame.merge_2x1(bT, \
                    bB, b)


print( A )
print( b )
ForwardSubstitution_unb( A, b )

print( 'updated b' )
print( b )

x[ 3 ] = b[ 3 ] / A[ 3,3 ]
x[ 2 ] = ( b[ 2 ] - A[ 2,3 ] * x[ 3 ] ) / A[ 2,2 ]
x[ 1 ] = ( b[ 1 ] - A[ 1,2 ] * x[ 2 ] - A[ 1,3 ] * x[ 3 ] ) / A[ 1,1 ]
x[ 0 ] = ( b[ 0 ] - A[ 0,1 ] * x[ 1 ] - A[ 0,2 ] * x[ 2 ]- A[ 0,3 ] * x[ 3 ] ) / A[ 0,0 ]

print( 'x = ' )
print( x )

print( 'x - xold' )
print( x - xold )