import multiprocessing
import random
import time
from pprint import pprint
from multiprocessing import Pool


def create_random_matrix(size):
    return [[random.random() for _ in range(size)] for __ in range(size)]

# def matrix_multiply_seq(matrix1: list[list[float]], matrix2: list[list[float]]) -> list[list[float]]:
#     """Выполняет последовательное перемножение двух матриц и возвращает результирующую матрицу"""
#     l1 = len(matrix1)
#     l2 = len(matrix2)
#     result = [[0 for _ in range(l2)] for __ in range(l1)]
#     for i in range(l1):
#         for j in range(l2):
#             for k in range(l2):
#                 result[i][j] += matrix1[i][k] * matrix2[k][j]
#     return result


matrix_result = [[0 for _ in range(1000)] for __ in range(1000)]

def matrix_multiply(args) -> None:
    """Перемножает строки от start_cnt до end_cnt, результат помещает в глобальную переменную matrix_result"""
    matrix1, matrix2, start_cnt, end_cnt = args
    for i in range(start_cnt, end_cnt):
        for j in range(len(matrix2)):
            for k in range(len(matrix2)):
                matrix_result[i][j] += matrix1[i - start_cnt][k] * matrix2[k][j]


def matrix_multiply_parralel(matrix1: list[list[float]], matrix2: list[list[float]], thread_count):
    """Выполняет парралеьное перемножение матриц"""
    l1 = len(matrix1)

    step = l1 // thread_count
    args = [(matrix1, matrix2, i, i + step) for i in range(0, l1, step)]
    args[-1] = (matrix1, matrix2, step * (l1 - 1), l1)  # Остаток на последний поток

    with Pool(processes=thread_count) as pool:
        pool.map(matrix_multiply, args)
        # pprint(matrix_result, compact=True)

def main():
    sizes = [100, 300, 500, 1000]
    num_threads = [2, 4, 6, 8, 12, 16, 20]
    print(f"cpu_count: {multiprocessing.cpu_count()}")

    for size in sizes:
        matrix1 = create_random_matrix(size)
        matrix2 = create_random_matrix(size)
        t0 = time.perf_counter()
        matrix_multiply((matrix1, matrix2, 0, len(matrix1)))
        # pprint(matrix_result[:size][:size], compact=True)
        print(f"Время последовательного перемножения матриц {size=:4}: \t\t\t\t{time.perf_counter() - t0:.3f}s")

        for threads in num_threads:
            start_time = time.perf_counter()
            matrix_multiply_parralel(matrix1, matrix2, threads)
            end_time = time.perf_counter()
            print(f"Время парралельного перемножения матриц {size=:4}, {threads=} : \t{end_time - start_time:.3f}")

        print("-" * 100)

if __name__ == '__main__':
    main()