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

//! Boundary condition definitions and application.
//!
//! This module provides comprehensive boundary condition types and efficient
//! application methods for finite element analysis.

pub mod contact;
pub mod dirichlet;
pub mod neumann;
pub mod robin;
pub mod thermal;

use crate::assembly::{AdvancedDofNumbering, DofComponent};
use crate::error::{BoundaryError, FeaResult};
use crate::mesh::{ElementId, Mesh, NodeId};
use nalgebra::{DVector, Vector3};
use std::collections::{HashMap, HashSet};

pub use contact::*;
pub use dirichlet::*;
pub use neumann::*;
pub use robin::*;
pub use thermal::*;

/// Time-dependent boundary condition function.
pub struct TimeFunction(pub Box<dyn Fn(f64) -> f64 + Send + Sync>);

impl std::fmt::Debug for TimeFunction {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("TimeFunction")
    }
}

impl std::ops::Deref for TimeFunction {
    type Target = dyn Fn(f64) -> f64 + Send + Sync;

    fn deref(&self) -> &Self::Target {
        &*self.0
    }
}

/// Spatial boundary condition function.
pub struct SpatialFunction(pub Box<dyn Fn(&Vector3<f64>) -> f64 + Send + Sync>);

impl std::fmt::Debug for SpatialFunction {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("SpatialFunction")
    }
}

impl std::ops::Deref for SpatialFunction {
    type Target = dyn Fn(&Vector3<f64>) -> f64 + Send + Sync;

    fn deref(&self) -> &Self::Target {
        &*self.0
    }
}

/// Combined time-spatial boundary condition function.
pub struct TimeSpatialFunction(pub Box<dyn Fn(f64, &Vector3<f64>) -> f64 + Send + Sync>);

impl std::fmt::Debug for TimeSpatialFunction {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("TimeSpatialFunction")
    }
}

impl std::ops::Deref for TimeSpatialFunction {
    type Target = dyn Fn(f64, &Vector3<f64>) -> f64 + Send + Sync;

    fn deref(&self) -> &Self::Target {
        &*self.0
    }
}

/// Boundary condition types supported by the FEA solver.
#[derive(Debug)]
pub enum BoundaryCondition {
    /// Dirichlet (essential) boundary condition
    Dirichlet(DirichletBC),
    /// Neumann (natural) boundary condition
    Neumann(NeumannBC),
    /// Robin (mixed) boundary condition
    Robin(RobinBC),
    /// Contact boundary condition
    Contact(ContactBC),
    /// Thermal boundary condition
    Thermal(ThermalBC),
}

impl BoundaryCondition {
    /// Get the nodes affected by this boundary condition.
    pub fn affected_nodes(&self) -> Vec<NodeId> {
        match self {
            Self::Dirichlet(bc) => bc.nodes.clone(),
            Self::Neumann(bc) => bc.nodes.clone(),
            Self::Robin(bc) => bc.nodes.clone(),
            Self::Contact(bc) => bc.slave_nodes.clone(),
            Self::Thermal(bc) => bc.nodes.clone(),
        }
    }

    /// Check if this boundary condition applies at the given time.
    pub fn is_active(&self, time: f64) -> bool {
        match self {
            Self::Dirichlet(bc) => bc.is_active(time),
            Self::Neumann(bc) => bc.is_active(time),
            Self::Robin(bc) => bc.is_active(time),
            Self::Contact(bc) => bc.is_active(time),
            Self::Thermal(bc) => bc.is_active(time),
        }
    }

    /// Apply this boundary condition to the global system.
    pub fn apply(
        &self,
        mesh: &Mesh,
        dof_numbering: &mut AdvancedDofNumbering,
        global_stiffness: &mut crate::assembly::SparseMatrix,
        global_force: &mut DVector<f64>,
        time: f64,
    ) -> FeaResult<()> {
        if !self.is_active(time) {
            return Ok(());
        }

        match self {
            Self::Dirichlet(bc) => {
                bc.apply(mesh, dof_numbering, global_stiffness, global_force, time)
            }
            Self::Neumann(bc) => {
                bc.apply(mesh, dof_numbering, global_stiffness, global_force, time)
            }
            Self::Robin(bc) => bc.apply(mesh, dof_numbering, global_stiffness, global_force, time),
            Self::Contact(bc) => {
                bc.apply(mesh, dof_numbering, global_stiffness, global_force, time)
            }
            Self::Thermal(bc) => {
                bc.apply(mesh, dof_numbering, global_stiffness, global_force, time)
            }
        }
    }
}

/// Collection of boundary conditions with management utilities.
#[derive(Debug)]
pub struct BoundaryConditionSet {
    /// All boundary conditions
    conditions: Vec<BoundaryCondition>,
    /// Node sets for organized boundary condition application
    node_sets: HashMap<String, HashSet<NodeId>>,
    /// Element sets for surface-based boundary conditions
    element_sets: HashMap<String, HashSet<ElementId>>,
    /// Time step information for time-dependent conditions
    time_info: TimeInfo,
}

/// Time stepping information for boundary conditions.
#[derive(Debug, Clone)]
pub struct TimeInfo {
    pub current_time: f64,
    pub time_step: f64,
    pub start_time: f64,
    pub end_time: f64,
}

impl Default for TimeInfo {
    fn default() -> Self {
        Self {
            current_time: 0.0,
            time_step: 0.1,
            start_time: 0.0,
            end_time: 1.0,
        }
    }
}

impl BoundaryConditionSet {
    /// Create a new boundary condition set.
    pub fn new() -> Self {
        Self {
            conditions: Vec::new(),
            node_sets: HashMap::new(),
            element_sets: HashMap::new(),
            time_info: TimeInfo::default(),
        }
    }

    /// Add a boundary condition to the set.
    pub fn add_condition(&mut self, condition: BoundaryCondition) {
        self.conditions.push(condition);
    }

    /// Add multiple boundary conditions.
    pub fn add_conditions(&mut self, conditions: Vec<BoundaryCondition>) {
        self.conditions.extend(conditions);
    }

    /// Get a reference to the conditions.
    pub fn conditions(&self) -> &[BoundaryCondition] {
        &self.conditions
    }

    /// Create a node set with a name.
    pub fn create_node_set(&mut self, name: String, nodes: HashSet<NodeId>) {
        self.node_sets.insert(name, nodes);
    }

    /// Create an element set with a name.
    pub fn create_element_set(&mut self, name: String, elements: HashSet<ElementId>) {
        self.element_sets.insert(name, elements);
    }

    /// Get nodes in a named set.
    pub fn get_node_set(&self, name: &str) -> Option<&HashSet<NodeId>> {
        self.node_sets.get(name)
    }

    /// Get elements in a named set.
    pub fn get_element_set(&self, name: &str) -> Option<&HashSet<ElementId>> {
        self.element_sets.get(name)
    }

    /// Apply all boundary conditions at the current time.
    pub fn apply_all(
        &self,
        mesh: &Mesh,
        dof_numbering: &mut AdvancedDofNumbering,
        global_stiffness: &mut crate::assembly::SparseMatrix,
        global_force: &mut DVector<f64>,
    ) -> FeaResult<()> {
        for condition in &self.conditions {
            condition.apply(
                mesh,
                dof_numbering,
                global_stiffness,
                global_force,
                self.time_info.current_time,
            )?;
        }
        Ok(())
    }

    /// Apply boundary conditions at a specific time.
    pub fn apply_at_time(
        &self,
        mesh: &Mesh,
        dof_numbering: &mut AdvancedDofNumbering,
        global_stiffness: &mut crate::assembly::SparseMatrix,
        global_force: &mut DVector<f64>,
        time: f64,
    ) -> FeaResult<()> {
        for condition in &self.conditions {
            condition.apply(mesh, dof_numbering, global_stiffness, global_force, time)?;
        }
        Ok(())
    }

    /// Update time information.
    pub fn set_time_info(&mut self, time_info: TimeInfo) {
        self.time_info = time_info;
    }

    /// Advance time by one step.
    pub fn advance_time(&mut self) {
        self.time_info.current_time += self.time_info.time_step;
    }

    /// Get current time.
    pub fn current_time(&self) -> f64 {
        self.time_info.current_time
    }

    /// Get all nodes affected by boundary conditions.
    pub fn affected_nodes(&self) -> HashSet<NodeId> {
        let mut nodes = HashSet::new();
        for condition in &self.conditions {
            nodes.extend(condition.affected_nodes());
        }
        nodes
    }

    /// Count active boundary conditions at current time.
    pub fn count_active_conditions(&self) -> usize {
        self.conditions
            .iter()
            .filter(|bc| bc.is_active(self.time_info.current_time))
            .count()
    }

    /// Get statistics about the boundary condition set.
    pub fn statistics(&self) -> BoundaryConditionStatistics {
        let mut dirichlet_count = 0;
        let mut neumann_count = 0;
        let mut robin_count = 0;
        let mut contact_count = 0;
        let mut thermal_count = 0;

        for condition in &self.conditions {
            match condition {
                BoundaryCondition::Dirichlet(_) => dirichlet_count += 1,
                BoundaryCondition::Neumann(_) => neumann_count += 1,
                BoundaryCondition::Robin(_) => robin_count += 1,
                BoundaryCondition::Contact(_) => contact_count += 1,
                BoundaryCondition::Thermal(_) => thermal_count += 1,
            }
        }

        BoundaryConditionStatistics {
            total_conditions: self.conditions.len(),
            active_conditions: self.count_active_conditions(),
            dirichlet_count,
            neumann_count,
            robin_count,
            contact_count,
            thermal_count,
            node_sets: self.node_sets.len(),
            element_sets: self.element_sets.len(),
            affected_nodes: self.affected_nodes().len(),
        }
    }

    /// Validate all boundary conditions.
    pub fn validate(&self, mesh: &Mesh) -> FeaResult<ValidationReport> {
        let mut errors = Vec::new();
        let mut warnings = Vec::new();

        // Check that all referenced nodes exist in the mesh
        for condition in &self.conditions {
            for node_id in condition.affected_nodes() {
                if !mesh.nodes.contains_key(&node_id) {
                    errors.push(ValidationError::NodeNotFound(node_id));
                }
            }
        }

        // Check for conflicting Dirichlet conditions
        let mut dirichlet_dofs = HashMap::new();
        for (index, condition) in self.conditions.iter().enumerate() {
            if let BoundaryCondition::Dirichlet(dirichlet_bc) = condition {
                for &node_id in &dirichlet_bc.nodes {
                    for &component in &dirichlet_bc.components {
                        let key = (node_id, component);
                        if let Some(existing_index) = dirichlet_dofs.get(&key) {
                            warnings.push(ValidationWarning::ConflictingDirichlet {
                                node_id,
                                component,
                                condition1: *existing_index,
                                condition2: index,
                            });
                        } else {
                            dirichlet_dofs.insert(key, index);
                        }
                    }
                }
            }
        }

        // Check node sets reference valid nodes
        for (set_name, nodes) in &self.node_sets {
            for &node_id in nodes {
                if !mesh.nodes.contains_key(&node_id) {
                    errors.push(ValidationError::NodeSetInvalidNode {
                        set_name: set_name.clone(),
                        node_id,
                    });
                }
            }
        }

        // Check element sets reference valid elements
        for (set_name, elements) in &self.element_sets {
            for &element_id in elements {
                if !mesh.elements.contains_key(&element_id) {
                    errors.push(ValidationError::ElementSetInvalidElement {
                        set_name: set_name.clone(),
                        element_id,
                    });
                }
            }
        }

        Ok(ValidationReport {
            is_valid: errors.is_empty(),
            errors,
            warnings,
        })
    }
}

impl Default for BoundaryConditionSet {
    fn default() -> Self {
        Self::new()
    }
}

/// Statistics about boundary conditions.
#[derive(Debug, Clone)]
pub struct BoundaryConditionStatistics {
    pub total_conditions: usize,
    pub active_conditions: usize,
    pub dirichlet_count: usize,
    pub neumann_count: usize,
    pub robin_count: usize,
    pub contact_count: usize,
    pub thermal_count: usize,
    pub node_sets: usize,
    pub element_sets: usize,
    pub affected_nodes: usize,
}

impl std::fmt::Display for BoundaryConditionStatistics {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "Boundary Condition Statistics:")?;
        writeln!(f, "  Total conditions: {}", self.total_conditions)?;
        writeln!(f, "  Active conditions: {}", self.active_conditions)?;
        writeln!(f, "  Dirichlet conditions: {}", self.dirichlet_count)?;
        writeln!(f, "  Neumann conditions: {}", self.neumann_count)?;
        writeln!(f, "  Robin conditions: {}", self.robin_count)?;
        writeln!(f, "  Contact conditions: {}", self.contact_count)?;
        writeln!(f, "  Thermal conditions: {}", self.thermal_count)?;
        writeln!(f, "  Node sets: {}", self.node_sets)?;
        writeln!(f, "  Element sets: {}", self.element_sets)?;
        writeln!(f, "  Affected nodes: {}", self.affected_nodes)?;
        Ok(())
    }
}

/// Validation errors for boundary conditions.
#[derive(Debug, Clone)]
pub enum ValidationError {
    NodeNotFound(NodeId),
    NodeSetInvalidNode {
        set_name: String,
        node_id: NodeId,
    },
    ElementSetInvalidElement {
        set_name: String,
        element_id: ElementId,
    },
}

/// Validation warnings for boundary conditions.
#[derive(Debug, Clone)]
pub enum ValidationWarning {
    ConflictingDirichlet {
        node_id: NodeId,
        component: DofComponent,
        condition1: usize,
        condition2: usize,
    },
}

/// Validation report for boundary conditions.
#[derive(Debug, Clone)]
pub struct ValidationReport {
    pub is_valid: bool,
    pub errors: Vec<ValidationError>,
    pub warnings: Vec<ValidationWarning>,
}

impl std::fmt::Display for ValidationReport {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "Boundary Condition Validation Report:")?;
        writeln!(f, "  Valid: {}", self.is_valid)?;
        writeln!(f, "  Errors: {}", self.errors.len())?;
        writeln!(f, "  Warnings: {}", self.warnings.len())?;

        if !self.errors.is_empty() {
            writeln!(f, "\nErrors:")?;
            for error in &self.errors {
                writeln!(f, "  {error:?}")?;
            }
        }

        if !self.warnings.is_empty() {
            writeln!(f, "\nWarnings:")?;
            for warning in &self.warnings {
                writeln!(f, "  {warning:?}")?;
            }
        }

        Ok(())
    }
}

/// Boundary condition application utilities.
pub struct BoundaryConditionApplicator;

impl BoundaryConditionApplicator {
    /// Apply Dirichlet boundary condition using elimination method.
    pub fn apply_dirichlet_elimination(
        dof: usize,
        value: f64,
        global_stiffness: &mut crate::assembly::SparseMatrix,
        global_force: &mut DVector<f64>,
    ) -> FeaResult<()> {
        global_stiffness.set_dirichlet(dof, value)?;
        global_force[dof] = value;
        Ok(())
    }

    /// Apply Neumann boundary condition (add to force vector).
    pub fn apply_neumann_force(
        dof: usize,
        force: f64,
        global_force: &mut DVector<f64>,
    ) -> FeaResult<()> {
        if dof < global_force.len() {
            global_force[dof] += force;
            Ok(())
        } else {
            Err(BoundaryError::DofOutOfBounds {
                dof,
                max_dof: global_force.len().saturating_sub(1),
            }
            .into())
        }
    }

    /// Apply Robin boundary condition (add to both stiffness and force).
    pub fn apply_robin(
        dof: usize,
        stiffness_contribution: f64,
        force_contribution: f64,
        global_stiffness: &mut crate::assembly::SparseMatrix,
        global_force: &mut DVector<f64>,
    ) -> FeaResult<()> {
        global_stiffness.add_entry(dof, dof, stiffness_contribution)?;
        global_force[dof] += force_contribution;
        Ok(())
    }
}

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

    #[test]
    fn test_boundary_condition_set_creation() {
        let bc_set = BoundaryConditionSet::new();
        assert_eq!(bc_set.conditions.len(), 0);
        assert_eq!(bc_set.node_sets.len(), 0);
    }

    #[test]
    fn test_time_info() {
        let time_info = TimeInfo::default();
        assert_eq!(time_info.current_time, 0.0);
        assert_eq!(time_info.time_step, 0.1);
    }

    #[test]
    fn test_node_set_creation() {
        let mut bc_set = BoundaryConditionSet::new();
        let mut nodes = HashSet::new();
        nodes.insert(NodeId(0));
        nodes.insert(NodeId(1));

        bc_set.create_node_set("fixed_nodes".to_string(), nodes.clone());
        assert_eq!(bc_set.get_node_set("fixed_nodes"), Some(&nodes));
    }

    #[test]
    fn test_boundary_condition_statistics() {
        let bc_set = BoundaryConditionSet::new();
        let stats = bc_set.statistics();

        assert_eq!(stats.total_conditions, 0);
        assert_eq!(stats.active_conditions, 0);
        assert_eq!(stats.dirichlet_count, 0);
    }

    #[test]
    fn test_validation_empty_set() {
        let rect = Rectangle::new(1.0, 1.0);
        let mesh = rect.generate_quad_mesh(1, 1, MaterialId(0)).unwrap();
        let bc_set = BoundaryConditionSet::new();

        let report = bc_set.validate(&mesh).unwrap();
        assert!(report.is_valid);
        assert_eq!(report.errors.len(), 0);
        assert_eq!(report.warnings.len(), 0);
    }

    #[test]
    fn test_boundary_condition_applicator() {
        use crate::assembly::SparseMatrix;

        let mut stiffness = SparseMatrix::new(3, 3);
        let mut force = DVector::zeros(3);

        // Test Dirichlet application
        BoundaryConditionApplicator::apply_dirichlet_elimination(
            0,
            1.0,
            &mut stiffness,
            &mut force,
        )
        .unwrap();
        assert_eq!(force[0], 1.0);

        // Test Neumann application
        BoundaryConditionApplicator::apply_neumann_force(1, 10.0, &mut force).unwrap();
        assert_eq!(force[1], 10.0);
    }
}