// Copyright (c) 2024 RustyTorch++ Team
// Licensed under the Apache License, Version 2.0

//! Time integration schemes for dynamic analysis.

use crate::assembly::SparseMatrix;
use crate::error::FeaResult;
use nalgebra::DVector;

/// Newmark time integration scheme.
#[derive(Debug)]
pub struct NewmarkIntegrator {
    /// Newmark beta parameter
    pub beta: f64,
    /// Newmark gamma parameter
    pub gamma: f64,
    /// Time step
    pub dt: f64,
}

impl NewmarkIntegrator {
    /// Create average acceleration method (beta=0.25, gamma=0.5)
    pub fn average_acceleration(dt: f64) -> Self {
        Self {
            beta: 0.25,
            gamma: 0.5,
            dt,
        }
    }

    /// Create linear acceleration method (beta=1/6, gamma=0.5)
    pub fn linear_acceleration(dt: f64) -> Self {
        Self {
            beta: 1.0 / 6.0,
            gamma: 0.5,
            dt,
        }
    }

    /// Perform one time step
    pub fn step(
        &self,
        mass: &SparseMatrix,
        damping: &SparseMatrix,
        stiffness: &SparseMatrix,
        force: &DVector<f64>,
        displacement: &mut DVector<f64>,
        velocity: &mut DVector<f64>,
        acceleration: &mut DVector<f64>,
    ) -> FeaResult<()> {
        let dt = self.dt;
        let beta = self.beta;
        let gamma = self.gamma;

        // Predictors
        let disp_pred = &*displacement + dt * &*velocity + (0.5 - beta) * dt * dt * &*acceleration;
        let vel_pred = &*velocity + (1.0 - gamma) * dt * &*acceleration;

        // Effective stiffness matrix
        let k_eff = stiffness.clone()
            + &(mass.clone() * (1.0 / (beta * dt * dt)))
            + &(damping.clone() * (gamma / (beta * dt)));

        // Effective force
        let f_eff = force
            + &mass.multiply_vector(&((1.0 / (beta * dt * dt)) * &disp_pred))?
            + &damping.multiply_vector(&((gamma / (beta * dt)) * &vel_pred))?;

        // Solve for displacement increment using appropriate solver
        let delta_disp = self.solve_system(&k_eff, &f_eff)?;

        // Correctors
        *displacement = disp_pred + beta * dt * dt * &delta_disp;
        *velocity = vel_pred + gamma * dt * &delta_disp;
        *acceleration = (1.0 / (beta * dt * dt))
            * (&delta_disp - dt * &*velocity - (0.5 - beta) * dt * dt * &*acceleration);

        Ok(())
    }

    /// Solve the linear system K*x = f using appropriate method
    fn solve_system(&self, k: &SparseMatrix, f: &DVector<f64>) -> FeaResult<DVector<f64>> {
        // Use direct solver for well-conditioned systems
        k.solve_vector(f)
    }
}

/// Central difference time integration.
#[derive(Debug)]
pub struct CentralDifferenceIntegrator {
    pub dt: f64,
}

impl CentralDifferenceIntegrator {
    pub fn new(dt: f64) -> Self {
        Self { dt }
    }

    pub fn step(
        &self,
        mass: &SparseMatrix,
        damping: &SparseMatrix,
        stiffness: &SparseMatrix,
        force: &DVector<f64>,
        displacement_prev: &DVector<f64>,
        displacement_curr: &DVector<f64>,
    ) -> FeaResult<DVector<f64>> {
        let dt2 = self.dt * self.dt;

        // M*a_n = F_n - C*v_n - K*u_n
        let velocity = (displacement_curr - displacement_prev) / self.dt;
        let rhs = force
            - &damping.multiply_vector(&velocity)?
            - &stiffness.multiply_vector(displacement_curr)?;
        let acceleration = mass.solve_vector(&rhs)?;

        // u_{n+1} = 2*u_n - u_{n-1} + dt^2*a_n
        let displacement_next =
            2.0 * displacement_curr - displacement_prev + &(dt2 * &acceleration);

        Ok(displacement_next)
    }
}