// Copyright (c) 2024 RustyTorch++ Team
// Licensed under the Apache License, Version 2.0

//! Nonlinear solvers for nonlinear finite element problems.

use super::{ConvergenceInfo, LinearSolver, NonlinearSolver, SolverOptions};
use crate::assembly::SparseMatrix;
use crate::error::FeaResult;
use nalgebra::DVector;
use std::time::Instant;

/// Newton-Raphson nonlinear solver.
#[derive(Debug)]
pub struct NewtonRaphson {
    /// Linear solver for Jacobian systems
    linear_solver: Box<dyn LinearSolver>,
}

impl NewtonRaphson {
    /// Create a new Newton-Raphson solver.
    pub fn new(linear_solver: Box<dyn LinearSolver>) -> Self {
        Self { linear_solver }
    }
}

impl NonlinearSolver for NewtonRaphson {
    fn solve_nonlinear<F, J>(
        &mut self,
        residual_function: F,
        jacobian_function: J,
        initial_guess: &DVector<f64>,
        options: &SolverOptions,
    ) -> FeaResult<(DVector<f64>, ConvergenceInfo)>
    where
        F: Fn(&DVector<f64>) -> FeaResult<DVector<f64>>,
        J: Fn(&DVector<f64>) -> FeaResult<SparseMatrix>,
    {
        let start_time = Instant::now();
        let mut info = ConvergenceInfo::new();
        let mut x = initial_guess.clone();

        for iter in 0..options.max_iterations {
            // Compute residual
            let residual = residual_function(&x)?;
            let residual_norm = residual.norm();
            info.add_residual(residual_norm);

            // Check convergence
            if iter == 0 {
                let initial_residual = residual_norm;
                if initial_residual < options.tolerance {
                    info.set_converged(0, residual_norm, 0.0);
                    break;
                }
            } else {
                let relative_residual = residual_norm / info.residual_history[0];
                if residual_norm < options.tolerance
                    || relative_residual < options.relative_tolerance
                {
                    info.set_converged(iter, residual_norm, relative_residual);
                    break;
                }
            }

            // Compute Jacobian
            let jacobian = jacobian_function(&x)?;

            // Solve Jacobian system: J * delta_x = -residual
            let neg_residual = -residual;
            let (delta_x, _) = self
                .linear_solver
                .solve(&jacobian, &neg_residual, options)?;

            // Update solution
            x += delta_x;
        }

        info.set_solve_time(start_time.elapsed());
        Ok((x, info))
    }

    fn name(&self) -> &'static str {
        "Newton-Raphson"
    }
}

/// Modified Newton solver (reuses Jacobian).
#[derive(Debug)]
pub struct ModifiedNewton {
    linear_solver: Box<dyn LinearSolver>,
    jacobian_reuse_count: usize,
}

impl ModifiedNewton {
    pub fn new(linear_solver: Box<dyn LinearSolver>) -> Self {
        Self {
            linear_solver,
            jacobian_reuse_count: 5,
        }
    }

    pub fn with_reuse_count(mut self, count: usize) -> Self {
        self.jacobian_reuse_count = count;
        self
    }
}

impl NonlinearSolver for ModifiedNewton {
    fn solve_nonlinear<F, J>(
        &mut self,
        residual_function: F,
        jacobian_function: J,
        initial_guess: &DVector<f64>,
        options: &SolverOptions,
    ) -> FeaResult<(DVector<f64>, ConvergenceInfo)>
    where
        F: Fn(&DVector<f64>) -> FeaResult<DVector<f64>>,
        J: Fn(&DVector<f64>) -> FeaResult<SparseMatrix>,
    {
        let start_time = Instant::now();
        let mut info = ConvergenceInfo::new();
        let mut x = initial_guess.clone();
        let mut cached_jacobian: Option<SparseMatrix> = None;

        for iter in 0..options.max_iterations {
            let residual = residual_function(&x)?;
            let residual_norm = residual.norm();
            info.add_residual(residual_norm);

            if iter == 0 {
                if residual_norm < options.tolerance {
                    info.set_converged(0, residual_norm, 0.0);
                    break;
                }
            } else {
                let relative_residual = residual_norm / info.residual_history[0];
                if residual_norm < options.tolerance
                    || relative_residual < options.relative_tolerance
                {
                    info.set_converged(iter, residual_norm, relative_residual);
                    break;
                }
            }

            // Reuse Jacobian for several iterations
            if cached_jacobian.is_none() || iter % self.jacobian_reuse_count == 0 {
                cached_jacobian = Some(jacobian_function(&x)?);
            }

            let jacobian = cached_jacobian.as_ref().unwrap();
            let neg_residual = -residual;
            let (delta_x, _) = self.linear_solver.solve(jacobian, &neg_residual, options)?;

            x += delta_x;
        }

        info.set_solve_time(start_time.elapsed());
        Ok((x, info))
    }

    fn name(&self) -> &'static str {
        "Modified Newton"
    }
}

/// Quasi-Newton solver with BFGS updates.
#[derive(Debug)]
pub struct QuasiNewton {
    linear_solver: Box<dyn LinearSolver>,
}

impl QuasiNewton {
    pub fn new(linear_solver: Box<dyn LinearSolver>) -> Self {
        Self { linear_solver }
    }
}

impl NonlinearSolver for QuasiNewton {
    fn solve_nonlinear<F, J>(
        &mut self,
        residual_function: F,
        jacobian_function: J,
        initial_guess: &DVector<f64>,
        options: &SolverOptions,
    ) -> FeaResult<(DVector<f64>, ConvergenceInfo)>
    where
        F: Fn(&DVector<f64>) -> FeaResult<DVector<f64>>,
        J: Fn(&DVector<f64>) -> FeaResult<SparseMatrix>,
    {
        let start_time = Instant::now();
        let mut info = ConvergenceInfo::new();
        let mut x = initial_guess.clone();

        // Use Newton-Raphson for first iteration to get initial Jacobian
        let residual = residual_function(&x)?;
        let initial_residual_norm = residual.norm();
        info.add_residual(initial_residual_norm);

        if initial_residual_norm < options.tolerance {
            info.set_converged(0, initial_residual_norm, 0.0);
            info.set_solve_time(start_time.elapsed());
            return Ok((x, info));
        }

        let mut jacobian = jacobian_function(&x)?;

        for iter in 0..options.max_iterations {
            let current_residual = residual_function(&x)?;
            let residual_norm = current_residual.norm();
            info.add_residual(residual_norm);

            let relative_residual = residual_norm / initial_residual_norm;
            if residual_norm < options.tolerance || relative_residual < options.relative_tolerance {
                info.set_converged(iter, residual_norm, relative_residual);
                break;
            }

            let neg_residual = -current_residual;
            let (delta_x, _) = self
                .linear_solver
                .solve(&jacobian, &neg_residual, options)?;

            x += &delta_x;

            // BFGS update would go here in a full implementation
            // For simplicity, we recompute the Jacobian every few iterations
            if iter % 5 == 0 {
                jacobian = jacobian_function(&x)?;
            }
        }

        info.set_solve_time(start_time.elapsed());
        Ok((x, info))
    }

    fn name(&self) -> &'static str {
        "Quasi-Newton (BFGS)"
    }
}

#[cfg(disabled)]
mod tests {
    use super::*;
    use crate::assembly::SparseMatrix;
    use crate::solvers::LuDirect;

    #[test]
    fn test_newton_raphson_creation() {
        let linear_solver = Box::new(LuDirect::new());
        let solver = NewtonRaphson::new(linear_solver);
        assert_eq!(solver.name(), "Newton-Raphson");
    }

    #[test]
    fn test_modified_newton_creation() {
        let linear_solver = Box::new(LuDirect::new());
        let solver = ModifiedNewton::new(linear_solver).with_reuse_count(3);
        assert_eq!(solver.name(), "Modified Newton");
        assert_eq!(solver.jacobian_reuse_count, 3);
    }

    #[test]
    fn test_quasi_newton_creation() {
        let linear_solver = Box::new(LuDirect::new());
        let solver = QuasiNewton::new(linear_solver);
        assert_eq!(solver.name(), "Quasi-Newton (BFGS)");
    }

    #[test]
    fn test_simple_nonlinear_solve() {
        let linear_solver = Box::new(LuDirect::new());
        let mut solver = NewtonRaphson::new(linear_solver);

        // Simple nonlinear system: f(x) = x^2 - 4 = 0, solution x = 2
        let residual_fn = |x: &DVector<f64>| -> FeaResult<DVector<f64>> {
            let mut r = DVector::zeros(1);
            r[0] = x[0] * x[0] - 4.0;
            Ok(r)
        };

        let jacobian_fn = |x: &DVector<f64>| -> FeaResult<SparseMatrix> {
            let mut j = SparseMatrix::new(1, 1);
            j.add_entry(0, 0, 2.0 * x[0]).unwrap();
            j.finalize().unwrap();
            Ok(j)
        };

        let initial_guess = DVector::from_vec(vec![1.0]);
        let options = SolverOptions::default();

        let result = solver.solve_nonlinear(residual_fn, jacobian_fn, &initial_guess, &options);
        assert!(result.is_ok());

        let (solution, info) = result.unwrap();
        assert!((solution[0] - 2.0).abs() < 0.1); // Should converge to x = 2
        assert!(info.converged);
    }
}