use std::time::Instant;

use ndarray::{Array1, Zip};
use ndarray_rand::{RandomExt, rand_distr::Uniform};


pub fn sigmoid_val(x: f32) -> f32 {
    return 1. / (1. + (-x).exp());
}

pub fn sigmoid_ref(x: &f32) -> f32 {
    return 1. / (1. + (-x).exp());
}

pub fn sigmoid_mut_ref(x: &mut f32) {
    *x = 1. / (1. + (-*x).exp());
}


pub fn sigmoid_vec(x: &Vec<f32>) -> Vec<f32> {
    return x.iter().map(|&v| sigmoid_val(v)).collect();
}


pub fn sigmoid_ndarr_map(x: Array1<f32>) -> Array1<f32> {
    return x.map(sigmoid_ref);
}

pub fn sigmoid_ndarr_mapv(x: Array1<f32>) -> Array1<f32> {
    return x.mapv(sigmoid_val);
}

pub fn sigmoid_ndarr_map_inplace(mut x: Array1<f32>) -> Array1<f32> {
    x.map_inplace(sigmoid_mut_ref);
    return x;
}

pub fn sigmoid_ndarr_mapv_inplace(mut x: Array1<f32>, virt: Option<&Array1<bool>>) -> Array1<f32> {
    if virt.is_some() {
        Zip::from(&mut x).and(virt.unwrap()).for_each(
            |elem, is_virt| *elem = if *is_virt { *elem } else { sigmoid_ref(elem) }
        );
    } else {
        x.mapv_inplace(sigmoid_val);
    }
    return x;
}

pub fn sigmoid_ndarr_par_map_inplace(mut x: Array1<f32>) -> Array1<f32> {
    x.par_map_inplace(sigmoid_mut_ref);
    return x;
}

pub fn sigmoid_ndarr_par_mapv_inplace(mut x: Array1<f32>) -> Array1<f32> {
    x.par_mapv_inplace(sigmoid_val);
    return x;
}


pub fn time_sigmoids(n: usize, arr_len: usize) {
    let arr: Array1<f32> = Array1::random(arr_len, Uniform::new(0., 1.));
    let vec: Vec<f32> = arr.to_vec();
    
    let t0 = Instant::now();
    for _ in 0..n { let _result = sigmoid_vec(&vec); }
    let elapsed = t0.elapsed();
    println!("sigmoid_vec took {:.5} seconds", elapsed.as_secs_f64());

    let mut arr_copy = arr.to_owned();

    let t0 = Instant::now();
    for _ in 0..n { arr_copy = sigmoid_ndarr_map(arr_copy); }
    let elapsed = t0.elapsed();
    println!("sigmoid_ndarr_map took {:.5} seconds", elapsed.as_secs_f64());

    arr_copy = arr.to_owned();
    let t0 = Instant::now();
    for _ in 0..n { arr_copy = sigmoid_ndarr_mapv(arr_copy); }
    let elapsed = t0.elapsed();
    println!("sigmoid_ndarr_mapv took {:.5} seconds", elapsed.as_secs_f64());

    arr_copy = arr.to_owned();
    let t0 = Instant::now();
    for _ in 0..n { arr_copy = sigmoid_ndarr_map_inplace(arr_copy); }
    let elapsed = t0.elapsed();
    println!("sigmoid_ndarr_map_inplace took {:.5} seconds", elapsed.as_secs_f64());

    arr_copy = arr.to_owned();
    let virt = ndarray::Array::from_elem(arr.raw_dim(), false);
    let t0 = Instant::now();
    for _ in 0..n { arr_copy = sigmoid_ndarr_mapv_inplace(arr_copy, Some(&virt)); }
    let elapsed = t0.elapsed();
    println!("sigmoid_ndarr_mapv_inplace took {:.5} seconds", elapsed.as_secs_f64());

    arr_copy = arr.to_owned();
    let t0 = Instant::now();
    for _ in 0..n { arr_copy = sigmoid_ndarr_par_map_inplace(arr_copy); }
    let elapsed = t0.elapsed();
    println!("sigmoid_ndarr_par_map_inplace took {:.5} seconds", elapsed.as_secs_f64());

    arr_copy = arr.to_owned();
    let t0 = Instant::now();
    for _ in 0..n { arr_copy = sigmoid_ndarr_par_mapv_inplace(arr_copy); }
    let elapsed = t0.elapsed();
    println!("sigmoid_ndarr_par_mapv_inplace took {:.5} seconds", elapsed.as_secs_f64());

}