//! Finite element model data structures.
//!
//! This module provides data structures for representing finite element models,
//! compatible with various FEM solvers including LS-DYNA and Abaqus.

use nalgebra::{Point3, Vector3};
use rtx_materials::{KelvinMaxwell, LinearElastic};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Element type enumeration.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ElementType {
    /// 4-node tetrahedral element.
    Tet4,
    /// 10-node tetrahedral element (quadratic).
    Tet10,
    /// 8-node hexahedral element.
    Hex8,
    /// 20-node hexahedral element (quadratic).
    Hex20,
    /// 4-node shell element.
    Shell4,
    /// 3-node triangular shell element.
    Tri3,
}

impl ElementType {
    /// Get the number of nodes for this element type.
    pub fn num_nodes(&self) -> usize {
        match self {
            ElementType::Tet4 => 4,
            ElementType::Tet10 => 10,
            ElementType::Hex8 => 8,
            ElementType::Hex20 => 20,
            ElementType::Shell4 => 4,
            ElementType::Tri3 => 3,
        }
    }

    /// Get the LS-DYNA element section type.
    pub fn lsdyna_section_type(&self) -> &'static str {
        match self {
            ElementType::Tet4 | ElementType::Tet10 => "SOLID",
            ElementType::Hex8 | ElementType::Hex20 => "SOLID",
            ElementType::Shell4 | ElementType::Tri3 => "SHELL",
        }
    }
}

/// Material definition for FEM export.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Material {
    /// Linear elastic material (MAT_001 in LS-DYNA).
    Elastic(LinearElastic),
    /// Kelvin-Maxwell viscoelastic material (MAT_076 in LS-DYNA).
    KelvinMaxwell(KelvinMaxwell),
    /// User-defined material with raw parameters.
    UserDefined {
        /// Material type identifier.
        mat_type: String,
        /// Material parameters.
        parameters: HashMap<String, f64>,
    },
}

impl Material {
    /// Get the LS-DYNA material type number.
    pub fn lsdyna_mat_type(&self) -> i32 {
        match self {
            Material::Elastic(_) => 1,
            Material::KelvinMaxwell(_) => 76,
            Material::UserDefined { .. } => 0,
        }
    }

    /// Get the density of the material.
    pub fn density(&self) -> Option<f64> {
        match self {
            Material::Elastic(e) => Some(e.density),
            Material::KelvinMaxwell(km) => Some(km.density),
            Material::UserDefined { parameters, .. } => parameters.get("density").copied(),
        }
    }
}

/// A finite element node.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Node {
    /// Node ID (1-based for FEM compatibility).
    pub id: u64,
    /// Node position in 3D space.
    pub position: Point3<f64>,
    /// Optional nodal displacement (for results).
    pub displacement: Option<Vector3<f64>>,
}

impl Node {
    /// Create a new node.
    pub fn new(id: u64, x: f64, y: f64, z: f64) -> Self {
        Self {
            id,
            position: Point3::new(x, y, z),
            displacement: None,
        }
    }
}

/// A finite element.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Element {
    /// Element ID (1-based for FEM compatibility).
    pub id: u64,
    /// Part ID this element belongs to.
    pub part_id: u64,
    /// Element type.
    pub element_type: ElementType,
    /// Node connectivity (1-based node IDs).
    pub nodes: Vec<u64>,
}

impl Element {
    /// Create a new element.
    pub fn new(id: u64, part_id: u64, element_type: ElementType, nodes: Vec<u64>) -> Self {
        Self {
            id,
            part_id,
            element_type,
            nodes,
        }
    }

    /// Create a 4-node tetrahedral element.
    pub fn tet4(id: u64, part_id: u64, n1: u64, n2: u64, n3: u64, n4: u64) -> Self {
        Self::new(id, part_id, ElementType::Tet4, vec![n1, n2, n3, n4])
    }

    /// Create a 10-node tetrahedral element.
    pub fn tet10(id: u64, part_id: u64, nodes: [u64; 10]) -> Self {
        Self::new(id, part_id, ElementType::Tet10, nodes.to_vec())
    }

    /// Create an 8-node hexahedral element.
    pub fn hex8(id: u64, part_id: u64, nodes: [u64; 8]) -> Self {
        Self::new(id, part_id, ElementType::Hex8, nodes.to_vec())
    }
}

/// A part definition grouping elements with a material.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Part {
    /// Part ID (1-based).
    pub id: u64,
    /// Part name/title.
    pub name: String,
    /// Section ID.
    pub section_id: u64,
    /// Material ID.
    pub material_id: u64,
    /// Element type for this part.
    pub element_type: ElementType,
}

impl Part {
    /// Create a new part.
    pub fn new(
        id: u64,
        name: impl Into<String>,
        section_id: u64,
        material_id: u64,
        element_type: ElementType,
    ) -> Self {
        Self {
            id,
            name: name.into(),
            section_id,
            material_id,
            element_type,
        }
    }
}

/// A node set for boundary conditions.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeSet {
    /// Node set ID.
    pub id: u64,
    /// Node set name.
    pub name: String,
    /// Node IDs in this set.
    pub nodes: Vec<u64>,
}

impl NodeSet {
    /// Create a new node set.
    pub fn new(id: u64, name: impl Into<String>, nodes: Vec<u64>) -> Self {
        Self {
            id,
            name: name.into(),
            nodes,
        }
    }
}

/// A complete finite element model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FEModel {
    /// Model title.
    pub title: String,
    /// Model description or study name.
    pub description: String,
    /// All nodes in the model.
    pub nodes: Vec<Node>,
    /// All elements in the model.
    pub elements: Vec<Element>,
    /// Part definitions.
    pub parts: Vec<Part>,
    /// Material definitions (keyed by material ID).
    pub materials: HashMap<u64, Material>,
    /// Node sets for boundary conditions.
    pub node_sets: Vec<NodeSet>,
    /// Units system identifier.
    pub units: String,
}

impl FEModel {
    /// Create a new empty FE model.
    pub fn new(title: impl Into<String>, description: impl Into<String>) -> Self {
        Self {
            title: title.into(),
            description: description.into(),
            nodes: Vec::new(),
            elements: Vec::new(),
            parts: Vec::new(),
            materials: HashMap::new(),
            node_sets: Vec::new(),
            units: "SI".to_string(),
        }
    }

    /// Add a node to the model.
    pub fn add_node(&mut self, node: Node) {
        self.nodes.push(node);
    }

    /// Add multiple nodes to the model.
    pub fn add_nodes(&mut self, nodes: impl IntoIterator<Item = Node>) {
        self.nodes.extend(nodes);
    }

    /// Add an element to the model.
    pub fn add_element(&mut self, element: Element) {
        self.elements.push(element);
    }

    /// Add multiple elements to the model.
    pub fn add_elements(&mut self, elements: impl IntoIterator<Item = Element>) {
        self.elements.extend(elements);
    }

    /// Add a part to the model.
    pub fn add_part(&mut self, part: Part) {
        self.parts.push(part);
    }

    /// Add a material to the model.
    pub fn add_material(&mut self, id: u64, material: Material) {
        self.materials.insert(id, material);
    }

    /// Add a node set to the model.
    pub fn add_node_set(&mut self, node_set: NodeSet) {
        self.node_sets.push(node_set);
    }

    /// Get the number of nodes.
    pub fn num_nodes(&self) -> usize {
        self.nodes.len()
    }

    /// Get the number of elements.
    pub fn num_elements(&self) -> usize {
        self.elements.len()
    }

    /// Get the number of parts.
    pub fn num_parts(&self) -> usize {
        self.parts.len()
    }

    /// Get a node by ID.
    pub fn get_node(&self, id: u64) -> Option<&Node> {
        self.nodes.iter().find(|n| n.id == id)
    }

    /// Get an element by ID.
    pub fn get_element(&self, id: u64) -> Option<&Element> {
        self.elements.iter().find(|e| e.id == id)
    }

    /// Get elements by part ID.
    pub fn get_elements_by_part(&self, part_id: u64) -> Vec<&Element> {
        self.elements
            .iter()
            .filter(|e| e.part_id == part_id)
            .collect()
    }

    /// Calculate the bounding box of the model.
    pub fn bounding_box(&self) -> Option<(Point3<f64>, Point3<f64>)> {
        if self.nodes.is_empty() {
            return None;
        }

        let mut min = self.nodes[0].position;
        let mut max = self.nodes[0].position;

        for node in &self.nodes {
            min.x = min.x.min(node.position.x);
            min.y = min.y.min(node.position.y);
            min.z = min.z.min(node.position.z);
            max.x = max.x.max(node.position.x);
            max.y = max.y.max(node.position.y);
            max.z = max.z.max(node.position.z);
        }

        Some((min, max))
    }

    /// Validate the model for completeness.
    pub fn validate(&self) -> Result<(), Vec<String>> {
        let mut errors = Vec::new();

        if self.nodes.is_empty() {
            errors.push("Model has no nodes".to_string());
        }

        if self.elements.is_empty() {
            errors.push("Model has no elements".to_string());
        }

        // Check that all element nodes exist
        let node_ids: std::collections::HashSet<u64> = self.nodes.iter().map(|n| n.id).collect();
        for element in &self.elements {
            for &node_id in &element.nodes {
                if !node_ids.contains(&node_id) {
                    errors.push(format!(
                        "Element {} references non-existent node {}",
                        element.id, node_id
                    ));
                }
            }
        }

        // Check that all element part IDs have corresponding parts
        let part_ids: std::collections::HashSet<u64> = self.parts.iter().map(|p| p.id).collect();
        for element in &self.elements {
            if !part_ids.contains(&element.part_id) {
                errors.push(format!(
                    "Element {} references non-existent part {}",
                    element.id, element.part_id
                ));
            }
        }

        // Check that all parts have corresponding materials
        for part in &self.parts {
            if !self.materials.contains_key(&part.material_id) {
                errors.push(format!(
                    "Part {} references non-existent material {}",
                    part.id, part.material_id
                ));
            }
        }

        if errors.is_empty() {
            Ok(())
        } else {
            Err(errors)
        }
    }
}

/// Builder for creating FE models.
#[derive(Debug)]
pub struct FEModelBuilder {
    model: FEModel,
    next_node_id: u64,
    next_element_id: u64,
    next_part_id: u64,
    next_material_id: u64,
    next_node_set_id: u64,
}

impl FEModelBuilder {
    /// Create a new model builder.
    pub fn new(title: impl Into<String>, description: impl Into<String>) -> Self {
        Self {
            model: FEModel::new(title, description),
            next_node_id: 1,
            next_element_id: 1,
            next_part_id: 1,
            next_material_id: 1,
            next_node_set_id: 1,
        }
    }

    /// Set the units system.
    pub fn units(mut self, units: impl Into<String>) -> Self {
        self.model.units = units.into();
        self
    }

    /// Add a node and return its ID.
    pub fn add_node(&mut self, x: f64, y: f64, z: f64) -> u64 {
        let id = self.next_node_id;
        self.next_node_id += 1;
        self.model.add_node(Node::new(id, x, y, z));
        id
    }

    /// Add multiple nodes from coordinate arrays.
    pub fn add_nodes_from_coords(&mut self, coords: &[[f64; 3]]) -> Vec<u64> {
        coords
            .iter()
            .map(|[x, y, z]| self.add_node(*x, *y, *z))
            .collect()
    }

    /// Add a material and return its ID.
    pub fn add_material(&mut self, material: Material) -> u64 {
        let id = self.next_material_id;
        self.next_material_id += 1;
        self.model.add_material(id, material);
        id
    }

    /// Add a part and return its ID.
    pub fn add_part(
        &mut self,
        name: impl Into<String>,
        material_id: u64,
        element_type: ElementType,
    ) -> u64 {
        let id = self.next_part_id;
        let section_id = id; // Use same ID for section
        self.next_part_id += 1;
        self.model
            .add_part(Part::new(id, name, section_id, material_id, element_type));
        id
    }

    /// Add an element and return its ID.
    pub fn add_element(&mut self, part_id: u64, element_type: ElementType, nodes: Vec<u64>) -> u64 {
        let id = self.next_element_id;
        self.next_element_id += 1;
        self.model
            .add_element(Element::new(id, part_id, element_type, nodes));
        id
    }

    /// Add a tetrahedral element and return its ID.
    pub fn add_tet4(&mut self, part_id: u64, n1: u64, n2: u64, n3: u64, n4: u64) -> u64 {
        self.add_element(part_id, ElementType::Tet4, vec![n1, n2, n3, n4])
    }

    /// Add a node set and return its ID.
    pub fn add_node_set(&mut self, name: impl Into<String>, nodes: Vec<u64>) -> u64 {
        let id = self.next_node_set_id;
        self.next_node_set_id += 1;
        self.model.add_node_set(NodeSet::new(id, name, nodes));
        id
    }

    /// Build the final model.
    pub fn build(self) -> FEModel {
        self.model
    }
}

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

    #[test]
    fn test_node_creation() {
        let node = Node::new(1, 1.0, 2.0, 3.0);
        assert_eq!(node.id, 1);
        assert_eq!(node.position.x, 1.0);
        assert_eq!(node.position.y, 2.0);
        assert_eq!(node.position.z, 3.0);
    }

    #[test]
    fn test_element_creation() {
        let elem = Element::tet4(1, 1, 1, 2, 3, 4);
        assert_eq!(elem.id, 1);
        assert_eq!(elem.part_id, 1);
        assert_eq!(elem.element_type, ElementType::Tet4);
        assert_eq!(elem.nodes.len(), 4);
    }

    #[test]
    fn test_model_builder() {
        let mut builder = FEModelBuilder::new("Test Model", "Test Study");

        // Add nodes
        let n1 = builder.add_node(0.0, 0.0, 0.0);
        let n2 = builder.add_node(1.0, 0.0, 0.0);
        let n3 = builder.add_node(0.5, 1.0, 0.0);
        let n4 = builder.add_node(0.5, 0.5, 1.0);

        // Add material
        let mat_id = builder.add_material(Material::Elastic(LinearElastic::brain_tissue()));

        // Add part
        let part_id = builder.add_part("Brain", mat_id, ElementType::Tet4);

        // Add element
        builder.add_tet4(part_id, n1, n2, n3, n4);

        let model = builder.build();

        assert_eq!(model.num_nodes(), 4);
        assert_eq!(model.num_elements(), 1);
        assert_eq!(model.num_parts(), 1);
        assert!(model.validate().is_ok());
    }

    #[test]
    fn test_bounding_box() {
        let mut builder = FEModelBuilder::new("Test", "");
        builder.add_node(-1.0, -2.0, -3.0);
        builder.add_node(1.0, 2.0, 3.0);

        let model = builder.build();
        let (min, max) = model.bounding_box().unwrap();

        assert_eq!(min.x, -1.0);
        assert_eq!(min.y, -2.0);
        assert_eq!(min.z, -3.0);
        assert_eq!(max.x, 1.0);
        assert_eq!(max.y, 2.0);
        assert_eq!(max.z, 3.0);
    }

    #[test]
    fn test_element_type_num_nodes() {
        assert_eq!(ElementType::Tet4.num_nodes(), 4);
        assert_eq!(ElementType::Tet10.num_nodes(), 10);
        assert_eq!(ElementType::Hex8.num_nodes(), 8);
    }
}