oxits 0.1.0

use crate::core::config::DistanceKind;

#[derive(Debug, Clone)]
pub struct RecurrencePlotConfig {
    pub dimension: usize,
    pub time_delay: usize,
    pub threshold: Option<f64>,
    pub percentage: Option<f64>,
    pub distance: DistanceKind,
}

impl RecurrencePlotConfig {
    pub fn new() -> Self {
        Self {
            dimension: 1,
            time_delay: 1,
            threshold: None,
            percentage: None,
            distance: DistanceKind::Squared,
        }
    }
}

impl Default for RecurrencePlotConfig {
    fn default() -> Self {
        Self::new()
    }
}

pub struct RecurrencePlot;

impl RecurrencePlot {
    /// Transform time series into recurrence plot images.
    /// Output shape: (n_samples, n_points, n_points)
    /// where n_points = n_timestamps - (dimension - 1) * time_delay
    pub fn transform(config: &RecurrencePlotConfig, x: &[Vec<f64>]) -> Vec<Vec<Vec<f64>>> {
        assert!(!x.is_empty(), "Input must have at least one sample");
        assert!(config.dimension >= 1, "dimension must be >= 1");
        assert!(config.time_delay >= 1, "time_delay must be >= 1");

        #[cfg(feature = "parallel")]
        {
            use rayon::prelude::*;
            return x
                .par_iter()
                .map(|sample| recurrence_plot_single(sample, config))
                .collect();
        }
        #[cfg(not(feature = "parallel"))]
        x.iter()
            .map(|sample| recurrence_plot_single(sample, config))
            .collect()
    }
}

/// Compute pairwise distance for a single element pair.
#[inline]
fn pairwise_dist(a: &[f64], b: &[f64], distance: DistanceKind) -> f64 {
    match distance {
        DistanceKind::Squared => a
            .iter()
            .zip(b.iter())
            .map(|(&x, &y)| {
                let d = x - y;
                d * d
            })
            .sum(),
        DistanceKind::Absolute => a.iter().zip(b.iter()).map(|(&x, &y)| (x - y).abs()).sum(),
    }
}

/// Fill a flat distance matrix for dim=1 (each "vector" is a single scalar).
fn fill_distances_dim1(sample: &[f64], n_points: usize, distance: DistanceKind, flat: &mut [f64]) {
    for i in 0..n_points {
        let si = sample[i];
        let row = &mut flat[i * n_points..(i + 1) * n_points];
        for (j, &sj) in sample[..n_points].iter().enumerate() {
            row[j] = match distance {
                DistanceKind::Squared => {
                    let d = si - sj;
                    d * d
                }
                DistanceKind::Absolute => (si - sj).abs(),
            };
        }
    }
}

/// Fill a flat distance matrix for dim>1 (delay-embedded vectors), exploiting symmetry.
fn fill_distances_general(
    vectors: &[Vec<f64>],
    n_points: usize,
    distance: DistanceKind,
    flat: &mut [f64],
) {
    for i in 0..n_points {
        for j in i + 1..n_points {
            let dist = pairwise_dist(&vectors[i], &vectors[j], distance);
            flat[i * n_points + j] = dist;
            flat[j * n_points + i] = dist;
        }
    }
}

/// Apply a fixed threshold: values <= threshold become 1.0, else 0.0.
fn apply_threshold(flat: &mut [f64], threshold: f64) {
    for d in flat.iter_mut() {
        *d = if *d <= threshold { 1.0 } else { 0.0 };
    }
}

/// Apply a percentage-based threshold using the k-th percentile of upper-triangle distances.
fn apply_percentage_threshold(flat: &mut [f64], n_points: usize, percentage: f64) {
    let n_pairs = n_points * (n_points - 1) / 2;
    let mut all_dists: Vec<f64> = Vec::with_capacity(n_pairs);
    for i in 0..n_points {
        for j in i + 1..n_points {
            all_dists.push(flat[i * n_points + j]);
        }
    }
    let idx = ((percentage / 100.0) * all_dists.len() as f64) as usize;
    let idx = idx.min(all_dists.len() - 1);
    let thresh = *all_dists
        .select_nth_unstable_by(idx, |a, b| a.partial_cmp(b).unwrap())
        .1;
    apply_threshold(flat, thresh);
}

fn recurrence_plot_single(sample: &[f64], config: &RecurrencePlotConfig) -> Vec<Vec<f64>> {
    let n = sample.len();
    let n_points = n - (config.dimension - 1) * config.time_delay;
    assert!(
        n_points >= 1,
        "Time series too short for given dimension and time_delay"
    );

    // Fast path: no threshold/percentage, use collect-per-row to avoid flat buffer
    if config.threshold.is_none() && config.percentage.is_none() {
        return if config.dimension == 1 {
            rp_dim1_direct(sample, n_points, config.distance)
        } else {
            let vectors = delay_embed(sample, n_points, config.dimension, config.time_delay);
            rp_general_direct(&vectors, n_points, config.distance)
        };
    }

    // Threshold/percentage path: need distances in a mutable buffer
    let mut flat = vec![0.0_f64; n_points * n_points];

    if config.dimension == 1 {
        fill_distances_dim1(sample, n_points, config.distance, &mut flat);
    } else {
        let vectors = delay_embed(sample, n_points, config.dimension, config.time_delay);
        fill_distances_general(&vectors, n_points, config.distance, &mut flat);
    }

    if let Some(threshold) = config.threshold {
        apply_threshold(&mut flat, threshold);
    } else if let Some(percentage) = config.percentage {
        apply_percentage_threshold(&mut flat, n_points, percentage);
    }

    flat.chunks_exact(n_points).map(|c| c.to_vec()).collect()
}

fn delay_embed(
    sample: &[f64],
    n_points: usize,
    dimension: usize,
    time_delay: usize,
) -> Vec<Vec<f64>> {
    (0..n_points)
        .map(|i| (0..dimension).map(|d| sample[i + d * time_delay]).collect())
        .collect()
}

/// Dim=1 fast path: row-by-row collect, no flat buffer needed.
#[inline]
fn rp_dim1_direct(sample: &[f64], n_points: usize, distance: DistanceKind) -> Vec<Vec<f64>> {
    (0..n_points)
        .map(|i| {
            let si = sample[i];
            match distance {
                DistanceKind::Squared => sample[..n_points]
                    .iter()
                    .map(|&sj| {
                        let d = si - sj;
                        d * d
                    })
                    .collect(),
                DistanceKind::Absolute => sample[..n_points]
                    .iter()
                    .map(|&sj| (si - sj).abs())
                    .collect(),
            }
        })
        .collect()
}

/// General path: row-by-row collect for delay-embedded vectors.
#[inline]
fn rp_general_direct(
    vectors: &[Vec<f64>],
    n_points: usize,
    distance: DistanceKind,
) -> Vec<Vec<f64>> {
    (0..n_points)
        .map(|i| {
            (0..n_points)
                .map(|j| match distance {
                    DistanceKind::Squared => vectors[i]
                        .iter()
                        .zip(vectors[j].iter())
                        .map(|(&a, &b)| {
                            let d = a - b;
                            d * d
                        })
                        .sum::<f64>(),
                    DistanceKind::Absolute => vectors[i]
                        .iter()
                        .zip(vectors[j].iter())
                        .map(|(&a, &b)| (a - b).abs())
                        .sum::<f64>(),
                })
                .collect()
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_rp_shape() {
        let config = RecurrencePlotConfig::new();
        let x = vec![vec![1.0, 2.0, 3.0, 4.0, 5.0]];
        let result = RecurrencePlot::transform(&config, &x);
        assert_eq!(result.len(), 1);
        assert_eq!(result[0].len(), 5);
        assert_eq!(result[0][0].len(), 5);
    }

    #[test]
    fn test_rp_symmetric() {
        let config = RecurrencePlotConfig::new();
        let x = vec![vec![0.0, 1.0, 0.0, -1.0, 0.0]];
        let result = RecurrencePlot::transform(&config, &x);
        for i in 0..5 {
            for j in 0..5 {
                assert!(
                    (result[0][i][j] - result[0][j][i]).abs() < 1e-10,
                    "RP should be symmetric"
                );
            }
        }
    }

    #[test]
    fn test_rp_diagonal_zero() {
        let config = RecurrencePlotConfig::new();
        let x = vec![vec![1.0, 2.0, 3.0]];
        let result = RecurrencePlot::transform(&config, &x);
        for i in 0..3 {
            assert!(result[0][i][i].abs() < 1e-10, "Diagonal should be 0");
        }
    }

    #[test]
    fn test_rp_with_threshold() {
        let config = RecurrencePlotConfig {
            threshold: Some(1.0),
            ..RecurrencePlotConfig::new()
        };
        let x = vec![vec![0.0, 0.5, 0.0, 2.0]];
        let result = RecurrencePlot::transform(&config, &x);
        // All values should be 0 or 1
        for row in &result[0] {
            for &v in row {
                assert!(v == 0.0 || v == 1.0, "Binary RP should be 0 or 1, got {v}");
            }
        }
    }

    #[test]
    fn test_rp_with_embedding() {
        let config = RecurrencePlotConfig {
            dimension: 2,
            time_delay: 1,
            ..RecurrencePlotConfig::new()
        };
        let x = vec![vec![1.0, 2.0, 3.0, 4.0, 5.0]];
        let result = RecurrencePlot::transform(&config, &x);
        // n_points = 5 - (2-1)*1 = 4
        assert_eq!(result[0].len(), 4);
        assert_eq!(result[0][0].len(), 4);
    }

    #[test]
    fn test_rp_default_config() {
        let config = RecurrencePlotConfig::default();
        assert_eq!(config.dimension, 1);
        assert_eq!(config.time_delay, 1);
    }

    #[test]
    fn test_rp_absolute_distance() {
        let config = RecurrencePlotConfig {
            distance: DistanceKind::Absolute,
            ..RecurrencePlotConfig::new()
        };
        let x = vec![vec![0.0, 3.0, 1.0]];
        let result = RecurrencePlot::transform(&config, &x);
        assert!((result[0][0][1] - 3.0).abs() < 1e-10); // |0 - 3| = 3
        assert!((result[0][0][2] - 1.0).abs() < 1e-10); // |0 - 1| = 1
    }

    #[test]
    fn test_rp_with_percentage() {
        let config = RecurrencePlotConfig {
            percentage: Some(50.0),
            ..RecurrencePlotConfig::new()
        };
        let x = vec![vec![0.0, 1.0, 2.0, 3.0]];
        let result = RecurrencePlot::transform(&config, &x);
        // All values should be 0 or 1
        for row in &result[0] {
            for &v in row {
                assert!(v == 0.0 || v == 1.0, "Binary RP, got {v}");
            }
        }
        // Diagonal should be 1 (distance 0 is always <= threshold)
        for i in 0..4 {
            assert!((result[0][i][i] - 1.0).abs() < 1e-10);
        }
    }

    #[test]
    fn test_rp_embedding_with_threshold() {
        let config = RecurrencePlotConfig {
            dimension: 2,
            time_delay: 1,
            threshold: Some(10.0),
            ..RecurrencePlotConfig::new()
        };
        let x = vec![vec![1.0, 2.0, 3.0, 4.0, 5.0]];
        let result = RecurrencePlot::transform(&config, &x);
        // n_points = 4, all values binary
        assert_eq!(result[0].len(), 4);
        for row in &result[0] {
            for &v in row {
                assert!(v == 0.0 || v == 1.0);
            }
        }
    }

    #[test]
    fn test_rp_embedding_with_percentage() {
        let config = RecurrencePlotConfig {
            dimension: 2,
            time_delay: 1,
            percentage: Some(50.0),
            ..RecurrencePlotConfig::new()
        };
        let x = vec![vec![0.0, 1.0, 0.0, -1.0, 0.0]];
        let result = RecurrencePlot::transform(&config, &x);
        assert_eq!(result[0].len(), 4);
        for row in &result[0] {
            for &v in row {
                assert!(v == 0.0 || v == 1.0);
            }
        }
    }

    #[test]
    fn test_rp_absolute_with_threshold() {
        let config = RecurrencePlotConfig {
            distance: DistanceKind::Absolute,
            threshold: Some(1.5),
            ..RecurrencePlotConfig::new()
        };
        let x = vec![vec![0.0, 1.0, 3.0]];
        let result = RecurrencePlot::transform(&config, &x);
        // |0-1|=1 <= 1.5 → 1; |0-3|=3 > 1.5 → 0
        assert!((result[0][0][1] - 1.0).abs() < 1e-10);
        assert!((result[0][0][2] - 0.0).abs() < 1e-10);
    }

    #[test]
    fn test_rp_embedding_absolute_direct() {
        // dim>1 direct path (no threshold/percentage) with absolute distance
        let config = RecurrencePlotConfig {
            dimension: 2,
            time_delay: 1,
            distance: DistanceKind::Absolute,
            ..RecurrencePlotConfig::new()
        };
        let x = vec![vec![1.0, 2.0, 3.0, 4.0]];
        let result = RecurrencePlot::transform(&config, &x);
        assert_eq!(result[0].len(), 3); // n_points = 4 - 1 = 3
                                        // Diagonal should be 0
        for i in 0..3 {
            assert!(result[0][i][i].abs() < 1e-10);
        }
    }

    #[test]
    fn test_rp_embedding_absolute_with_threshold() {
        // dim>1 threshold path with absolute distance
        let config = RecurrencePlotConfig {
            dimension: 2,
            time_delay: 1,
            distance: DistanceKind::Absolute,
            threshold: Some(5.0),
            ..RecurrencePlotConfig::new()
        };
        let x = vec![vec![0.0, 1.0, 2.0, 3.0, 4.0]];
        let result = RecurrencePlot::transform(&config, &x);
        assert_eq!(result[0].len(), 4);
        for row in &result[0] {
            for &v in row {
                assert!(v == 0.0 || v == 1.0);
            }
        }
    }
}