import numpy as np

num_addorsub=0
num_mul=0
num_assign=0

def matrix_add(matrix_a, matrix_b):
    '''
    :param matrix_a:
    :param matrix_b:
    :return:matrix_c=matrix_a+matrix_b
    '''
    rows = len(matrix_a) # get numbers of rows
    columns = len(matrix_a[0]) # get numbers of cols
    matrix_c = [list() for i in range(rows)] # build matrix 2d list
    for i in range(rows):
        for j in range(columns):
            matrix_c_temp = matrix_a[i][j] + matrix_b[i][j]
            global num_addorsub,num_assign
            num_addorsub = num_addorsub + 1
            num_assign = num_assign + 1
            matrix_c[i].append(matrix_c_temp)
    return matrix_c


def matrix_minus(matrix_a, matrix_b):
    '''
    :param matrix_a:
    :param matrix_b:
    :return:matrix_c=matrix_a-matrix_b
    '''
    rows = len(matrix_a)
    columns = len(matrix_a[0])
    matrix_c = [list() for i in range(rows)]
    for i in range(rows):
        for j in range(columns):
            matrix_c_temp = matrix_a[i][j] - matrix_b[i][j]
            global num_addorsub,num_assign
            num_addorsub = num_addorsub + 1
            num_assign=num_assign + 1
            matrix_c[i].append(matrix_c_temp)
    return matrix_c


def matrix_divide(matrix_a, row, column):
    '''
    :param matrix_a:
    :param row:
    :param column:
    :return: matrix_b=matrix_a(row,column) to divide matrix_a
    '''
    rows = len(matrix_a)
    columns = len(matrix_a[0])
    matrix_b = [list() for i in range(rows // 2)]
    k = 0
    for i in range((row - 1) * rows // 2, row * rows // 2):
        for j in range((column - 1) * columns // 2, column * columns // 2):
            matrix_c_temp = matrix_a[i][j]
            matrix_b[k].append(matrix_c_temp)
        k += 1
    return matrix_b


def matrix_merge(matrix_11, matrix_12, matrix_21, matrix_22):
    '''
    :param matrix_11:
    :param matrix_12:
    :param matrix_21:
    :param matrix_22:
    :return:mariix merged by 4 parts above
    '''
    length = len(matrix_11)
    matrix_all = [list() for i in range(length * 2)]  # build a matrix of double rows
    for i in range(length):
        # for each row. matrix_all list contain row of matrix_11 and matrix_12
        matrix_all[i] = matrix_11[i] + matrix_12[i]
    for j in range(length):
        # for each row. matrix_all list contain row of matrix_21 and matrix_22
        matrix_all[length + j] = matrix_21[j] + matrix_22[j]
    return matrix_all
 

def strassen(matrix_a, matrix_b):
    '''
    :param matrix_a:
    :param matrix_b:
    :return:matrix_a * matrix_b
    '''
    row_a = len(matrix_a)
    col_a = len(matrix_a[0])
    row_b = len(matrix_b)
    col_b = len(matrix_b[0])
    if col_a != row_b:
        print('matrix_a and matrix_b can not be multiplied')
        return
    global num_mul,num_addorsub
    if row_a == 1 or col_a == 1 or row_b == 1 or col_b == 1:
        matrix_all = [list() for i in range(row_a)]
        for i in range(row_a):
            for j in range(col_b):
                matrix_all_temp = 0
                for k in range(col_a):
                    matrix_all_temp += matrix_a[i][k] * matrix_b[k][j]
                    num_mul = num_mul + 1
                    num_addorsub = num_addorsub + 1
                matrix_all[i].append(matrix_all_temp)
    else:
        # 10 first parts of computing
        s1 = matrix_minus((matrix_divide(matrix_b, 1, 2)), (matrix_divide(matrix_b, 2, 2)))
        s2 = matrix_add((matrix_divide(matrix_a, 1, 1)), (matrix_divide(matrix_a, 1, 2)))
        s3 = matrix_add((matrix_divide(matrix_a, 2, 1)), (matrix_divide(matrix_a, 2, 2)))
        s4 = matrix_minus((matrix_divide(matrix_b, 2, 1)), (matrix_divide(matrix_b, 1, 1)))
        s5 = matrix_add((matrix_divide(matrix_a, 1, 1)), (matrix_divide(matrix_a, 2, 2)))
        s6 = matrix_add((matrix_divide(matrix_b, 1, 1)), (matrix_divide(matrix_b, 2, 2)))
        s7 = matrix_minus((matrix_divide(matrix_a, 1, 2)), (matrix_divide(matrix_a, 2, 2)))
        s8 = matrix_add((matrix_divide(matrix_b, 2, 1)), (matrix_divide(matrix_b, 2, 2)))
        s9 = matrix_minus((matrix_divide(matrix_a, 1, 1)), (matrix_divide(matrix_a, 2, 1)))
        s10 = matrix_add((matrix_divide(matrix_b, 1, 1)), (matrix_divide(matrix_b, 1, 2)))
        # 7 second parts of computing
        p1 = strassen(matrix_divide(matrix_a, 1, 1), s1)
        p2 = strassen(s2, matrix_divide(matrix_b, 2, 2))
        p3 = strassen(s3, matrix_divide(matrix_b, 1, 1))
        p4 = strassen(matrix_divide(matrix_a, 2, 2), s4)
        p5 = strassen(s5, s6)
        p6 = strassen(s7, s8)
        p7 = strassen(s9, s10)
        # 4 final parts of result
        c11 = matrix_add(matrix_add(p5, p4), matrix_minus(p6, p2))
        c12 = matrix_add(p1, p2)
        c21 = matrix_add(p3, p4)
        c22 = matrix_minus(matrix_add(p5, p1), matrix_add(p3, p7))
        matrix_all = matrix_merge(c11, c12, c21, c22)
        global num_assign
        num_assign =num_assign+22
    return matrix_all


def main():

    # statistical data
    A = np.random.random_integers(-5, 5, size=(256, 64))
    print("\nRandom Matrix A:\n", A)
    B = np.random.random_integers(-5, 5, size=(64, 128))
    print("\nRandom Matrix B:\n", B)

    C_verification=np.dot(A,B)

    result = strassen(A, B)
    print("\n A*B Result of matrixs by generate randomly\n",np.array(result))

    print("\nfrequency of add/sub",num_addorsub)
    print("frequency of assign", num_assign)
    print("frequency of mul", num_mul)

    if (C_verification==result).all():
        print("\nCorrect")
    else:
        print("\nWrong")

if __name__ == '__main__':
    main()

# frequency of add/sub 4499186
# frequency of assign 3989370
# frequency of mul 941192
# 2097152