scirs2-integrate 0.4.1

//! Vector-valued integration
//!
//! This module provides integration methods for vector-valued functions.
//! These methods are useful when you need to integrate a function that
//! returns arrays rather than scalar values.

use scirs2_core::ndarray::Array1;
use std::cmp::Ordering;
use std::collections::BinaryHeap;
use std::f64::consts::PI;
use std::fmt;

use crate::error::{IntegrateError, IntegrateResult};

/// Result type for vector-valued integration
#[derive(Clone, Debug)]
pub struct QuadVecResult<T> {
    /// The integral estimate
    pub integral: Array1<T>,
    /// The error estimate
    pub error: Array1<T>,
    /// Number of function evaluations
    pub nfev: usize,
    /// Number of integration subintervals used
    pub nintervals: usize,
    /// Whether the integration converged successfully
    pub success: bool,
}

impl<T: fmt::Display> fmt::Display for QuadVecResult<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "QuadVecResult(\n  integral=[{:}],\n  error=[{:}],\n  nfev={},\n  nintervals={},\n  success={}\n)",
            self.integral
                .iter()
                .map(|v| format!("{v}"))
                .collect::<Vec<_>>()
                .join(", "),
            self.error
                .iter()
                .map(|v| format!("{v}"))
                .collect::<Vec<_>>()
                .join(", "),
            self.nfev,
            self.nintervals,
            self.success
        )
    }
}

/// Options for quad_vec integration
#[derive(Clone, Debug)]
pub struct QuadVecOptions {
    /// Absolute tolerance
    pub epsabs: f64,
    /// Relative tolerance
    pub epsrel: f64,
    /// Norm to use for error estimation
    pub norm: NormType,
    /// Maximum number of subintervals
    pub limit: usize,
    /// Quadrature rule to use
    pub rule: QuadRule,
    /// Additional points where the integrand should be sampled
    pub points: Option<Vec<f64>>,
}

impl Default for QuadVecOptions {
    fn default() -> Self {
        Self {
            epsabs: 1e-10,
            epsrel: 1e-8,
            norm: NormType::L2,
            limit: 50,
            rule: QuadRule::GK21,
            points: None,
        }
    }
}

/// Type of norm to use for error estimation
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum NormType {
    /// Maximum absolute value
    Max,
    /// Euclidean (L2) norm
    L2,
}

/// Quadrature rule to use
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum QuadRule {
    /// 15-point Gauss-Kronrod rule
    GK15,
    /// 21-point Gauss-Kronrod rule
    GK21,
    /// Composite trapezoidal rule
    Trapezoid,
}

/// Subinterval for adaptive quadrature
#[derive(Clone, Debug)]
struct Subinterval {
    /// Left endpoint
    a: f64,
    /// Right endpoint
    b: f64,
    /// Integral estimate on this subinterval
    integral: Array1<f64>,
    /// Error estimate on this subinterval
    error: Array1<f64>,
    /// Norm of the error estimate (priority for subdivision)
    error_norm: f64,
}

impl PartialEq for Subinterval {
    fn eq(&self, other: &Self) -> bool {
        self.error_norm == other.error_norm
    }
}

impl Eq for Subinterval {}

impl PartialOrd for Subinterval {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Subinterval {
    fn cmp(&self, other: &Self) -> Ordering {
        // We want a max heap, so reverse the ordering
        other
            .error_norm
            .partial_cmp(&self.error_norm)
            .unwrap_or(Ordering::Equal)
    }
}

/// Compute a norm for error estimation
#[allow(dead_code)]
fn compute_norm(array: &Array1<f64>, normtype: NormType) -> f64 {
    match normtype {
        NormType::Max => {
            let mut max_abs = 0.0;
            for &val in array.iter() {
                let abs_val = val.abs();
                if abs_val > max_abs {
                    max_abs = abs_val;
                }
            }
            max_abs
        }
        NormType::L2 => {
            let mut sum_squares: f64 = 0.0;
            for &val in array.iter() {
                sum_squares += val * val;
            }
            sum_squares.sqrt()
        }
    }
}

/// Integration of a vector-valued function.
///
/// This function is similar to `quad`, but for functions that return
/// arrays (vectors) rather than scalars.
///
/// # Parameters
///
/// * `f` - Function to integrate. Should take a float and return an array.
/// * `a` - Lower bound of integration.
/// * `b` - Upper bound of integration.
/// * `options` - Integration options (optional).
///
/// # Returns
///
/// A result containing the integral estimate, error, and other information.
///
/// # Examples
///
/// ```
/// use scirs2_core::ndarray::{Array1, arr1};
/// use scirs2_integrate::quad_vec::{quad_vec, QuadVecOptions};
///
/// // Integrate a function that returns a 2D vector
/// let f = |x: f64| arr1(&[x.sin(), x.cos()]);
/// let result = quad_vec(f, 0.0, std::f64::consts::PI, None).expect("Operation failed");
///
/// // Result should be approximately [2.0, 0.0]
/// assert!((result.integral[0] - 2.0).abs() < 1e-10);
/// assert!(result.integral[1].abs() < 1e-10);
/// ```
#[allow(dead_code)]
pub fn quad_vec<F>(
    f: F,
    a: f64,
    b: f64,
    options: Option<QuadVecOptions>,
) -> IntegrateResult<QuadVecResult<f64>>
where
    F: Fn(f64) -> Array1<f64>,
{
    let options = options.unwrap_or_default();

    // Validate inputs
    if !a.is_finite() || !b.is_finite() {
        return Err(IntegrateError::ValueError(
            "Integration limits must be finite".to_string(),
        ));
    }

    // Check if interval is effectively zero
    if (b - a).abs() <= f64::EPSILON * a.abs().max(b.abs()) {
        // Evaluate at midpoint to get vector dimension
        let fval = f((a + b) / 2.0);
        let zeros = Array1::zeros(fval.len());

        return Ok(QuadVecResult {
            integral: zeros.clone(),
            error: zeros,
            nfev: 1,
            nintervals: 0,
            success: true,
        });
    }

    // Determine initial intervals
    let intervals = if let Some(ref points) = options.points {
        // Start with user-supplied breakpoints
        let mut sorted_points: Vec<f64> = points.clone();
        sorted_points.sort_by(|a, b| a.partial_cmp(b).unwrap_or(Ordering::Equal));

        // Filter out points outside [a, b] and remove duplicates
        let mut filtered_points: Vec<f64> = Vec::new();
        for &point in sorted_points.iter() {
            if point > a
                && point < b
                && (filtered_points.is_empty()
                    || (point - filtered_points.last().expect("Operation failed")).abs()
                        > f64::EPSILON)
            {
                filtered_points.push(point);
            }
        }

        // Create initial intervals
        let mut intervals: Vec<(f64, f64)> = Vec::new();

        if filtered_points.is_empty() {
            intervals.push((a, b));
        } else {
            intervals.push((a, filtered_points[0]));

            for i in 0..filtered_points.len() - 1 {
                intervals.push((filtered_points[i], filtered_points[i + 1]));
            }

            intervals.push((*filtered_points.last().expect("Operation failed"), b));
        }

        intervals
    } else {
        // Just use the whole interval
        vec![(a, b)]
    };

    // Evaluate function once to determine output size
    let fval = f((intervals[0].0 + intervals[0].1) / 2.0);
    let output_size = fval.len();

    // Initialize priority queue for adaptive subdivision
    let mut subintervals = BinaryHeap::new();
    let mut nfev = 1; // We've already evaluated f once

    // Process initial intervals
    for (a_i, b_i) in intervals {
        let (integral, error, evals) = evaluate_interval(&f, a_i, b_i, output_size, options.rule)?;

        nfev += evals;

        let error_norm = compute_norm(&error, options.norm);

        subintervals.push(Subinterval {
            a: a_i,
            b: b_i,
            integral,
            error,
            error_norm,
        });
    }

    // Adaptive subdivision
    while subintervals.len() < options.limit {
        // Get interval with largest error
        let interval = match subintervals.pop() {
            Some(i) => i,
            None => break, // Shouldn't happen, but just in case
        };

        // Check if we've reached desired accuracy
        let total_integral = get_total(&subintervals, &interval, |i| &i.integral);
        let total_error = get_total(&subintervals, &interval, |i| &i.error);

        let error_norm = compute_norm(&total_error, options.norm);
        let abs_tol = options.epsabs;
        let rel_tol = options.epsrel * compute_norm(&total_integral, options.norm);

        if error_norm <= abs_tol || error_norm <= rel_tol {
            // Add the interval back, we'll compute the final result later
            subintervals.push(interval);
            break;
        }

        // Split the interval
        let mid = (interval.a + interval.b) / 2.0;

        // Evaluate on the two halves
        let (left_integral, left_error, left_evals) =
            evaluate_interval(&f, interval.a, mid, output_size, options.rule)?;

        let (right_integral, right_error, right_evals) =
            evaluate_interval(&f, mid, interval.b, output_size, options.rule)?;

        nfev += left_evals + right_evals;

        // Create new intervals
        let left_error_norm = compute_norm(&left_error, options.norm);
        let right_error_norm = compute_norm(&right_error, options.norm);

        subintervals.push(Subinterval {
            a: interval.a,
            b: mid,
            integral: left_integral,
            error: left_error,
            error_norm: left_error_norm,
        });

        subintervals.push(Subinterval {
            a: mid,
            b: interval.b,
            integral: right_integral,
            error: right_error,
            error_norm: right_error_norm,
        });
    }

    // Compute final result
    let interval_vec: Vec<Subinterval> = subintervals.into_vec();
    let mut total_integral = Array1::zeros(output_size);
    let mut total_error = Array1::zeros(output_size);

    for interval in &interval_vec {
        for (i, &val) in interval.integral.iter().enumerate() {
            total_integral[i] += val;
        }

        for (i, &val) in interval.error.iter().enumerate() {
            total_error[i] += val;
        }
    }

    // Check for convergence
    let error_norm = compute_norm(&total_error, options.norm);
    let abs_tol = options.epsabs;
    let rel_tol = options.epsrel * compute_norm(&total_integral, options.norm);

    let success = error_norm <= abs_tol || error_norm <= rel_tol;

    Ok(QuadVecResult {
        integral: total_integral,
        error: total_error,
        nfev,
        nintervals: interval_vec.len(),
        success,
    })
}

/// Compute a property of all intervals combined
#[allow(dead_code)]
fn get_total<F, T>(heap: &BinaryHeap<Subinterval>, extra: &Subinterval, extract: F) -> Array1<T>
where
    F: Fn(&Subinterval) -> &Array1<T>,
    T: Clone + scirs2_core::numeric::Zero,
{
    let mut result = extract(extra).clone();

    for interval in heap.iter() {
        let property = extract(interval);

        for (i, val) in property.iter().enumerate() {
            result[i] = result[i].clone() + val.clone();
        }
    }

    result
}

/// Evaluate the integral on a specific interval
#[allow(dead_code)]
fn evaluate_interval<F>(
    f: &F,
    a: f64,
    b: f64,
    output_size: usize,
    rule: QuadRule,
) -> IntegrateResult<(Array1<f64>, Array1<f64>, usize)>
where
    F: Fn(f64) -> Array1<f64>,
{
    match rule {
        QuadRule::GK15 => {
            // Gauss-Kronrod 15-point rule (7-point Gauss, 15-point Kronrod)
            // Points and weights from SciPy
            let nodes = [
                -0.9914553711208126f64,
                -0.9491079123427585,
                -0.8648644233597691,
                -0.7415311855993944,
                -0.5860872354676911,
                -0.4058451513773972,
                -0.2077849550078985,
                0.0,
                0.2077849550078985,
                0.4058451513773972,
                0.5860872354676911,
                0.7415311855993944,
                0.8648644233597691,
                0.9491079123427585,
                0.9914553711208126,
            ];

            let weights_k = [
                0.022935322010529224f64,
                0.063_092_092_629_978_56,
                0.10479001032225018,
                0.14065325971552592,
                0.169_004_726_639_267_9,
                0.190_350_578_064_785_4,
                0.20443294007529889,
                0.20948214108472782,
                0.20443294007529889,
                0.190_350_578_064_785_4,
                0.169_004_726_639_267_9,
                0.14065325971552592,
                0.10479001032225018,
                0.063_092_092_629_978_56,
                0.022935322010529224,
            ];

            // Abscissae for the 7-point Gauss rule (odd indices of xgk)
            let weights_g = [
                0.129_484_966_168_869_7_f64,
                0.27970539148927664,
                0.381_830_050_505_118_9,
                0.417_959_183_673_469_4,
                0.381_830_050_505_118_9,
                0.27970539148927664,
                0.129_484_966_168_869_7,
            ];

            evaluate_rule(f, a, b, output_size, &nodes, &weights_g, &weights_k)
        }
        QuadRule::GK21 => {
            // Gauss-Kronrod 21-point rule (10-point Gauss, 21-point Kronrod)
            // Points and weights from SciPy
            let nodes = [
                -0.9956571630258081f64,
                -0.9739065285171717,
                -0.9301574913557082,
                -0.8650633666889845,
                -0.7808177265864169,
                -0.6794095682990244,
                -0.5627571346686047,
                -0.4333953941292472,
                -0.2943928627014602,
                -0.1488743389816312,
                0.0,
                0.1488743389816312,
                0.2943928627014602,
                0.4333953941292472,
                0.5627571346686047,
                0.6794095682990244,
                0.7808177265864169,
                0.8650633666889845,
                0.9301574913557082,
                0.9739065285171717,
                0.9956571630258081,
            ];

            let weights_k = [
                0.011694638867371874f64,
                0.032558162307964725,
                0.054755896574351995,
                0.075_039_674_810_919_96,
                0.093_125_454_583_697_6,
                0.109_387_158_802_297_64,
                0.123_491_976_262_065_84,
                0.134_709_217_311_473_34,
                0.142_775_938_577_060_09,
                0.147_739_104_901_338_49,
                0.149_445_554_002_916_9,
                0.147_739_104_901_338_49,
                0.142_775_938_577_060_09,
                0.134_709_217_311_473_34,
                0.123_491_976_262_065_84,
                0.109_387_158_802_297_64,
                0.093_125_454_583_697_6,
                0.075_039_674_810_919_96,
                0.054755896574351995,
                0.032558162307964725,
                0.011694638867371874,
            ];

            // Abscissae for the 10-point Gauss rule (every other point)
            let weights_g = [
                0.066_671_344_308_688_14f64,
                0.149_451_349_150_580_6,
                0.219_086_362_515_982_04,
                0.269_266_719_309_996_36,
                0.295_524_224_714_752_9,
                0.295_524_224_714_752_9,
                0.269_266_719_309_996_36,
                0.219_086_362_515_982_04,
                0.149_451_349_150_580_6,
                0.066_671_344_308_688_14,
            ];

            evaluate_rule(f, a, b, output_size, &nodes, &weights_g, &weights_k)
        }
        QuadRule::Trapezoid => {
            // Simple trapezoid rule with 15 points
            let n = 15;
            let mut integral = Array1::zeros(output_size);
            let mut error = Array1::zeros(output_size);

            let h = (b - a) / (n as f64 - 1.0);
            let fa = f(a);
            let fb = f(b);

            // Add endpoints with half weight
            for (i, (&fa_i, &fb_i)) in fa.iter().zip(fb.iter()).enumerate() {
                integral[i] = 0.5 * (fa_i + fb_i);
            }

            // Add interior points
            for j in 1..n - 1 {
                let x = a + (j as f64) * h;
                let fx = f(x);

                for (i, &fx_i) in fx.iter().enumerate() {
                    integral[i] += fx_i;
                }
            }

            // Scale by h
            for i in 0..output_size {
                integral[i] *= h;

                // Crude error estimate
                error[i] = 1e-2 * integral[i].abs();
            }

            Ok((integral, error, n))
        }
    }
}

/// Evaluate a Gauss-Kronrod rule on an interval
#[allow(dead_code)]
fn evaluate_rule<F>(
    f: &F,
    a: f64,
    b: f64,
    output_size: usize,
    nodes: &[f64],
    weights_g: &[f64],
    weights_k: &[f64],
) -> IntegrateResult<(Array1<f64>, Array1<f64>, usize)>
where
    F: Fn(f64) -> Array1<f64>,
{
    let _n = nodes.len();

    let mut integral_k = Array1::zeros(output_size);
    let mut integral_g = Array1::zeros(output_size);

    // Map nodes to [a, b]
    let mid = (a + b) / 2.0;
    let half_length = (b - a) / 2.0;

    let mut nfev = 0;

    // For GK rules, Gauss points are at odd indices (1, 3, 5, ...)
    let mut gauss_idx = 0;

    // Evaluate at all Kronrod points
    for (i, &node) in nodes.iter().enumerate() {
        let x = mid + half_length * node;
        let fx = f(x);
        nfev += 1;

        // Add to Kronrod integral
        for (j, &fx_j) in fx.iter().enumerate() {
            integral_k[j] += weights_k[i] * fx_j;
        }

        // Check if this is also a Gauss point
        // For GK15: Gauss points are at indices 1, 3, 5, 7, 9, 11, 13
        // For GK21: Gauss points are at indices 1, 3, 5, 7, 9, 11, 13, 15, 17, 19
        if i % 2 == 1 && gauss_idx < weights_g.len() {
            for (j, &fx_j) in fx.iter().enumerate() {
                integral_g[j] += weights_g[gauss_idx] * fx_j;
            }
            gauss_idx += 1;
        }
    }

    // Scale by half-length
    integral_k *= half_length;
    integral_g *= half_length;

    // Compute error estimate
    // Error is estimated as (200 * |I_k - I_g|)^1.5
    let mut error = Array1::zeros(output_size);
    for i in 0..output_size {
        let diff = (integral_k[i] - integral_g[i]).abs();
        error[i] = (200.0 * diff).powf(1.5_f64);
    }

    Ok((integral_k, error, nfev))
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_abs_diff_eq;
    use scirs2_core::ndarray::arr1;

    #[test]
    fn test_simple_integral() {
        // Integrate [x, x^2] from 0 to 1
        let f = |x: f64| arr1(&[x, x * x]);
        let result = quad_vec(f, 0.0, 1.0, None).expect("Operation failed");

        assert_abs_diff_eq!(result.integral[0], 0.5, epsilon = 1e-10);
        assert_abs_diff_eq!(result.integral[1], 1.0 / 3.0, epsilon = 1e-10);
        assert!(result.success);
    }

    #[test]
    fn test_trig_functions() {
        // Integrate [sin(x), cos(x)] from 0 to π
        let f = |x: f64| arr1(&[x.sin(), x.cos()]);
        let result = quad_vec(f, 0.0, PI, None).expect("Operation failed");

        assert_abs_diff_eq!(result.integral[0], 2.0, epsilon = 1e-10);
        assert_abs_diff_eq!(result.integral[1], 0.0, epsilon = 1e-10);
        assert!(result.success);
    }

    #[test]
    fn test_with_breakpoints() {
        // Integrate [x, x^2] from 0 to 2 with a breakpoint at x=1
        let f = |x: f64| arr1(&[x, x * x]);

        let options = QuadVecOptions {
            points: Some(vec![1.0]),
            ..Default::default()
        };

        let result = quad_vec(f, 0.0, 2.0, Some(options)).expect("Operation failed");

        assert_abs_diff_eq!(result.integral[0], 2.0, epsilon = 1e-10);
        assert_abs_diff_eq!(result.integral[1], 8.0 / 3.0, epsilon = 1e-10);
        assert!(result.success);
    }

    #[test]
    fn test_different_rules() {
        // Test with different quadrature rules
        let f = |x: f64| arr1(&[x.sin()]);

        let options_gk15 = QuadVecOptions {
            rule: QuadRule::GK15,
            ..Default::default()
        };

        let options_gk21 = QuadVecOptions {
            rule: QuadRule::GK21,
            ..Default::default()
        };

        let options_trapezoid = QuadVecOptions {
            rule: QuadRule::Trapezoid,
            ..Default::default()
        };

        let result_gk15 = quad_vec(f, 0.0, PI, Some(options_gk15)).expect("Operation failed");
        let result_gk21 = quad_vec(f, 0.0, PI, Some(options_gk21)).expect("Operation failed");
        let result_trapezoid =
            quad_vec(f, 0.0, PI, Some(options_trapezoid)).expect("Operation failed");

        assert_abs_diff_eq!(result_gk15.integral[0], 2.0, epsilon = 1e-10);
        assert_abs_diff_eq!(result_gk21.integral[0], 2.0, epsilon = 1e-10);
        assert_abs_diff_eq!(result_trapezoid.integral[0], 2.0, epsilon = 2e-3); // Lower precision for trapezoid
    }

    #[test]
    fn test_error_norms() {
        // Test Max norm
        let arr = arr1(&[1.0, -2.0, 0.5]);
        let max_norm = compute_norm(&arr, NormType::Max);
        assert_abs_diff_eq!(max_norm, 2.0, epsilon = 1e-10);

        // Test L2 norm
        let l2_norm = compute_norm(&arr, NormType::L2);
        assert_abs_diff_eq!(
            l2_norm,
            (1.0f64 * 1.0 + 2.0 * 2.0 + 0.5 * 0.5).sqrt(),
            epsilon = 1e-10
        );
    }
}