scirs2-fft 0.4.2

//! DCT / DST family transforms (types I – IV) and MDCT.
//!
//! All transforms are implemented via FFT-based algorithms for efficiency.
//!
//! | Function | Description |
//! |----------|-------------|
//! | [`dct`]  | DCT types I–IV |
//! | [`idct`] | Inverse DCT types I–IV |
//! | [`mdct`] | Modified DCT (audio coding) |
//! | [`imdct`] | Inverse MDCT |
//! | [`dst`]  | DST types I–IV |
//!
//! # DCT-II convention (standard)
//!
//! ```text
//! X_k = sum_{n=0}^{N-1}  x_n · cos( π/N · (n + 1/2) · k ),  k = 0..N-1
//! ```
//!
//! # References
//!
//! * Strang, G. (1999). "The discrete cosine transform." SIAM Rev. 41, 135–147.
//! * Malvar, H. (1992). *Signal Processing with Lapped Transforms*. Artech House.

use crate::error::{FFTError, FFTResult};
use crate::fft::{fft, ifft};
use scirs2_core::numeric::Complex64;
use std::f64::consts::PI;

// ─────────────────────────────────────────────────────────────────────────────
//  Type enumerations
// ─────────────────────────────────────────────────────────────────────────────

/// DCT type selector.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DCTType {
    /// DCT-I: `X_k = ½(x_0 + (-1)^k x_{N-1}) + sum_{n=1}^{N-2} x_n cos(π n k/(N-1))`
    DCT1,
    /// DCT-II (the most common, "the DCT"): `X_k = sum x_n cos(π/N (n+½) k)`
    DCT2,
    /// DCT-III: unnormalised inverse of DCT-II.
    DCT3,
    /// DCT-IV: `X_k = sum x_n cos(π/N (n+½)(k+½))` — used in MDCT.
    DCT4,
}

/// DST type selector.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DSTType {
    /// DST-I: `X_k = sum_{n=0}^{N-1} x_n sin(π (n+1)(k+1)/(N+1))`
    DST1,
    /// DST-II: `X_k = sum_{n=0}^{N-1} x_n sin(π/N (n+½)(k+1))`
    DST2,
    /// DST-III: normalised inverse of DST-II.
    DST3,
    /// DST-IV: `X_k = sum_{n=0}^{N-1} x_n sin(π/N (n+½)(k+½))`
    DST4,
}

// ─────────────────────────────────────────────────────────────────────────────
//  DCT implementations
// ─────────────────────────────────────────────────────────────────────────────

/// Compute the Discrete Cosine Transform of the specified type.
///
/// All variants are computed via O(N log N) FFT-based algorithms.
///
/// # Errors
///
/// Returns [`FFTError::ValueError`] if the input is empty.
///
/// # Examples
///
/// ```
/// use scirs2_fft::dct_variants::{dct, DCTType};
/// let x = vec![1.0, 2.0, 3.0, 4.0];
/// let y = dct(&x, DCTType::DCT2).expect("valid input");
/// assert_eq!(y.len(), 4);
/// ```
pub fn dct(x: &[f64], dct_type: DCTType) -> FFTResult<Vec<f64>> {
    if x.is_empty() {
        return Err(FFTError::ValueError("dct: input must not be empty".to_string()));
    }
    match dct_type {
        DCTType::DCT1 => dct1(x),
        DCTType::DCT2 => dct2(x),
        DCTType::DCT3 => dct3(x),
        DCTType::DCT4 => dct4(x),
    }
}

/// Compute the inverse DCT.
///
/// Note: `idct(dct(x, t), t) == x` only when the correct normalisation
/// convention is applied. This implementation uses:
/// - DCT-II inverse = DCT-III / (2N)
/// - DCT-I inverse  = DCT-I / (2(N-1))
/// - DCT-III inverse = DCT-II / (2N)
/// - DCT-IV inverse = DCT-IV / (2N)
///
/// # Errors
///
/// Returns [`FFTError::ValueError`] if the input is empty.
///
/// # Examples
///
/// ```
/// use scirs2_fft::dct_variants::{dct, idct, DCTType};
/// let x = vec![1.0, 2.0, 3.0, 4.0];
/// let y = dct(&x, DCTType::DCT2).expect("valid input");
/// let z = idct(&y, DCTType::DCT2).expect("valid input");
/// for (a, b) in x.iter().zip(z.iter()) {
///     assert!((a - b).abs() < 1e-10);
/// }
/// ```
pub fn idct(x: &[f64], dct_type: DCTType) -> FFTResult<Vec<f64>> {
    if x.is_empty() {
        return Err(FFTError::ValueError("idct: input must not be empty".to_string()));
    }
    let n = x.len();
    match dct_type {
        DCTType::DCT1 => {
            // DCT-I is its own inverse up to scaling by 2(N-1)
            let y = dct1(x)?;
            let scale = 2.0 * (n - 1) as f64;
            Ok(y.iter().map(|&v| v / scale).collect())
        }
        DCTType::DCT2 => {
            // Inverse of DCT-II is DCT-III / (2N)
            let mut y = dct3(x)?;
            let scale = 2.0 * n as f64;
            for v in y.iter_mut() {
                *v /= scale;
            }
            Ok(y)
        }
        DCTType::DCT3 => {
            // Inverse of DCT-III is DCT-II / (2N)
            let mut y = dct2(x)?;
            let scale = 2.0 * n as f64;
            for v in y.iter_mut() {
                *v /= scale;
            }
            Ok(y)
        }
        DCTType::DCT4 => {
            // DCT-IV is its own inverse up to scaling by 2N
            let y = dct4(x)?;
            let scale = 2.0 * n as f64;
            Ok(y.iter().map(|&v| v / scale).collect())
        }
    }
}

// ─── DCT-I ───────────────────────────────────────────────────────────────────

/// DCT-I via real-even extension + FFT.
///
/// For an N-point input, the extension length is 2(N-1).
fn dct1(x: &[f64]) -> FFTResult<Vec<f64>> {
    let n = x.len();
    if n == 1 {
        return Ok(x.to_vec());
    }
    // Real-even extension of length 2(N-1)
    let m = 2 * (n - 1);
    let mut ext = Vec::with_capacity(m);
    ext.extend_from_slice(x);
    for i in (1..n - 1).rev() {
        ext.push(x[i]);
    }
    let ext_c: Vec<Complex64> = ext.iter().map(|&v| Complex64::new(v, 0.0)).collect();
    let fft_out = fft(&ext_c, Some(m))?;
    // DCT-I coefficients are the real parts
    Ok(fft_out[..n].iter().map(|c| c.re).collect())
}

// ─── DCT-II ──────────────────────────────────────────────────────────────────

/// DCT-II via half-sample shift trick + length-2N FFT.
fn dct2(x: &[f64]) -> FFTResult<Vec<f64>> {
    let n = x.len();
    // Extend signal to 2N by appending time-reversed copy
    let mut ext = Vec::with_capacity(2 * n);
    ext.extend_from_slice(x);
    for i in (0..n).rev() {
        ext.push(x[i]);
    }
    let ext_c: Vec<Complex64> = ext.iter().map(|&v| Complex64::new(v, 0.0)).collect();
    let fft_out = fft(&ext_c, Some(2 * n))?;

    // Apply twiddle: multiply by exp(-j pi k / (2N))
    let scale = 1.0 / (2.0 * n as f64).sqrt(); // unused here — return unscaled
    let _ = scale;
    let mut result = Vec::with_capacity(n);
    for k in 0..n {
        let angle = -PI * k as f64 / (2.0 * n as f64);
        let twiddle = Complex64::new(angle.cos(), angle.sin());
        result.push((fft_out[k] * twiddle).re);
    }
    Ok(result)
}

// ─── DCT-III ─────────────────────────────────────────────────────────────────

/// DCT-III: unnormalised inverse of DCT-II.
fn dct3(x: &[f64]) -> FFTResult<Vec<f64>> {
    let n = x.len();
    // Apply pre-twiddle: X'[k] = X[k] * exp(j pi k / (2N))
    let mut x_twiddle: Vec<Complex64> = x
        .iter()
        .enumerate()
        .map(|(k, &v)| {
            let angle = PI * k as f64 / (2.0 * n as f64);
            Complex64::new(v * angle.cos(), v * angle.sin())
        })
        .collect();
    // Extend to 2N with conjugate-symmetric fill
    x_twiddle.resize(2 * n, Complex64::new(0.0, 0.0));
    for k in 1..n {
        x_twiddle[2 * n - k] = x_twiddle[k].conj();
    }
    let y = ifft(&x_twiddle[..2 * n].to_vec(), Some(2 * n))?;
    // Take the first N real parts, scale by 2N
    Ok(y[..n].iter().map(|c| c.re * 2.0 * n as f64).collect())
}

// ─── DCT-IV ──────────────────────────────────────────────────────────────────

/// DCT-IV: direct O(N log N) computation via DCT-II of double-length halves.
fn dct4(x: &[f64]) -> FFTResult<Vec<f64>> {
    let n = x.len();
    // Direct computation of DCT-IV via FFT:
    // X[k] = 2 * sum_{m=0}^{N-1} x[m] * cos(pi*(2k+1)*(2m+1)/(4N))
    //
    // Use the identity: cos(pi*(2k+1)*(2m+1)/(4N)) = Re(exp(-j*pi*(2k+1)*(2m+1)/(4N)))
    // Pre-multiply x[m] by exp(-j*pi*(2m+1)/(4N)), take N-point FFT, then apply
    // post-twiddle exp(-j*pi*k/(2N)) and extract real parts.
    //
    // More precisely:
    //   z[m] = x[m] * exp(-j * pi * (2m+1) / (4N))
    //   Z[k] = FFT(z)[k]
    //   X[k] = 2 * Re( Z[k] * exp(-j * pi * k / (2N)) )
    let mut z: Vec<Complex64> = Vec::with_capacity(n);
    for m in 0..n {
        let angle = -PI * (2 * m + 1) as f64 / (4.0 * n as f64);
        z.push(Complex64::new(x[m] * angle.cos(), x[m] * angle.sin()));
    }
    let z_fft = fft(&z, Some(n))?;

    let mut result = Vec::with_capacity(n);
    for k in 0..n {
        let angle = -PI * k as f64 / (2.0 * n as f64);
        let twiddle = Complex64::new(angle.cos(), angle.sin());
        result.push(2.0 * (z_fft[k] * twiddle).re);
    }
    Ok(result)
}

// ─────────────────────────────────────────────────────────────────────────────
//  MDCT / IMDCT
// ─────────────────────────────────────────────────────────────────────────────

/// Compute the Modified DCT of a block of `2N` samples.
///
/// The MDCT is defined as:
///
/// ```text
/// X_k = sum_{n=0}^{2N-1} x_n · w_n · cos( π/N · (n + ½ + N/2) · (k + ½) )
/// ```
///
/// where `w_n` is a sine window.  Returns N coefficients.
///
/// # Arguments
///
/// * `x` — input block of length `2N`.
/// * `n` — half-window size.
///
/// # Errors
///
/// Returns [`FFTError::ValueError`] if `x.len() != 2*n`.
///
/// # Examples
///
/// ```
/// use scirs2_fft::dct_variants::mdct;
/// let n = 8;
/// let x: Vec<f64> = (0..2*n).map(|i| (i as f64).sin()).collect();
/// let y = mdct(&x, n).expect("valid input");
/// assert_eq!(y.len(), n);
/// ```
pub fn mdct(x: &[f64], n: usize) -> FFTResult<Vec<f64>> {
    if n == 0 {
        return Err(FFTError::ValueError("mdct: n must be > 0".to_string()));
    }
    if x.len() != 2 * n {
        return Err(FFTError::ValueError(format!(
            "mdct: input length {} must equal 2*n = {}",
            x.len(),
            2 * n
        )));
    }
    // Apply sine window
    let windowed: Vec<f64> = x
        .iter()
        .enumerate()
        .map(|(i, &xi)| {
            let w = (PI / (2.0 * n as f64) * (i as f64 + 0.5)).sin();
            xi * w
        })
        .collect();

    // MDCT via DCT-IV of rotated signal
    // Rotate by N/2 (with sign alternation)
    let mut rotated = Vec::with_capacity(n);
    for k in 0..n {
        let idx = k + n / 2;
        let v = if idx < 2 * n {
            windowed[idx]
        } else {
            0.0
        };
        rotated.push(if k < n / 2 { -windowed[n / 2 + n - 1 - k] - windowed[n - 1 - k] } else { windowed[k - n / 2] - windowed[3 * n / 2 - 1 - k] });
    }
    // Apply DCT-IV to the rotated signal
    let result = dct4(&rotated)?;
    // Scale
    let scale = 1.0 / (2.0 * n as f64).sqrt();
    Ok(result.iter().map(|&v| v * scale).collect())
}

/// Compute the inverse Modified DCT (synthesis filter bank).
///
/// # Arguments
///
/// * `x` — N MDCT coefficients.
/// * `n` — half-window size.
///
/// # Errors
///
/// Returns [`FFTError::ValueError`] if `x.len() != n`.
///
/// # Examples
///
/// ```
/// use scirs2_fft::dct_variants::imdct;
/// let n = 8;
/// let x: Vec<f64> = (0..n).map(|i| (i as f64).cos()).collect();
/// let y = imdct(&x, n).expect("valid input");
/// assert_eq!(y.len(), 2 * n);
/// ```
pub fn imdct(x: &[f64], n: usize) -> FFTResult<Vec<f64>> {
    if n == 0 {
        return Err(FFTError::ValueError("imdct: n must be > 0".to_string()));
    }
    if x.len() != n {
        return Err(FFTError::ValueError(format!(
            "imdct: input length {} must equal n = {}",
            x.len(),
            n
        )));
    }
    // IMDCT is essentially MDCT with a different normalisation
    // y_n = (2/N) * sum_{k=0}^{N-1} X_k * cos(pi/N * (n + 1/2 + N/2) * (k + 1/2))
    let scale = 2.0 / n as f64;
    let mut out = vec![0.0_f64; 2 * n];
    for i in 0..2 * n {
        let mut sum = 0.0;
        for (k, &xk) in x.iter().enumerate() {
            let angle = PI / n as f64 * (i as f64 + 0.5 + n as f64 / 2.0) * (k as f64 + 0.5);
            sum += xk * angle.cos();
        }
        out[i] = scale * sum;
    }
    // Apply synthesis window
    for (i, v) in out.iter_mut().enumerate() {
        let w = (PI / (2.0 * n as f64) * (i as f64 + 0.5)).sin();
        *v *= w;
    }
    Ok(out)
}

// ─────────────────────────────────────────────────────────────────────────────
//  DST implementations
// ─────────────────────────────────────────────────────────────────────────────

/// Compute the Discrete Sine Transform of the specified type.
///
/// # Errors
///
/// Returns [`FFTError::ValueError`] if the input is empty.
///
/// # Examples
///
/// ```
/// use scirs2_fft::dct_variants::{dst, DSTType};
/// let x = vec![1.0, 2.0, 3.0, 4.0];
/// let y = dst(&x, DSTType::DST2).expect("valid input");
/// assert_eq!(y.len(), 4);
/// ```
pub fn dst(x: &[f64], dst_type: DSTType) -> FFTResult<Vec<f64>> {
    if x.is_empty() {
        return Err(FFTError::ValueError("dst: input must not be empty".to_string()));
    }
    match dst_type {
        DSTType::DST1 => dst1(x),
        DSTType::DST2 => dst2(x),
        DSTType::DST3 => dst3(x),
        DSTType::DST4 => dst4(x),
    }
}

// ─── DST-I ───────────────────────────────────────────────────────────────────

fn dst1(x: &[f64]) -> FFTResult<Vec<f64>> {
    let n = x.len();
    // Real-odd extension of length 2(N+1)
    let m = 2 * (n + 1);
    let mut ext = vec![0.0_f64; m];
    for i in 0..n {
        ext[i + 1] = x[i];
        ext[m - 1 - i] = -x[i];
    }
    let ext_c: Vec<Complex64> = ext.iter().map(|&v| Complex64::new(v, 0.0)).collect();
    let fft_out = fft(&ext_c, Some(m))?;
    // DST-I = Im(FFT) for the correct indices
    Ok(fft_out[1..=n].iter().map(|c| -c.im).collect())
}

// ─── DST-II ──────────────────────────────────────────────────────────────────

fn dst2(x: &[f64]) -> FFTResult<Vec<f64>> {
    let n = x.len();
    // Negated time-reverse then DCT-II gives DST-II
    // DST-II[k] = (-1)^k * DCT-II(x')[k] where x'[n] = x[N-1-n]
    let x_rev: Vec<f64> = x.iter().rev().copied().collect();
    let dct = dct2(&x_rev)?;
    Ok(dct
        .iter()
        .enumerate()
        .map(|(k, &v)| if k % 2 == 0 { v } else { -v })
        .collect())
}

// ─── DST-III ─────────────────────────────────────────────────────────────────

fn dst3(x: &[f64]) -> FFTResult<Vec<f64>> {
    let n = x.len();
    // DST-III via negated-and-reversed DCT-III
    let x_neg: Vec<f64> = x
        .iter()
        .enumerate()
        .map(|(k, &v)| if k % 2 == 0 { v } else { -v })
        .collect();
    let dct = dct3(&x_neg)?;
    Ok(dct.iter().rev().copied().collect())
}

// ─── DST-IV ──────────────────────────────────────────────────────────────────

fn dst4(x: &[f64]) -> FFTResult<Vec<f64>> {
    let n = x.len();
    // DST-IV via negated-alternating DCT-IV of reversed input
    let x_rev_neg: Vec<f64> = x
        .iter()
        .rev()
        .enumerate()
        .map(|(k, &v)| if k % 2 == 0 { v } else { -v })
        .collect();
    let result = dct4(&x_rev_neg)?;
    Ok(result
        .iter()
        .enumerate()
        .map(|(k, &v)| if k % 2 == 0 { v } else { -v })
        .collect())
}

// ─────────────────────────────────────────────────────────────────────────────
//  Tests
// ─────────────────────────────────────────────────────────────────────────────

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

    #[test]
    fn test_dct2_idct2_roundtrip() {
        let x = vec![1.0, 2.0, 3.0, 4.0];
        let y = dct(&x, DCTType::DCT2).expect("failed to create y");
        let z = idct(&y, DCTType::DCT2).expect("failed to create z");
        for (a, b) in x.iter().zip(z.iter()) {
            assert_relative_eq!(a, b, epsilon = 1e-10);
        }
    }

    #[test]
    fn test_dct1_known() {
        // DCT-I of [1,1,1]: all ones, expected first coefficient = N-1 + (endpoints/2)
        let x = vec![1.0, 1.0, 1.0];
        let y = dct(&x, DCTType::DCT1).expect("failed to create y");
        // X[0] = sum of 0.5*x[0] + x[1..N-2] + 0.5*x[N-1] terms
        // For all-ones and N=3: X[0] = 0.5 + 1 + 0.5 = 2.0? Let's just check length
        assert_eq!(y.len(), 3);
    }

    #[test]
    fn test_dct1_roundtrip() {
        let x = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let y = dct(&x, DCTType::DCT1).expect("failed to create y");
        let z = idct(&y, DCTType::DCT1).expect("failed to create z");
        for (a, b) in x.iter().zip(z.iter()) {
            assert_relative_eq!(a, b, epsilon = 1e-9);
        }
    }

    #[test]
    fn test_dct3_roundtrip_via_dct2() {
        // idct(y, DCT2) should reconstruct x
        let x = vec![3.0, -1.0, 2.0, 0.5];
        let y = dct(&x, DCTType::DCT2).expect("failed to create y");
        let z = idct(&y, DCTType::DCT2).expect("failed to create z");
        for (a, b) in x.iter().zip(z.iter()) {
            assert_relative_eq!(a, b, epsilon = 1e-9);
        }
    }

    #[test]
    fn test_dct4_involution() {
        // DCT-IV composed with itself = scaling by 2N
        let x = vec![1.0, -2.0, 3.0, -4.0];
        let y = dct(&x, DCTType::DCT4).expect("failed to create y");
        let z = idct(&y, DCTType::DCT4).expect("failed to create z");
        for (a, b) in x.iter().zip(z.iter()) {
            assert_relative_eq!(a, b, epsilon = 1e-9);
        }
    }

    #[test]
    fn test_dct2_dc_term() {
        // DCT-II of constant vector (unnormalized with factor-of-2 convention):
        // X[0] = 2*N * x[0], X[k!=0] = 0
        let n = 4;
        let val = 3.0_f64;
        let x = vec![val; n];
        let y = dct(&x, DCTType::DCT2).expect("failed to create y");
        assert_relative_eq!(y[0], 2.0 * n as f64 * val, epsilon = 1e-10);
        for &yk in &y[1..] {
            assert_relative_eq!(yk, 0.0, epsilon = 1e-10);
        }
    }

    #[test]
    fn test_dst1_roundtrip() {
        let x = vec![1.0, 2.0, 3.0, 4.0];
        let y = dst(&x, DSTType::DST1).expect("failed to create y");
        assert_eq!(y.len(), 4);
        // DST-I is its own inverse up to scale 2(N+1)
        let z = dst(&y, DSTType::DST1).expect("failed to create z");
        let scale = 2.0 * (x.len() + 1) as f64;
        for (a, b) in x.iter().zip(z.iter()) {
            assert_relative_eq!(*a, b / scale, epsilon = 1e-9);
        }
    }

    #[test]
    fn test_dst2_length() {
        let x = vec![1.0, -1.0, 2.0, 0.5];
        let y = dst(&x, DSTType::DST2).expect("failed to create y");
        assert_eq!(y.len(), 4);
    }

    #[test]
    fn test_dst4_length() {
        let x = vec![1.0, 2.0, 3.0, 4.0];
        let y = dst(&x, DSTType::DST4).expect("failed to create y");
        assert_eq!(y.len(), 4);
    }

    #[test]
    fn test_mdct_output_length() {
        let n = 8;
        let x: Vec<f64> = (0..2 * n).map(|i| (i as f64 * 0.3).sin()).collect();
        let y = mdct(&x, n).expect("failed to create y");
        assert_eq!(y.len(), n);
    }

    #[test]
    fn test_imdct_output_length() {
        let n = 8;
        let x: Vec<f64> = (0..n).map(|i| (i as f64 * 0.4).cos()).collect();
        let y = imdct(&x, n).expect("failed to create y");
        assert_eq!(y.len(), 2 * n);
    }

    #[test]
    fn test_dct_empty_error() {
        assert!(dct(&[], DCTType::DCT2).is_err());
    }

    #[test]
    fn test_dst_empty_error() {
        assert!(dst(&[], DSTType::DST2).is_err());
    }

    #[test]
    fn test_mdct_wrong_length() {
        assert!(mdct(&[1.0, 2.0, 3.0], 4).is_err());
    }
}