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

//! Contact boundary conditions for mechanical contact problems.

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

/// Contact boundary condition for node-to-node or node-to-surface contact.
#[derive(Debug)]
pub struct ContactBC {
    /// Master nodes (contact target)
    pub master_nodes: Vec<NodeId>,
    /// Slave nodes (contacting surface)
    pub slave_nodes: Vec<NodeId>,
    /// Contact stiffness (penalty parameter)
    pub contact_stiffness: f64,
    /// Friction coefficient
    pub friction_coefficient: f64,
    /// Initial gap
    pub initial_gap: f64,
    /// Contact tolerance
    pub tolerance: f64,
    /// Time range when active
    pub time_range: Option<(f64, f64)>,
    /// Contact detection method
    pub detection_method: ContactDetectionMethod,
    /// Contact enforcement method
    pub enforcement_method: ContactEnforcementMethod,
}

/// Contact detection methods.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ContactDetectionMethod {
    /// Node-to-node contact
    NodeToNode,
    /// Node-to-surface contact
    NodeToSurface,
    /// Surface-to-surface contact
    SurfaceToSurface,
}

/// Contact enforcement methods.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ContactEnforcementMethod {
    /// Penalty method
    Penalty,
    /// Lagrange multiplier method
    LagrangeMultiplier,
    /// Augmented Lagrangian method
    AugmentedLagrangian,
}

impl ContactBC {
    /// Create a new contact boundary condition.
    pub fn new(
        master_nodes: Vec<NodeId>,
        slave_nodes: Vec<NodeId>,
        contact_stiffness: f64,
    ) -> Self {
        Self {
            master_nodes,
            slave_nodes,
            contact_stiffness,
            friction_coefficient: 0.0,
            initial_gap: 0.0,
            tolerance: 1e-6,
            time_range: None,
            detection_method: ContactDetectionMethod::NodeToNode,
            enforcement_method: ContactEnforcementMethod::Penalty,
        }
    }

    /// Set friction coefficient.
    pub fn with_friction(mut self, friction_coefficient: f64) -> Self {
        self.friction_coefficient = friction_coefficient;
        self
    }

    /// Set initial gap.
    pub fn with_gap(mut self, gap: f64) -> Self {
        self.initial_gap = gap;
        self
    }

    /// Set contact tolerance.
    pub fn with_tolerance(mut self, tolerance: f64) -> Self {
        self.tolerance = tolerance;
        self
    }

    /// Set detection method.
    pub fn with_detection_method(mut self, method: ContactDetectionMethod) -> Self {
        self.detection_method = method;
        self
    }

    /// Set enforcement method.
    pub fn with_enforcement_method(mut self, method: ContactEnforcementMethod) -> Self {
        self.enforcement_method = method;
        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
        }
    }

    /// Apply contact 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 contact boundary condition with current displacement solution.
    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(());
        }

        match self.enforcement_method {
            ContactEnforcementMethod::Penalty => self.apply_penalty_method(
                mesh,
                dof_numbering,
                global_stiffness,
                global_force,
                current_solution,
            ),
            ContactEnforcementMethod::LagrangeMultiplier => self.apply_lagrange_method(
                mesh,
                dof_numbering,
                global_stiffness,
                global_force,
                current_solution,
            ),
            ContactEnforcementMethod::AugmentedLagrangian => self
                .apply_augmented_lagrangian_method(
                    mesh,
                    dof_numbering,
                    global_stiffness,
                    global_force,
                    current_solution,
                ),
        }
    }

    /// Apply penalty method for contact enforcement.
    fn apply_penalty_method(
        &self,
        mesh: &Mesh,
        dof_numbering: &mut AdvancedDofNumbering,
        global_stiffness: &mut SparseMatrix,
        global_force: &mut DVector<f64>,
        current_solution: Option<&DVector<f64>>,
    ) -> FeaResult<()> {
        match self.detection_method {
            ContactDetectionMethod::NodeToNode => self.apply_node_to_node_penalty(
                mesh,
                dof_numbering,
                global_stiffness,
                global_force,
                current_solution,
            ),
            ContactDetectionMethod::NodeToSurface => self.apply_node_to_surface_penalty(
                mesh,
                dof_numbering,
                global_stiffness,
                global_force,
                current_solution,
            ),
            ContactDetectionMethod::SurfaceToSurface => {
                // Simplified implementation
                self.apply_node_to_node_penalty(
                    mesh,
                    dof_numbering,
                    global_stiffness,
                    global_force,
                    current_solution,
                )
            }
        }
    }

    /// Apply node-to-node contact penalty method.
    fn apply_node_to_node_penalty(
        &self,
        mesh: &Mesh,
        dof_numbering: &mut AdvancedDofNumbering,
        global_stiffness: &mut SparseMatrix,
        global_force: &mut DVector<f64>,
        current_solution: Option<&DVector<f64>>,
    ) -> FeaResult<()> {
        // Simple node-to-node contact
        let min_nodes = self.master_nodes.len().min(self.slave_nodes.len());

        for i in 0..min_nodes {
            let master_node = self.master_nodes[i];
            let slave_node = self.slave_nodes[i];

            // Get node positions
            let master_pos = mesh
                .nodes
                .get(&master_node)
                .ok_or(BoundaryError::NodeNotFound {
                    node_id: master_node.0,
                })?
                .coordinates
                .clone();
            let slave_pos = mesh
                .nodes
                .get(&slave_node)
                .ok_or(BoundaryError::NodeNotFound {
                    node_id: slave_node.0,
                })?
                .coordinates
                .clone();

            // Calculate gap (assuming contact in Z direction)
            let gap = if let Some(solution) = current_solution {
                // Get current displacements
                let master_disp =
                    self.get_node_displacement(master_node, dof_numbering, solution)?;
                let slave_disp = self.get_node_displacement(slave_node, dof_numbering, solution)?;

                let current_master_pos = master_pos + master_disp;
                let current_slave_pos = slave_pos + slave_disp;

                (current_slave_pos[2] - current_master_pos[2]) - self.initial_gap
            } else {
                (slave_pos[2] - master_pos[2]) - self.initial_gap
            };

            // Apply contact forces if in contact (gap <= 0)
            if gap <= self.tolerance {
                let penetration = -gap.min(0.0);
                let contact_force = self.contact_stiffness * penetration;

                // Apply normal contact force
                if let (Some(master_dof), Some(slave_dof)) = (
                    dof_numbering.get_dof(master_node, DofComponent::DisplacementZ),
                    dof_numbering.get_dof(slave_node, DofComponent::DisplacementZ),
                ) {
                    // Add penalty stiffness
                    global_stiffness.add_entry(master_dof, master_dof, self.contact_stiffness)?;
                    global_stiffness.add_entry(slave_dof, slave_dof, self.contact_stiffness)?;
                    global_stiffness.add_entry(master_dof, slave_dof, -self.contact_stiffness)?;
                    global_stiffness.add_entry(slave_dof, master_dof, -self.contact_stiffness)?;

                    // Add contact forces
                    global_force[master_dof] += contact_force;
                    global_force[slave_dof] -= contact_force;
                }

                // Apply friction forces if friction coefficient > 0
                if self.friction_coefficient > 0.0 {
                    self.apply_friction_forces(
                        master_node,
                        slave_node,
                        contact_force,
                        dof_numbering,
                        global_stiffness,
                        global_force,
                        current_solution,
                    )?;
                }
            }
        }

        Ok(())
    }

    /// Apply node-to-surface contact penalty method.
    fn apply_node_to_surface_penalty(
        &self,
        _mesh: &Mesh,
        _dof_numbering: &mut AdvancedDofNumbering,
        _global_stiffness: &mut SparseMatrix,
        _global_force: &mut DVector<f64>,
        _current_solution: Option<&DVector<f64>>,
    ) -> FeaResult<()> {
        // Simplified implementation - would require surface projection algorithms
        Ok(())
    }

    /// Apply Lagrange multiplier method.
    fn apply_lagrange_method(
        &self,
        _mesh: &Mesh,
        _dof_numbering: &mut AdvancedDofNumbering,
        _global_stiffness: &mut SparseMatrix,
        _global_force: &mut DVector<f64>,
        _current_solution: Option<&DVector<f64>>,
    ) -> FeaResult<()> {
        // Would require augmenting the system with Lagrange multiplier DOFs
        Ok(())
    }

    /// Apply augmented Lagrangian method.
    fn apply_augmented_lagrangian_method(
        &self,
        _mesh: &Mesh,
        _dof_numbering: &mut AdvancedDofNumbering,
        _global_stiffness: &mut SparseMatrix,
        _global_force: &mut DVector<f64>,
        _current_solution: Option<&DVector<f64>>,
    ) -> FeaResult<()> {
        // Combination of penalty and Lagrange multiplier methods
        Ok(())
    }

    /// Apply friction forces.
    fn apply_friction_forces(
        &self,
        master_node: NodeId,
        slave_node: NodeId,
        normal_force: f64,
        dof_numbering: &mut AdvancedDofNumbering,
        _global_stiffness: &mut SparseMatrix,
        global_force: &mut DVector<f64>,
        current_solution: Option<&DVector<f64>>,
    ) -> FeaResult<()> {
        if let Some(solution) = current_solution {
            // Get relative tangential velocity/displacement
            let master_disp = self.get_node_displacement(master_node, dof_numbering, solution)?;
            let slave_disp = self.get_node_displacement(slave_node, dof_numbering, solution)?;

            let relative_disp = slave_disp - master_disp;
            let tangential_disp = Vector3::new(relative_disp.x, relative_disp.y, 0.0);
            let tangential_magnitude = tangential_disp.norm();

            if tangential_magnitude > self.tolerance {
                let friction_force = self.friction_coefficient * normal_force;
                let friction_direction = tangential_disp / tangential_magnitude;

                // Apply friction forces in X and Y directions
                for (component, direction_component) in [
                    (DofComponent::DisplacementX, friction_direction.x),
                    (DofComponent::DisplacementY, friction_direction.y),
                ] {
                    if let (Some(master_dof), Some(slave_dof)) = (
                        dof_numbering.get_dof(master_node, component),
                        dof_numbering.get_dof(slave_node, component),
                    ) {
                        let friction_component = friction_force * direction_component;
                        global_force[master_dof] += friction_component;
                        global_force[slave_dof] -= friction_component;
                    }
                }
            }
        }

        Ok(())
    }

    /// Get node displacement from solution vector.
    fn get_node_displacement(
        &self,
        node_id: NodeId,
        dof_numbering: &AdvancedDofNumbering,
        solution: &DVector<f64>,
    ) -> FeaResult<Vector3<f64>> {
        let mut displacement = Vector3::zeros();

        for (i, component) in [
            DofComponent::DisplacementX,
            DofComponent::DisplacementY,
            DofComponent::DisplacementZ,
        ]
        .iter()
        .enumerate()
        {
            if let Some(dof) = dof_numbering.get_dof(node_id, *component)
                && dof < solution.len()
            {
                displacement[i] = solution[dof];
            }
        }

        Ok(displacement)
    }
}

/// Contact patterns for common contact scenarios.
pub struct ContactPatterns;

impl ContactPatterns {
    /// Create contact between two surfaces.
    pub fn surface_to_surface_contact(
        master_nodes: Vec<NodeId>,
        slave_nodes: Vec<NodeId>,
        contact_stiffness: f64,
        friction_coefficient: f64,
    ) -> ContactBC {
        ContactBC::new(master_nodes, slave_nodes, contact_stiffness)
            .with_friction(friction_coefficient)
            .with_detection_method(ContactDetectionMethod::SurfaceToSurface)
    }

    /// Create self-contact for a single surface.
    pub fn self_contact(
        nodes: Vec<NodeId>,
        contact_stiffness: f64,
        friction_coefficient: f64,
    ) -> ContactBC {
        ContactBC::new(nodes.clone(), nodes, contact_stiffness)
            .with_friction(friction_coefficient)
            .with_detection_method(ContactDetectionMethod::NodeToNode)
    }

    /// Create rigid contact (infinite stiffness approximation).
    pub fn rigid_contact(master_nodes: Vec<NodeId>, slave_nodes: Vec<NodeId>) -> ContactBC {
        ContactBC::new(master_nodes, slave_nodes, 1e12) // Very high stiffness
            .with_enforcement_method(ContactEnforcementMethod::Penalty)
    }

    /// Create frictionless contact.
    pub fn frictionless_contact(
        master_nodes: Vec<NodeId>,
        slave_nodes: Vec<NodeId>,
        contact_stiffness: f64,
    ) -> ContactBC {
        ContactBC::new(master_nodes, slave_nodes, contact_stiffness).with_friction(0.0)
    }

    /// Create stick contact (no sliding).
    pub fn stick_contact(
        master_nodes: Vec<NodeId>,
        slave_nodes: Vec<NodeId>,
        contact_stiffness: f64,
    ) -> ContactBC {
        ContactBC::new(master_nodes, slave_nodes, contact_stiffness).with_friction(1e6)
        // Very high friction coefficient
    }
}

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

    #[test]
    fn test_contact_bc_creation() {
        let bc = ContactBC::new(vec![NodeId(0), NodeId(1)], vec![NodeId(2), NodeId(3)], 1e6);

        assert_eq!(bc.master_nodes, vec![NodeId(0), NodeId(1)]);
        assert_eq!(bc.slave_nodes, vec![NodeId(2), NodeId(3)]);
        assert_eq!(bc.contact_stiffness, 1e6);
        assert_eq!(bc.friction_coefficient, 0.0);
    }

    #[test]
    fn test_contact_bc_with_friction() {
        let bc = ContactBC::new(vec![NodeId(0)], vec![NodeId(1)], 1e5).with_friction(0.3);

        assert_eq!(bc.friction_coefficient, 0.3);
    }

    #[test]
    fn test_contact_detection_methods() {
        let methods = [
            ContactDetectionMethod::NodeToNode,
            ContactDetectionMethod::NodeToSurface,
            ContactDetectionMethod::SurfaceToSurface,
        ];
        assert_eq!(methods.len(), 3);
    }

    #[test]
    fn test_contact_enforcement_methods() {
        let methods = [
            ContactEnforcementMethod::Penalty,
            ContactEnforcementMethod::LagrangeMultiplier,
            ContactEnforcementMethod::AugmentedLagrangian,
        ];
        assert_eq!(methods.len(), 3);
    }

    #[test]
    fn test_contact_patterns_surface_to_surface() {
        let bc = ContactPatterns::surface_to_surface_contact(
            vec![NodeId(0), NodeId(1)],
            vec![NodeId(2), NodeId(3)],
            1e6,
            0.3,
        );

        assert_eq!(bc.friction_coefficient, 0.3);
        assert_eq!(
            bc.detection_method,
            ContactDetectionMethod::SurfaceToSurface
        );
    }

    #[test]
    fn test_contact_patterns_self_contact() {
        let nodes = vec![NodeId(0), NodeId(1), NodeId(2)];
        let bc = ContactPatterns::self_contact(nodes.clone(), 1e5, 0.2);

        assert_eq!(bc.master_nodes, nodes);
        assert_eq!(bc.slave_nodes, nodes);
        assert_eq!(bc.friction_coefficient, 0.2);
    }

    #[test]
    fn test_contact_patterns_rigid_contact() {
        let bc = ContactPatterns::rigid_contact(vec![NodeId(0)], vec![NodeId(1)]);

        assert_eq!(bc.contact_stiffness, 1e12);
        assert_eq!(bc.enforcement_method, ContactEnforcementMethod::Penalty);
    }

    #[test]
    fn test_contact_patterns_frictionless() {
        let bc = ContactPatterns::frictionless_contact(vec![NodeId(0)], vec![NodeId(1)], 5e5);

        assert_eq!(bc.friction_coefficient, 0.0);
        assert_eq!(bc.contact_stiffness, 5e5);
    }

    #[test]
    fn test_contact_patterns_stick_contact() {
        let bc = ContactPatterns::stick_contact(vec![NodeId(0)], vec![NodeId(1)], 1e6);

        assert_eq!(bc.friction_coefficient, 1e6);
    }
}