rustytorch/crates/specialized/rtx-fea/src/boundary/robin.rs

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

//! Robin (mixed) boundary conditions combining Dirichlet and Neumann effects.

use super::{BoundaryConditionApplicator, TimeFunction};
use crate::assembly::{AdvancedDofNumbering, DofComponent, SparseMatrix};
use crate::error::{BoundaryError, FeaResult};
use crate::mesh::{Mesh, NodeId};
use nalgebra::DVector;

/// Robin boundary condition types for different physics.
#[derive(Debug)]
pub enum RobinType {
    /// Linear spring: k * u + c = force
    LinearSpring {
        spring_constant: f64,
        reference_displacement: f64,
    },
    /// Convective heat transfer: h * (T - T_ambient) = q
    ConvectiveHeatTransfer {
        heat_transfer_coefficient: f64,
        ambient_temperature: f64,
    },
    /// Radiation heat transfer: σε * (T^4 - T_ambient^4) = q
    RadiationHeatTransfer {
        stefan_boltzmann: f64,
        emissivity: f64,
        ambient_temperature: f64,
    },
    /// Damped spring: c * du/dt + k * u = force
    DampedSpring {
        spring_constant: f64,
        damping_coefficient: f64,
        reference_displacement: f64,
    },
    /// Custom Robin condition: α * u + β * du/dn = γ
    Custom {
        alpha: f64,
        beta: f64,
        gamma: f64,
        time_function: Option<TimeFunction>,
    },
}

impl RobinType {
    /// Get stiffness contribution for this Robin condition.
    pub fn get_stiffness_contribution(&self, time: f64, current_value: f64) -> f64 {
        match self {
            Self::LinearSpring {
                spring_constant, ..
            } => *spring_constant,
            Self::ConvectiveHeatTransfer {
                heat_transfer_coefficient,
                ..
            } => *heat_transfer_coefficient,
            Self::RadiationHeatTransfer {
                stefan_boltzmann,
                emissivity,
                ambient_temperature: _,
            } => {
                // Linearized radiation: d/dT(σε(T^4 - T_amb^4)) ≈ 4σεT^3
                4.0 * stefan_boltzmann * emissivity * current_value.powi(3)
            }
            Self::DampedSpring {
                spring_constant, ..
            } => *spring_constant,
            Self::Custom {
                alpha,
                time_function,
                ..
            } => {
                if let Some(func) = time_function {
                    alpha * func(time)
                } else {
                    *alpha
                }
            }
        }
    }

    /// Get force contribution for this Robin condition.
    pub fn get_force_contribution(&self, time: f64, current_value: f64) -> f64 {
        match self {
            Self::LinearSpring {
                spring_constant,
                reference_displacement,
            } => spring_constant * reference_displacement,
            Self::ConvectiveHeatTransfer {
                heat_transfer_coefficient,
                ambient_temperature,
            } => heat_transfer_coefficient * ambient_temperature,
            Self::RadiationHeatTransfer {
                stefan_boltzmann,
                emissivity,
                ambient_temperature,
            } => {
                // Linearized: σε * T_amb^4 + 3σεT^3 * T_amb
                stefan_boltzmann
                    * emissivity
                    * (ambient_temperature.powi(4)
                        + 3.0 * current_value.powi(3) * ambient_temperature)
            }
            Self::DampedSpring {
                spring_constant,
                reference_displacement,
                ..
            } => spring_constant * reference_displacement,
            Self::Custom {
                beta,
                gamma,
                time_function,
                ..
            } => {
                if let Some(func) = time_function {
                    beta * func(time) + gamma
                } else {
                    *gamma
                }
            }
        }
    }

    /// Create a simple spring condition.
    pub fn spring(spring_constant: f64, reference_displacement: f64) -> Self {
        Self::LinearSpring {
            spring_constant,
            reference_displacement,
        }
    }

    /// Create a convective heat transfer condition.
    pub fn convection(heat_transfer_coefficient: f64, ambient_temperature: f64) -> Self {
        Self::ConvectiveHeatTransfer {
            heat_transfer_coefficient,
            ambient_temperature,
        }
    }

    /// Create a radiation heat transfer condition.
    pub fn radiation(emissivity: f64, ambient_temperature: f64) -> Self {
        Self::RadiationHeatTransfer {
            stefan_boltzmann: 5.67e-8, // Stefan-Boltzmann constant
            emissivity,
            ambient_temperature,
        }
    }

    /// Create a damped spring condition.
    pub fn damped_spring(
        spring_constant: f64,
        damping_coefficient: f64,
        reference_displacement: f64,
    ) -> Self {
        Self::DampedSpring {
            spring_constant,
            damping_coefficient,
            reference_displacement,
        }
    }
}

/// Robin boundary condition.
#[derive(Debug)]
pub struct RobinBC {
    /// Nodes where this condition applies
    pub nodes: Vec<NodeId>,
    /// DOF components affected
    pub components: Vec<DofComponent>,
    /// Robin condition type
    pub condition_type: RobinType,
    /// Time range when this condition is active
    pub time_range: Option<(f64, f64)>,
    /// Scaling factor
    pub scaling_factor: f64,
    /// Previous time step values for time derivatives
    pub previous_values: Option<Vec<f64>>,
    /// Previous time for derivative calculation
    pub previous_time: Option<f64>,
}

impl RobinBC {
    /// Create a new Robin boundary condition.
    pub fn new(
        nodes: Vec<NodeId>,
        components: Vec<DofComponent>,
        condition_type: RobinType,
    ) -> Self {
        Self {
            nodes,
            components,
            condition_type,
            time_range: None,
            scaling_factor: 1.0,
            previous_values: None,
            previous_time: None,
        }
    }

    /// Create a spring boundary condition.
    pub fn spring_support(
        nodes: Vec<NodeId>,
        components: Vec<DofComponent>,
        spring_constant: f64,
        reference_displacement: f64,
    ) -> Self {
        Self::new(
            nodes,
            components,
            RobinType::spring(spring_constant, reference_displacement),
        )
    }

    /// Create a convective heat transfer boundary condition.
    pub fn convective_boundary(
        _nodes: Vec<NodeId>,
        heat_transfer_coefficient: f64,
        ambient_temperature: f64,
    ) -> Self {
        Self::new(
            vec![],
            vec![DofComponent::Temperature],
            RobinType::convection(heat_transfer_coefficient, ambient_temperature),
        )
    }

    /// Set time range for this boundary condition.
    pub fn with_time_range(mut self, start_time: f64, end_time: f64) -> Self {
        self.time_range = Some((start_time, end_time));
        self
    }

    /// Set scaling factor.
    pub fn with_scaling(mut self, factor: f64) -> Self {
        self.scaling_factor = factor;
        self
    }

    /// Check if this boundary condition is active at the given time.
    pub fn is_active(&self, time: f64) -> bool {
        if let Some((start_time, end_time)) = self.time_range {
            time >= start_time && time <= end_time
        } else {
            true
        }
    }

    /// Update previous values for time derivative calculation.
    pub fn update_previous_values(&mut self, values: Vec<f64>, time: f64) {
        self.previous_values = Some(values);
        self.previous_time = Some(time);
    }

    /// Get time derivative for damped conditions.
    fn get_time_derivative(&self, current_values: &[f64], dt: f64) -> Vec<f64> {
        if let (Some(prev_values), Some(_)) = (&self.previous_values, self.previous_time) {
            current_values
                .iter()
                .zip(prev_values.iter())
                .map(|(curr, prev)| (curr - prev) / dt)
                .collect()
        } else {
            vec![0.0; current_values.len()]
        }
    }

    /// Apply this Robin boundary condition to the global system.
    pub fn apply(
        &self,
        mesh: &Mesh,
        dof_numbering: &mut AdvancedDofNumbering,
        global_stiffness: &mut SparseMatrix,
        global_force: &mut DVector<f64>,
        time: f64,
    ) -> FeaResult<()> {
        self.apply_with_solution(
            mesh,
            dof_numbering,
            global_stiffness,
            global_force,
            time,
            None,
        )
    }

    /// Apply Robin boundary condition with current solution values.
    pub fn apply_with_solution(
        &self,
        mesh: &Mesh,
        dof_numbering: &mut AdvancedDofNumbering,
        global_stiffness: &mut SparseMatrix,
        global_force: &mut DVector<f64>,
        time: f64,
        current_solution: Option<&DVector<f64>>,
    ) -> FeaResult<()> {
        if !self.is_active(time) {
            return Ok(());
        }

        for &node_id in &self.nodes {
            // Get node position
            let node = mesh
                .nodes
                .get(&node_id)
                .ok_or(BoundaryError::NodeNotFound { node_id: node_id.0 })?;
            let _node_position = node.position();

            // Apply to each component
            for &component in &self.components {
                if let Some(dof) = dof_numbering.get_dof(node_id, component) {
                    // Get current value for nonlinear Robin conditions
                    let current_value = if let Some(solution) = current_solution {
                        if dof < solution.len() {
                            solution[dof]
                        } else {
                            0.0
                        }
                    } else {
                        0.0
                    };

                    // Get stiffness and force contributions
                    let stiffness_contribution = self
                        .condition_type
                        .get_stiffness_contribution(time, current_value)
                        * self.scaling_factor;
                    let force_contribution = self
                        .condition_type
                        .get_force_contribution(time, current_value)
                        * self.scaling_factor;

                    // Apply Robin condition
                    BoundaryConditionApplicator::apply_robin(
                        dof,
                        stiffness_contribution,
                        force_contribution,
                        global_stiffness,
                        global_force,
                    )?;
                }
            }
        }

        Ok(())
    }

    /// Apply damped Robin condition (requires previous time step data).
    pub fn apply_damped(
        &mut self,
        mesh: &Mesh,
        dof_numbering: &mut AdvancedDofNumbering,
        global_stiffness: &mut SparseMatrix,
        global_force: &mut DVector<f64>,
        time: f64,
        current_solution: &DVector<f64>,
        dt: f64,
    ) -> FeaResult<()> {
        if !self.is_active(time) {
            return Ok(());
        }

        // Handle damped spring conditions
        if let RobinType::DampedSpring {
            spring_constant,
            damping_coefficient,
            reference_displacement,
        } = &self.condition_type
        {
            // Collect current values
            let mut current_values = Vec::new();
            let mut dof_indices = Vec::new();

            for &node_id in &self.nodes {
                for &component in &self.components {
                    if let Some(dof) = dof_numbering.get_dof(node_id, component)
                        && dof < current_solution.len()
                    {
                        current_values.push(current_solution[dof]);
                        dof_indices.push(dof);
                    }
                }
            }

            // Get time derivatives
            let time_derivatives = self.get_time_derivative(&current_values, dt);

            // Apply damped Robin condition
            for (i, &dof) in dof_indices.iter().enumerate() {
                let velocity = if i < time_derivatives.len() {
                    time_derivatives[i]
                } else {
                    0.0
                };

                let _current_displacement = if i < current_values.len() {
                    current_values[i]
                } else {
                    0.0
                };

                // Stiffness: k + c/dt (implicit time integration)
                let stiffness_contribution =
                    (spring_constant + damping_coefficient / dt) * self.scaling_factor;

                // Force: k * u_ref + c * v_prev / dt
                let force_contribution = (spring_constant * reference_displacement
                    + damping_coefficient * velocity / dt)
                    * self.scaling_factor;

                BoundaryConditionApplicator::apply_robin(
                    dof,
                    stiffness_contribution,
                    force_contribution,
                    global_stiffness,
                    global_force,
                )?;
            }

            // Update previous values
            self.update_previous_values(current_values, time);
        } else {
            // For non-damped conditions, use regular apply
            self.apply_with_solution(
                mesh,
                dof_numbering,
                global_stiffness,
                global_force,
                time,
                Some(current_solution),
            )?;
        }

        Ok(())
    }

    /// Get all DOFs affected by this boundary condition.
    pub fn get_affected_dofs(&self, dof_numbering: &AdvancedDofNumbering) -> Vec<usize> {
        let mut dofs = Vec::new();

        for &node_id in &self.nodes {
            for &component in &self.components {
                if let Some(dof) = dof_numbering.get_dof(node_id, component) {
                    dofs.push(dof);
                }
            }
        }

        dofs.sort_unstable();
        dofs
    }
}

/// Collection of common Robin boundary condition patterns.
pub struct RobinPatterns;

impl RobinPatterns {
    /// Create elastic foundation (Winkler foundation).
    pub fn elastic_foundation(nodes: Vec<NodeId>, foundation_modulus: f64) -> RobinBC {
        RobinBC::spring_support(
            nodes,
            vec![DofComponent::DisplacementY], // Assuming Y is vertical
            foundation_modulus,
            0.0,
        )
    }

    /// Create viscous damper.
    pub fn viscous_damper(
        nodes: Vec<NodeId>,
        components: Vec<DofComponent>,
        damping_coefficient: f64,
    ) -> RobinBC {
        RobinBC::new(
            nodes,
            components,
            RobinType::damped_spring(0.0, damping_coefficient, 0.0),
        )
    }

    /// Create spring-damper system.
    pub fn spring_damper(
        nodes: Vec<NodeId>,
        components: Vec<DofComponent>,
        spring_constant: f64,
        damping_coefficient: f64,
        reference_displacement: f64,
    ) -> RobinBC {
        RobinBC::new(
            nodes,
            components,
            RobinType::damped_spring(spring_constant, damping_coefficient, reference_displacement),
        )
    }

    /// Create thermal convection boundary.
    pub fn thermal_convection(
        nodes: Vec<NodeId>,
        heat_transfer_coefficient: f64,
        ambient_temperature: f64,
    ) -> RobinBC {
        RobinBC::convective_boundary(nodes, heat_transfer_coefficient, ambient_temperature)
    }

    /// Create thermal radiation boundary.
    pub fn thermal_radiation(
        nodes: Vec<NodeId>,
        emissivity: f64,
        ambient_temperature: f64,
    ) -> RobinBC {
        RobinBC::new(
            nodes,
            vec![DofComponent::Temperature],
            RobinType::radiation(emissivity, ambient_temperature),
        )
    }

    /// Create acoustic impedance boundary.
    pub fn acoustic_impedance(nodes: Vec<NodeId>, impedance: f64) -> RobinBC {
        RobinBC::new(
            nodes,
            vec![DofComponent::Pressure],
            RobinType::Custom {
                alpha: impedance,
                beta: 1.0,
                gamma: 0.0,
                time_function: None,
            },
        )
    }

    /// Create contact spring (penalty method).
    pub fn contact_spring(nodes: Vec<NodeId>, contact_stiffness: f64, gap: f64) -> RobinBC {
        RobinBC::spring_support(
            nodes,
            vec![DofComponent::DisplacementZ], // Assuming Z is contact direction
            contact_stiffness,
            gap,
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mesh::{MaterialId, geometry::Rectangle};

    #[test]
    fn test_robin_type_spring() {
        let robin = RobinType::spring(1000.0, 0.1);

        let stiffness = robin.get_stiffness_contribution(0.0, 0.0);
        let force = robin.get_force_contribution(0.0, 0.0);

        assert_eq!(stiffness, 1000.0);
        assert_eq!(force, 100.0); // 1000.0 * 0.1
    }

    #[test]
    fn test_robin_type_convection() {
        let robin = RobinType::convection(25.0, 300.0);

        let stiffness = robin.get_stiffness_contribution(0.0, 350.0);
        let force = robin.get_force_contribution(0.0, 350.0);

        assert_eq!(stiffness, 25.0);
        assert_eq!(force, 7500.0); // 25.0 * 300.0
    }

    #[test]
    fn test_robin_type_radiation() {
        let robin = RobinType::radiation(0.8, 300.0);

        let stiffness = robin.get_stiffness_contribution(0.0, 350.0);
        assert!(stiffness > 0.0); // Should be positive

        let force = robin.get_force_contribution(0.0, 350.0);
        assert!(force > 0.0); // Should be positive for radiation
    }

    #[test]
    fn test_robin_bc_creation() {
        let bc = RobinBC::spring_support(
            vec![NodeId(0), NodeId(1)],
            vec![DofComponent::DisplacementY],
            1000.0,
            0.0,
        );

        assert_eq!(bc.nodes, vec![NodeId(0), NodeId(1)]);
        assert_eq!(bc.components, vec![DofComponent::DisplacementY]);
        assert_eq!(bc.scaling_factor, 1.0);
    }

    #[test]
    fn test_robin_bc_with_time_range() {
        let bc = RobinBC::spring_support(
            vec![NodeId(0)],
            vec![DofComponent::DisplacementX],
            500.0,
            0.05,
        )
        .with_time_range(1.0, 5.0);

        assert!(!bc.is_active(0.5));
        assert!(bc.is_active(3.0));
        assert!(!bc.is_active(6.0));
    }

    #[test]
    fn test_robin_patterns_elastic_foundation() {
        let bc = RobinPatterns::elastic_foundation(vec![NodeId(0), NodeId(1)], 50000.0);

        assert_eq!(bc.nodes, vec![NodeId(0), NodeId(1)]);
        assert_eq!(bc.components, vec![DofComponent::DisplacementY]);
    }

    #[test]
    fn test_robin_patterns_spring_damper() {
        let bc = RobinPatterns::spring_damper(
            vec![NodeId(0)],
            vec![DofComponent::DisplacementX],
            1000.0,
            50.0,
            0.0,
        );

        if let RobinType::DampedSpring {
            spring_constant,
            damping_coefficient,
            reference_displacement,
        } = bc.condition_type
        {
            assert_eq!(spring_constant, 1000.0);
            assert_eq!(damping_coefficient, 50.0);
            assert_eq!(reference_displacement, 0.0);
        } else {
            panic!("Expected DampedSpring type");
        }
    }

    #[test]
    fn test_robin_patterns_thermal_convection() {
        let bc = RobinPatterns::thermal_convection(vec![NodeId(0)], 15.0, 293.15);

        assert_eq!(bc.components, vec![DofComponent::Temperature]);
    }

    #[test]
    fn test_time_derivative_calculation() {
        let mut bc = RobinBC::spring_support(
            vec![NodeId(0)],
            vec![DofComponent::DisplacementX],
            1000.0,
            0.0,
        );

        // Set previous values
        bc.update_previous_values(vec![0.1], 1.0);

        // Calculate derivatives
        let current_values = vec![0.15];
        let dt = 0.1;
        let derivatives = bc.get_time_derivative(&current_values, dt);

        assert_eq!(derivatives.len(), 1);
        assert!((derivatives[0] - 0.5).abs() < 1e-10); // (0.15 - 0.1) / 0.1 = 0.5
    }

    #[test]
    fn test_robin_type_custom() {
        let robin = RobinType::Custom {
            alpha: 2.0,
            beta: 3.0,
            gamma: 5.0,
            time_function: None,
        };

        let stiffness = robin.get_stiffness_contribution(0.0, 0.0);
        let force = robin.get_force_contribution(0.0, 0.0);

        assert_eq!(stiffness, 2.0);
        assert_eq!(force, 5.0);
    }
}