from concurrent.futures import ThreadPoolExecutor
import time
import numpy as np
import argparse

# Последовательное умножение
def matrix_multiply_sequential(A, B):
    n = len(A)
    C = [[0] * n for _ in range(n)]
    
    # Транспонируем матрицу B для оптимального доступа по строкам
    B_T = [[B[j][i] for j in range(n)] for i in range(n)]
    
    for i in range(n):
        A_row = A[i]
        for j in range(n):
            B_col = B_T[j]
            sum_ij = 0
            for k in range(n):
                sum_ij += A_row[k] * B_col[k]
            C[i][j] = sum_ij
    
    return C

# Вычисляет подматрицу C
def worker(A, B, C, start_row, end_row):
    n = len(A)
    for i in range(start_row, end_row):
        for j in range(n):
            sum_ij = 0
            for k in range(n):
                sum_ij += A[i][k] * B[k][j]
            C[i][j] = sum_ij

# Параллельное умножение матриц
def matrix_multiply_parallel(A, B, num_threads):
    n = len(A)
    C = [[0] * n for _ in range(n)]
    
    # Разбиваем строки между потоками
    rows_per_thread = n // num_threads
    extra_rows = n % num_threads
    row_splits = [rows_per_thread + (1 if i < extra_rows else 0) for i in range(num_threads)]
    row_indices = [sum(row_splits[:i]) for i in range(num_threads + 1)]

    # Параллельно выполняем умножение подматриц
    with ThreadPoolExecutor(max_workers=num_threads) as executor:
        futures = [
            executor.submit(worker, A, B, C, row_indices[i], row_indices[i+1])
            for i in range(num_threads)
        ]
        for future in futures:
            future.result()
    
    return C

def benchmark(matrix_sizes, num_threads_list):
    for size in matrix_sizes:
        A = [[1] * size for _ in range(size)]  # Матрицы, заполненные единицами для упрощения теста
        B = [[1] * size for _ in range(size)]
        print(f"\nРазмер матриц: {size}x{size}")
        
        # Бенчмарк для последовательного алгоритма
        start_time = time.time()
        C_seq = matrix_multiply_sequential(A, B)
        sequential_time = time.time() - start_time
        print(f"Последовательное умножение заняло: {sequential_time:.4f} секунд")

        # Бенчмарк для параллельного алгоритма
        for num_threads in num_threads_list:
            start_time = time.time()
            C_par = matrix_multiply_parallel(A, B, num_threads)
            parallel_time = time.time() - start_time
            print(f"Параллельное умножение с {num_threads} потоками заняло: {parallel_time:.4f} секунд")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Запуск бенчмарков для умножения матриц.")
    parser.add_argument(
        "--threads",
        type=int,
        nargs="+",
        required=True,
        help="Список количества потоков для параллельного алгоритма (например, --threads 1 2 4 8)"
    )
    args = parser.parse_args()
    
    # Задание размеров матриц для тестов
    matrix_sizes = [100, 300, 500]
    num_threads_list = args.threads  # Получаем список потоков из аргументов командной строки

    # Запуск бенчмарка
    benchmark(matrix_sizes, num_threads_list)