//! Newton-homotopy continuation solver for robust cold starts.
//!
//! This module provides [`HomotopyConfig`], a globally-convergent continuation
//! solver that improves on a naive cold start without requiring a database of
//! previous solutions (the IPM BOLT `GLBL`/`iPRVS` approach) or any manual
//! tuning.
//!
//! # The Newton homotopy
//!
//! Given the target system `F(x) = 0` and an arbitrary initial guess `x₀`, define
//! the homotopy
//!
//! ```text
//! H(x, λ) = F(x) − (1 − λ) · F(x₀),   λ ∈ [0, 1]
//! ```
//!
//! At `λ = 0`, `H(x₀, 0) = F(x₀) − F(x₀) = 0`, so the initial guess is an
//! **exact** solution of the deformed system. At `λ = 1`, `H(x, 1) = F(x)`, the
//! real system. The solver walks `λ` from 0 to 1, solving `H(·, λ) = 0` with an
//! inner Newton iteration at each step and using the previous converged point as
//! the next initial guess.
//!
//! # Why it reuses the analytic Jacobian unchanged
//!
//! The subtracted term `(1 − λ)·F(x₀)` is **constant in `x`**, so
//!
//! ```text
//! ∂H/∂x = ∂F/∂x = J(x)
//! ```
//!
//! The inner Newton step therefore uses the exact, component-wise analytic
//! Jacobian assembled by [`System::assemble_jacobian`] — no finite differences
//! and no changes to any component are required. This keeps Entropyk's
//! structural advantage over finite-difference solvers (IPM eKINSOL) while adding
//! cold-start robustness. A `use_numerical_jacobian` flag is provided for parity
//! with [`crate::strategies::NewtonConfig`] when a component's analytic Jacobian
//! is unavailable.
//!
//! # Adaptive step control
//!
//! The `λ` increment starts at `1 / initial_steps` and is **halved** whenever an
//! inner Newton solve fails to converge, retrying from the last good `λ`. On
//! success the increment is gently grown again. This predictor–corrector scheme
//! automatically takes small steps through difficult regions (phase boundaries,
//! stiff correlations) and large steps through easy ones.

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

use crate::jacobian::JacobianMatrix;
use crate::metadata::SimulationMetadata;
use crate::solver::{
    apply_newton_step, dominant_residual, ConvergedState, ConvergenceDiagnostics,
    ConvergenceStatus, IterationDiagnostics, Solver, SolverError, SolverType,
};
use crate::system::System;
use entropyk_components::JacobianBuilder;

/// Configuration for the Newton-homotopy continuation solver.
///
/// Solves `F(x) = 0` by continuation on the homotopy
/// `H(x, λ) = F(x) − (1 − λ)·F(x₀)` from `λ = 0` (where `x₀` is exact) to
/// `λ = 1` (the real system).
#[derive(Debug, Clone, PartialEq)]
pub struct HomotopyConfig {
    /// Initial number of `λ` subdivisions. The starting step is `1 / initial_steps`.
    /// Default: 10.
    pub initial_steps: usize,
    /// Maximum Newton iterations allowed for each inner `λ` solve. Default: 50.
    pub inner_max_iterations: usize,
    /// Convergence tolerance (L2 residual norm) for each inner `λ` solve.
    /// Default: 1e-8.
    pub inner_tolerance: f64,
    /// Final convergence tolerance (L2 residual norm) checked at `λ = 1`.
    /// Default: 1e-6.
    pub tolerance: f64,
    /// Smallest allowed `λ` increment before the solver gives up. Default: 1e-4.
    pub min_lambda_step: f64,
    /// Divergence guard: inner residual norm above this aborts the inner solve.
    /// Default: 1e12.
    pub divergence_threshold: f64,
    /// Use a finite-difference Jacobian instead of the analytic one. Default: false.
    pub use_numerical_jacobian: bool,
    /// Relative step for the finite-difference Jacobian. Default: 1e-5.
    pub numerical_epsilon: f64,
    /// Optional overall time budget.
    pub timeout: Option<Duration>,
    /// Initial guess `x₀`. When `None`, a zero vector is used.
    pub initial_state: Option<Vec<f64>>,
}

impl Default for HomotopyConfig {
    fn default() -> Self {
        Self {
            initial_steps: 10,
            inner_max_iterations: 50,
            inner_tolerance: 1e-8,
            tolerance: 1e-6,
            min_lambda_step: 1e-4,
            divergence_threshold: 1e12,
            use_numerical_jacobian: false,
            numerical_epsilon: 1e-5,
            timeout: None,
            initial_state: None,
        }
    }
}

impl HomotopyConfig {
    /// Sets the initial guess `x₀` for the continuation.
    pub fn with_initial_state(mut self, state: Vec<f64>) -> Self {
        self.initial_state = Some(state);
        self
    }

    /// Sets the initial number of `λ` subdivisions.
    pub fn with_initial_steps(mut self, steps: usize) -> Self {
        self.initial_steps = steps.max(1);
        self
    }

    /// Selects the finite-difference Jacobian (analytic is the default).
    pub fn with_numerical_jacobian(mut self, enabled: bool) -> Self {
        self.use_numerical_jacobian = enabled;
        self
    }

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

    fn failure_diagnostics(
        &self,
        iterations: usize,
        final_residual: f64,
        residuals: &[f64],
        elapsed_ms: u64,
    ) -> Option<ConvergenceDiagnostics> {
        if iterations == 0 {
            return None;
        }

        let (max_residual_index, max_residual) = dominant_residual(residuals);
        let mut diagnostics = ConvergenceDiagnostics::with_capacity(1);
        diagnostics.iterations = iterations;
        diagnostics.final_residual = final_residual;
        diagnostics.best_residual = final_residual;
        diagnostics.converged = false;
        diagnostics.timing_ms = elapsed_ms;
        diagnostics.final_solver = Some(SolverType::Homotopy);
        diagnostics.push_iteration(IterationDiagnostics {
            iteration: iterations,
            residual_norm: final_residual,
            delta_norm: 0.0,
            alpha: Some(1.0),
            jacobian_frozen: false,
            jacobian_condition: None,
            max_residual_index,
            max_residual,
            recoverable_events: 0,
        });
        Some(diagnostics)
    }

    /// Runs the inner Newton iteration for `H(x, λ) = F(x) − (1 − λ)·r0 = 0`.
    ///
    /// Mutates `state` in place. Returns `Ok(iterations)` if the inner system
    /// converged below `inner_tolerance`, or `Err(())` if it diverged, the
    /// Jacobian was singular, or the iteration budget was exhausted. On failure
    /// the caller restores the previous good state.
    #[allow(clippy::too_many_arguments)]
    fn inner_newton(
        &self,
        system: &mut System,
        state: &mut [f64],
        r0: &[f64],
        lambda: f64,
        clipping_mask: &[Option<(f64, f64)>],
        residuals: &mut Vec<f64>,
        residuals_h: &mut Vec<f64>,
        jacobian: &mut JacobianMatrix,
        jacobian_builder: &mut JacobianBuilder,
    ) -> Result<usize, Option<SolverError>> {
        let offset = 1.0 - lambda;

        for k in 0..self.inner_max_iterations {
            // Evaluate F(x) and form the homotopy residual H = F − (1 − λ)·r0.
            // Recoverable domain violations shrink λ (KINSOL convention, Story
            // 1.3); fatal errors abort the whole homotopy immediately instead
            // of burning the λ budget on a bug.
            if let Err(e) = system.compute_residuals(state, residuals) {
                return Err(if e.is_recoverable() {
                    None
                } else {
                    Some(SolverError::InvalidSystem {
                        message: format!("Failed to compute residuals: {:?}", e),
                    })
                });
            }
            for i in 0..residuals.len() {
                residuals_h[i] = residuals[i] - offset * r0[i];
            }

            let norm = Self::residual_norm(residuals_h.as_slice());
            if norm < self.inner_tolerance {
                return Ok(k);
            }
            if !norm.is_finite() || norm > self.divergence_threshold {
                return Err(None);
            }

            // ∂H/∂x = ∂F/∂x, so the Jacobian of F is used unchanged.
            if self.use_numerical_jacobian {
                let eps = self.numerical_epsilon;
                let compute = |s: &[f64], r: &mut [f64]| {
                    let s_vec = s.to_vec();
                    let mut r_vec = vec![0.0; r.len()];
                    let res = system.compute_residuals(&s_vec, &mut r_vec);
                    r.copy_from_slice(&r_vec);
                    res.map(|_| ()).map_err(|e| format!("{:?}", e))
                };
                match JacobianMatrix::numerical(compute, state, residuals.as_slice(), eps) {
                    Ok(jm) => jacobian.as_matrix_mut().copy_from(jm.as_matrix()),
                    Err(_) => return Err(None),
                }
            } else {
                jacobian_builder.clear();
                if let Err(e) = system.assemble_jacobian(state, jacobian_builder) {
                    return Err(if e.is_recoverable() {
                        None
                    } else {
                        Some(SolverError::InvalidSystem {
                            message: format!("Failed to assemble Jacobian: {:?}", e),
                        })
                    });
                }
                jacobian.update_from_builder(jacobian_builder.entries());
            }

            // Solve J·Δx = −H (the solve routine negates the supplied residual).
            let delta = match jacobian.solve(residuals_h.as_slice()) {
                Some(d) => d,
                None => return Err(None),
            };

            apply_newton_step(state, &delta, clipping_mask, 1.0);
        }

        // Final convergence check after the last step.
        if let Err(e) = system.compute_residuals(state, residuals) {
            return Err(if e.is_recoverable() {
                None
            } else {
                Some(SolverError::InvalidSystem {
                    message: format!("Failed to compute residuals: {:?}", e),
                })
            });
        }
        for i in 0..residuals.len() {
            residuals_h[i] = residuals[i] - offset * r0[i];
        }
        if Self::residual_norm(residuals_h.as_slice()) < self.inner_tolerance {
            Ok(self.inner_max_iterations)
        } else {
            Err(None)
        }
    }
}

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

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

        // Working buffers (allocated once, reused across every λ step).
        // A caller-supplied initial guess MUST match the system size: silently
        // substituting zeros would hide a caller bug behind an opaque later failure.
        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![0.0; n_equations];
        let mut residuals_h = vec![0.0; n_equations];
        let mut jacobian = JacobianMatrix::zeros(n_equations, n_state);
        let mut jacobian_builder = JacobianBuilder::new();
        let mut state_saved = vec![0.0; n_state];

        let clipping_mask: Vec<Option<(f64, f64)>> = (0..n_state)
            .map(|i| system.get_bounds_for_state_index(i))
            .collect();

        // r0 = F(x0). By construction H(x0, 0) = 0, so x0 is exact at λ = 0.
        system
            .compute_residuals(&state, &mut residuals)
            .map_err(|e| SolverError::InvalidSystem {
                message: format!("Failed to compute initial residuals: {:?}", e),
            })?;
        let r0 = residuals.clone();
        let initial_norm = Self::residual_norm(&r0);

        // F(x0) must be finite for the deformation H(x,λ)=F(x)-(1-λ)F(x0) to be
        // well defined. A non-finite r0 (e.g. a zero (P,h) cold start hitting the
        // fluid backend) would make every continuation step doomed; fail early
        // with an actionable message instead of running the whole loop.
        if !initial_norm.is_finite() {
            return Err(SolverError::InvalidSystem {
                message: "Initial residual F(x0) is non-finite; the initial guess is infeasible \
                          for the fluid backend (provide a physical initial_state)"
                    .to_string(),
            });
        }

        // Already solved? Skip the continuation entirely.
        if initial_norm < self.tolerance {
            return Ok(ConvergedState::new(
                state,
                0,
                initial_norm,
                ConvergenceStatus::Converged,
                SimulationMetadata::new(system.input_hash()),
            ));
        }

        let max_step = 4.0 / self.initial_steps.max(1) as f64;
        // Guard against a non-positive min step (e.g. struct-literal misconfig),
        // which would otherwise let dlambda shrink forever and hang the solver.
        let min_lambda_step = self.min_lambda_step.max(1e-12);
        let mut lambda = 0.0_f64;
        let mut dlambda = 1.0 / self.initial_steps.max(1) as f64;
        let mut total_iterations = 0usize;

        while lambda < 1.0 {
            if let Some(timeout) = self.timeout {
                if start_time.elapsed() > timeout {
                    let compute_ok = system.compute_residuals(&state, &mut residuals).is_ok();
                    let final_residual = if compute_ok {
                        Self::residual_norm(&residuals)
                    } else {
                        f64::INFINITY
                    };
                    let diagnostics = self.failure_diagnostics(
                        total_iterations,
                        final_residual,
                        if compute_ok { &residuals } else { &[] },
                        start_time.elapsed().as_millis() as u64,
                    );
                    return Err(SolverError::Timeout {
                        timeout_ms: timeout.as_millis() as u64,
                    }
                    .with_optional_diagnostics(diagnostics));
                }
            }

            let target = (lambda + dlambda).min(1.0);
            state_saved.copy_from_slice(&state);

            match self.inner_newton(
                system,
                &mut state,
                &r0,
                target,
                &clipping_mask,
                &mut residuals,
                &mut residuals_h,
                &mut jacobian,
                &mut jacobian_builder,
            ) {
                Ok(iters) => {
                    total_iterations += iters;
                    lambda = target;
                    // Step succeeded: gently grow the increment for the next step.
                    dlambda = (dlambda * 1.5).min(max_step);
                }
                Err(None) => {
                    // Step failed (recoverable): restore and halve the increment, then retry.
                    state.copy_from_slice(&state_saved);
                    dlambda *= 0.5;
                    if dlambda < min_lambda_step {
                        let compute_ok = system.compute_residuals(&state, &mut residuals).is_ok();
                        let final_residual = if compute_ok {
                            Self::residual_norm(&residuals)
                        } else {
                            f64::INFINITY
                        };
                        let diagnostics = self.failure_diagnostics(
                            total_iterations,
                            final_residual,
                            if compute_ok { &residuals } else { &[] },
                            start_time.elapsed().as_millis() as u64,
                        );
                        return Err(SolverError::NonConvergence {
                            iterations: total_iterations,
                            final_residual,
                        }
                        .with_optional_diagnostics(diagnostics));
                    }
                }
                Err(Some(fatal)) => {
                    // Fatal evaluation error: do not burn the λ budget on it.
                    return Err(fatal);
                }
            }
        }

        // At λ = 1, H == F: verify the real system is actually solved.
        system
            .compute_residuals(&state, &mut residuals)
            .map_err(|e| SolverError::InvalidSystem {
                message: format!("Failed to compute final residuals: {:?}", e),
            })?;
        let final_norm = Self::residual_norm(&residuals);

        if final_norm < self.tolerance {
            let status = if !system.saturated_variables().is_empty() {
                ConvergenceStatus::ControlSaturation
            } else {
                ConvergenceStatus::Converged
            };
            Ok(ConvergedState::new(
                state,
                total_iterations,
                final_norm,
                status,
                SimulationMetadata::new(system.input_hash()),
            ))
        } else {
            let diagnostics = self.failure_diagnostics(
                total_iterations,
                final_norm,
                &residuals,
                start_time.elapsed().as_millis() as u64,
            );
            Err(SolverError::NonConvergence {
                iterations: total_iterations,
                final_residual: final_norm,
            }
            .with_optional_diagnostics(diagnostics))
        }
    }

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

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

    #[test]
    fn test_homotopy_default_config() {
        let cfg = HomotopyConfig::default();
        assert_eq!(cfg.initial_steps, 10);
        assert!(!cfg.use_numerical_jacobian);
        assert!((cfg.tolerance - 1e-6).abs() < 1e-15);
    }

    #[test]
    fn test_homotopy_builders() {
        let cfg = HomotopyConfig::default()
            .with_initial_steps(20)
            .with_numerical_jacobian(true)
            .with_initial_state(vec![1.0, 2.0]);
        assert_eq!(cfg.initial_steps, 20);
        assert!(cfg.use_numerical_jacobian);
        assert_eq!(cfg.initial_state, Some(vec![1.0, 2.0]));
    }

    #[test]
    fn test_homotopy_initial_steps_floor_is_one() {
        let cfg = HomotopyConfig::default().with_initial_steps(0);
        assert_eq!(cfg.initial_steps, 1);
    }

    #[test]
    fn test_homotopy_residual_norm() {
        assert!((HomotopyConfig::residual_norm(&[3.0, 4.0]) - 5.0).abs() < 1e-12);
    }

    #[test]
    fn test_homotopy_empty_system_errors() {
        let mut system = System::new();
        system.finalize().unwrap();
        let mut solver = HomotopyConfig::default();
        assert!(solver.solve(&mut system).is_err());
    }

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