import flame
import laff as laff

def Tmvmult_ln_unb_var1(L, x, y):

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

    xT, \
    xB  = flame.part_2x1(x, \
                         0, 'TOP')

    yT, \
    yB  = flame.part_2x1(y, \
                         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')

        x0,   \
        chi1, \
        x2    = flame.repart_2x1_to_3x1(xT, \
                                        xB, \
                                        1, 'BOTTOM')

        y0,   \
        psi1, \
        y2    = flame.repart_2x1_to_3x1(yT, \
                                        yB, \
                                        1, 'BOTTOM')

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

        laff.dots( l10t, x0, psi1 )
        laff.dots( lambda11, chi1, psi1 )

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

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

        xT, \
        xB  = flame.cont_with_3x1_to_2x1(x0,   \
                                         chi1, \
                                         x2,   \
                                         'TOP')

        yT, \
        yB  = flame.cont_with_3x1_to_2x1(y0,   \
                                         psi1, \
                                         y2,   \
                                         'TOP')

    flame.merge_2x1(yT, \
                    yB, y)



from numpy import random
from numpy import matrix

L = matrix( random.rand( 4,4 ) )
x = matrix( random.rand( 4,1 ) )
y = matrix( random.rand( 4,1 ) )
yold = matrix( random.rand( 4,1 ) )

# Notice that L is not lower triangular.  We will only use the lower triangular part.

print( 'L before =' )
print( L )

print( 'x before =' )
print( x )

print( 'y before =' )
print( y )

import numpy as np
laff.copy( y, yold )   # save the original vector y

Tmvmult_ln_unb_var1( L, x, y )

print( 'y after =' )
print( y )

print( 'y - ( np.tril( L ) * x + yold ) = ' ) #np.tril makes the matrix lower triangular
print( y - ( np.tril( L ) * x + yold ) )

import flame
import laff as laff

def Tmvmult_ln_unb_var2(L, x, y):

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

    xT, \
    xB  = flame.part_2x1(x, \
                         0, 'TOP')

    yT, \
    yB  = flame.part_2x1(y, \
                         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')

        x0,   \
        chi1, \
        x2    = flame.repart_2x1_to_3x1(xT, \
                                        xB, \
                                        1, 'BOTTOM')

        y0,   \
        psi1, \
        y2    = flame.repart_2x1_to_3x1(yT, \
                                        yB, \
                                        1, 'BOTTOM')

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

        laff.axpy( chi1, lambda11, psi1 )
        laff.axpy( chi1, l21, y2 )

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

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

        xT, \
        xB  = flame.cont_with_3x1_to_2x1(x0,   \
                                         chi1, \
                                         x2,   \
                                         'TOP')

        yT, \
        yB  = flame.cont_with_3x1_to_2x1(y0,   \
                                         psi1, \
                                         y2,   \
                                         'TOP')

    flame.merge_2x1(yT, \
                    yB, y)



from numpy import random
from numpy import matrix

L = matrix( random.rand( 4,4 ) )
x = matrix( random.rand( 4,1 ) )
y = matrix( random.rand( 4,1 ) )
yold = matrix( random.rand( 4,1 ) )

# L is not lower triangular.  We will only use the lower triangular part.

print( 'L before =' )
print( L )

print( 'x before =' )
print( x )

print( 'y before =' )
print( y )

laff.copy( y, yold )   # save the original vector y

Tmvmult_ln_unb_var2( L, x, y )

print( 'y after =' )
print( y )

print( 'y - ( np.tril( L ) * x + yold ) = ' ) #np.tril makes the matrix lower triangular
print( y - ( np.tril( L ) * x + yold ) )