imgal 0.3.1

use std::collections::HashMap;

use ndarray::{
    Array2, Array3, ArrayBase, ArrayView1, ArrayView2, AsArray, Axis, Ix1, Ix3, ViewRepr, Zip, s,
    stack,
};
use rayon::prelude::*;

use crate::integration::midpoint;
use crate::parameter::omega;
use crate::prelude::*;

/// Compute the real and imaginary (G, S) coordinates of a 3D decay image.
///
/// # Description
///
/// Computes the real (G) and imaginary (S) components using normalized sine
/// and cosine Fourier transforms:
///
/// ```text
/// G = ∫(I(t) * cos(nωt) * dt) / ∫(I(t) * dt)
/// S = ∫(I(t) * sin(nωt) * dt) / ∫(I(t) * dt)
/// ```
///
/// # Arguments
///
/// * `data`: The input 3D decay image.
/// * `period`: The period (*i.e.* time interval).
/// * `harmonic`: The harmonic value. If `None`, then `harmonic = 1.0`.
/// * `axis`: The decay or lifetime axis. If `None`, then `axis = 2`.
/// * `threads`: The requested number of threads to use for parallel execution.
///   If `None` or `Some(1)` sequential execution is used. If `Some(0)`, then
///   the maximum available parallelism is used. Thread counts are clamped to
///   the systems maximum.
///
/// # Returns
///
/// * `Ok(Array3<f64>)`: The real and imaginary coordinates as a 3D
///   (row, col, ch) image, where G and S are indexed at `0` and `1`
///   respectively on the *channel* axis.
/// * `Err(ImgalError)`: If `axis >= 3`.
pub fn gs_image<'a, T, A>(
    data: A,
    period: f64,
    mask: Option<ArrayView2<bool>>,
    harmonic: Option<f64>,
    axis: Option<usize>,
    threads: Option<usize>,
) -> Result<Array3<f64>, ImgalError>
where
    A: AsArray<'a, T, Ix3>,
    T: 'a + AsNumeric,
{
    let axis = axis.unwrap_or(2);
    if axis >= 3 {
        return Err(ImgalError::InvalidAxis {
            axis_idx: axis,
            dim_len: 3,
        });
    }
    let data: ArrayBase<ViewRepr<&'a T>, Ix3> = data.into();
    let h = harmonic.unwrap_or(1.0);
    let w = omega(period);
    let n: usize = data.len_of(Axis(axis));
    let dt: f64 = period / n as f64;
    let h_w_dt: f64 = h * w * dt;
    let mut w_cos_buf: Vec<f64> = Vec::with_capacity(n);
    let mut w_sin_buf: Vec<f64> = Vec::with_capacity(n);
    let mut shape = data.shape().to_vec();
    shape.remove(axis);
    let mut g_arr = Array2::<f64>::zeros((shape[0], shape[1]));
    let mut s_arr = Array2::<f64>::zeros((shape[0], shape[1]));
    for i in 0..n {
        w_cos_buf.push(f64::cos(h_w_dt * (i as f64)));
        w_sin_buf.push(f64::sin(h_w_dt * (i as f64)));
    }
    let lanes = data.lanes(Axis(axis));
    let gs_calc = |ln: ArrayView1<T>, g: &mut f64, s: &mut f64| {
        let mut iv = 0.0;
        let mut gv = 0.0;
        let mut sv = 0.0;
        ln.iter()
            .zip(w_cos_buf.iter())
            .zip(w_sin_buf.iter())
            .for_each(|((v, cosv), sinv)| {
                let vf: f64 = (*v).to_f64();
                iv += vf;
                gv += vf * cosv;
                sv += vf * sinv;
            });
        iv *= dt;
        gv *= dt;
        sv *= dt;
        *g = gv / iv;
        *s = sv / iv;
    };
    let gs_msk_calc = |ln: ArrayView1<T>, m: &bool, g: &mut f64, s: &mut f64| {
        if *m {
            let mut iv = 0.0;
            let mut gv = 0.0;
            let mut sv = 0.0;
            ln.iter()
                .zip(w_cos_buf.iter())
                .zip(w_sin_buf.iter())
                .for_each(|((v, cosv), sinv)| {
                    let vf: f64 = (*v).to_f64();
                    iv += vf;
                    gv += vf * cosv;
                    sv += vf * sinv;
                });
            iv *= dt;
            gv *= dt;
            sv *= dt;
            *g = gv / iv;
            *s = sv / iv;
        } else {
            *g = 0.0;
            *s = 0.0;
        }
    };
    if let Some(msk) = mask {
        par!(threads,
            seq_exp: Zip::from(lanes).and(msk).and(&mut g_arr).and(&mut s_arr)
                .for_each(&gs_msk_calc),
            par_exp: Zip::from(lanes).and(msk).and(&mut g_arr).and(&mut s_arr)
                .par_for_each(&gs_msk_calc));
    } else {
        par!(threads,
            seq_exp: Zip::from(lanes).and(&mut g_arr).and(&mut s_arr)
                .for_each(&gs_calc),
            par_exp: Zip::from(lanes).and(&mut g_arr).and(&mut s_arr)
                .par_for_each(&gs_calc));
    }
    Ok(stack(Axis(2), &[g_arr.view(), s_arr.view()]).unwrap())
}

/// Compute the real and imaginary (G, S) coordinates of a HashMap of ROI point
/// clouds
///
/// # Description
///
/// Computes the real and imaginary (G, S) coordinates each ROI point cloud
/// stored in a HashMap.
///
/// ```text
/// G = ∫(I(t) * cos(nωt) * dt) / ∫(I(t) * dt)
/// S = ∫(I(t) * sin(nωt) * dt) / ∫(I(t) * dt)
/// ```
///
/// # Arguments
///
/// * `data`: The input decay 3D image.
/// * `period`: The period (*i.e.* time interval).
/// * `rois`: A HashMap of point clouds representing Regions of Interests
///   (ROIs). 2D ROIs are expected.
/// * `harmonic`: The harmonic value. If `None`, then `harmonic = 1.0`.
/// * `axis`: The decay or lifetime axis. If `None`, then `axis = 2`.
/// * `threads`: The requested number of threads to use for parallel execution.
///   If `None` or `Some(1)` sequential execution is used. If `Some(0)`, then
///   the maximum available parallelism is used. Thread counts are clamped to
///   the systems maximum.
///
/// # Returns
///
/// * `Ok(HashMap<u64, Array2<f64>>)`: A HashMap where the keys are the ROI
///   labels and values are the G and S values computed at each point in the
///   input ROI point cloud. Each computed ROI point cloud has shape `(p, 2)`,
///   where `p` is the number of points.
/// * `Err(ImgalError)`: If `axis >= 3`.
pub fn gs_roi<'a, T, A>(
    data: A,
    period: f64,
    rois: &HashMap<u64, Array2<usize>>,
    harmonic: Option<f64>,
    axis: Option<usize>,
    threads: Option<usize>,
) -> Result<HashMap<u64, Array2<f64>>, ImgalError>
where
    A: AsArray<'a, T, Ix3>,
    T: 'a + AsNumeric,
{
    let data: ArrayBase<ViewRepr<&'a T>, Ix3> = data.into();
    let axis = axis.unwrap_or(2);
    if axis >= 3 {
        return Err(ImgalError::InvalidAxis {
            axis_idx: axis,
            dim_len: 3,
        });
    }
    let vec_to_arr = |k: u64, v: Vec<Vec<f64>>| {
        let arr = Array2::from_shape_vec((v.len(), v[0].len()), v.into_iter().flatten().collect())
            .expect("Failed to reshape ROI point cloud into an Array2<f64>.");
        (k, arr)
    };
    let roi_gs_calc_seq = || {
        let mut cloud_map: HashMap<u64, Vec<Vec<f64>>> = HashMap::new();
        rois.iter().for_each(|(&k, v)| {
            let roi_coords = v.lanes(Axis(1));
            roi_coords.into_iter().for_each(|p| {
                let row = p[0];
                let col = p[1];
                let ln = match axis {
                    0 => data.slice(s![.., row, col]),
                    1 => data.slice(s![row, .., col]),
                    _ => data.slice(s![row, col, ..]),
                };
                let g = real_coord(ln, period, harmonic, None);
                let s = imaginary_coord(ln, period, harmonic, None);
                cloud_map.entry(k).or_default().push(vec![g, s]);
            });
        });
        cloud_map
    };
    let roi_gs_calc_par = || {
        rois.into_par_iter()
            .fold(
                HashMap::new,
                |mut map: HashMap<u64, Vec<Vec<f64>>>, (&k, v)| {
                    let roi_coords = v.lanes(Axis(1));
                    roi_coords.into_iter().for_each(|p| {
                        let row = p[0];
                        let col = p[1];
                        let ln = match axis {
                            0 => data.slice(s![.., row, col]),
                            1 => data.slice(s![row, .., col]),
                            _ => data.slice(s![row, col, ..]),
                        };
                        let g = real_coord(ln, period, harmonic, None);
                        let s = imaginary_coord(ln, period, harmonic, None);
                        map.entry(k).or_default().push(vec![g, s]);
                    });
                    map
                },
            )
            .reduce(HashMap::new, |mut map_a, map_b| {
                map_b.into_iter().for_each(|(k, mut v)| {
                    map_a.entry(k).or_insert_with(Vec::new).append(&mut v);
                });
                map_a
            })
    };
    let cloud_map = par!(threads,
        seq_exp: roi_gs_calc_seq(),
        par_exp: roi_gs_calc_par());
    Ok(cloud_map
        .into_iter()
        .map(|(k, v)| vec_to_arr(k, v))
        .collect())
}

/// Compute the imaginary (S) component of a 1D decay array.
///
/// # Description
///
/// Computes the imaginary (S) component is calculated using the normalized sine
/// Fourier transform:
///
/// ```text
/// S = ∫(I(t) * sin(nωt) * dt) / ∫(I(t) * dt)
/// ```
///
/// Where `n` and `ω` are harmonic and omega values respectively.
///
/// # Arguments
///
/// * `data`: The input 1D decay array.
/// * `period`: The period (*i.e.* time interval).
/// * `harmonic`: The harmonic value. If `None`, then `harmonic = 1.0`.
/// * `threads`: The requested number of threads to use for parallel execution.
///   If `None` or `Some(1)` sequential execution is used. If `Some(0)`, then
///   the maximum available parallelism is used. Thread counts are clamped to
///   the systems maximum.
///
/// # Returns
///
/// * `f64`: The imaginary component, S.
#[inline]
pub fn imaginary_coord<'a, T, A>(
    data: A,
    period: f64,
    harmonic: Option<f64>,
    threads: Option<usize>,
) -> f64
where
    A: AsArray<'a, T, Ix1>,
    T: 'a + AsNumeric,
{
    let data: ArrayBase<ViewRepr<&'a T>, Ix1> = data.into();
    let h = harmonic.unwrap_or(1.0);
    let w = omega(period);
    let n = data.len();
    let dt = period / (n as f64);
    let h_w_dt = h * w * dt;
    let buf: Vec<f64> = (0..n)
        .map(|i| data[i].to_f64() * (h_w_dt * i as f64).sin())
        .collect();
    midpoint(&buf, Some(dt), threads) / midpoint(data, Some(dt), threads)
}

/// Compute the real (G) component of a 1D decay array.
///
/// # Description
///
/// Computes the real (G) component is calculated using the normalized cosine
/// Fourier transform:
///
/// ```text
/// G = ∫(I(t) * cos(nωt) * dt) / ∫(I(t) * dt)
/// ```
///
/// Where `n` and `ω` are harmonic and omega values respectively.
///
/// # Arguments
///
/// * `data`: The 1D decay array.
/// * `period`: The period, (*i.e.* time interval).
/// * `harmonic`: The harmonic value. If `None`, then `harmonic = 1.0`.
/// * `threads`: The requested number of threads to use for parallel execution.
///   If `None` or `Some(1)` sequential execution is used. If `Some(0)`, then
///   the maximum available parallelism is used. Thread counts are clamped to
///   the systems maximum.
///
/// # Returns
///
/// * `f64`: The real component, G.
#[inline]
pub fn real_coord<'a, T, A>(
    data: A,
    period: f64,
    harmonic: Option<f64>,
    threads: Option<usize>,
) -> f64
where
    A: AsArray<'a, T, Ix1>,
    T: 'a + AsNumeric,
{
    let data: ArrayBase<ViewRepr<&'a T>, Ix1> = data.into();
    let h = harmonic.unwrap_or(1.0);
    let w = omega(period);
    let n = data.len();
    let dt = period / (n as f64);
    let h_w_dt = h * w * dt;
    let buf: Vec<f64> = (0..n)
        .map(|i| data[i].to_f64() * (h_w_dt * i as f64).cos())
        .collect();
    midpoint(&buf, Some(dt), threads) / midpoint(data, Some(dt), threads)
}