//! Newton-Raphson solver implementation.
//!
//! Provides [`NewtonConfig`] which implements the Newton-Raphson method for
//! solving systems of non-linear equations with quadratic convergence.

use std::time::{Duration, Instant};

use crate::criteria::ConvergenceCriteria;
use crate::jacobian::JacobianMatrix;
use crate::metadata::SimulationMetadata;
use crate::solver::{
    apply_newton_step, dominant_residual, ConvergedState, ConvergenceDiagnostics,
    ConvergenceStatus, IterationDiagnostics, JacobianFreezingConfig, Solver, SolverError,
    SolverType, TimeoutConfig, VerboseConfig,
};
use crate::system::System;
use entropyk_components::{ComponentError, JacobianBuilder};
use entropyk_solver_core::{DomainViolation, LinearSolver};

/// Configuration for the Newton-Raphson solver.
///
/// Solves F(x) = 0 by iterating: x_{k+1} = x_k - α·J^{-1}·r(x_k)
/// where J is the Jacobian matrix and α is the step length.
#[derive(Debug, Clone, PartialEq)]
pub struct NewtonConfig {
    /// Maximum iterations before declaring non-convergence. Default: 100.
    pub max_iterations: usize,
    /// Convergence tolerance (L2 norm). Default: 1e-6.
    pub tolerance: f64,
    /// Enable Armijo line-search. Default: false.
    pub line_search: bool,
    /// Optional time budget.
    pub timeout: Option<Duration>,
    /// Use numerical Jacobian (finite differences). Default: false.
    pub use_numerical_jacobian: bool,
    /// Armijo condition constant. Default: 1e-4.
    pub line_search_armijo_c: f64,
    /// Max backtracking iterations. Default: 20.
    pub line_search_max_backtracks: usize,
    /// Divergence threshold. Default: 1e10.
    pub divergence_threshold: f64,
    /// Timeout behavior configuration.
    pub timeout_config: TimeoutConfig,
    /// Previous state for ZOH fallback.
    pub previous_state: Option<Vec<f64>>,
    /// Residual for previous_state.
    pub previous_residual: Option<f64>,
    /// Smart initial state for cold-start.
    pub initial_state: Option<Vec<f64>>,
    /// Multi-circuit convergence criteria.
    pub convergence_criteria: Option<ConvergenceCriteria>,
    /// Jacobian-freezing optimization.
    pub jacobian_freezing: Option<JacobianFreezingConfig>,
    /// Verbose mode configuration for diagnostics.
    pub verbose_config: VerboseConfig,
}

impl Default for NewtonConfig {
    fn default() -> Self {
        Self {
            max_iterations: 100,
            tolerance: 1e-6,
            line_search: false,
            timeout: None,
            use_numerical_jacobian: false,
            line_search_armijo_c: 1e-4,
            line_search_max_backtracks: 20,
            divergence_threshold: 1e10,
            timeout_config: TimeoutConfig::default(),
            previous_state: None,
            previous_residual: None,
            initial_state: None,
            convergence_criteria: None,
            jacobian_freezing: None,
            verbose_config: VerboseConfig::default(),
        }
    }
}

impl NewtonConfig {
    /// Sets the initial state for cold-start solving.
    pub fn with_initial_state(mut self, state: Vec<f64>) -> Self {
        self.initial_state = Some(state);
        self
    }

    /// Sets multi-circuit convergence criteria.
    pub fn with_convergence_criteria(mut self, criteria: ConvergenceCriteria) -> Self {
        self.convergence_criteria = Some(criteria);
        self
    }

    /// Enables Jacobian-freezing optimization.
    pub fn with_jacobian_freezing(mut self, config: JacobianFreezingConfig) -> Self {
        self.jacobian_freezing = Some(config);
        self
    }

    /// Enables verbose mode for diagnostics.
    pub fn with_verbose(mut self, config: VerboseConfig) -> Self {
        self.verbose_config = config;
        self
    }

    /// Computes the L2 norm of the residual vector.
    fn residual_norm(residuals: &[f64]) -> f64 {
        residuals.iter().map(|r| r * r).sum::<f64>().sqrt()
    }

    /// Handles timeout based on configuration.
    fn handle_timeout(
        &self,
        best_state: &[f64],
        best_residual: f64,
        iterations: usize,
        timeout: Duration,
        system: &System,
    ) -> Result<ConvergedState, SolverError> {
        if !self.timeout_config.return_best_state_on_timeout {
            return Err(SolverError::Timeout {
                timeout_ms: timeout.as_millis() as u64,
            });
        }

        if self.timeout_config.zoh_fallback {
            if let Some(ref prev_state) = self.previous_state {
                let residual = self.previous_residual.unwrap_or(best_residual);
                tracing::info!(iterations, residual, "ZOH fallback");
                return Ok(ConvergedState::new(
                    prev_state.clone(),
                    iterations,
                    residual,
                    ConvergenceStatus::TimedOutWithBestState,
                    SimulationMetadata::new(system.input_hash()),
                ));
            }
        }

        tracing::info!(iterations, best_residual, "Returning best state on timeout");
        Ok(ConvergedState::new(
            best_state.to_vec(),
            iterations,
            best_residual,
            ConvergenceStatus::TimedOutWithBestState,
            SimulationMetadata::new(system.input_hash()),
        ))
    }

    /// Checks for divergence based on residual growth.
    fn check_divergence(
        &self,
        current_norm: f64,
        previous_norm: f64,
        divergence_count: &mut usize,
    ) -> Option<SolverError> {
        if current_norm > self.divergence_threshold {
            return Some(SolverError::Divergence {
                reason: format!(
                    "Residual {} exceeds threshold {}",
                    current_norm, self.divergence_threshold
                ),
            });
        }

        if current_norm > previous_norm {
            *divergence_count += 1;
            if *divergence_count >= 3 {
                return Some(SolverError::Divergence {
                    reason: format!(
                        "Residual increased 3x: {:.6e} → {:.6e}",
                        previous_norm, current_norm
                    ),
                });
            }
        } else {
            *divergence_count = 0;
        }

        None
    }

    /// Performs Armijo line search. Returns Some(alpha) if valid step found.
    /// hot path. `state_copy` and `new_residuals` must have appropriate lengths.
    ///
    /// Recoverable component domain violations (KINSOL `> 0`, see
    /// [`ComponentError::is_recoverable`]) shrink the step and are counted in
    /// `recoverable_events` for the iteration diagnostics; a fatal evaluation
    /// error aborts the search immediately instead of burning backtracks.
    #[allow(clippy::too_many_arguments)]
    fn line_search(
        &self,
        system: &System,
        state: &mut Vec<f64>,
        delta: &[f64],
        _residuals: &[f64],
        current_norm: f64,
        state_copy: &mut [f64],
        new_residuals: &mut Vec<f64>,
        clipping_mask: &[Option<(f64, f64)>],
        recoverable_events: &mut u32,
    ) -> Result<Option<f64>, SolverError> {
        let mut alpha: f64 = 1.0;
        state_copy.copy_from_slice(state);
        let gradient_dot_delta = -current_norm;

        for _backtrack in 0..self.line_search_max_backtracks {
            apply_newton_step(state, delta, clipping_mask, alpha);

            if let Err(e) = system.compute_residuals(state, new_residuals) {
                state.copy_from_slice(state_copy);
                if e.is_recoverable() {
                    // Recoverable domain violation (KINSOL > 0): shrink and retry.
                    alpha *= 0.5;
                    *recoverable_events += 1;
                    continue;
                }
                // Fatal evaluation error: abort without burning backtracks.
                return Err(SolverError::InvalidSystem {
                    message: format!("Failed to compute residuals: {:?}", e),
                });
            }

            let new_norm = Self::residual_norm(new_residuals);
            if new_norm <= current_norm + self.line_search_armijo_c * alpha * gradient_dot_delta {
                tracing::debug!(
                    alpha,
                    old_norm = current_norm,
                    new_norm,
                    "Line search accepted"
                );
                return Ok(Some(alpha));
            }

            state.copy_from_slice(state_copy);
            alpha *= 0.5;
        }

        tracing::warn!(
            "Line search failed after {} backtracks",
            self.line_search_max_backtracks
        );
        Ok(None)
    }

    fn finalize_failure_diagnostics(
        &self,
        mut diagnostics: Option<ConvergenceDiagnostics>,
        iterations: usize,
        final_residual: f64,
        best_residual: f64,
        elapsed_ms: u64,
        jacobian_condition_final: Option<f64>,
        final_state: Option<Vec<f64>>,
    ) -> Option<ConvergenceDiagnostics> {
        if let Some(ref mut diag) = diagnostics {
            diag.iterations = iterations;
            diag.final_residual = final_residual;
            diag.best_residual = best_residual;
            diag.converged = false;
            diag.timing_ms = elapsed_ms;
            diag.jacobian_condition_final = jacobian_condition_final;
            diag.final_solver = Some(SolverType::NewtonRaphson);

            if self.verbose_config.dump_final_state {
                diag.final_state = final_state;
                let json_output = diag.dump_diagnostics(self.verbose_config.output_format);
                tracing::warn!(
                    iterations,
                    final_residual,
                    "Non-convergence diagnostics:\n{}",
                    json_output
                );
            }
        }
        diagnostics
    }
}

impl Solver for NewtonConfig {
    fn solve(&mut self, system: &mut System) -> Result<ConvergedState, SolverError> {
        let start_time = Instant::now();

        // Initialize diagnostics collection if verbose mode enabled
        let verbose_enabled = self.verbose_config.enabled && self.verbose_config.is_any_enabled();
        let mut diagnostics = if verbose_enabled {
            Some(ConvergenceDiagnostics::with_capacity(self.max_iterations))
        } else {
            None
        };

        tracing::info!(
            max_iterations = self.max_iterations,
            tolerance = self.tolerance,
            line_search = self.line_search,
            verbose = verbose_enabled,
            "Newton-Raphson solver starting"
        );

        let n_state = system.full_state_vector_len();
        let n_equations: usize = system
            .traverse_for_jacobian()
            .map(|(_, c, _)| c.n_equations())
            .sum::<usize>()
            + system.constraints().count()
            + system.coupling_residual_count()
            + 2 * system.saturated_controller_count()
            + system.mass_flow_closure_count();

        if n_state == 0 || n_equations == 0 {
            return Err(SolverError::InvalidSystem {
                message: "Empty system has no state variables or equations".to_string(),
            });
        }

        // Pre-allocate all buffers. A caller-supplied initial state MUST match
        // the full state length: a debug_assert would abort (violating zero-panic)
        // and a silent zeros fallback would solve a different problem. Fail cleanly.
        let mut state: Vec<f64> = match self.initial_state.as_ref() {
            Some(s) if s.len() == n_state => s.clone(),
            Some(s) => {
                return Err(SolverError::InvalidSystem {
                    message: format!(
                        "initial_state length {} does not match system state length {}",
                        s.len(),
                        n_state
                    ),
                });
            }
            None => vec![0.0; n_state],
        };
        let mut residuals: Vec<f64> = vec![0.0; n_equations];
        let mut jacobian_builder = JacobianBuilder::new();
        let mut divergence_count: usize = 0;
        let mut previous_norm: f64;
        let mut state_copy: Vec<f64> = vec![0.0; n_state]; // Pre-allocated for line search
        let mut new_residuals: Vec<f64> = vec![0.0; n_equations]; // Pre-allocated for line search
        let mut prev_iteration_state: Vec<f64> = vec![0.0; n_state]; // For convergence delta check

        // Pre-allocate best-state tracking buffer (Story 4.5 - AC: #5)
        let mut best_state: Vec<f64> = vec![0.0; n_state];
        let mut best_residual: f64;

        // Jacobian-freezing tracking state
        let mut jacobian_matrix = JacobianMatrix::zeros(n_equations, n_state);
        let mut frozen_count: usize = 0;
        let mut force_recompute: bool = true;

        // LinearSolver backend (FR3/Story 1.5): nalgebra by default, faer via
        // config/env dispatch. Factorizes at each fresh assembly and reuses
        // the factorization (with cached scalings) across frozen iterations.
        let mut linear_backend =
            crate::linear::make_linear_backend(None).map_err(|e| SolverError::InvalidSystem {
                message: format!("Failed to select linear backend: {e}"),
            })?;
        linear_backend
            .set_problem(n_equations, n_state)
            .map_err(|e| SolverError::InvalidSystem {
                message: format!("Failed to initialize linear backend: {e}"),
            })?;
        let mut delta_buf: Vec<f64> = vec![0.0; n_state];

        // Cached condition number (for verbose mode when Jacobian frozen)
        let mut cached_condition: Option<f64> = None;

        // Pre-compute clipping mask
        let clipping_mask: Vec<Option<(f64, f64)>> = (0..n_state)
            .map(|i| system.get_solver_bounds_for_state_index(i))
            .collect();

        // Initial residual computation. A recoverable domain violation at the
        // cold start is a typed outcome; fatal errors stay InvalidSystem.
        system
            .compute_residuals(&state, &mut residuals)
            .map_err(|e| match e {
                ComponentError::DomainViolation(violation) => {
                    SolverError::DomainViolation(violation)
                }
                fatal => SolverError::InvalidSystem {
                    message: format!("Failed to compute initial residuals: {:?}", fatal),
                },
            })?;

        let mut current_norm = Self::residual_norm(&residuals);
        best_state.copy_from_slice(&state);
        best_residual = current_norm;

        // Optional per-equation residual dump (debugging aid, controlled by env
        // var so production runs are unaffected). Prints the 8 largest-magnitude
        // residuals at the cold start so we can pinpoint which component / state
        // slot dominates the residual norm — exactly the NLPD "bottleneck
        // equation" diagnostic but at the solver entry point.
        if std::env::var("ENTROPYK_DUMP_RESIDUALS").is_ok() {
            let mut indexed: Vec<(usize, f64)> = residuals
                .iter()
                .enumerate()
                .map(|(i, &r)| (i, r.abs()))
                .collect();
            indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
            tracing::info!(
                residual_norm = current_norm,
                top_residuals = ?indexed
                    .iter()
                    .take(8)
                    .map(|(i, v)| format!("[{i}]={v:.4e}"))
                    .collect::<Vec<_>>()
                    .join(" "),
                "Initial residual breakdown (top 8 by magnitude)"
            );
        }

        tracing::debug!(iteration = 0, residual_norm = current_norm, "Initial state");

        // Check if already converged
        if current_norm < self.tolerance {
            let status = if !system.saturated_variables().is_empty() {
                ConvergenceStatus::ControlSaturation
            } else {
                ConvergenceStatus::Converged
            };

            if let Some(ref criteria) = self.convergence_criteria {
                let report = criteria.check(&state, None, &residuals, system);
                if report.is_globally_converged() {
                    tracing::info!(
                        iterations = 0,
                        final_residual = current_norm,
                        "Converged at initial state (criteria)"
                    );
                    return Ok(ConvergedState::with_report(
                        state,
                        0,
                        current_norm,
                        status,
                        report,
                        SimulationMetadata::new(system.input_hash()),
                    ));
                }
            } else {
                tracing::info!(
                    iterations = 0,
                    final_residual = current_norm,
                    "Converged at initial state"
                );
                return Ok(ConvergedState::new(
                    state,
                    0,
                    current_norm,
                    status,
                    SimulationMetadata::new(system.input_hash()),
                ));
            }
        }

        // Main Newton-Raphson iteration loop
        for iteration in 1..=self.max_iterations {
            prev_iteration_state.copy_from_slice(&state);

            // Check timeout
            if let Some(timeout) = self.timeout {
                if start_time.elapsed() > timeout {
                    tracing::info!(iteration, elapsed_ms = ?start_time.elapsed(), best_residual, "Solver timed out");
                    let failure_diagnostics = self.finalize_failure_diagnostics(
                        diagnostics.take(),
                        iteration - 1,
                        current_norm,
                        best_residual,
                        start_time.elapsed().as_millis() as u64,
                        cached_condition,
                        Some(state.clone()),
                    );
                    return self
                        .handle_timeout(&best_state, best_residual, iteration - 1, timeout, system)
                        .map_err(|err| err.with_optional_diagnostics(failure_diagnostics));
                }
            }

            // Jacobian Assembly / Freeze Decision
            let should_recompute = if let Some(ref freeze_cfg) = self.jacobian_freezing {
                if force_recompute {
                    true
                } else if frozen_count >= freeze_cfg.max_frozen_iters {
                    tracing::debug!(iteration, frozen_count, "Jacobian freeze limit reached");
                    true
                } else {
                    false
                }
            } else {
                true
            };

            let jacobian_frozen_this_iter = !should_recompute;

            if should_recompute {
                // Fresh Jacobian assembly (in-place update)
                jacobian_builder.clear();
                if self.use_numerical_jacobian {
                    // Numerical Jacobian via finite differences
                    let compute_residuals_fn = |s: &[f64], r: &mut [f64]| {
                        let s_vec = s.to_vec();
                        let mut r_vec = vec![0.0; r.len()];
                        let result = system.compute_residuals(&s_vec, &mut r_vec);
                        r.copy_from_slice(&r_vec);
                        result.map(|_| ()).map_err(|e| format!("{:?}", e))
                    };
                    let jm =
                        JacobianMatrix::numerical(compute_residuals_fn, &state, &residuals, 1e-5)
                            .map_err(|e| SolverError::InvalidSystem {
                            message: format!("Failed to compute numerical Jacobian: {}", e),
                        })?;
                    jacobian_matrix.as_matrix_mut().copy_from(jm.as_matrix());
                } else {
                    system
                        .assemble_jacobian(&state, &mut jacobian_builder)
                        .map_err(|e| SolverError::InvalidSystem {
                            message: format!("Failed to assemble Jacobian: {:?}", e),
                        })?;
                    jacobian_matrix.update_from_builder(jacobian_builder.entries());
                };

                // Feed the LinearSolver backend with the freshly assembled
                // matrix (assembles + factorizes; frozen iterations skip this).
                linear_backend
                    .set_matrix(jacobian_matrix.as_matrix())
                    .map_err(|e| SolverError::InvalidSystem {
                        message: format!("Failed to factorize Jacobian: {e}"),
                    })?;

                frozen_count = 0;
                force_recompute = false;

                // Compute and cache condition number if verbose mode enabled
                if verbose_enabled && self.verbose_config.log_jacobian_condition {
                    let cond = jacobian_matrix.estimate_condition_number();
                    cached_condition = cond;
                    if let Some(c) = cond {
                        tracing::info!(
                            iteration,
                            condition_number = c,
                            "Jacobian condition number"
                        );
                        if c > 1e10 {
                            tracing::warn!(
                                iteration,
                                condition_number = c,
                                "Ill-conditioned Jacobian detected (κ > 1e10)"
                            );
                        }
                    }
                }

                tracing::debug!(iteration, "Fresh Jacobian computed");
            } else {
                frozen_count += 1;
                tracing::debug!(iteration, frozen_count, "Reusing frozen Jacobian");
            }

            // Solve J·Δx = -r through the LinearSolver backend (FR3). The
            // factorization is fresh when assembled above and reused when the
            // Jacobian is frozen. Non-square systems keep the legacy
            // least-squares path.
            let delta = if n_equations != n_state {
                match jacobian_matrix.solve(&residuals) {
                    Some(d) => d,
                    None => {
                        let failure_diagnostics = self.finalize_failure_diagnostics(
                            diagnostics.take(),
                            iteration,
                            current_norm,
                            best_residual,
                            start_time.elapsed().as_millis() as u64,
                            cached_condition,
                            Some(state.clone()),
                        );
                        return Err(SolverError::Divergence {
                            reason: "Jacobian is singular".to_string(),
                        }
                        .with_optional_diagnostics(failure_diagnostics));
                    }
                }
            } else {
                for (d, r) in delta_buf.iter_mut().zip(residuals.iter()) {
                    *d = -*r;
                }
                match linear_backend.solve_in_place(&mut delta_buf) {
                    Ok(()) => delta_buf.clone(),
                    Err(_) => {
                        let failure_diagnostics = self.finalize_failure_diagnostics(
                            diagnostics.take(),
                            iteration,
                            current_norm,
                            best_residual,
                            start_time.elapsed().as_millis() as u64,
                            cached_condition,
                            Some(state.clone()),
                        );
                        return Err(SolverError::Divergence {
                            reason: "Jacobian is singular".to_string(),
                        }
                        .with_optional_diagnostics(failure_diagnostics));
                    }
                }
            };

            // Apply step with optional line search
            let mut recoverable_events = 0u32;
            let alpha = if self.line_search {
                match self.line_search(
                    system,
                    &mut state,
                    &delta,
                    &residuals,
                    current_norm,
                    &mut state_copy,
                    &mut new_residuals,
                    &clipping_mask,
                    &mut recoverable_events,
                )? {
                    Some(a) => a,
                    None => {
                        let failure_diagnostics = self.finalize_failure_diagnostics(
                            diagnostics.take(),
                            iteration,
                            current_norm,
                            best_residual,
                            start_time.elapsed().as_millis() as u64,
                            cached_condition,
                            Some(state.clone()),
                        );
                        // Backtracks exhausted on recoverable domain violations:
                        // a typed outcome, not an untyped stall (AC #2).
                        if recoverable_events > 0 {
                            return Err(SolverError::DomainViolation(DomainViolation {
                                component: None,
                                detail:
                                    "line search exhausted after recoverable domain violation(s)"
                                        .into(),
                            })
                            .with_optional_diagnostics(failure_diagnostics));
                        }
                        return Err(SolverError::Divergence {
                            reason: "Line search failed".to_string(),
                        }
                        .with_optional_diagnostics(failure_diagnostics));
                    }
                }
            } else {
                apply_newton_step(&mut state, &delta, &clipping_mask, 1.0);
                1.0
            };

            system
                .compute_residuals(&state, &mut residuals)
                .map_err(|e| match e {
                    ComponentError::DomainViolation(violation) => {
                        SolverError::DomainViolation(violation)
                    }
                    fatal => SolverError::InvalidSystem {
                        message: format!("Failed to compute residuals: {:?}", fatal),
                    },
                })?;

            previous_norm = current_norm;
            current_norm = Self::residual_norm(&residuals);

            // Compute delta norm for diagnostics
            let delta_norm: f64 = state
                .iter()
                .zip(prev_iteration_state.iter())
                .map(|(s, p)| (s - p).powi(2))
                .sum::<f64>()
                .sqrt();

            if current_norm < best_residual {
                best_state.copy_from_slice(&state);
                best_residual = current_norm;
                tracing::debug!(iteration, best_residual, "Best state updated");
            }

            // Jacobian-freeze feedback
            if let Some(ref freeze_cfg) = self.jacobian_freezing {
                if previous_norm > 0.0
                    && current_norm / previous_norm >= (1.0 - freeze_cfg.threshold)
                {
                    if frozen_count > 0 || !force_recompute {
                        tracing::debug!(
                            iteration,
                            current_norm,
                            previous_norm,
                            "Unfreezing Jacobian"
                        );
                    }
                    force_recompute = true;
                    frozen_count = 0;
                }
            }

            // Verbose mode: Log iteration residuals
            if verbose_enabled && self.verbose_config.log_residuals {
                tracing::info!(
                    iteration,
                    residual_norm = current_norm,
                    delta_norm = delta_norm,
                    alpha = alpha,
                    jacobian_frozen = jacobian_frozen_this_iter,
                    "Newton iteration"
                );
            }

            // Collect iteration diagnostics
            if let Some(ref mut diag) = diagnostics {
                let (max_residual_index, max_residual) = dominant_residual(&residuals);
                diag.push_iteration(IterationDiagnostics {
                    iteration,
                    residual_norm: current_norm,
                    delta_norm,
                    alpha: Some(alpha),
                    jacobian_frozen: jacobian_frozen_this_iter,
                    jacobian_condition: cached_condition,
                    max_residual_index,
                    max_residual,
                    recoverable_events,
                });
            }

            tracing::debug!(
                iteration,
                residual_norm = current_norm,
                alpha,
                "Newton iteration complete"
            );

            // Check convergence
            let converged = if let Some(ref criteria) = self.convergence_criteria {
                let report =
                    criteria.check(&state, Some(&prev_iteration_state), &residuals, system);
                if report.is_globally_converged() {
                    let status = if !system.saturated_variables().is_empty() {
                        ConvergenceStatus::ControlSaturation
                    } else {
                        ConvergenceStatus::Converged
                    };

                    // Finalize diagnostics
                    if let Some(ref mut diag) = diagnostics {
                        diag.iterations = iteration;
                        diag.final_residual = current_norm;
                        diag.best_residual = best_residual;
                        diag.converged = true;
                        diag.timing_ms = start_time.elapsed().as_millis() as u64;
                        diag.jacobian_condition_final = cached_condition;
                        diag.final_solver = Some(SolverType::NewtonRaphson);

                        if self.verbose_config.log_residuals {
                            tracing::info!("{}", diag.summary());
                        }
                    }

                    tracing::info!(
                        iterations = iteration,
                        final_residual = current_norm,
                        "Converged (criteria)"
                    );
                    let result = ConvergedState::with_report(
                        state,
                        iteration,
                        current_norm,
                        status,
                        report,
                        SimulationMetadata::new(system.input_hash()),
                    );
                    return Ok(if let Some(d) = diagnostics {
                        ConvergedState {
                            diagnostics: Some(d),
                            ..result
                        }
                    } else {
                        result
                    });
                }
                false
            } else {
                current_norm < self.tolerance
            };

            if converged {
                let status = if !system.saturated_variables().is_empty() {
                    ConvergenceStatus::ControlSaturation
                } else {
                    ConvergenceStatus::Converged
                };

                // Finalize diagnostics
                if let Some(ref mut diag) = diagnostics {
                    diag.iterations = iteration;
                    diag.final_residual = current_norm;
                    diag.best_residual = best_residual;
                    diag.converged = true;
                    diag.timing_ms = start_time.elapsed().as_millis() as u64;
                    diag.jacobian_condition_final = cached_condition;
                    diag.final_solver = Some(SolverType::NewtonRaphson);

                    if self.verbose_config.log_residuals {
                        tracing::info!("{}", diag.summary());
                    }
                }

                tracing::info!(
                    iterations = iteration,
                    final_residual = current_norm,
                    "Converged"
                );
                let result = ConvergedState::new(
                    state,
                    iteration,
                    current_norm,
                    status,
                    SimulationMetadata::new(system.input_hash()),
                );
                return Ok(if let Some(d) = diagnostics {
                    ConvergedState {
                        diagnostics: Some(d),
                        ..result
                    }
                } else {
                    result
                });
            }

            if let Some(err) =
                self.check_divergence(current_norm, previous_norm, &mut divergence_count)
            {
                tracing::warn!(
                    iteration,
                    residual_norm = current_norm,
                    "Divergence detected"
                );
                let failure_diagnostics = self.finalize_failure_diagnostics(
                    diagnostics.take(),
                    iteration,
                    current_norm,
                    best_residual,
                    start_time.elapsed().as_millis() as u64,
                    cached_condition,
                    Some(state.clone()),
                );
                return Err(err.with_optional_diagnostics(failure_diagnostics));
            }
        }

        // Non-convergence: dump diagnostics if enabled
        let failure_diagnostics = self.finalize_failure_diagnostics(
            diagnostics.take(),
            self.max_iterations,
            current_norm,
            best_residual,
            start_time.elapsed().as_millis() as u64,
            cached_condition,
            Some(state.clone()),
        );

        tracing::warn!(
            max_iterations = self.max_iterations,
            final_residual = current_norm,
            "Did not converge"
        );
        Err(SolverError::NonConvergence {
            iterations: self.max_iterations,
            final_residual: current_norm,
        }
        .with_optional_diagnostics(failure_diagnostics))
    }

    fn with_timeout(mut self, timeout: Duration) -> Self {
        self.timeout = Some(timeout);
        self
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::solver::Solver;
    use crate::system::System;
    use std::time::Duration;

    #[test]
    fn test_newton_config_with_timeout() {
        let cfg = NewtonConfig::default().with_timeout(Duration::from_millis(100));
        assert_eq!(cfg.timeout, Some(Duration::from_millis(100)));
    }

    #[test]
    fn test_newton_config_default() {
        let cfg = NewtonConfig::default();
        assert_eq!(cfg.max_iterations, 100);
        assert!(cfg.tolerance > 0.0 && cfg.tolerance < 1e-3);
    }

    #[test]
    fn test_newton_solver_trait_object() {
        let mut boxed: Box<dyn Solver> = Box::new(NewtonConfig::default());
        let mut system = System::new();
        system.finalize().unwrap();
        assert!(boxed.solve(&mut system).is_err());
    }
}
