import numpy as np
import time
from multiprocessing import Pool, cpu_count


def generate_matrix(size, value_range=(0, 200)):
    return np.random.randint(value_range[0], value_range[1], (size, size))


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


def worker_multiply(args):
    i_range, matrix_a, matrix_b, size = args
    result_part = np.zeros((len(i_range), size), dtype=int)
    for idx, i in enumerate(i_range):
        for j in range(size):
            for k in range(size):
                result_part[idx][j] += matrix_a[i][k] * matrix_b[k][j]
    return result_part


def multiply_matrices_parallel(matrix_a, matrix_b, num_threads):
    size = len(matrix_a)
    step = size // num_threads
    ranges = [range(i, min(i + step, size)) for i in range(0, size, step)]

    with Pool(processes=num_threads) as pool:
        results = pool.map(worker_multiply, [(i_range, matrix_a, matrix_b, size) for i_range in ranges])

    return np.vstack(results)

def benchmark(matrix_a, matrix_b, num_threads):
    print(f"\nМатрицы размера {len(matrix_a)}x{len(matrix_a)}:")

    start = time.time()
    result_seq = multiply_matrices_sequential(matrix_a, matrix_b)
    sequential_time = time.time() - start
    print(f"Последовательное умножение заняло: {sequential_time:.4f} секунд")

    start = time.time()
    result_par = multiply_matrices_parallel(matrix_a, matrix_b, num_threads)
    parallel_time = time.time() - start
    print(f"Параллельное умножение ({num_threads} потоков) заняло: {parallel_time:.4f} секунд")

    assert np.array_equal(result_seq, result_par), "Ошибка: результаты не совпадают!"
    print("Результаты совпадают.")

    return sequential_time, parallel_time


if __name__ == "__main__":
    sizes = [100, 300, 500]
    num_threads = min(cpu_count(), 4)

    for size in sizes:
        matrix_a = generate_matrix(size)
        matrix_b = generate_matrix(size)
        benchmark(matrix_a, matrix_b, num_threads)