import flame
import laff as laff

def Mvmult_n_unb_var1B(A, x, y):

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

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

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

        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( a10t, x0, psi1 )
        laff.dots( alpha11, chi1, psi1 )
        laff.dots( a12t, x2, psi1 )

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

        ATL, ATR, \
        ABL, ABR  = flame.cont_with_3x3_to_2x2(A00,  a01,     A02,  \
                                               a10t, alpha11, a12t, \
                                               A20,  a21,     A22,  \
                                               '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

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

print( 'A before =' )
print( A )

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

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

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

Mvmult_n_unb_var1B( A, x, y )

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

print( 'y - ( A * x + yold ) = ' )
print( y - ( A * x + yold ) )

import flame
import laff as laff

def Mvmult_n_unb_var2B(A, x, y):

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

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

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

        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, a01, y0 )
        laff.axpy( chi1, alpha11, psi1 )
        laff.axpy( chi1, a21, y2 )

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

        ATL, ATR, \
        ABL, ABR  = flame.cont_with_3x3_to_2x2(A00,  a01,     A02,  \
                                               a10t, alpha11, a12t, \
                                               A20,  a21,     A22,  \
                                               '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

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

print( 'A before =' )
print( A )

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

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

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

Mvmult_n_unb_var2B( A, x, y )

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

print( 'y - ( A * x + yold ) = ' )
print( y - ( A * x + yold ) )