//! Pseudo-transient continuation (PTC / Ψtc) globalization.
//!
//! When Newton-with-line-search stagnates at a local minimum of ‖F‖ — the
//! classic failure on fully-coupled vapor-compression systems with two-phase
//! pressure drops at high EXV opening — PTC walks the physical relaxation
//! dynamics instead of minimizing ‖F‖ along a line. Each iteration is one
//! linearized implicit Euler step of the fictitious dynamics `x′ = −F(x)`:
//!
//! ```text
//!   (δ⁻¹·I + J(x))·s = −F(x),   x ← x + s
//! ```
//!
//! with the timestep δ grown by the Switched Evolution Relaxation (SER) rule
//! as the residual shrinks. Small δ far from the solution reconditions the
//! near-singular momentum block (orifice ↔ two-phase ΔP) and bounds the step;
//! δ → ∞ recovers exact Newton with quadratic terminal convergence.
//!
//! References: Kelley & Keyes, *Convergence Analysis of Pseudo-transient
//! Continuation*, SIAM J. Numer. Anal. 35:508–523 (1998); Mulder & Van Leer
//! (1985) for the SER rule; PETSc `TSPSEUDO` for the incremental SER variant.

use crate::jacobian::JacobianMatrix;
use crate::metadata::SimulationMetadata;
use crate::solver::{apply_newton_step, ConvergedState, ConvergenceStatus, Solver, SolverError};
use crate::system::System;
use entropyk_components::JacobianBuilder;
use entropyk_solver_core::LinearSolver;
use std::time::{Duration, Instant};

/// Configuration for pseudo-transient continuation.
///
/// Used as a globalization stage of [`super::FallbackSolver`]: only invoked
/// when the primary Newton/Picard stages fail, so it never perturbs problems
/// that already converge.
#[derive(Debug, Clone, PartialEq)]
pub struct PtcConfig {
    /// Maximum PTC iterations.
    pub max_iterations: usize,
    /// Convergence tolerance on the residual L2 norm.
    pub tolerance: f64,
    /// Initial timestep δ₀ (in scaled-residual units). Small δ₀ regularizes
    /// the first steps; too small wastes iterations in the induction phase.
    pub delta0: f64,
    /// SER growth factor (> 1). Controls how fast δ grows as ‖F‖ shrinks.
    pub growth: f64,
    /// Maximum timestep before the diagonal shift is negligible (pure Newton).
    pub delta_max: f64,
    /// Timestep shrink factor on rejected steps (< 1).
    pub shrink: f64,
    /// Minimum timestep: if δ falls below this after repeated rejections the
    /// method has failed (detectable failure mode, Kelley-Keyes Thm 2.3).
    pub delta_min: f64,
    /// Optional initial state (cold start shared with the primary solvers).
    pub initial_state: Option<Vec<f64>>,
    /// Optional timeout.
    pub timeout: Option<Duration>,
}

impl Default for PtcConfig {
    fn default() -> Self {
        Self {
            max_iterations: 200,
            tolerance: 1e-6,
            delta0: 1e-3,
            growth: 1.5,
            delta_max: 1e9,
            shrink: 0.25,
            delta_min: 1e-12,
            initial_state: None,
            timeout: None,
        }
    }
}

impl PtcConfig {
    /// L2 residual norm, consistent with the other strategies.
    fn residual_norm(residuals: &[f64]) -> f64 {
        let norm = residuals.iter().map(|r| r * r).sum::<f64>().sqrt();
        if norm.is_finite() {
            norm
        } else {
            f64::MAX
        }
    }
}

impl Solver for PtcConfig {
    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(),
            });
        }

        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 saved_state: Vec<f64> = vec![0.0; n_state];
        let mut jacobian_builder = JacobianBuilder::new();
        let mut jacobian_matrix = JacobianMatrix::zeros(n_equations, n_state);

        // LinearSolver backend (FR3): factorizes at each fresh assembly;
        // bit-identical to JacobianMatrix::solve.
        let mut linear_backend = crate::linear::NalgebraLuSolver::new();
        linear_backend
            .set_problem(n_equations, n_state)
            .map_err(|e| SolverError::InvalidSystem {
                message: format!("Failed to initialize linear backend: {e}"),
            })?;
        let mut step_buf: Vec<f64> = vec![0.0; n_state];

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

        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);
        let mut prev_norm = current_norm;

        tracing::info!(
            max_iterations = self.max_iterations,
            tolerance = self.tolerance,
            delta0 = self.delta0,
            residual_norm = current_norm,
            "Pseudo-transient continuation starting"
        );

        if current_norm < self.tolerance {
            return Ok(ConvergedState::new(
                state,
                0,
                current_norm,
                ConvergenceStatus::Converged,
                SimulationMetadata::new(system.input_hash()),
            ));
        }

        let mut delta = self.delta0;

        for iteration in 1..=self.max_iterations {
            if let Some(timeout) = self.timeout {
                if start_time.elapsed() > timeout {
                    tracing::info!(iteration, "PTC timed out");
                    return Err(SolverError::Timeout {
                        timeout_ms: timeout.as_millis() as u64,
                    });
                }
            }

            // Fresh Jacobian each iteration (δ changes every step anyway).
            jacobian_builder.clear();
            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());

            // Diagonal shift: (δ⁻¹·I + J)·s = −F.
            let shift = 1.0 / delta;
            {
                let n = jacobian_matrix.as_matrix().nrows();
                let m = jacobian_matrix.as_matrix_mut();
                for i in 0..n {
                    m[(i, i)] += shift;
                }
            }

            // Solve through the LinearSolver backend (FR3); non-square keeps
            // the legacy least-squares path.
            let step = if n_equations == n_state {
                linear_backend
                    .set_matrix(jacobian_matrix.as_matrix())
                    .map_err(|e| SolverError::InvalidSystem {
                        message: format!("Failed to factorize Jacobian: {e}"),
                    })?;
                for (d, r) in step_buf.iter_mut().zip(residuals.iter()) {
                    *d = -*r;
                }
                linear_backend
                    .solve_in_place(&mut step_buf)
                    .ok()
                    .map(|_| step_buf.clone())
            } else {
                jacobian_matrix.solve(&residuals)
            };
            let Some(step) = step else {
                // Singular even with the shift: shrink and retry without
                // consuming an iteration of state progress.
                delta *= self.shrink;
                if delta < self.delta_min {
                    return Err(SolverError::NonConvergence {
                        iterations: iteration - 1,
                        final_residual: current_norm,
                    });
                }
                continue;
            };

            saved_state.copy_from_slice(&state);
            // No line search: the δ⁻¹ term already bounds the step.
            apply_newton_step(&mut state, &step, &clipping_mask, 1.0);

            if let Err(e) = system.compute_residuals(&state, &mut residuals) {
                if !e.is_recoverable() {
                    // Fatal evaluation error: abort immediately, do not burn
                    // the shrink budget on it.
                    return Err(SolverError::InvalidSystem {
                        message: format!("Failed to compute residuals: {:?}", e),
                    });
                }
                // Recoverable domain violation (KINSOL > 0): reject and shrink
                // (residual undefined).
                state.copy_from_slice(&saved_state);
                delta *= self.shrink;
                if delta < self.delta_min {
                    return Err(SolverError::NonConvergence {
                        iterations: iteration - 1,
                        final_residual: current_norm,
                    });
                }
                continue;
            }
            current_norm = Self::residual_norm(&residuals);

            tracing::debug!(
                iteration,
                residual_norm = current_norm,
                delta,
                "PTC iteration"
            );

            if current_norm < self.tolerance {
                tracing::info!(
                    iterations = iteration,
                    final_residual = current_norm,
                    "PTC converged"
                );
                return Ok(ConvergedState::new(
                    state,
                    iteration,
                    current_norm,
                    ConvergenceStatus::Converged,
                    SimulationMetadata::new(system.input_hash()),
                ));
            }

            if current_norm > prev_norm {
                // Reject: restore and cut the timestep.
                state.copy_from_slice(&saved_state);
                delta *= self.shrink;
                tracing::debug!(iteration, delta, "PTC step rejected, shrinking timestep");
                if delta < self.delta_min {
                    tracing::warn!(
                        iteration,
                        final_residual = current_norm,
                        "PTC timestep collapsed — globalization failed"
                    );
                    return Err(SolverError::NonConvergence {
                        iterations: iteration - 1,
                        final_residual: current_norm,
                    });
                }
                continue;
            }

            // SER incremental rule (PETSc form), bounded by delta_max.
            delta = (self.growth * delta * prev_norm / current_norm).min(self.delta_max);
            prev_norm = current_norm;
        }

        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::*;

    #[test]
    fn test_ptc_config_defaults() {
        let cfg = PtcConfig::default();
        assert_eq!(cfg.max_iterations, 200);
        assert!(cfg.delta0 > 0.0 && cfg.delta0 < 1.0);
        assert!(cfg.growth > 1.0);
        assert!(cfg.delta_max > cfg.delta0);
        assert!(cfg.shrink > 0.0 && cfg.shrink < 1.0);
        assert!(cfg.delta_min > 0.0 && cfg.delta_min < cfg.delta0);
    }

    #[test]
    fn test_ptc_rejects_empty_system() {
        let mut system = System::new();
        system.finalize().unwrap();
        let mut solver = PtcConfig::default();
        let result = solver.solve(&mut system);
        assert!(matches!(result, Err(SolverError::InvalidSystem { .. })));
    }
}