//! Solver trait abstraction and core types for thermodynamic system solving.
//!
//! Provides the `Solver` trait (object-safe interface) and `SolverStrategy` enum
//! (zero-cost static dispatch) for solver strategies.

use std::time::Duration;
use thiserror::Error;

use crate::criteria::ConvergenceReport;
use crate::metadata::SimulationMetadata;
use crate::system::System;

// Re-export strategies and their types
pub use crate::strategies::{
    FallbackConfig, FallbackSolver, NewtonConfig, PicardConfig, SolverStrategy,
};

// ─────────────────────────────────────────────────────────────────────────────
// Error types
// ─────────────────────────────────────────────────────────────────────────────

/// Errors that can occur during solving.
///
/// All variants carry enough context for the caller to diagnose the failure
/// and decide on a recovery strategy (e.g. retry with a different solver,
/// return best-known state, or propagate to the user).
#[derive(Error, Debug, Clone, PartialEq)]
pub enum SolverError {
    /// The solver exhausted its iteration budget without converging.
    ///
    /// The `final_residual` is the $\ell_2$ norm of the residual vector at the
    /// last iteration. Use this to assess how close the solver got.
    #[error(
        "Solver did not converge after {iterations} iterations \
         (final residual norm: {final_residual:.3e})"
    )]
    NonConvergence {
        /// Number of iterations performed before giving up.
        iterations: usize,
        /// $\ell_2$ norm of the residual vector at the last iteration.
        final_residual: f64,
    },

    /// The solver exceeded its configured time budget.
    ///
    /// The `timeout_ms` field records the budget that was exceeded.
    /// Story 4.5 (Time-Budgeted Solving) will return the best-known state
    /// alongside this error.
    #[error("Solver timed out after {timeout_ms} ms")]
    Timeout {
        /// The configured timeout in milliseconds.
        timeout_ms: u64,
    },

    /// The iteration sequence diverged (residual norm is growing unboundedly).
    ///
    /// This typically indicates a poor initial guess or an ill-conditioned system.
    #[error("Solver diverged: {reason}")]
    Divergence {
        /// Human-readable description of why divergence was detected.
        reason: String,
    },

    /// The system description is invalid and cannot be solved.
    ///
    /// Examples: empty system, singular Jacobian at the initial point,
    /// inconsistent constraints.
    #[error("Invalid system: {message}")]
    InvalidSystem {
        /// Human-readable description of the system defect.
        message: String,
    },

    /// Post-solve validation failed (e.g., mass or energy balance violation).
    #[error("Validation failed: mass error = {mass_error:.3e} kg/s, energy error = {energy_error:.3e} W")]
    Validation {
        /// Mass balance error in kg/s
        mass_error: f64,
        /// Energy balance error in W
        energy_error: f64,
    },
}

// ─────────────────────────────────────────────────────────────────────────────
// Result types
// ─────────────────────────────────────────────────────────────────────────────

/// Convergence status of a completed solve.
#[derive(Debug, Clone, PartialEq)]
pub enum ConvergenceStatus {
    /// The solver converged within tolerance and iteration budget.
    Converged,
    /// The solver stopped due to timeout but returned the best-known state.
    /// (Used by Story 4.5 — Time-Budgeted Solving.)
    TimedOutWithBestState,
    /// The solver converged, but one or more control variables saturated at bounds.
    /// (Used by Story 5.2 — Bounded Control Variables)
    ControlSaturation,
}

/// The result of a successful (or best-effort) solve.
///
/// Contains the final state vector and convergence metadata.
#[derive(Debug, Clone)]
pub struct ConvergedState {
    /// Final state vector $\mathbf{x}^*$ at convergence.
    ///
    /// Layout mirrors `System::state_vector_len()`:
    /// `[P_edge0, h_edge0, P_edge1, h_edge1, ...]`.
    pub state: Vec<f64>,

    /// Number of iterations performed.
    pub iterations: usize,

    /// $\ell_2$ norm of the residual vector at the final state.
    pub final_residual: f64,

    /// Whether the solver fully converged or returned a best-effort state.
    pub status: ConvergenceStatus,

    /// Optional multi-circuit convergence report (Story 4.7).
    ///
    /// `Some(report)` when [`ConvergenceCriteria`] was set on the solver.
    /// `None` when using the raw tolerance check (backward-compatible default).
    pub convergence_report: Option<ConvergenceReport>,

    /// Traceability metadata for reproducibility.
    pub metadata: SimulationMetadata,
}

impl ConvergedState {
    /// Creates a new `ConvergedState` without a convergence report (backward-compat).
    pub fn new(
        state: Vec<f64>,
        iterations: usize,
        final_residual: f64,
        status: ConvergenceStatus,
        metadata: SimulationMetadata,
    ) -> Self {
        Self {
            state,
            iterations,
            final_residual,
            status,
            convergence_report: None,
            metadata,
        }
    }

    /// Creates a `ConvergedState` with an attached convergence report.
    pub fn with_report(
        state: Vec<f64>,
        iterations: usize,
        final_residual: f64,
        status: ConvergenceStatus,
        report: ConvergenceReport,
        metadata: SimulationMetadata,
    ) -> Self {
        Self {
            state,
            iterations,
            final_residual,
            status,
            convergence_report: Some(report),
            metadata,
        }
    }

    /// Returns `true` if the solver fully converged (not just best-effort).
    pub fn is_converged(&self) -> bool {
        matches!(
            self.status,
            ConvergenceStatus::Converged | ConvergenceStatus::ControlSaturation
        )
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Solver trait
// ─────────────────────────────────────────────────────────────────────────────

/// Object-safe interface for all solver strategies.
///
/// Implementors must be able to solve a [`System`] of equations and optionally
/// respect a time budget. The trait is intentionally minimal and object-safe so
/// that user-provided custom solvers can be used via `Box<dyn Solver>`.
///
/// # Object Safety
///
/// This trait is object-safe:
/// - No generic methods.
/// - No `Self` in return types.
/// - `with_timeout` takes `self` by value — **not** object-safe as written.
///
/// > **Note:** `with_timeout` is excluded from the object-safe vtable. Use the
/// > concrete type (e.g. `NewtonConfig`) or `SolverStrategy` to configure
/// > timeouts, then box the result.
///
/// # Example
///
/// ```rust,no_run
/// use entropyk_solver::solver::Solver;
/// use entropyk_solver::system::System;
///
/// fn run(solver: &mut dyn Solver, system: &mut System) {
///     match solver.solve(system) {
///         Ok(state) => println!("Converged in {} iters", state.iterations),
///         Err(e) => eprintln!("Solver failed: {}", e),
///     }
/// }
/// ```
pub trait Solver {
    /// Solve the system of equations.
    ///
    /// Mutates `system` state in-place during iteration. On success, returns a
    /// [`ConvergedState`] with the final state vector and convergence metadata.
    ///
    /// # Errors
    ///
    /// - [`SolverError::NonConvergence`] — max iterations exceeded.
    /// - [`SolverError::Timeout`] — time budget exceeded.
    /// - [`SolverError::Divergence`] — residual norm growing unboundedly.
    /// - [`SolverError::InvalidSystem`] — system is ill-formed.
    fn solve(&mut self, system: &mut System) -> Result<ConvergedState, SolverError>;

    /// Configure a time budget for this solver (builder pattern).
    ///
    /// If the solver exceeds `timeout`, it stops and returns
    /// [`SolverError::Timeout`]. Actual enforcement is implemented in
    /// Stories 4.2 / 4.3; this method stores the configuration.
    ///
    /// # Note on Object Safety
    ///
    /// This method is bounded by `where Self: Sized` and is therefore excluded
    /// from the vtable. It **cannot** be called through a `dyn Solver` reference.
    /// Configure the timeout on the concrete type before boxing.
    fn with_timeout(self, timeout: Duration) -> Self
    where
        Self: Sized;
}

// ─────────────────────────────────────────────────────────────────────────────
// Timeout Configuration (Story 4.5)
// ─────────────────────────────────────────────────────────────────────────────

/// Configuration for timeout behavior in solvers.
///
/// This struct controls how solvers handle timeout conditions, enabling
/// graceful degradation for HIL (Hardware-in-the-Loop) scenarios where
/// real-time constraints must never be violated.
///
/// # Example
///
/// ```rust
/// use entropyk_solver::solver::TimeoutConfig;
///
/// // Default: return best state on timeout (graceful degradation)
/// let config = TimeoutConfig::default();
/// assert!(config.return_best_state_on_timeout);
/// assert!(!config.zoh_fallback);
///
/// // HIL scenario with Zero-Order Hold fallback
/// let hil_config = TimeoutConfig {
///     zoh_fallback: true,
///     ..Default::default()
/// };
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct TimeoutConfig {
    /// Return best-known state on timeout instead of error.
    ///
    /// When `true` (default), the solver returns `Ok(ConvergedState)` with
    /// `status: TimedOutWithBestState` on timeout, containing the best state
    /// encountered during iteration (lowest residual norm).
    ///
    /// When `false`, the solver returns `Err(SolverError::Timeout)` on timeout.
    pub return_best_state_on_timeout: bool,

    /// On timeout, return previous state (ZOH) instead of current best.
    ///
    /// Zero-Order Hold (ZOH) fallback is useful for HIL scenarios where the
    /// previous known-good state should be used when a new solution cannot be
    /// found within the time budget.
    ///
    /// Requires `previous_state` to be set before solving.
    /// Default: `false`.
    pub zoh_fallback: bool,
}

impl Default for TimeoutConfig {
    fn default() -> Self {
        Self {
            return_best_state_on_timeout: true,
            zoh_fallback: false,
        }
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Jacobian Freezing Configuration (Story 4.8)
// ─────────────────────────────────────────────────────────────────────────────

/// Configuration for Jacobian-freezing optimization.
///
/// When enabled, the Newton-Raphson solver reuses the previously computed
/// Jacobian matrix for up to `max_frozen_iters` consecutive iterations,
/// provided the residual norm is still decreasing. This avoids expensive
/// Jacobian assembly and can reduce per-iteration CPU time by up to ~80%.
///
/// # Auto-disable on divergence
///
/// If the residual norm *increases* while a frozen Jacobian is being used,
/// the solver immediately forces a fresh Jacobian computation on the next
/// iteration and resets the frozen-iteration counter.
///
/// # Example
///
/// ```rust
/// use entropyk_solver::solver::{NewtonConfig, JacobianFreezingConfig};
///
/// let config = NewtonConfig::default()
///     .with_jacobian_freezing(JacobianFreezingConfig {
///         max_frozen_iters: 3,
///         threshold: 0.1,
///     });
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct JacobianFreezingConfig {
    /// Maximum number of consecutive iterations the Jacobian may be reused
    /// without recomputing.
    ///
    /// After this many frozen iterations the solver forces a fresh assembly,
    /// even if the residual is still decreasing.  Default: 3.
    pub max_frozen_iters: usize,

    /// Residual-norm ratio threshold below which freezing is considered safe.
    ///
    /// Freezing is only attempted when
    /// `current_norm / previous_norm < (1.0 - threshold)`,
    /// ensuring that convergence is still progressing sufficiently.
    /// Default: 0.1 (i.e., at least a 10 % residual decrease per step).
    pub threshold: f64,
}

impl Default for JacobianFreezingConfig {
    fn default() -> Self {
        Self {
            max_frozen_iters: 3,
            threshold: 0.1,
        }
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Helper functions
// ─────────────────────────────────────────────────────────────────────────────

/// Applies a Newton step to the state vector, clamping bounded variables.
///
/// Update formula: x_new = clamp(x_old + alpha * delta)
pub(crate) fn apply_newton_step(
    state: &mut [f64],
    delta: &[f64],
    clipping_mask: &[Option<(f64, f64)>],
    alpha: f64,
) {
    for (i, s) in state.iter_mut().enumerate() {
        let proposed = *s + alpha * delta[i];
        *s = match &clipping_mask[i] {
            Some((min, max)) => proposed.clamp(*min, *max),
            None => proposed,
        };
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Tests
// ─────────────────────────────────────────────────────────────────────────────

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

    /// Verify that `Box<dyn Solver>` compiles (trait is object-safe for `solve`).
    #[test]
    fn test_solver_trait_object_safety() {
        fn accepts_dyn_solver(_solver: &mut dyn Solver) {}
        let mut newton = NewtonConfig::default();
        accepts_dyn_solver(&mut newton);
    }

    /// Verify that `Box<dyn Solver>` can be constructed from concrete types.
    #[test]
    fn test_box_dyn_solver_compiles() {
        let _boxed: Box<dyn Solver> = Box::new(NewtonConfig::default());
    }

    /// Verify `SolverError::NonConvergence` Display message.
    #[test]
    fn test_error_non_convergence_display() {
        let err = SolverError::NonConvergence {
            iterations: 42,
            final_residual: 1.23e-3,
        };
        assert!(err.to_string().contains("42"));
    }

    /// Verify `ConvergedState` fields are accessible and `is_converged()` works.
    #[test]
    fn test_converged_state_fields_accessible() {
        let state = ConvergedState::new(
            vec![1.0, 2.0, 3.0],
            15,
            1e-8,
            ConvergenceStatus::Converged,
            SimulationMetadata::new("mock_hash".to_string()),
        );
        assert_eq!(state.state, vec![1.0, 2.0, 3.0]);
        assert!(state.is_converged());
    }

    /// Verify `ConvergedState` with `TimedOutWithBestState` status.
    #[test]
    fn test_converged_state_timed_out_not_converged() {
        let state = ConvergedState::new(
            vec![0.0],
            50,
            1e-3,
            ConvergenceStatus::TimedOutWithBestState,
            SimulationMetadata::new("mock_hash".to_string()),
        );
        assert!(!state.is_converged());
    }

    /// Verify `Box<dyn Solver>` can call `solve()` through the vtable.
    #[test]
    fn test_box_dyn_solver_can_call_solve() {
        let mut boxed: Box<dyn Solver> = Box::new(NewtonConfig::default());
        let mut system = System::new();
        system.finalize().unwrap();
        let result = boxed.solve(&mut system);
        assert!(result.is_err(), "dyn Solver::solve must dispatch correctly");
    }

    // ── Story 5.6: Control Variable Step Clipping Tests ───────────────────────

    #[test]
    fn test_bounded_variable_clipped_at_max() {
        let mut state = vec![0.5];
        let delta = vec![2.0];
        let mask = vec![Some((0.0, 1.0))];
        apply_newton_step(&mut state, &delta, &mask, 1.0);
        assert_eq!(state[0], 1.0, "Should be clipped to max bound");
    }

    #[test]
    fn test_bounded_variable_clipped_at_min() {
        let mut state = vec![0.5];
        let delta = vec![-2.0];
        let mask = vec![Some((0.0, 1.0))];
        apply_newton_step(&mut state, &delta, &mask, 1.0);
        assert_eq!(state[0], 0.0, "Should be clipped to min bound");
    }

    #[test]
    fn test_edge_states_not_clipped() {
        let mut state = vec![0.5, 10.0];
        let delta = vec![-2.0, 50.0];
        let mask = vec![Some((0.0, 1.0)), None];
        apply_newton_step(&mut state, &delta, &mask, 1.0);
        assert_eq!(state[0], 0.0, "Bounded variable should be clipped");
        assert_eq!(state[1], 60.0, "Unbounded variable should NOT be clipped");
    }
}
