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

//! Neumann (natural) boundary conditions for force and traction application.

use super::{BoundaryConditionApplicator, SpatialFunction, TimeFunction, TimeSpatialFunction};
use crate::assembly::{AdvancedDofNumbering, DofComponent, SparseMatrix};
use crate::elements::{FiniteElement, StandardFiniteElement};
use crate::error::{BoundaryError, FeaResult};
use crate::mesh::{Element, ElementId, Mesh, NodeId};
use nalgebra::{DMatrix, DVector, Vector3};

/// Neumann boundary condition types.
#[derive(Debug)]
pub enum NeumannType {
    /// Fixed force/traction value
    Fixed(f64),
    /// Time-dependent force/traction
    TimeDependent(TimeFunction),
    /// Spatially varying force/traction
    Spatial(SpatialFunction),
    /// Time and spatially varying force/traction
    TimeSpatial(TimeSpatialFunction),
    /// Pressure loading (normal to surface)
    Pressure {
        pressure: f64,
        time_function: Option<TimeFunction>,
    },
    /// Distributed load along elements
    DistributedLoad {
        load_per_length: f64,
        direction: Vector3<f64>,
    },
    /// Body force (volume loading)
    BodyForce {
        force_density: Vector3<f64>,
        time_function: Option<TimeFunction>,
    },
}

impl NeumannType {
    /// Evaluate the force/traction at given time and position.
    pub fn evaluate(&self, time: f64, position: &Vector3<f64>) -> f64 {
        match self {
            Self::Fixed(value) => *value,
            Self::TimeDependent(func) => func(time),
            Self::Spatial(func) => func(position),
            Self::TimeSpatial(func) => func(time, position),
            Self::Pressure {
                pressure,
                time_function,
            } => {
                let base_pressure = *pressure;
                if let Some(func) = time_function {
                    base_pressure * func(time)
                } else {
                    base_pressure
                }
            }
            Self::DistributedLoad {
                load_per_length, ..
            } => *load_per_length,
            Self::BodyForce { .. } => 0.0, // Handled separately
        }
    }

    /// Get force direction vector.
    pub fn get_direction(&self) -> Vector3<f64> {
        match self {
            Self::DistributedLoad { direction, .. } => *direction,
            Self::BodyForce { force_density, .. } => *force_density,
            _ => Vector3::new(1.0, 0.0, 0.0), // Default direction
        }
    }

    /// Create a harmonic force function.
    pub fn harmonic(amplitude: f64, frequency: f64, phase: f64) -> Self {
        Self::TimeDependent(TimeFunction(Box::new(move |t| {
            amplitude * (frequency * t + phase).sin()
        })))
    }

    /// Create a ramp loading function.
    pub fn ramp_load(start_time: f64, end_time: f64, start_load: f64, end_load: f64) -> Self {
        Self::TimeDependent(TimeFunction(Box::new(move |t| {
            if t <= start_time {
                start_load
            } else if t >= end_time {
                end_load
            } else {
                start_load + (end_load - start_load) * (t - start_time) / (end_time - start_time)
            }
        })))
    }

    /// Create a step loading function.
    pub fn step_load(step_time: f64, before_load: f64, after_load: f64) -> Self {
        Self::TimeDependent(TimeFunction(Box::new(move |t| {
            if t < step_time {
                before_load
            } else {
                after_load
            }
        })))
    }
}

/// Neumann boundary condition for nodal forces.
#[derive(Debug)]
pub struct NeumannBC {
    /// Nodes where this condition applies
    pub nodes: Vec<NodeId>,
    /// DOF components affected
    pub components: Vec<DofComponent>,
    /// Force/traction type
    pub condition_type: NeumannType,
    /// Time range when this condition is active
    pub time_range: Option<(f64, f64)>,
    /// Load ramping factor
    pub ramping_factor: f64,
    /// Whether to distribute load equally among nodes
    pub distribute_equally: bool,
}

impl NeumannBC {
    /// Create a fixed force boundary condition.
    pub fn fixed_force(nodes: Vec<NodeId>, components: Vec<DofComponent>, force: f64) -> Self {
        Self {
            nodes,
            components,
            condition_type: NeumannType::Fixed(force),
            time_range: None,
            ramping_factor: 1.0,
            distribute_equally: true,
        }
    }

    /// Create a time-dependent force boundary condition.
    pub fn time_dependent_force(
        nodes: Vec<NodeId>,
        components: Vec<DofComponent>,
        time_function: TimeFunction,
    ) -> Self {
        Self {
            nodes,
            components,
            condition_type: NeumannType::TimeDependent(time_function),
            time_range: None,
            ramping_factor: 1.0,
            distribute_equally: true,
        }
    }

    /// Create a pressure loading condition.
    pub fn pressure_load(
        nodes: Vec<NodeId>,
        pressure: f64,
        time_function: Option<TimeFunction>,
    ) -> Self {
        Self {
            nodes,
            components: vec![DofComponent::DisplacementZ], // Assuming Z is normal
            condition_type: NeumannType::Pressure {
                pressure,
                time_function,
            },
            time_range: None,
            ramping_factor: 1.0,
            distribute_equally: true,
        }
    }

    /// Create a distributed load condition.
    pub fn distributed_load(
        nodes: Vec<NodeId>,
        load_per_length: f64,
        direction: Vector3<f64>,
    ) -> Self {
        let components = if direction.x.abs() > 1e-15 {
            vec![DofComponent::DisplacementX]
        } else if direction.y.abs() > 1e-15 {
            vec![DofComponent::DisplacementY]
        } else {
            vec![DofComponent::DisplacementZ]
        };

        Self {
            nodes,
            components,
            condition_type: NeumannType::DistributedLoad {
                load_per_length,
                direction,
            },
            time_range: None,
            ramping_factor: 1.0,
            distribute_equally: true,
        }
    }

    /// 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 ramping factor.
    pub fn with_ramping(mut self, factor: f64) -> Self {
        self.ramping_factor = factor;
        self
    }

    /// Set load distribution method.
    pub fn with_distribution(mut self, distribute_equally: bool) -> Self {
        self.distribute_equally = distribute_equally;
        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
        }
    }

    /// Get the force value at given time and node.
    pub fn get_force(&self, time: f64, node_position: &Vector3<f64>) -> f64 {
        let base_force = self.condition_type.evaluate(time, node_position);
        let ramped_force = base_force * self.ramping_factor;

        if self.distribute_equally && self.nodes.len() > 1 {
            ramped_force / self.nodes.len() as f64
        } else {
            ramped_force
        }
    }

    /// Apply this Neumann 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<()> {
        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();

            // Get force value
            let force_value = self.get_force(time, &node_position);

            // Apply to each component
            for &component in &self.components {
                if let Some(dof) = dof_numbering.get_dof(node_id, component) {
                    BoundaryConditionApplicator::apply_neumann_force(
                        dof,
                        force_value,
                        global_force,
                    )?;
                }
            }
        }

        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
    }
}

/// Surface traction boundary condition for element faces.
#[derive(Debug)]
pub struct SurfaceTractionBC {
    /// Elements where traction is applied
    pub elements: Vec<ElementId>,
    /// Face indices for each element
    pub face_indices: Vec<usize>,
    /// Traction value
    pub traction: Vector3<f64>,
    /// Time function for traction
    pub time_function: Option<TimeFunction>,
    /// Time range when active
    pub time_range: Option<(f64, f64)>,
}

impl SurfaceTractionBC {
    /// Create a new surface traction boundary condition.
    pub fn new(elements: Vec<ElementId>, face_indices: Vec<usize>, traction: Vector3<f64>) -> Self {
        Self {
            elements,
            face_indices,
            traction,
            time_function: None,
            time_range: None,
        }
    }

    /// Set time function for traction.
    pub fn with_time_function(mut self, time_function: TimeFunction) -> Self {
        self.time_function = Some(time_function);
        self
    }

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

    /// Check if active at 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
        }
    }

    /// Get traction at time.
    pub fn get_traction(&self, time: f64) -> Vector3<f64> {
        if let Some(ref time_func) = self.time_function {
            self.traction * time_func(time)
        } else {
            self.traction
        }
    }

    /// Apply surface traction to global system.
    pub fn apply(
        &self,
        mesh: &Mesh,
        dof_numbering: &mut AdvancedDofNumbering,
        _global_stiffness: &mut SparseMatrix,
        global_force: &mut DVector<f64>,
        time: f64,
    ) -> FeaResult<()> {
        if !self.is_active(time) {
            return Ok(());
        }

        let traction = self.get_traction(time);

        for (elem_idx, &element_id) in self.elements.iter().enumerate() {
            let element = mesh
                .elements
                .get(&element_id)
                .ok_or(BoundaryError::ElementNotFound {
                    element_id: element_id.0,
                })?;

            let face_index = if elem_idx < self.face_indices.len() {
                self.face_indices[elem_idx]
            } else {
                0 // Default to first face
            };

            // Get element nodes
            let element_nodes: Result<Vec<_>, _> = element
                .nodes
                .iter()
                .map(|&node_id| {
                    mesh.nodes
                        .get(&node_id)
                        .ok_or(BoundaryError::NodeNotFound { node_id: node_id.0 })
                })
                .collect();
            let element_nodes = element_nodes?;

            // Create finite element
            let finite_element = StandardFiniteElement::new(
                element.element_type,
                element_nodes.iter().map(|n| n.position()).collect(),
            );

            // Compute equivalent nodal forces from surface traction
            let nodal_forces =
                self.compute_equivalent_nodal_forces(&finite_element, face_index, &traction)?;

            // Add to global force vector
            for (local_node_idx, &node_id) in element.nodes.iter().enumerate() {
                let force = nodal_forces.row(local_node_idx);

                // Apply force components
                for (comp_idx, &component) in [
                    DofComponent::DisplacementX,
                    DofComponent::DisplacementY,
                    DofComponent::DisplacementZ,
                ]
                .iter()
                .enumerate()
                {
                    if let Some(dof) = dof_numbering.get_dof(node_id, component)
                        && comp_idx < force.len()
                    {
                        BoundaryConditionApplicator::apply_neumann_force(
                            dof,
                            force[comp_idx],
                            global_force,
                        )?;
                    }
                }
            }
        }

        Ok(())
    }

    /// Compute equivalent nodal forces from surface traction.
    fn compute_equivalent_nodal_forces(
        &self,
        finite_element: &dyn FiniteElement,
        _face_index: usize,
        traction: &Vector3<f64>,
    ) -> FeaResult<DMatrix<f64>> {
        // Compute nodal forces using surface integration over element face
        let num_nodes = finite_element.num_nodes();
        let mut nodal_forces = DMatrix::zeros(num_nodes, 3);

        // TODO: Implement face integration methods on FiniteElement trait
        // For now, return placeholder nodal forces
        // These methods would need to be added to the FiniteElement trait:
        // - get_face_nodes(face_index)
        // - get_face_quadrature_points(face_index)
        // - evaluate_face_shape_functions(face_index, coords)
        // - compute_face_jacobian(face_index, coords)

        // Placeholder implementation - distribute traction equally to all nodes
        let traction_per_node = traction / (num_nodes as f64);
        for i in 0..num_nodes {
            nodal_forces[(i, 0)] = traction_per_node.x;
            nodal_forces[(i, 1)] = traction_per_node.y;
            nodal_forces[(i, 2)] = traction_per_node.z;
        }

        Ok(nodal_forces)
    }
}

/// Body force boundary condition for volume loading.
#[derive(Debug)]
pub struct BodyForceBC {
    /// Elements where body force is applied
    pub elements: Vec<ElementId>,
    /// Body force density (force per unit volume)
    pub force_density: Vector3<f64>,
    /// Time function for body force
    pub time_function: Option<TimeFunction>,
    /// Time range when active
    pub time_range: Option<(f64, f64)>,
}

impl BodyForceBC {
    /// Create a new body force boundary condition.
    pub fn new(elements: Vec<ElementId>, force_density: Vector3<f64>) -> Self {
        Self {
            elements,
            force_density,
            time_function: None,
            time_range: None,
        }
    }

    /// Set time function.
    pub fn with_time_function(mut self, time_function: TimeFunction) -> Self {
        self.time_function = Some(time_function);
        self
    }

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

    /// Check if active.
    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
        }
    }

    /// Get body force at time.
    pub fn get_body_force(&self, time: f64) -> Vector3<f64> {
        if let Some(ref time_func) = self.time_function {
            self.force_density * time_func(time)
        } else {
            self.force_density
        }
    }

    /// Apply body force to global system.
    pub fn apply(
        &self,
        mesh: &Mesh,
        dof_numbering: &mut AdvancedDofNumbering,
        _global_stiffness: &mut SparseMatrix,
        global_force: &mut DVector<f64>,
        time: f64,
    ) -> FeaResult<()> {
        if !self.is_active(time) {
            return Ok(());
        }

        let body_force = self.get_body_force(time);

        for &element_id in &self.elements {
            let element = mesh
                .elements
                .get(&element_id)
                .ok_or(BoundaryError::ElementNotFound {
                    element_id: element_id.0,
                })?;

            // Compute actual element volume using quadrature integration
            let element_volume = self.compute_element_volume(mesh, element)?;

            // Distribute body force equally among nodes
            let force_per_node = body_force * element_volume / element.nodes.len() as f64;

            for &node_id in &element.nodes {
                for (comp_idx, &component) in [
                    DofComponent::DisplacementX,
                    DofComponent::DisplacementY,
                    DofComponent::DisplacementZ,
                ]
                .iter()
                .enumerate()
                {
                    if let Some(dof) = dof_numbering.get_dof(node_id, component) {
                        BoundaryConditionApplicator::apply_neumann_force(
                            dof,
                            force_per_node[comp_idx],
                            global_force,
                        )?;
                    }
                }
            }
        }

        Ok(())
    }

    /// Compute element volume using quadrature integration.
    fn compute_element_volume(&self, mesh: &Mesh, element: &Element) -> FeaResult<f64> {
        let element_nodes: Result<Vec<_>, _> = element
            .nodes
            .iter()
            .map(|&node_id| {
                mesh.nodes
                    .get(&node_id)
                    .ok_or(BoundaryError::NodeNotFound { node_id: node_id.0 })
            })
            .collect();
        let element_nodes = element_nodes?;

        let finite_element = StandardFiniteElement::new(
            element.element_type,
            element_nodes.iter().map(|n| n.position()).collect(),
        );
        let quadrature_points = finite_element.get_volume_quadrature_points()?;

        let mut volume = 0.0;
        for quad_point in quadrature_points {
            let jacobian = finite_element.compute_jacobian(&quad_point.coords)?;
            volume += quad_point.weight * jacobian.determinant().abs();
        }

        Ok(volume)
    }
}

/// Collection of common Neumann boundary condition patterns.
pub struct NeumannPatterns;

impl NeumannPatterns {
    /// Create concentrated load at a point.
    pub fn concentrated_load(node: NodeId, component: DofComponent, force: f64) -> NeumannBC {
        NeumannBC::fixed_force(vec![node], vec![component], force)
    }

    /// Create distributed load along a line of nodes.
    pub fn distributed_line_load(
        nodes: Vec<NodeId>,
        component: DofComponent,
        total_load: f64,
    ) -> NeumannBC {
        NeumannBC::fixed_force(nodes, vec![component], total_load).with_distribution(true)
    }

    /// Create pressure load on surface.
    pub fn pressure_load(nodes: Vec<NodeId>, pressure: f64) -> NeumannBC {
        NeumannBC::pressure_load(nodes, pressure, None)
    }

    /// Create time-varying concentrated load.
    pub fn harmonic_load(
        node: NodeId,
        component: DofComponent,
        amplitude: f64,
        frequency: f64,
    ) -> NeumannBC {
        NeumannBC::time_dependent_force(
            vec![node],
            vec![component],
            TimeFunction(Box::new(move |t| {
                amplitude * (2.0 * std::f64::consts::PI * frequency * t).sin()
            })),
        )
    }

    /// Create gravity load.
    pub fn gravity_load(elements: Vec<ElementId>, density: f64, gravity: f64) -> BodyForceBC {
        BodyForceBC::new(elements, Vector3::new(0.0, -density * gravity, 0.0))
    }

    /// Create thermal load.
    pub fn thermal_load(
        elements: Vec<ElementId>,
        thermal_expansion: f64,
        temperature_change: f64,
        elastic_modulus: f64,
    ) -> BodyForceBC {
        let thermal_stress = elastic_modulus * thermal_expansion * temperature_change;
        BodyForceBC::new(
            elements,
            Vector3::new(thermal_stress, thermal_stress, thermal_stress),
        )
    }
}

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

    #[test]
    fn test_neumann_type_fixed() {
        let neumann = NeumannType::Fixed(10.0);
        let position = Vector3::new(1.0, 2.0, 3.0);
        assert_eq!(neumann.evaluate(0.0, &position), 10.0);
        assert_eq!(neumann.evaluate(5.0, &position), 10.0);
    }

    #[test]
    fn test_neumann_type_harmonic() {
        let neumann = NeumannType::harmonic(5.0, 1.0, 0.0);
        let position = Vector3::new(0.0, 0.0, 0.0);

        assert!((neumann.evaluate(0.0, &position) - 0.0).abs() < 1e-10);
        assert!((neumann.evaluate(std::f64::consts::PI / 2.0, &position) - 5.0).abs() < 1e-10);
    }

    #[test]
    fn test_neumann_bc_creation() {
        let bc = NeumannBC::fixed_force(vec![NodeId(0)], vec![DofComponent::DisplacementX], 100.0);

        assert_eq!(bc.nodes, vec![NodeId(0)]);
        assert_eq!(bc.components, vec![DofComponent::DisplacementX]);
        assert!(bc.distribute_equally);
    }

    #[test]
    fn test_pressure_load() {
        let bc = NeumannBC::pressure_load(vec![NodeId(0), NodeId(1)], 1000.0, None);

        let position = Vector3::new(0.0, 0.0, 0.0);
        let force = bc.get_force(0.0, &position);
        assert_eq!(force, 500.0); // Distributed equally between 2 nodes
    }

    #[test]
    fn test_time_range() {
        let bc = NeumannBC::fixed_force(vec![NodeId(0)], vec![DofComponent::DisplacementY], 50.0)
            .with_time_range(1.0, 3.0);

        assert!(!bc.is_active(0.5));
        assert!(bc.is_active(2.0));
        assert!(!bc.is_active(4.0));
    }

    #[test]
    fn test_surface_traction_bc() {
        let bc = SurfaceTractionBC::new(vec![ElementId(0)], vec![0], Vector3::new(1.0, 0.0, 0.0));

        assert_eq!(bc.elements, vec![ElementId(0)]);
        assert_eq!(bc.traction, Vector3::new(1.0, 0.0, 0.0));
        assert!(bc.is_active(0.0));
    }

    #[test]
    fn test_body_force_bc() {
        let bc = BodyForceBC::new(
            vec![ElementId(0), ElementId(1)],
            Vector3::new(0.0, -9.81, 0.0),
        );

        assert_eq!(bc.elements, vec![ElementId(0), ElementId(1)]);
        assert_eq!(bc.force_density, Vector3::new(0.0, -9.81, 0.0));
        assert!(bc.is_active(0.0));
    }

    #[test]
    fn test_neumann_patterns_concentrated_load() {
        let bc = NeumannPatterns::concentrated_load(NodeId(5), DofComponent::DisplacementZ, 250.0);

        assert_eq!(bc.nodes, vec![NodeId(5)]);
        assert_eq!(bc.components, vec![DofComponent::DisplacementZ]);
    }

    #[test]
    fn test_neumann_patterns_gravity_load() {
        let elements = vec![ElementId(0), ElementId(1)];
        let bc = NeumannPatterns::gravity_load(elements.clone(), 7850.0, 9.81);

        assert_eq!(bc.elements, elements);
        assert!((bc.force_density.y + 7850.0 * 9.81).abs() < 1e-10);
    }

    #[test]
    fn test_ramp_load() {
        let neumann = NeumannType::ramp_load(1.0, 3.0, 0.0, 100.0);
        let position = Vector3::new(0.0, 0.0, 0.0);

        assert_eq!(neumann.evaluate(0.0, &position), 0.0);
        assert_eq!(neumann.evaluate(1.0, &position), 0.0);
        assert_eq!(neumann.evaluate(2.0, &position), 50.0);
        assert_eq!(neumann.evaluate(3.0, &position), 100.0);
        assert_eq!(neumann.evaluate(4.0, &position), 100.0);
    }

    #[test]
    fn test_step_load() {
        let neumann = NeumannType::step_load(2.0, 0.0, 75.0);
        let position = Vector3::new(0.0, 0.0, 0.0);

        assert_eq!(neumann.evaluate(1.0, &position), 0.0);
        assert_eq!(neumann.evaluate(2.0, &position), 75.0);
        assert_eq!(neumann.evaluate(3.0, &position), 75.0);
    }
}