#![allow(dead_code)]

extern crate blas_src;

use std::env;
use std::fs::File;
use std::path::Path;
use std::time::Instant;

use csv;
use glob::glob;
use ndarray::{Array, Array1, Array2};

mod sigmoids;  // different custom implementations of the sigmoid function
use sigmoids::{sigmoid_ndarr_mapv_inplace as sigmoid};


const DIR_DATA: &str = "data";


fn csv_row_to_float_vec(row: &csv::StringRecord) -> Vec<f32> {
    return row.iter().map(
        |item| {item.parse::<f32>().expect("Not a float")}
    ).collect();
}


fn csv_to_float_vec(file_path: &Path) -> (Vec<f32>, usize) {
    let mut output: Vec<f32> = Vec::new();
    let mut num_columns: usize = 0;
    let file = File::open(file_path).expect("Can't open file");
    let mut reader = csv::ReaderBuilder::new().has_headers(false).from_reader(file);
    for result in reader.records() {
        let row: Vec<f32> = csv_row_to_float_vec(&result.expect("No csv record"));
        if num_columns == 0 {
            num_columns = row.len();
        }
        output.extend(row);
    }
    return (output, num_columns);
}

fn csv_to_array1(file_path: &Path) -> Array1<f32> {
    let (vec, _) = csv_to_float_vec(&file_path);
    return Array::from_shape_vec((vec.len(), ), vec).unwrap();
}

fn csv_to_array2(file_path: &Path) -> Array2<f32> {
    let (vec, num_colums) = csv_to_float_vec(&file_path);
    return Array2::from_shape_vec((vec.len() / num_colums, num_colums), vec).unwrap();
}

fn load_test_data() -> (Array1<f32>, Vec<(Array2<f32>, Array1<f32>)>) {
    let test_data_dir = Path::new(".").join(DIR_DATA);
    let inputs: Array1<f32> = csv_to_array1(&test_data_dir.join("inputs.csv"));
    let weights_glob = test_data_dir.join("weights*.csv");
    let biases_glob = test_data_dir.join("biases*.csv");
    let weights_iter = glob(weights_glob.to_str().unwrap()).unwrap().map(
        |path_result| {csv_to_array2(&path_result.unwrap())}
    );
    let biases_iter = glob(biases_glob.to_str().unwrap()).unwrap().map(
        |path_result| {csv_to_array1(&path_result.unwrap())}
    );
    return (inputs, weights_iter.zip(biases_iter).collect());
}


fn layer_func<F>(input_vector: &Array1<f32>, weight_matrix: &Array2<f32>,
                 bias_vector: &Array1<f32>, activation: &F,
                 virtual_vector: Option<&Array1<bool>>) -> Array1<f32>
where F: Fn(Array1<f32>, Option<&Array1<bool>>) -> Array1<f32> {
    return activation(weight_matrix.dot(input_vector) + bias_vector, virtual_vector)
}


fn feed_forward<F>(x: &Array1<f32>, layers: &[(Array2<f32>, Array1<f32>)], activation: &F) -> Array1<f32>
where F: Fn(Array1<f32>, Option<&Array1<bool>>) -> Array1<f32> {
    let mut y = x.to_owned();
    for (w, b) in layers {
        y = layer_func(&y, w, b, activation, Some(&Array::from_elem(b.raw_dim(), false)))
    }
    return y;
}

fn time_feed_forward(n: usize) {
    let (inp, layers) = load_test_data();
    let t0 = Instant::now();
    for _ in 0..n { feed_forward(&inp, &layers, &sigmoid); }
    let elapsed = t0.elapsed();
    println!("{:.5}", elapsed.as_secs_f64());
}


fn main() {
    let args: Vec<String> = env::args().collect();
    let n: usize = args[1].parse::<usize>().expect("Not an integer");
    time_feed_forward(n);
}