oxits 0.1.0

use crate::core::config::SsaGrouping;
use realfft::RealFftPlanner;
use std::sync::Arc;

#[derive(Debug, Clone)]
pub struct SsaConfig {
    pub window_size: usize,
    pub groups: SsaGrouping,
    pub lower_frequency_bound: f64,
    pub lower_frequency_contribution: f64,
}

impl SsaConfig {
    pub fn new(window_size: usize) -> Self {
        Self {
            window_size,
            groups: SsaGrouping::None,
            lower_frequency_bound: 0.075,
            lower_frequency_contribution: 0.85,
        }
    }
}

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

pub struct Ssa;

impl Ssa {
    /// Decompose time series using Singular Spectrum Analysis.
    ///
    /// Output shape: (n_samples, n_components, n_timestamps)
    /// where n_components depends on `groups`:
    ///   - None: window_size components
    ///   - Auto: 3 components (trend, seasonal, residual)
    ///   - Explicit(groups): groups.len() components
    pub fn transform(config: &SsaConfig, x: &[Vec<f64>]) -> Vec<Vec<Vec<f64>>> {
        assert!(!x.is_empty(), "Input must have at least one sample");
        let n_timestamps = x[0].len();
        assert!(
            config.window_size >= 2 && config.window_size <= n_timestamps,
            "window_size must be in [2, n_timestamps]"
        );

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

        #[cfg(not(feature = "parallel"))]
        x.iter().map(|sample| ssa_single(sample, config)).collect()
    }
}

/// Build the trajectory matrix (w x k) from a time series by lagged embedding.
fn build_trajectory(sample: &[f64], w: usize, k: usize) -> Vec<Vec<f64>> {
    let mut trajectory = vec![vec![0.0; k]; w];
    for i in 0..w {
        trajectory[i][..k].copy_from_slice(&sample[i..(k + i)]);
    }
    trajectory
}

/// Compute X * X^T (symmetric w x w matrix) from trajectory matrix.
fn compute_xxt(trajectory: &[Vec<f64>], w: usize) -> Vec<Vec<f64>> {
    let mut xxt = vec![vec![0.0; w]; w];
    for i in 0..w {
        for j in i..w {
            let dot: f64 = trajectory[i]
                .iter()
                .zip(trajectory[j].iter())
                .map(|(&a, &b)| a * b)
                .sum();
            xxt[i][j] = dot;
            xxt[j][i] = dot;
        }
    }
    xxt
}

/// Compute elementary matrices: X_elem[i] = u_i * (u_i^T * X) for each eigenvector.
fn compute_elementary_matrices(
    trajectory: &[Vec<f64>],
    eigenvectors: &[Vec<f64>],
    w: usize,
    k: usize,
) -> Vec<Vec<Vec<f64>>> {
    (0..w)
        .map(|i| {
            let u: Vec<f64> = (0..w).map(|r| eigenvectors[r][i]).collect();
            let ut_x: Vec<f64> = (0..k)
                .map(|j| (0..w).map(|r| u[r] * trajectory[r][j]).sum::<f64>())
                .collect();
            (0..w)
                .map(|r| ut_x.iter().map(|&v| u[r] * v).collect())
                .collect()
        })
        .collect()
}

/// Sum elementary matrices within each group to produce grouped (w x k) matrices.
fn sum_groups(
    x_elem: &[Vec<Vec<f64>>],
    groups: &[Vec<usize>],
    w: usize,
    k: usize,
) -> Vec<Vec<Vec<f64>>> {
    groups
        .iter()
        .map(|group| {
            let mut mat = vec![vec![0.0; k]; w];
            for &idx in group {
                if idx < x_elem.len() {
                    for r in 0..w {
                        for c in 0..k {
                            mat[r][c] += x_elem[idx][r][c];
                        }
                    }
                }
            }
            mat
        })
        .collect()
}

fn ssa_single(sample: &[f64], config: &SsaConfig) -> Vec<Vec<f64>> {
    let n = sample.len();
    let w = config.window_size;
    let k = n - w + 1;

    let trajectory = build_trajectory(sample, w, k);
    let xxt = compute_xxt(&trajectory, w);
    let (eigenvalues, eigenvectors) = symmetric_eigen(&xxt);
    let x_elem = compute_elementary_matrices(&trajectory, &eigenvectors, w, k);

    let groups = match &config.groups {
        SsaGrouping::None => (0..w).map(|i| vec![i]).collect::<Vec<Vec<usize>>>(),
        SsaGrouping::Auto => auto_grouping(
            &eigenvalues,
            &eigenvectors,
            w,
            config.lower_frequency_bound,
            config.lower_frequency_contribution,
        ),
        SsaGrouping::Explicit(groups) => groups.clone(),
    };

    let grouped = sum_groups(&x_elem, &groups, w, k);
    grouped
        .iter()
        .map(|mat| diagonal_averaging(mat, n))
        .collect()
}

/// Diagonal averaging (Hankelization) to convert a (w × k) matrix back to
/// a time series of length n = w + k - 1.
fn diagonal_averaging(mat: &[Vec<f64>], n: usize) -> Vec<f64> {
    let w = mat.len();
    let k = mat[0].len();
    let l_star = w.min(k);
    let k_star = w.max(k);

    // If w > k, we need to transpose conceptually
    let transposed = w > k;

    let mut result = vec![0.0; n];
    for s in 0..n {
        let mut sum = 0.0;
        let mut count = 0;

        if s < l_star - 1 {
            // First segment
            for m in 0..=s {
                if transposed {
                    sum += mat[s - m][m];
                } else {
                    sum += mat[m][s - m];
                }
                count += 1;
            }
        } else if s < k_star {
            // Middle segment
            for m in 0..l_star {
                if transposed {
                    sum += mat[s - m][m];
                } else {
                    sum += mat[m][s - m];
                }
                count += 1;
            }
        } else {
            // Last segment
            for m in (s - k_star + 1)..=(n - k_star) {
                if transposed {
                    sum += mat[s - m][m];
                } else {
                    sum += mat[m][s - m];
                }
                count += 1;
            }
        }

        result[s] = sum / count as f64;
    }

    result
}

/// Automatic grouping using periodogram analysis of eigenvectors.
/// Returns 3 groups: trend, seasonal, residual.
fn auto_grouping(
    _eigenvalues: &[f64],
    eigenvectors: &[Vec<f64>],
    window_size: usize,
    lower_frequency_bound: f64,
    lower_frequency_contribution: f64,
) -> Vec<Vec<usize>> {
    let c = lower_frequency_contribution;

    // Compute frequency axis
    let n_freq = window_size / 2 + 1;
    let freqs: Vec<f64> = (0..n_freq).map(|i| i as f64 / window_size as f64).collect();

    // Find frequency index for trend boundary
    let idx_trend = freqs
        .iter()
        .rposition(|&f| f < lower_frequency_bound)
        .unwrap_or(0);
    let idx_resid = n_freq / 2;

    // Create FFT plan once for all eigenvectors (all same length = window_size)
    let mut planner = RealFftPlanner::<f64>::new();
    let fft = planner.plan_fft_forward(window_size);

    let mut trend_indices = Vec::new();
    let mut season_indices = Vec::new();
    let mut resid_indices = Vec::new();

    for comp in 0..window_size {
        // Extract eigenvector
        let v: Vec<f64> = (0..window_size).map(|r| eigenvectors[r][comp]).collect();

        // Compute periodogram using shared FFT plan
        let pxx = compute_periodogram(&v, &fft);

        // Compute cumulative sum
        let cumsum: Vec<f64> = pxx
            .iter()
            .scan(0.0, |acc, &x| {
                *acc += x;
                Some(*acc)
            })
            .collect();

        let total = *cumsum.last().unwrap();
        if total == 0.0 {
            resid_indices.push(comp);
            continue;
        }

        let trend_ratio = cumsum[idx_trend] / total;
        let resid_ratio = if idx_resid < cumsum.len() {
            cumsum[idx_resid] / total
        } else {
            1.0
        };

        if trend_ratio > c {
            trend_indices.push(comp);
        } else if resid_ratio < c {
            resid_indices.push(comp);
        } else {
            season_indices.push(comp);
        }
    }

    vec![trend_indices, season_indices, resid_indices]
}

/// Compute periodogram of a signal using FFT.
/// Returns power spectral density with appropriate scaling.
fn compute_periodogram(signal: &[f64], fft: &Arc<dyn realfft::RealToComplex<f64>>) -> Vec<f64> {
    let n = signal.len();
    if n == 0 {
        return Vec::new();
    }

    let mut data = signal.to_vec();
    let mut spectrum = fft.make_output_vec();
    fft.process(&mut data, &mut spectrum).unwrap();

    // Compute power spectral density (ortho normalization)
    let norm = n as f64;
    let n_freq = spectrum.len();
    let mut pxx: Vec<f64> = spectrum
        .iter()
        .map(|c| (c.re * c.re + c.im * c.im) / norm)
        .collect();

    // Double the non-DC, non-Nyquist components
    if n.is_multiple_of(2) {
        for p in pxx.iter_mut().take(n_freq - 1).skip(1) {
            *p *= 2.0;
        }
    } else {
        for p in pxx.iter_mut().skip(1) {
            *p *= 2.0;
        }
    }

    pxx
}

/// Find the largest off-diagonal element in a symmetric matrix.
/// Returns (max_abs_value, row, col) with row < col.
fn find_max_off_diagonal(mat: &[Vec<f64>], n: usize) -> (f64, usize, usize) {
    let mut max_val = 0.0_f64;
    let mut p = 0;
    let mut q = 1;
    for i in 0..n {
        for j in (i + 1)..n {
            if mat[i][j].abs() > max_val {
                max_val = mat[i][j].abs();
                p = i;
                q = j;
            }
        }
    }
    (max_val, p, q)
}

/// Compute the Jacobi rotation angle for pivot (p, q).
fn jacobi_rotation_angle(mat: &[Vec<f64>], p: usize, q: usize) -> f64 {
    if (mat[q][q] - mat[p][p]).abs() < 1e-30 {
        std::f64::consts::FRAC_PI_4
    } else {
        0.5 * (2.0 * mat[p][q] / (mat[p][p] - mat[q][q])).atan()
    }
}

/// Apply a Jacobi rotation at pivot (p, q) to the symmetric matrix in-place.
fn apply_jacobi_rotation(
    mat: &mut [Vec<f64>],
    n: usize,
    p: usize,
    q: usize,
    cos_t: f64,
    sin_t: f64,
) {
    let old_pp = mat[p][p];
    let old_qq = mat[q][q];
    let old_pq = mat[p][q];

    for i in 0..n {
        if i != p && i != q {
            let old_ip = mat[i][p];
            let old_iq = mat[i][q];
            mat[i][p] = cos_t * old_ip + sin_t * old_iq;
            mat[p][i] = mat[i][p];
            mat[i][q] = -sin_t * old_ip + cos_t * old_iq;
            mat[q][i] = mat[i][q];
        }
    }
    mat[p][p] = cos_t * cos_t * old_pp + 2.0 * sin_t * cos_t * old_pq + sin_t * sin_t * old_qq;
    mat[q][q] = sin_t * sin_t * old_pp - 2.0 * sin_t * cos_t * old_pq + cos_t * cos_t * old_qq;
    mat[p][q] = 0.0;
    mat[q][p] = 0.0;
}

/// Rotate eigenvector columns p and q by (cos_t, sin_t).
fn rotate_eigenvectors(v: &mut [Vec<f64>], n: usize, p: usize, q: usize, cos_t: f64, sin_t: f64) {
    for i in 0..n {
        let vp = v[i][p];
        let vq = v[i][q];
        v[i][p] = cos_t * vp + sin_t * vq;
        v[i][q] = -sin_t * vp + cos_t * vq;
    }
}

/// Sort eigenvalues descending and reorder eigenvector columns to match.
fn sort_eigen_pairs(mat: &[Vec<f64>], v: &[Vec<f64>], n: usize) -> (Vec<f64>, Vec<Vec<f64>>) {
    let mut eigen_pairs: Vec<(f64, usize)> = (0..n).map(|i| (mat[i][i], i)).collect();
    eigen_pairs.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap());

    let eigenvalues: Vec<f64> = eigen_pairs.iter().map(|&(val, _)| val).collect();
    let eigenvectors: Vec<Vec<f64>> = (0..n)
        .map(|row| eigen_pairs.iter().map(|&(_, col)| v[row][col]).collect())
        .collect();
    (eigenvalues, eigenvectors)
}

/// Symmetric eigendecomposition using Jacobi iteration.
/// Returns (eigenvalues, eigenvectors) with eigenvalues sorted descending.
/// `eigenvectors[i][j]` = j-th component of i-th eigenvector (column-major).
fn symmetric_eigen(a: &[Vec<f64>]) -> (Vec<f64>, Vec<Vec<f64>>) {
    let n = a.len();
    let mut mat: Vec<Vec<f64>> = a.to_vec();

    let mut v = vec![vec![0.0; n]; n];
    for i in 0..n {
        v[i][i] = 1.0;
    }

    let max_iter = 100 * n * n;
    let tol = 1e-12;

    for _ in 0..max_iter {
        let (max_val, p, q) = find_max_off_diagonal(&mat, n);
        if max_val < tol {
            break;
        }

        let theta = jacobi_rotation_angle(&mat, p, q);
        let (cos_t, sin_t) = (theta.cos(), theta.sin());

        apply_jacobi_rotation(&mut mat, n, p, q, cos_t, sin_t);
        rotate_eigenvectors(&mut v, n, p, q, cos_t, sin_t);
    }

    sort_eigen_pairs(&mat, &v, n)
}

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

    #[test]
    fn test_ssa_basic_no_grouping() {
        let config = SsaConfig::new(3);
        let x = vec![vec![1.0, 2.0, 3.0, 4.0, 5.0]];
        let result = Ssa::transform(&config, &x);
        assert_eq!(result.len(), 1);
        // window_size=3 => 3 components
        assert_eq!(result[0].len(), 3);
        // Each component has n_timestamps length
        assert_eq!(result[0][0].len(), 5);
    }

    #[test]
    fn test_ssa_reconstruction() {
        // Sum of all components should reconstruct original signal
        let config = SsaConfig::new(3);
        let x = vec![vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]];
        let result = Ssa::transform(&config, &x);

        let n = x[0].len();
        let mut reconstructed = vec![0.0; n];
        for component in &result[0] {
            for (i, &v) in component.iter().enumerate() {
                reconstructed[i] += v;
            }
        }

        for (i, (&orig, &recon)) in x[0].iter().zip(reconstructed.iter()).enumerate() {
            assert!(
                (orig - recon).abs() < 1e-8,
                "Reconstruction failed at index {i}: {orig} != {recon}"
            );
        }
    }

    #[test]
    fn test_ssa_multiple_samples() {
        let config = SsaConfig::new(3);
        let x = vec![vec![1.0, 2.0, 3.0, 4.0, 5.0], vec![5.0, 4.0, 3.0, 2.0, 1.0]];
        let result = Ssa::transform(&config, &x);
        assert_eq!(result.len(), 2);
    }

    #[test]
    fn test_ssa_explicit_groups() {
        let config = SsaConfig {
            window_size: 4,
            groups: SsaGrouping::Explicit(vec![vec![0, 1], vec![2, 3]]),
            ..SsaConfig::default()
        };
        let x = vec![vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]];
        let result = Ssa::transform(&config, &x);
        assert_eq!(result[0].len(), 2); // 2 groups

        // Reconstruction should still work
        let n = x[0].len();
        let mut reconstructed = vec![0.0; n];
        for component in &result[0] {
            for (i, &v) in component.iter().enumerate() {
                reconstructed[i] += v;
            }
        }
        for (i, (&orig, &recon)) in x[0].iter().zip(reconstructed.iter()).enumerate() {
            assert!(
                (orig - recon).abs() < 1e-8,
                "Reconstruction failed at index {i}: {orig} != {recon}"
            );
        }
    }

    #[test]
    fn test_ssa_auto_groups() {
        let config = SsaConfig {
            window_size: 5,
            groups: SsaGrouping::Auto,
            ..SsaConfig::default()
        };
        let x = vec![vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]];
        let result = Ssa::transform(&config, &x);
        assert_eq!(result[0].len(), 3); // trend, seasonal, residual

        // Reconstruction
        let n = x[0].len();
        let mut reconstructed = vec![0.0; n];
        for component in &result[0] {
            for (i, &v) in component.iter().enumerate() {
                reconstructed[i] += v;
            }
        }
        for (i, (&orig, &recon)) in x[0].iter().zip(reconstructed.iter()).enumerate() {
            assert!(
                (orig - recon).abs() < 1e-6,
                "Reconstruction failed at index {i}: {orig} != {recon}"
            );
        }
    }

    #[test]
    fn test_diagonal_averaging_simple() {
        // 2x3 matrix, should produce length 4 time series
        let mat = vec![vec![0.0, 1.0, 2.0], vec![3.0, 4.0, 5.0]];
        let result = diagonal_averaging(&mat, 4);
        assert_eq!(result.len(), 4);
        // Anti-diagonal averages:
        // s=0: mat[0][0] = 0
        // s=1: (mat[1][0] + mat[0][1]) / 2 = (3+1)/2 = 2
        // s=2: (mat[1][1] + mat[0][2]) / 2 = (4+2)/2 = 3
        // s=3: mat[1][2] = 5
        assert!((result[0] - 0.0).abs() < 1e-10);
        assert!((result[1] - 2.0).abs() < 1e-10);
        assert!((result[2] - 3.0).abs() < 1e-10);
        assert!((result[3] - 5.0).abs() < 1e-10);
    }

    #[test]
    fn test_symmetric_eigen() {
        // Simple 2x2 symmetric matrix
        let a = vec![vec![2.0, 1.0], vec![1.0, 2.0]];
        let (vals, vecs) = symmetric_eigen(&a);
        // Eigenvalues should be 3 and 1
        assert!((vals[0] - 3.0).abs() < 1e-8);
        assert!((vals[1] - 1.0).abs() < 1e-8);
        // Eigenvectors should be orthogonal
        let dot: f64 = vecs[0][0] * vecs[0][1] + vecs[1][0] * vecs[1][1];
        assert!(dot.abs() < 1e-8);
    }
}