//! 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, ConvergedState, ConvergenceStatus, JacobianFreezingConfig, Solver,
    SolverError, TimeoutConfig,
};
use crate::system::System;
use entropyk_components::JacobianBuilder;

/// 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>,
}

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,
        }
    }
}

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
    }

    /// 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.
    #[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)>],
    ) -> Option<f64> {
        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 system.compute_residuals(state, new_residuals).is_err() {
                state.copy_from_slice(state_copy);
                alpha *= 0.5;
                continue;
            }

            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 Some(alpha);
            }

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

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

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

        tracing::info!(
            max_iterations = self.max_iterations,
            tolerance = self.tolerance,
            line_search = self.line_search,
            "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();

        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
        let mut state: Vec<f64> = self
            .initial_state
            .as_ref()
            .map(|s| {
                debug_assert_eq!(s.len(), n_state, "initial_state length mismatch");
                if s.len() == n_state { s.clone() } else { vec![0.0; n_state] }
            })
            .unwrap_or_else(|| 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;

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

        // Initial residual computation
        system
            .compute_residuals(&state, &mut residuals)
            .map_err(|e| SolverError::InvalidSystem {
                message: format!("Failed to compute initial residuals: {:?}", e),
            })?;

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

        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");
                    return self.handle_timeout(&best_state, best_residual, iteration - 1, timeout, system);
                }
            }

            // 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
            };

            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());
                };

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

            // Solve J·Δx = -r
            let delta = match jacobian_matrix.solve(&residuals) {
                Some(d) => d,
                None => {
                    return Err(SolverError::Divergence {
                        reason: "Jacobian is singular".to_string(),
                    });
                }
            };

            // Apply step with optional line search
            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,
                ) {
                    Some(a) => a,
                    None => {
                        return Err(SolverError::Divergence {
                            reason: "Line search failed".to_string(),
                        });
                    }
                }
            } else {
                apply_newton_step(&mut state, &delta, &clipping_mask, 1.0);
                1.0
            };

            system.compute_residuals(&state, &mut residuals)
                .map_err(|e| SolverError::InvalidSystem {
                    message: format!("Failed to compute residuals: {:?}", e),
                })?;

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

            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;
                }
            }

            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
                    };
                    tracing::info!(iterations = iteration, final_residual = current_norm, "Converged (criteria)");
                    return Ok(ConvergedState::with_report(
                        state, iteration, current_norm, status, report, SimulationMetadata::new(system.input_hash()),
                    ));
                }
                false
            } else {
                current_norm < self.tolerance
            };

            if converged {
                let status = if !system.saturated_variables().is_empty() {
                    ConvergenceStatus::ControlSaturation
                } else {
                    ConvergenceStatus::Converged
                };
                tracing::info!(iterations = iteration, final_residual = current_norm, "Converged");
                return Ok(ConvergedState::new(
                    state, iteration, current_norm, status, SimulationMetadata::new(system.input_hash()),
                ));
            }

            if let Some(err) = self.check_divergence(current_norm, previous_norm, &mut divergence_count) {
                tracing::warn!(iteration, residual_norm = current_norm, "Divergence detected");
                return Err(err);
            }
        }

        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,
        })
    }

    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());
    }
}