use std::thread;
use std::sync::mpsc;
use std::time;

// Квадратная матрица хранится в куче как последовательность всех её элементов, при этом на этапе компиляции всегда известен её размер
struct Matrix<const N: usize> {
    inner: Vec<[i32; N]>
}

impl<const N: usize> Matrix<N> {
    fn new() -> Matrix<N> {
        Matrix { inner: vec![[0; N]; N] }
    }
    
    // Псевдослучайная генерация данных для матрицы
    fn random(seed: i32) -> Matrix<N> {
        let mut matrix = Matrix::new();

        for i in 0..N {
            for j in 0..N {
                matrix.inner[i][j] = i as i32 + j as i32 - seed;
            }
        }

        matrix
    }
}

// Перемножение матриц с помощью произвольного количества потоков, размер матриц и количество потоков известно на этапе компиляции
fn multiply<const N: usize, const M: usize>(f: &Matrix<N>, s: &Matrix<N>) -> Matrix<N> {
    let mut result = Matrix::<N>::new();
    let part_size = N / M;
    let last_part_size = N - (N / M) * (M - 1);
    let (sender, receiver) = mpsc::channel();
    
    // Конструкция scope гарантирует, что созданные внутри потоки завершат работу там же и не буду существовать дольше, чем сами данные
    thread::scope(move |scope| {
        for n_part in 1..M {
            let sender = sender.clone();
            scope.spawn(move || {
                for row in (if n_part != 0 {n_part - 1} else {n_part}) * part_size..n_part * part_size {
                    let mut data = Box::new([0; N]);
                    for col in 0..N {
                        data[col] = f.inner[row].iter().zip(s.inner.iter().map(|row| row[col])).map(|(v1, v2)| v1 * v2).sum()
                    }
                    sender.send((row, data)).unwrap();
                }
            });
        }
        // В случае одного потока или если осталась часть матрицы, которую не удалось передать в поток
        if last_part_size != 0 {
            for row in N-last_part_size..N {
                let mut data = Box::new([0; N]);
                for col in 0..N {
                    data[col] = f.inner[row].iter().zip(s.inner.iter().map(|row| row[col])).map(|(v1, v2)| v1 * v2).sum()
                }
                sender.send((row, data)).unwrap();
            }
        }
    });
    
    // Потоки отдают результаты в виде сообщений в mpsc канале, при этом передаётся указатель на данные, а не сами данные
    while let Ok((row, data)) = receiver.recv() {
        for col in 0..N {
            result.inner[row][col] = data[col];
        }
    }

    result
}

// Замер времени выполнения, при этом размеры матриц и количество потоков известны на этапе компиляции
fn benchmark<const N: usize, const M: usize>() {
    let first = Matrix::<N>::random(1);
    let second = Matrix::<N>::random(2);

    let start_time = time::Instant::now();
    multiply::<N, M>(&first, &second);
    println!("{0:>4}x{0:<4} ({1:>2} потоков): {2} сек", N, M, (time::Instant::now() - start_time).as_secs_f64())
}

fn main() {
    benchmark::<100, 1>();
    benchmark::<300, 1>();
    benchmark::<500, 1>();
    benchmark::<1000, 1>();
    benchmark::<1200, 1>();

    benchmark::<100, 4>();
    benchmark::<300, 4>();
    benchmark::<500, 4>();
    benchmark::<1000, 4>();
    benchmark::<1200, 4>();

    benchmark::<100, 8>();
    benchmark::<300, 8>();
    benchmark::<500, 8>();
    benchmark::<1000, 8>();
    benchmark::<1200, 8>();
}