#![feature(core_intrinsics)]

use std::{env, intrinsics::fadd_fast, time::Instant};

fn matrix_new(n: usize) -> Vec<f64> {
    let n = n as isize;
    let tmp = 1.0 / n as f64 / n as f64;
    (0..n)
        .flat_map(|i| (0..n).map(move |j| tmp * (i - j) as f64 * (i + j) as f64))
        .collect()
}

fn matrix_transpose(n: usize, src: &[f64], dst: &mut [f64]) {
    for (i, dst_row) in dst.chunks_exact_mut(n).enumerate() {
        dst_row
            .iter_mut()
            .enumerate()
            .for_each(|(j, t)| *t = src[j * n + i]);
    }
}

fn matrix_mul(n: usize, a: &[f64], b: &[f64], c: &mut [f64]) {
    let mut t = vec![0.0; n * n];
    matrix_transpose(n, b, &mut t);

    for (c_row, a_row) in c.chunks_exact_mut(n).zip(a.chunks_exact(n)) {
        for (t_col, c) in t.chunks_exact(n).zip(c_row) {
            *c = inner_product(*c, a_row, t_col);
        }
    }
}

fn inner_product(init: f64, a: &[f64], b: &[f64]) -> f64 {
    a.iter()
        .zip(b)
        .fold(init, |acc, (&x, &y)| unsafe { fadd_fast(acc, x * y) })
}

fn main() -> Result<(), &'static str> {
    let n: usize = env::args()
        .nth(1)
        .ok_or("usage: <app> N")
        .and_then(|n| n.parse().map_err(|_| "N is not a number"))?;

    let a = matrix_new(n);
    let b = matrix_new(n);
    let mut c = vec![0.0; n * n];

    let start_time = Instant::now();
    matrix_mul(n, &a, &b, &mut c);
    let time = start_time.elapsed();

    println!("time = {:?}", time);
    println!("n = {}", n);
    println!("result = {}", c[n * (n / 2) + n / 2]);

    Ok(())
}