temporal-neural-solver 0.1.2

//! Real solver integration - simplified version
//! This would use the actual sublinear solver if it compiled properly

use ndarray::{Array1, Array2};
use std::time::Instant;

/// Simplified sparse matrix for demonstration
pub struct SparseMatrix {
    pub rows: usize,
    pub cols: usize,
    pub values: Vec<(usize, usize, f64)>,
}

impl SparseMatrix {
    pub fn from_triplets(triplets: Vec<(usize, usize, f64)>, rows: usize, cols: usize) -> Self {
        SparseMatrix {
            rows,
            cols,
            values: triplets,
        }
    }

    /// Matrix-vector multiplication
    pub fn multiply(&self, x: &[f64]) -> Vec<f64> {
        let mut result = vec![0.0; self.rows];
        for (i, j, val) in &self.values {
            if *j < x.len() {
                result[*i] += val * x[*j];
            }
        }
        result
    }
}

/// Real Neumann series solver implementation
pub struct NeumannSolver {
    max_iterations: usize,
    tolerance: f64,
}

impl NeumannSolver {
    pub fn new(max_iterations: usize, tolerance: f64) -> Self {
        Self {
            max_iterations,
            tolerance,
        }
    }

    /// Solve Ax = b using Neumann series expansion
    /// (I - M)^(-1) = I + M + M^2 + M^3 + ...
    pub fn solve(&self, a: &SparseMatrix, b: &[f64]) -> SolverResult {
        let start = Instant::now();
        let n = b.len();

        // Initial guess x = b
        let mut x = b.to_vec();
        let mut residual = vec![0.0; n];
        let mut iterations = 0;

        // Jacobi preconditioner (diagonal scaling)
        let mut diagonal = vec![1.0; n];
        for (i, j, val) in &a.values {
            if i == j {
                diagonal[*i] = *val;
            }
        }

        // Iterate: x_{k+1} = b + M * x_k where M = I - D^{-1}A
        for iter in 0..self.max_iterations {
            // Compute residual = b - Ax
            let ax = a.multiply(&x);
            for i in 0..n {
                residual[i] = b[i] - ax[i];
            }

            // Check convergence
            let residual_norm: f64 = residual.iter().map(|r| r * r).sum::<f64>().sqrt();
            if residual_norm < self.tolerance {
                iterations = iter + 1;
                break;
            }

            // Update x = x + D^{-1} * residual (Jacobi step)
            for i in 0..n {
                if diagonal[i].abs() > 1e-10 {
                    x[i] += residual[i] / diagonal[i];
                }
            }

            iterations = iter + 1;
        }

        // Final residual calculation
        let ax_final = a.multiply(&x);
        let final_residual: Vec<f64> = (0..n).map(|i| b[i] - ax_final[i]).collect();
        let residual_norm = final_residual.iter().map(|r| r * r).sum::<f64>().sqrt();

        SolverResult {
            solution: x,
            residual_norm,
            iterations,
            time_elapsed: start.elapsed(),
        }
    }
}

pub struct SolverResult {
    pub solution: Vec<f64>,
    pub residual_norm: f64,
    pub iterations: usize,
    pub time_elapsed: std::time::Duration,
}

/// Forward push solver for graph-based systems
pub struct ForwardPushSolver {
    epsilon: f64,
    max_iterations: usize,
}

impl ForwardPushSolver {
    pub fn new(epsilon: f64, max_iterations: usize) -> Self {
        Self {
            epsilon,
            max_iterations,
        }
    }

    /// Forward push algorithm for PageRank-style problems
    pub fn solve(&self, adjacency: &Array2<f32>, teleport: &Array1<f32>) -> Array1<f32> {
        let n = adjacency.shape()[0];
        let mut estimate = Array1::zeros(n);
        let mut residual = teleport.clone();

        for _ in 0..self.max_iterations {
            // Find node with largest residual
            let (max_idx, &max_residual) = residual
                .iter()
                .enumerate()
                .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
                .unwrap();

            if max_residual < self.epsilon as f32 {
                break;
            }

            // Push residual forward
            estimate[max_idx] += residual[max_idx];

            // Distribute to neighbors
            let out_degree: f32 = (0..n).map(|j| adjacency[[max_idx, j]]).sum();
            if out_degree > 0.0 {
                for j in 0..n {
                    if adjacency[[max_idx, j]] > 0.0 {
                        residual[j] += 0.85 * residual[max_idx] * adjacency[[max_idx, j]] / out_degree;
                    }
                }
            }

            residual[max_idx] = 0.0;
        }

        estimate
    }
}

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

    #[test]
    fn test_neumann_solver() {
        // Create a simple diagonally dominant system
        // [2, -1, 0]   [1]
        // [-1, 2, -1] * x = [0]
        // [0, -1, 2]   [1]
        let matrix = SparseMatrix::from_triplets(
            vec![
                (0, 0, 2.0), (0, 1, -1.0),
                (1, 0, -1.0), (1, 1, 2.0), (1, 2, -1.0),
                (2, 1, -1.0), (2, 2, 2.0),
            ],
            3,
            3,
        );

        let b = vec![1.0, 0.0, 1.0];

        let solver = NeumannSolver::new(100, 1e-6);
        let result = solver.solve(&matrix, &b);

        println!("Solution: {:?}", result.solution);
        println!("Iterations: {}", result.iterations);
        println!("Residual norm: {}", result.residual_norm);
        println!("Time: {:?}", result.time_elapsed);

        // Check that solution is reasonable
        assert!(result.residual_norm < 1e-5);
        assert!(result.iterations < 100);
    }

    #[test]
    fn test_forward_push() {
        let mut adjacency = Array2::zeros((3, 3));
        adjacency[[0, 1]] = 1.0;
        adjacency[[1, 2]] = 1.0;
        adjacency[[2, 0]] = 1.0;

        let teleport = Array1::from_vec(vec![0.33, 0.33, 0.34]);

        let solver = ForwardPushSolver::new(1e-6, 100);
        let result = solver.solve(&adjacency, &teleport);

        println!("PageRank scores: {:?}", result);
        assert!(result.sum() > 0.0);
    }
}