//! Intelligent fallback solver implementation.
//!
//! This module provides the [`FallbackSolver`] which implements an intelligent
//! fallback strategy between Newton-Raphson and Sequential Substitution (Picard).
//!
//! # Strategy
//!
//! The fallback solver implements the following algorithm:
//!
//! 1. Start with Newton-Raphson (quadratic convergence)
//! 2. If Newton diverges, switch to Picard (more robust)
//! 3. If Picard stabilizes (residual < threshold), try returning to Newton
//! 4. If max switches reached, stay on current solver permanently
//!
//! # Features
//!
//! - Automatic fallback from Newton to Picard on divergence
//! - Return to Newton when Picard stabilizes the solution
//! - Maximum switch limit to prevent infinite oscillation
//! - Time-budgeted solving with graceful degradation (Story 4.5)
//! - Smart initialization support (Story 4.6)
//! - Multi-circuit convergence criteria (Story 4.7)

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

use crate::criteria::ConvergenceCriteria;
use crate::metadata::SimulationMetadata;
use crate::solver::{ConvergedState, ConvergenceStatus, Solver, SolverError};
use crate::system::System;

use super::{NewtonConfig, PicardConfig};

/// Configuration for the intelligent fallback solver.
///
/// The fallback solver starts with Newton-Raphson (quadratic convergence) and
/// automatically switches to Sequential Substitution (Picard) if Newton diverges.
/// It can return to Newton when Picard stabilizes the solution.
///
/// # Example
///
/// ```rust
/// use entropyk_solver::solver::{FallbackConfig, FallbackSolver, Solver};
/// use std::time::Duration;
///
/// let config = FallbackConfig {
///     fallback_enabled: true,
///     return_to_newton_threshold: 1e-3,
///     max_fallback_switches: 2,
/// };
///
/// let solver = FallbackSolver::new(config)
///     .with_timeout(Duration::from_secs(1));
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct FallbackConfig {
    /// Enable automatic fallback from Newton to Picard on divergence.
    ///
    /// When `true` (default), the solver switches to Picard if Newton diverges.
    /// When `false`, the solver runs pure Newton or Picard without fallback.
    pub fallback_enabled: bool,

    /// Residual norm threshold for returning to Newton from Picard.
    ///
    /// When Picard reduces the residual below this threshold, the solver
    /// attempts to return to Newton for faster convergence.
    /// Default: $10^{-3}$.
    pub return_to_newton_threshold: f64,

    /// Maximum number of solver switches before staying on current solver.
    ///
    /// Prevents infinite oscillation between Newton and Picard.
    /// Default: 2.
    pub max_fallback_switches: usize,
}

impl Default for FallbackConfig {
    fn default() -> Self {
        Self {
            fallback_enabled: true,
            return_to_newton_threshold: 1e-3,
            max_fallback_switches: 2,
        }
    }
}

/// Tracks which solver is currently active in the fallback loop.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CurrentSolver {
    Newton,
    Picard,
}

/// Internal state for the fallback solver.
struct FallbackState {
    current_solver: CurrentSolver,
    switch_count: usize,
    /// Whether we've permanently committed to Picard (after max switches or Newton re-divergence)
    committed_to_picard: bool,
    /// Best state encountered across all solver invocations (Story 4.5 - AC: #4)
    best_state: Option<Vec<f64>>,
    /// Best residual norm across all solver invocations (Story 4.5 - AC: #4)
    best_residual: Option<f64>,
}

impl FallbackState {
    fn new() -> Self {
        Self {
            current_solver: CurrentSolver::Newton,
            switch_count: 0,
            committed_to_picard: false,
            best_state: None,
            best_residual: None,
        }
    }

    /// Update best state if the given residual is better (Story 4.5 - AC: #4).
    fn update_best_state(&mut self, state: &[f64], residual: f64) {
        if self.best_residual.is_none() || residual < self.best_residual.unwrap() {
            self.best_state = Some(state.to_vec());
            self.best_residual = Some(residual);
        }
    }
}

/// Intelligent fallback solver that switches between Newton-Raphson and Picard.
///
/// The fallback solver implements the following algorithm:
///
/// 1. Start with Newton-Raphson (quadratic convergence)
/// 2. If Newton diverges, switch to Picard (more robust)
/// 3. If Picard stabilizes (residual < threshold), try returning to Newton
/// 4. If max switches reached, stay on current solver permanently
///
/// # Timeout Handling
///
/// The timeout applies to the total solving time across all solver switches.
/// Each solver inherits the remaining time budget.
///
/// # Pre-Allocated Buffers
///
/// All buffers are pre-allocated once before the fallback loop to avoid
/// heap allocation during solver switches (NFR4).
#[derive(Debug, Clone)]
pub struct FallbackSolver {
    /// Fallback behavior configuration.
    pub config: FallbackConfig,
    /// Newton-Raphson configuration.
    pub newton_config: NewtonConfig,
    /// Sequential Substitution (Picard) configuration.
    pub picard_config: PicardConfig,
}

impl FallbackSolver {
    /// Creates a new fallback solver with the given configuration.
    pub fn new(config: FallbackConfig) -> Self {
        Self {
            config,
            newton_config: NewtonConfig::default(),
            picard_config: PicardConfig::default(),
        }
    }

    /// Creates a fallback solver with default configuration.
    pub fn default_solver() -> Self {
        Self::new(FallbackConfig::default())
    }

    /// Sets custom Newton-Raphson configuration.
    pub fn with_newton_config(mut self, config: NewtonConfig) -> Self {
        self.newton_config = config;
        self
    }

    /// Sets custom Picard configuration.
    pub fn with_picard_config(mut self, config: PicardConfig) -> Self {
        self.picard_config = config;
        self
    }

    /// Sets the initial state for cold-start solving (Story 4.6 — builder pattern).
    ///
    /// Delegates to both `newton_config` and `picard_config` so the initial state
    /// is used regardless of which solver is active in the fallback loop.
    pub fn with_initial_state(mut self, state: Vec<f64>) -> Self {
        self.newton_config.initial_state = Some(state.clone());
        self.picard_config.initial_state = Some(state);
        self
    }

    /// Sets multi-circuit convergence criteria (Story 4.7 — builder pattern).
    ///
    /// Delegates to both `newton_config` and `picard_config` so criteria are
    /// applied regardless of which solver is active in the fallback loop.
    pub fn with_convergence_criteria(mut self, criteria: ConvergenceCriteria) -> Self {
        self.newton_config.convergence_criteria = Some(criteria.clone());
        self.picard_config.convergence_criteria = Some(criteria);
        self
    }

    /// Main fallback solving loop.
    ///
    /// Implements the intelligent fallback algorithm:
    /// - Start with Newton-Raphson
    /// - Switch to Picard on Newton divergence
    /// - Return to Newton when Picard stabilizes (if under switch limit and residual below threshold)
    /// - Stay on Picard permanently after max switches or if Newton re-diverges
    fn solve_with_fallback(
        &mut self,
        system: &mut System,
        start_time: Instant,
        timeout: Option<Duration>,
    ) -> Result<ConvergedState, SolverError> {
        let mut state = FallbackState::new();

        // Pre-configure solver configs once
        let mut newton_cfg = self.newton_config.clone();
        let mut picard_cfg = self.picard_config.clone();

        loop {
            // Check remaining time budget
            let remaining = timeout.map(|t| t.saturating_sub(start_time.elapsed()));

            // Check for timeout before running solver
            if let Some(remaining_time) = remaining {
                if remaining_time.is_zero() {
                    return Err(SolverError::Timeout {
                        timeout_ms: timeout.unwrap().as_millis() as u64,
                    });
                }
            }

            // Run current solver with remaining time
            newton_cfg.timeout = remaining;
            picard_cfg.timeout = remaining;

            let result = match state.current_solver {
                CurrentSolver::Newton => newton_cfg.solve(system),
                CurrentSolver::Picard => picard_cfg.solve(system),
            };

            match result {
                Ok(converged) => {
                    // Update best state tracking (Story 4.5 - AC: #4)
                    state.update_best_state(&converged.state, converged.final_residual);

                    tracing::info!(
                        solver = match state.current_solver {
                            CurrentSolver::Newton => "NewtonRaphson",
                            CurrentSolver::Picard => "Picard",
                        },
                        iterations = converged.iterations,
                        final_residual = converged.final_residual,
                        switch_count = state.switch_count,
                        "Fallback solver converged"
                    );
                    return Ok(converged);
                }
                Err(SolverError::Timeout { timeout_ms }) => {
                    // Story 4.5 - AC: #4: Return best state on timeout if available
                    if let (Some(best_state), Some(best_residual)) =
                        (state.best_state.clone(), state.best_residual)
                    {
                        tracing::info!(
                            best_residual = best_residual,
                            "Returning best state across all solver invocations on timeout"
                        );
                        return Ok(ConvergedState::new(
                            best_state,
                            0, // iterations not tracked across switches
                            best_residual,
                            ConvergenceStatus::TimedOutWithBestState,
                            SimulationMetadata::new(system.input_hash()),
                        ));
                    }
                    return Err(SolverError::Timeout { timeout_ms });
                }
                Err(SolverError::Divergence { ref reason }) => {
                    // Handle divergence based on current solver and state
                    if !self.config.fallback_enabled {
                        tracing::info!(
                            solver = match state.current_solver {
                                CurrentSolver::Newton => "NewtonRaphson",
                                CurrentSolver::Picard => "Picard",
                            },
                            reason = reason,
                            "Divergence detected, fallback disabled"
                        );
                        return result;
                    }

                    match state.current_solver {
                        CurrentSolver::Newton => {
                            // Newton diverged - switch to Picard (stay there permanently after max switches)
                            if state.switch_count >= self.config.max_fallback_switches {
                                // Max switches reached - commit to Picard permanently
                                state.committed_to_picard = true;
                                state.current_solver = CurrentSolver::Picard;
                                tracing::info!(
                                    switch_count = state.switch_count,
                                    max_switches = self.config.max_fallback_switches,
                                    "Max switches reached, committing to Picard permanently"
                                );
                            } else {
                                // Switch to Picard
                                state.switch_count += 1;
                                state.current_solver = CurrentSolver::Picard;
                                tracing::warn!(
                                    switch_count = state.switch_count,
                                    reason = reason,
                                    "Newton diverged, switching to Picard"
                                );
                            }
                            // Continue loop with Picard
                        }
                        CurrentSolver::Picard => {
                            // Picard diverged - if we were trying Newton again, commit to Picard permanently
                            if state.switch_count > 0 && !state.committed_to_picard {
                                state.committed_to_picard = true;
                                tracing::info!(
                                    switch_count = state.switch_count,
                                    reason = reason,
                                    "Newton re-diverged after return from Picard, staying on Picard permanently"
                                );
                                // Stay on Picard and try again
                            } else {
                                // Picard diverged with no return attempt - no more fallbacks available
                                tracing::warn!(
                                    reason = reason,
                                    "Picard diverged, no more fallbacks available"
                                );
                                return result;
                            }
                        }
                    }
                }
                Err(SolverError::NonConvergence {
                    iterations,
                    final_residual,
                }) => {
                    // Non-convergence: check if we should try the other solver
                    if !self.config.fallback_enabled {
                        return Err(SolverError::NonConvergence {
                            iterations,
                            final_residual,
                        });
                    }

                    match state.current_solver {
                        CurrentSolver::Newton => {
                            // Newton didn't converge - try Picard
                            if state.switch_count >= self.config.max_fallback_switches {
                                // Max switches reached - commit to Picard permanently
                                state.committed_to_picard = true;
                                state.current_solver = CurrentSolver::Picard;
                                tracing::info!(
                                    switch_count = state.switch_count,
                                    "Max switches reached, committing to Picard permanently"
                                );
                            } else {
                                state.switch_count += 1;
                                state.current_solver = CurrentSolver::Picard;
                                tracing::info!(
                                    switch_count = state.switch_count,
                                    iterations = iterations,
                                    final_residual = final_residual,
                                    "Newton did not converge, switching to Picard"
                                );
                            }
                            // Continue loop with Picard
                        }
                        CurrentSolver::Picard => {
                            // Picard didn't converge - check if we should try Newton
                            if state.committed_to_picard
                                || state.switch_count >= self.config.max_fallback_switches
                            {
                                tracing::info!(
                                    iterations = iterations,
                                    final_residual = final_residual,
                                    "Picard did not converge, no more fallbacks"
                                );
                                return Err(SolverError::NonConvergence {
                                    iterations,
                                    final_residual,
                                });
                            }

                            // Check if residual is low enough to try Newton
                            if final_residual < self.config.return_to_newton_threshold {
                                state.switch_count += 1;
                                state.current_solver = CurrentSolver::Newton;
                                tracing::info!(
                                    switch_count = state.switch_count,
                                    final_residual = final_residual,
                                    threshold = self.config.return_to_newton_threshold,
                                    "Picard stabilized, attempting Newton return"
                                );
                                // Continue loop with Newton
                            } else {
                                // Stay on Picard and keep trying
                                tracing::debug!(
                                    final_residual = final_residual,
                                    threshold = self.config.return_to_newton_threshold,
                                    "Picard not yet stabilized, aborting"
                                );
                                return Err(SolverError::NonConvergence {
                                    iterations,
                                    final_residual,
                                });
                            }
                        }
                    }
                }
                Err(other) => {
                    // InvalidSystem or other errors - propagate immediately
                    return Err(other);
                }
            }
        }
    }
}

impl Solver for FallbackSolver {
    fn solve(&mut self, system: &mut System) -> Result<ConvergedState, SolverError> {
        let start_time = Instant::now();
        let timeout = self.newton_config.timeout.or(self.picard_config.timeout);

        tracing::info!(
            fallback_enabled = self.config.fallback_enabled,
            return_to_newton_threshold = self.config.return_to_newton_threshold,
            max_fallback_switches = self.config.max_fallback_switches,
            "Fallback solver starting"
        );

        if self.config.fallback_enabled {
            self.solve_with_fallback(system, start_time, timeout)
        } else {
            // Fallback disabled - run pure Newton
            self.newton_config.solve(system)
        }
    }

    fn with_timeout(mut self, timeout: Duration) -> Self {
        self.newton_config.timeout = Some(timeout);
        self.picard_config.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_fallback_config_defaults() {
        let cfg = FallbackConfig::default();
        assert!(cfg.fallback_enabled);
        assert!((cfg.return_to_newton_threshold - 1e-3).abs() < 1e-15);
        assert_eq!(cfg.max_fallback_switches, 2);
    }

    #[test]
    fn test_fallback_solver_new() {
        let config = FallbackConfig {
            fallback_enabled: false,
            return_to_newton_threshold: 5e-4,
            max_fallback_switches: 3,
        };
        let solver = FallbackSolver::new(config.clone());
        assert_eq!(solver.config, config);
    }

    #[test]
    fn test_fallback_solver_with_timeout() {
        let timeout = Duration::from_millis(500);
        let solver = FallbackSolver::default_solver().with_timeout(timeout);
        assert_eq!(solver.newton_config.timeout, Some(timeout));
        assert_eq!(solver.picard_config.timeout, Some(timeout));
    }

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