gspx 0.1.2

use ndarray::{Array1, Array2, ArrayView1, ArrayView2, s};

use crate::{
    Laplacian,
    error::GspError,
    kernel::{ChebyKernel, estimate_spectral_bound},
};

/// Sampling strategy for Chebyshev fitting.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Sampling {
    /// Chebyshev nodes (recommended default).
    Chebyshev,
    /// Uniform linear spacing over the spectrum.
    Linear,
    /// Quadratically biased spacing toward low eigenvalues.
    Quadratic,
    /// Logarithmically spaced samples.
    Logarithmic,
}

/// Configuration for [`ChebyModel`] fitting.
#[derive(Debug, Clone)]
pub struct ChebyFitOptions {
    /// Number of sample points used by non-Chebyshev sampling modes.
    pub n_samples: Option<usize>,
    /// Sampling strategy for fitting.
    pub sampling: Sampling,
    /// Lower spectral bound for fitting.
    pub min_lambda: f64,
    /// Relative trimming tolerance for coefficient truncation.
    pub rtol: f64,
    /// Enables adaptive order growth until `target_error` is reached.
    pub adaptive: bool,
    /// Maximum adaptive order.
    pub max_order: usize,
    /// Relative error target for adaptive fitting.
    pub target_error: f64,
}

impl Default for ChebyFitOptions {
    fn default() -> Self {
        Self {
            n_samples: None,
            sampling: Sampling::Chebyshev,
            min_lambda: 0.0,
            rtol: 1e-12,
            adaptive: false,
            max_order: 500,
            target_error: 1e-10,
        }
    }
}

/// Chebyshev polynomial approximation model.
#[derive(Debug, Clone)]
pub struct ChebyModel {
    /// Coefficient matrix with shape `(order + 1, n_channels)`.
    pub coefficients: Array2<f64>,
    /// Upper spectrum bound used for scaling into the Chebyshev domain.
    pub spectrum_bound: f64,
    /// Lower spectrum bound used for scaling into the Chebyshev domain.
    pub min_lambda: f64,
}

impl ChebyModel {
    /// Fits a Chebyshev approximation for `f` on `[min_lambda, spectrum_bound]`.
    ///
    /// # Errors
    /// Returns an error when dimensions or bounds are invalid.
    pub fn fit<F>(
        f: F,
        order: usize,
        spectrum_bound: f64,
        opts: ChebyFitOptions,
    ) -> Result<Self, GspError>
    where
        F: Fn(&Array1<f64>) -> Array2<f64>,
    {
        if order < 1 {
            return Err(GspError::InvalidKernel("order must be >= 1".to_string()));
        }
        if spectrum_bound <= opts.min_lambda {
            return Err(GspError::InvalidKernel(
                "spectrum_bound must be greater than min_lambda".to_string(),
            ));
        }

        if opts.adaptive {
            return Self::adaptive_fit(f, order, spectrum_bound, opts);
        }

        let coeffs = compute_cheby_coefficients(
            &f,
            order,
            spectrum_bound,
            opts.min_lambda,
            opts.n_samples,
            opts.sampling,
        )?;
        let coeffs = truncate_coeffs(coeffs, opts.rtol);
        Ok(Self {
            coefficients: coeffs,
            spectrum_bound,
            min_lambda: opts.min_lambda,
        })
    }

    /// Builds a validated [`ChebyKernel`] fitted for a specific Laplacian.
    ///
    /// # Errors
    /// Returns an error when spectral-bound estimation or fitting fails.
    pub fn kernel_on_laplacian<F>(
        l: &Laplacian,
        f: F,
        order: usize,
        opts: ChebyFitOptions,
    ) -> Result<ChebyKernel, GspError>
    where
        F: Fn(&Array1<f64>) -> Array2<f64>,
    {
        let bound = estimate_spectral_bound(l)?;
        let model = Self::fit(f, order, bound, opts)?;
        Ok(ChebyKernel {
            coefficients: model.coefficients,
            spectrum_bound: model.spectrum_bound,
            min_lambda: model.min_lambda,
        })
    }

    /// Evaluates the fitted Chebyshev model at sample points `x`.
    ///
    /// # Errors
    /// Returns an error when kernel parameters are invalid.
    pub fn evaluate(&self, x: ArrayView1<f64>) -> Result<Array2<f64>, GspError> {
        let kernel = ChebyKernel {
            coefficients: self.coefficients.clone(),
            spectrum_bound: self.spectrum_bound,
            min_lambda: self.min_lambda,
        };
        kernel.evaluate(x)
    }

    fn adaptive_fit<F>(
        f: F,
        start_order: usize,
        spectrum_bound: f64,
        opts: ChebyFitOptions,
    ) -> Result<Self, GspError>
    where
        F: Fn(&Array1<f64>) -> Array2<f64>,
    {
        let test_x = Array1::linspace(opts.min_lambda, spectrum_bound, 1000);
        let f_exact = f(&test_x);
        let mut order = start_order.max(8);
        loop {
            let coeffs = compute_cheby_coefficients(
                &f,
                order,
                spectrum_bound,
                opts.min_lambda,
                opts.n_samples,
                opts.sampling,
            )?;
            let coeffs = ensure_2d_real(coeffs)?;
            let model = Self {
                coefficients: coeffs.clone(),
                spectrum_bound,
                min_lambda: opts.min_lambda,
            };
            let y = model.evaluate(test_x.view())?;
            let rel = (&f_exact - &y)
                .mapv(|v| v.abs())
                .iter()
                .zip(f_exact.iter())
                .map(|(e, y0)| e / y0.abs().max(1e-15))
                .fold(0.0f64, |a, b| a.max(b));
            if rel <= opts.target_error || order >= opts.max_order {
                return Ok(Self {
                    coefficients: truncate_coeffs(coeffs, opts.rtol),
                    spectrum_bound,
                    min_lambda: opts.min_lambda,
                });
            }
            order = ((order as f64) * 1.5).floor() as usize + 1;
            order = order.min(opts.max_order);
        }
    }
}

fn ensure_2d_real(mut coefficients: Array2<f64>) -> Result<Array2<f64>, GspError> {
    if coefficients.ndim() != 2 {
        return Err(GspError::Dimensions("expected 2D array".to_string()));
    }
    if coefficients.nrows() == 0 || coefficients.ncols() == 0 {
        return Err(GspError::Dimensions(
            "coefficient matrix must be non-empty".to_string(),
        ));
    }
    if !coefficients.is_standard_layout() {
        coefficients = coefficients.to_owned();
    }
    Ok(coefficients)
}

fn truncate_coeffs(coefficients: Array2<f64>, rtol: f64) -> Array2<f64> {
    let threshold = coefficients
        .iter()
        .map(|value| value.abs())
        .fold(0.0f64, |current_max, value| current_max.max(value))
        * rtol;
    let row_maxima: Vec<f64> = coefficients
        .rows()
        .into_iter()
        .map(|row| {
            row.iter()
                .fold(0.0f64, |current_max, value| current_max.max(value.abs()))
        })
        .collect();
    let mut last_nonzero_row = 1usize;
    for (index, row_maximum) in row_maxima.into_iter().enumerate() {
        if row_maximum > threshold {
            last_nonzero_row = index + 1;
        }
    }
    coefficients.slice(s![0..last_nonzero_row, ..]).to_owned()
}

fn compute_cheby_coefficients<F>(
    f: &F,
    order: usize,
    spectrum_bound: f64,
    min_lambda: f64,
    n_samples: Option<usize>,
    sampling: Sampling,
) -> Result<Array2<f64>, GspError>
where
    F: Fn(&Array1<f64>) -> Array2<f64>,
{
    let lambda_range = spectrum_bound - min_lambda;
    let lambda_mid = (spectrum_bound + min_lambda) / 2.0;

    if sampling == Sampling::Chebyshev {
        let sample_count = order + 1;
        let cheby_indices = Array1::from((0..=order).map(|value| value as f64).collect::<Vec<_>>());
        let cheby_nodes =
            cheby_indices.mapv(|index| (std::f64::consts::PI * index / (order as f64)).cos());
        let sample_x = cheby_nodes.mapv(|x| lambda_mid + 0.5 * lambda_range * x);
        let sample_values = f(&sample_x);
        let n_channels = sample_values.ncols();
        let mut coefficients = Array2::<f64>::zeros((sample_count, n_channels));
        let mut endpoint_weights = Array1::<f64>::ones(sample_count);
        endpoint_weights[0] = 0.5;
        endpoint_weights[sample_count - 1] = 0.5;

        for coeff_index in 0..=order {
            let cheby_basis = cheby_indices.mapv(|index| {
                ((coeff_index as f64) * std::f64::consts::PI * index / (order as f64)).cos()
            });
            let scale = if coeff_index == 0 || coeff_index == order {
                1.0 / (order as f64)
            } else {
                2.0 / (order as f64)
            };
            for channel in 0..n_channels {
                let mut sum = 0.0;
                for sample_index in 0..sample_count {
                    sum += endpoint_weights[sample_index]
                        * sample_values[[sample_index, channel]]
                        * cheby_basis[sample_index];
                }
                coefficients[[coeff_index, channel]] = scale * sum;
            }
        }
        return Ok(coefficients);
    }

    let sample_count = n_samples.unwrap_or((4 * (order + 1)).max(1000));
    let interpolation_grid = Array1::linspace(0.0, 1.0, sample_count);
    let sample_x = match sampling {
        Sampling::Linear => interpolation_grid.mapv(|t| min_lambda + lambda_range * t),
        Sampling::Quadratic => interpolation_grid.mapv(|t| min_lambda + lambda_range * t * t),
        Sampling::Logarithmic => {
            let eps = (min_lambda * 0.001).max(1e-10);
            interpolation_grid.mapv(|t| {
                ((min_lambda + eps).ln() + t * (spectrum_bound / (min_lambda + eps)).ln()).exp()
            })
        }
        Sampling::Chebyshev => unreachable!(),
    };
    let x_scaled = sample_x.mapv(|x| 2.0 * (x - min_lambda) / lambda_range - 1.0);
    let sample_values = f(&sample_x);
    let mut cheby_matrix = Array2::<f64>::zeros((sample_count, order + 1));
    cheby_matrix.column_mut(0).fill(1.0);
    if order >= 1 {
        cheby_matrix.column_mut(1).assign(&x_scaled);
    }
    for degree in 2..=order {
        for sample_index in 0..sample_count {
            cheby_matrix[[sample_index, degree]] =
                2.0 * x_scaled[sample_index] * cheby_matrix[[sample_index, degree - 1]]
                    - cheby_matrix[[sample_index, degree - 2]];
        }
    }

    // Chebyshev-weighted least squares: w(x) = 1/sqrt(1 - x^2) is the natural
    // orthogonality weight for Chebyshev polynomials, critical for conditioning.
    let weights = x_scaled.mapv(|x| 1.0 / (1.0 - (x * x).min(0.9999)).sqrt());
    for sample_index in 0..sample_count {
        let w = weights[sample_index];
        for degree in 0..=order {
            cheby_matrix[[sample_index, degree]] *= w;
        }
    }
    let weighted_values = &sample_values * &weights.insert_axis(ndarray::Axis(1));
    real_lstsq(cheby_matrix.view(), weighted_values.view())
}

fn real_lstsq(a: ArrayView2<f64>, b: ArrayView2<f64>) -> Result<Array2<f64>, GspError> {
    let (m, n) = a.dim();
    if b.nrows() != m {
        return Err(GspError::Dimensions(format!(
            "least squares row mismatch: A is {m}x{n}, B has {} rows",
            b.nrows()
        )));
    }
    let p = b.ncols();
    let a_data = a.iter().copied().collect::<Vec<_>>();
    let b_data = b.iter().copied().collect::<Vec<_>>();
    let adm = nalgebra::DMatrix::<f64>::from_row_slice(m, n, &a_data);
    let bdm = nalgebra::DMatrix::<f64>::from_row_slice(m, p, &b_data);
    let svd = adm.svd(true, true);
    let x = svd
        .solve(&bdm, 1e-12)
        .map_err(|_| GspError::Factorization("real least-squares solve failed".to_string()))?;
    let mut out = Array2::<f64>::zeros((n, p));
    for row in 0..n {
        for col in 0..p {
            out[[row, col]] = x[(row, col)];
        }
    }
    Ok(out)
}