import numpy as np
import laff
import flame

# Applying the LU Factorization to a matrix is the process of
# computing a unit lower triangular matrix, L, and upper 
# triangular matrix, U, such that A = L U.  To avoid nasty fractions
# in these caclulations, we create our matrix A from two matrices, L and U, 
# whose elements we know to be integer valued.

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 )

# Much later, we are also going to solve A x = y.  Here we create that y:
x2 = np.matrix( '1;\
                 2;\
                 -2;\
                 1' )
y = A * x2
print( 'y = ' )
print( y )

import flame
import laff

def LU_unb_var5(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
LU_unb_var5( A )


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

import flame

def Ltrsv_unb_var1(L, b):

    LTL, LTR, \
    LBL, LBR  = flame.part_2x2(L, \
                               0, 0, 'TL')

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

    while LTL.shape[0] < L.shape[0]:

        L00,  l01,     L02,  \
        l10t, lambda11, l12t, \
        L20,  l21,     L22   = flame.repart_2x2_to_3x3(LTL, LTR, \
                                                       LBL, LBR, \
                                                       1, 1, 'BR')

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

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

        laff.axpy( -beta1, l21, b2 ) # b2 := b2 - beta1*l21

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

        LTL, LTR, \
        LBL, LBR  = flame.cont_with_3x3_to_2x2(L00,  l01,      L02,  \
                                               l10t, lambda11, l12t, \
                                               L20,  l21,      L22,  \
                                               'TL')

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

    flame.merge_2x1(bT, \
                    bB, b)


print( A )
print( b )
Ltrsv_unb_var1( 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 )

import flame
import laff as laff

def Utrsv_unb_var1(U, b):

    UTL, UTR, \
    UBL, UBR  = flame.part_2x2(U, \
                               0, 0, 'BR')

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

    while UBR.shape[0] < U.shape[0]:

        U00,  u01,       U02,  \
        u10t, upsilon11, u12t, \
        U20,  u21,       U22   = flame.repart_2x2_to_3x3(UTL, UTR, \
                                                         UBL, UBR, \
                                                         1, 1, 'TL')

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

        #------------------------------------------------------------#
        
        laff.dots( -u12t, b2, beta1 ) # beta1 := beta1 - u21 * b2
        laff.invscal( upsilon11, beta1 ) # beta1 := beta1/upsilon11

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

        UTL, UTR, \
        UBL, UBR  = flame.cont_with_3x3_to_2x2(U00,  u01,       U02,  \
                                               u10t, upsilon11, u12t, \
                                               U20,  u21,       U22,  \
                                               'BR')

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

    flame.merge_2x1(bT, \
                    bB, b)



# just to be sure, let's start over.  We'll recreate A, x, and b, run all the routines, and
# then compare the updated b to the original vector x.

A = L * U
b = A * x

LU_unb_var5( A )
Ltrsv_unb_var1( A, b )
Utrsv_unb_var1( A, b )

print( 'updated b' )
print( b )
print( 'original x' )
print( x )
print( 'b - x' )
print( b - x )


def Solve( A, b ):
    
    # insert appropriate calls to routines you have written here!
    LU_unb_var5( A )
    Ltrsv_unb_var1( A, b )
    Utrsv_unb_var1( A, b )
    
    

# just to be sure, let's start over.  We'll recreate A, x, and b, run all the routines, and
# then compare the updated b to the original vector x.

A = L * U
b = A * x

Solve( A, b )

print( 'updated b' )
print( b )
print( 'original x' )
print( x )
print( 'b - x' )
print( b - x )


# insert appropriate calls here.

Ltrsv_unb_var1( A, y )
Utrsv_unb_var1( A, y )

print( 'y = ' )
print( y )

print( 'x2 - y =' )
print( x2 - y )