import numpy as np
import time
import multiprocessing


def sequential_matrix_multiply(matrix_a, matrix_b):
    result = np.zeros((len(matrix_a), len(matrix_b[0])))
    for i in range(len(matrix_a)):
        for j in range(len(matrix_b[0])):
            for k in range(len(matrix_b)):
                result[i][j] += matrix_a[i][k] * matrix_b[k][j]
    return result


def parallel_matrix_multiply_worker(args):
    matrix_a, matrix_b, row_start, row_end, result = args
    local_result = np.zeros((row_end - row_start, len(matrix_b[0])))
    for i in range(row_start, row_end):
        for j in range(len(matrix_b[0])):
            for k in range(len(matrix_b)):
                local_result[i - row_start][j] += matrix_a[i][k] * matrix_b[k][j]
    result.extend(local_result)


def parallel_matrix_multiply(matrix_a, matrix_b, num_processes=2):
    num_rows_a = len(matrix_a)
    chunk_size = num_rows_a // num_processes
    processes = []
    manager = multiprocessing.Manager()
    result = manager.list()

    for i in range(num_processes):
        row_start = i * chunk_size
        row_end = (i + 1) * chunk_size if i < num_processes - 1 else num_rows_a
        process_args = (matrix_a, matrix_b, row_start, row_end, result)
        process = multiprocessing.Process(target=parallel_matrix_multiply_worker, args=(process_args,))
        processes.append(process)

    for process in processes:
        process.start()

    for process in processes:
        process.join()

    return np.vstack(result)


def run_test(matrix_size, num_processes=2):
    matrix_a = np.random.rand(matrix_size, matrix_size)
    matrix_b = np.random.rand(matrix_size, matrix_size)

    start_time = time.time()
    result_sequential = sequential_matrix_multiply(matrix_a, matrix_b)
    sequential_time = time.time() - start_time
    print(f"Последовательноe умножение заняло ({matrix_size}x{matrix_size}): {sequential_time} секунд")

    start_time = time.time()
    result_parallel = parallel_matrix_multiply(matrix_a, matrix_b, num_processes)
    parallel_time = time.time() - start_time
    print(
        f"Параллельное умножение матрицы ({matrix_size}x{matrix_size}) с {num_processes} потоками заняло: {parallel_time} секунд")
    print("========================================")

# Тесты для матриц размером 100x100, 300x300 и 500x500 с разным числом процессов

# Бенчмарки для матриц размером 100, 300, 500 строк
if __name__ == '__main__':
    run_test(100, num_processes=2)
    run_test(100, num_processes=4)
    run_test(300, num_processes=2)
    run_test(300, num_processes=4)
    run_test(500, num_processes=2)
    run_test(500, num_processes=4)