//! Mesh generation functions for common geometries

use crate::error::FeaResult;
use nalgebra::DVector;

use super::element_types::ElementType;
use super::elements::{Element, MaterialId};
use super::mesh_core::{MaterialProperties, Mesh};
use super::nodes::{Node, NodeId};

impl Mesh {
    /// Generate a rectangular 2D mesh.
    pub fn generate_rectangle(width: f64, height: f64, nx: usize, ny: usize) -> FeaResult<Self> {
        let mut mesh = Self::new(2)?;

        let dx = width / (nx - 1) as f64;
        let dy = height / (ny - 1) as f64;

        // Create nodes
        let mut node_grid = vec![vec![NodeId(0); nx]; ny];
        for j in 0..ny {
            for i in 0..nx {
                let x = i as f64 * dx;
                let y = j as f64 * dy;
                let node = Node::new_2d(x, y);
                node_grid[j][i] = mesh.add_node(node);
            }
        }

        // Create quad elements
        let mat_id = MaterialId(0);
        for j in 0..ny - 1 {
            for i in 0..nx - 1 {
                let n1 = node_grid[j][i];
                let n2 = node_grid[j][i + 1];
                let n3 = node_grid[j + 1][i + 1];
                let n4 = node_grid[j + 1][i];

                let element = Element::new(ElementType::Quad4, vec![n1, n2, n3, n4], mat_id)?;
                mesh.add_element(element)?;
            }
        }

        // Add default material
        mesh.materials.insert(
            mat_id,
            MaterialProperties {
                name: "default".to_string(),
                youngs_modulus: 200e9,
                poissons_ratio: 0.3,
                density: 7850.0,
                thermal_conductivity: None,
                specific_heat: None,
                thermal_expansion: None,
            },
        );

        Ok(mesh)
    }

    /// Generate a box-shaped 3D mesh.
    pub fn generate_box(
        width: f64,
        height: f64,
        depth: f64,
        nx: usize,
        ny: usize,
        nz: usize,
    ) -> FeaResult<Self> {
        let mut mesh = Self::new(3)?;

        let dx = width / (nx - 1) as f64;
        let dy = height / (ny - 1) as f64;
        let dz = depth / (nz - 1) as f64;

        // Create nodes
        let mut node_grid = vec![vec![vec![NodeId(0); nx]; ny]; nz];
        for k in 0..nz {
            for j in 0..ny {
                for i in 0..nx {
                    let x = i as f64 * dx;
                    let y = j as f64 * dy;
                    let z = k as f64 * dz;
                    let node = Node::new_3d(x, y, z);
                    node_grid[k][j][i] = mesh.add_node(node);
                }
            }
        }

        // Create hex elements
        let mat_id = MaterialId(0);
        for k in 0..nz - 1 {
            for j in 0..ny - 1 {
                for i in 0..nx - 1 {
                    let n1 = node_grid[k][j][i];
                    let n2 = node_grid[k][j][i + 1];
                    let n3 = node_grid[k][j + 1][i + 1];
                    let n4 = node_grid[k][j + 1][i];
                    let n5 = node_grid[k + 1][j][i];
                    let n6 = node_grid[k + 1][j][i + 1];
                    let n7 = node_grid[k + 1][j + 1][i + 1];
                    let n8 = node_grid[k + 1][j + 1][i];

                    let element = Element::new(
                        ElementType::Hex8,
                        vec![n1, n2, n3, n4, n5, n6, n7, n8],
                        mat_id,
                    )?;
                    mesh.add_element(element)?;
                }
            }
        }

        // Add default material
        mesh.materials.insert(
            mat_id,
            MaterialProperties {
                name: "default".to_string(),
                youngs_modulus: 200e9,
                poissons_ratio: 0.3,
                density: 7850.0,
                thermal_conductivity: None,
                specific_heat: None,
                thermal_expansion: None,
            },
        );

        Ok(mesh)
    }

    /// Generate an icosphere mesh (geodesic sphere).
    pub fn generate_icosphere(radius: f64, subdivisions: usize) -> FeaResult<Self> {
        let mut mesh = Self::new(3)?;

        // Golden ratio
        let phi = f64::midpoint(1.0, 5.0_f64.sqrt());

        // Initial icosahedron vertices (normalized to unit sphere)
        let t = 1.0 / (1.0 + phi * phi).sqrt();
        let s = phi * t;

        let initial_vertices = vec![
            DVector::from_vec(vec![-t, s, 0.0]),
            DVector::from_vec(vec![t, s, 0.0]),
            DVector::from_vec(vec![-t, -s, 0.0]),
            DVector::from_vec(vec![t, -s, 0.0]),
            DVector::from_vec(vec![0.0, -t, s]),
            DVector::from_vec(vec![0.0, t, s]),
            DVector::from_vec(vec![0.0, -t, -s]),
            DVector::from_vec(vec![0.0, t, -s]),
            DVector::from_vec(vec![s, 0.0, -t]),
            DVector::from_vec(vec![s, 0.0, t]),
            DVector::from_vec(vec![-s, 0.0, -t]),
            DVector::from_vec(vec![-s, 0.0, t]),
        ];

        // Scale to desired radius
        let mut vertices: Vec<DVector<f64>> = initial_vertices
            .into_iter()
            .map(|v| {
                let normalized = &v / v.norm();
                normalized * radius
            })
            .collect();

        // Initial icosahedron faces (triangles)
        let mut faces = vec![
            (0, 11, 5),
            (0, 5, 1),
            (0, 1, 7),
            (0, 7, 10),
            (0, 10, 11),
            (1, 5, 9),
            (5, 11, 4),
            (11, 10, 2),
            (10, 7, 6),
            (7, 1, 8),
            (3, 9, 4),
            (3, 4, 2),
            (3, 2, 6),
            (3, 6, 8),
            (3, 8, 9),
            (4, 9, 5),
            (2, 4, 11),
            (6, 2, 10),
            (8, 6, 7),
            (9, 8, 1),
        ];

        // Subdivide faces
        for _ in 0..subdivisions {
            let mut new_faces = Vec::new();
            let mut edge_midpoints = std::collections::HashMap::new();

            for &(v0, v1, v2) in &faces {
                // Get or create edge midpoints
                let m01 = *edge_midpoints
                    .entry((v0.min(v1), v0.max(v1)))
                    .or_insert_with(|| {
                        let mid = (&vertices[v0] + &vertices[v1]) * 0.5;
                        let normalized = &mid / mid.norm();
                        let scaled = normalized * radius;
                        vertices.push(scaled);
                        vertices.len() - 1
                    });

                let m12 = *edge_midpoints
                    .entry((v1.min(v2), v1.max(v2)))
                    .or_insert_with(|| {
                        let mid = (&vertices[v1] + &vertices[v2]) * 0.5;
                        let normalized = &mid / mid.norm();
                        let scaled = normalized * radius;
                        vertices.push(scaled);
                        vertices.len() - 1
                    });

                let m20 = *edge_midpoints
                    .entry((v2.min(v0), v2.max(v0)))
                    .or_insert_with(|| {
                        let mid = (&vertices[v2] + &vertices[v0]) * 0.5;
                        let normalized = &mid / mid.norm();
                        let scaled = normalized * radius;
                        vertices.push(scaled);
                        vertices.len() - 1
                    });

                // Create 4 new triangles
                new_faces.push((v0, m01, m20));
                new_faces.push((v1, m12, m01));
                new_faces.push((v2, m20, m12));
                new_faces.push((m01, m12, m20));
            }

            faces = new_faces;
        }

        // Add vertices as nodes
        let mut node_ids = Vec::new();
        for vertex in vertices {
            let node = Node::new(vertex);
            node_ids.push(mesh.add_node(node));
        }

        // Add faces as triangular elements
        let mat_id = MaterialId(0);
        for (v0, v1, v2) in faces {
            let element = Element::new(
                ElementType::Tri3,
                vec![node_ids[v0], node_ids[v1], node_ids[v2]],
                mat_id,
            )?;
            mesh.add_element(element)?;
        }

        // Add default material
        mesh.materials.insert(
            mat_id,
            MaterialProperties {
                name: "sphere".to_string(),
                youngs_modulus: 200e9,
                poissons_ratio: 0.3,
                density: 7850.0,
                thermal_conductivity: None,
                specific_heat: None,
                thermal_expansion: None,
            },
        );

        Ok(mesh)
    }

    /// Generate a UV sphere mesh.
    pub fn generate_uv_sphere(
        radius: f64,
        n_latitude: usize,
        n_longitude: usize,
    ) -> FeaResult<Self> {
        let mut mesh = Self::new(3)?;

        // Create nodes
        let mut node_grid = vec![vec![None; n_longitude]; n_latitude + 1];

        // Add top pole
        let top_node = mesh.add_node(Node::new_3d(0.0, 0.0, radius));

        // Add middle latitude circles
        for i in 1..n_latitude {
            let theta = std::f64::consts::PI * (i as f64) / (n_latitude as f64);
            let sin_theta = theta.sin();
            let cos_theta = theta.cos();

            for j in 0..n_longitude {
                let phi = 2.0 * std::f64::consts::PI * (j as f64) / (n_longitude as f64);
                let x = radius * sin_theta * phi.cos();
                let y = radius * sin_theta * phi.sin();
                let z = radius * cos_theta;

                let node = Node::new_3d(x, y, z);
                node_grid[i][j] = Some(mesh.add_node(node));
            }
        }

        // Add bottom pole
        let bottom_node = mesh.add_node(Node::new_3d(0.0, 0.0, -radius));

        // Create elements
        let mat_id = MaterialId(0);

        // Top cap (triangles connecting to north pole)
        for j in 0..n_longitude {
            let j_next = (j + 1) % n_longitude;
            let n1 = top_node;
            let n2 = node_grid[1][j].unwrap();
            let n3 = node_grid[1][j_next].unwrap();

            let element = Element::new(ElementType::Tri3, vec![n1, n2, n3], mat_id)?;
            mesh.add_element(element)?;
        }

        // Middle bands (quads)
        for i in 1..n_latitude - 1 {
            for j in 0..n_longitude {
                let j_next = (j + 1) % n_longitude;
                let n1 = node_grid[i][j].unwrap();
                let n2 = node_grid[i][j_next].unwrap();
                let n3 = node_grid[i + 1][j_next].unwrap();
                let n4 = node_grid[i + 1][j].unwrap();

                let element = Element::new(ElementType::Quad4, vec![n1, n2, n3, n4], mat_id)?;
                mesh.add_element(element)?;
            }
        }

        // Bottom cap (triangles connecting to south pole)
        for j in 0..n_longitude {
            let j_next = (j + 1) % n_longitude;
            let n1 = node_grid[n_latitude - 1][j].unwrap();
            let n2 = node_grid[n_latitude - 1][j_next].unwrap();
            let n3 = bottom_node;

            let element = Element::new(ElementType::Tri3, vec![n1, n2, n3], mat_id)?;
            mesh.add_element(element)?;
        }

        // Add default material
        mesh.materials.insert(
            mat_id,
            MaterialProperties {
                name: "sphere".to_string(),
                youngs_modulus: 200e9,
                poissons_ratio: 0.3,
                density: 7850.0,
                thermal_conductivity: None,
                specific_heat: None,
                thermal_expansion: None,
            },
        );

        Ok(mesh)
    }

    /// Refine a sphere mesh while maintaining spherical shape.
    pub fn refine_sphere(&mut self, radius: f64) -> FeaResult<()> {
        // Store original elements
        let original_elements: Vec<_> = self.elements.clone().into_iter().collect();

        // Clear elements (keep nodes)
        self.elements.clear();

        // Track edge midpoints
        let mut edge_midpoints = std::collections::HashMap::new();

        // Helper to get or create spherical midpoint
        let get_spherical_midpoint =
            |n1: NodeId,
             n2: NodeId,
             edge_midpoints: &mut std::collections::HashMap<(NodeId, NodeId), NodeId>,
             mesh: &mut Self|
             -> NodeId {
                let edge_key = if n1 < n2 { (n1, n2) } else { (n2, n1) };

                if let Some(&mid_id) = edge_midpoints.get(&edge_key) {
                    return mid_id;
                }

                // Create midpoint on sphere surface
                let node1 = &mesh.nodes[&n1];
                let node2 = &mesh.nodes[&n2];
                let midpoint = (&node1.coordinates + &node2.coordinates) * 0.5;

                // Project to sphere surface
                let normalized = &midpoint / midpoint.norm();
                let spherical_point = normalized * radius;

                let mid_node = Node::new(spherical_point);
                let mid_id = mesh.add_node(mid_node);
                edge_midpoints.insert(edge_key, mid_id);
                mid_id
            };

        // Process each element
        for (_elem_id, element) in original_elements {
            let nodes = &element.nodes;
            let mat_id = element.material_id;

            match element.element_type {
                ElementType::Tri3 => {
                    // Subdivide triangle into 4 triangles
                    let n0 = nodes[0];
                    let n1 = nodes[1];
                    let n2 = nodes[2];

                    let m01 = get_spherical_midpoint(n0, n1, &mut edge_midpoints, self);
                    let m12 = get_spherical_midpoint(n1, n2, &mut edge_midpoints, self);
                    let m20 = get_spherical_midpoint(n2, n0, &mut edge_midpoints, self);

                    // Create 4 new triangles
                    self.add_element(Element::new(ElementType::Tri3, vec![n0, m01, m20], mat_id)?)?;
                    self.add_element(Element::new(ElementType::Tri3, vec![m01, n1, m12], mat_id)?)?;
                    self.add_element(Element::new(ElementType::Tri3, vec![m20, m12, n2], mat_id)?)?;
                    self.add_element(Element::new(
                        ElementType::Tri3,
                        vec![m01, m12, m20],
                        mat_id,
                    )?)?;
                }

                ElementType::Quad4 => {
                    // Subdivide quad into 4 quads
                    let n0 = nodes[0];
                    let n1 = nodes[1];
                    let n2 = nodes[2];
                    let n3 = nodes[3];

                    let m01 = get_spherical_midpoint(n0, n1, &mut edge_midpoints, self);
                    let m12 = get_spherical_midpoint(n1, n2, &mut edge_midpoints, self);
                    let m23 = get_spherical_midpoint(n2, n3, &mut edge_midpoints, self);
                    let m30 = get_spherical_midpoint(n3, n0, &mut edge_midpoints, self);

                    // Create center point on sphere
                    let node0 = &self.nodes[&n0];
                    let node2 = &self.nodes[&n2];
                    let center_coords = (&node0.coordinates + &node2.coordinates) * 0.5;
                    let normalized = &center_coords / center_coords.norm();
                    let spherical_center = normalized * radius;
                    let center = self.add_node(Node::new(spherical_center));

                    // Create 4 new quads
                    self.add_element(Element::new(
                        ElementType::Quad4,
                        vec![n0, m01, center, m30],
                        mat_id,
                    )?)?;
                    self.add_element(Element::new(
                        ElementType::Quad4,
                        vec![m01, n1, m12, center],
                        mat_id,
                    )?)?;
                    self.add_element(Element::new(
                        ElementType::Quad4,
                        vec![center, m12, n2, m23],
                        mat_id,
                    )?)?;
                    self.add_element(Element::new(
                        ElementType::Quad4,
                        vec![m30, center, m23, n3],
                        mat_id,
                    )?)?;
                }

                _ => {
                    // For other element types, just copy them back
                    self.add_element(element)?;
                }
            }
        }

        Ok(())
    }

    /// Generate a cylinder mesh.
    pub fn generate_cylinder(
        radius: f64,
        height: f64,
        n_radial: usize,
        n_axial: usize,
    ) -> FeaResult<Self> {
        let mut mesh = Self::new(3)?;

        // Create nodes
        let mut node_layers = Vec::new();
        for k in 0..=n_axial {
            let z = (k as f64 / n_axial as f64) * height;
            let mut layer = Vec::new();

            for i in 0..n_radial {
                let theta = 2.0 * std::f64::consts::PI * (i as f64) / (n_radial as f64);
                let x = radius * theta.cos();
                let y = radius * theta.sin();

                layer.push(mesh.add_node(Node::new_3d(x, y, z)));
            }
            node_layers.push(layer);
        }

        // Create elements
        let mat_id = MaterialId(0);
        for k in 0..n_axial {
            if n_radial == 3 {
                // Special case for triangular cross-section - create one wedge per layer
                let n1 = node_layers[k][0]; // Bottom triangle vertex 1
                let n2 = node_layers[k][1]; // Bottom triangle vertex 2
                let n3 = node_layers[k][2]; // Bottom triangle vertex 3
                let n4 = node_layers[k + 1][0]; // Top triangle vertex 1
                let n5 = node_layers[k + 1][1]; // Top triangle vertex 2
                let n6 = node_layers[k + 1][2]; // Top triangle vertex 3

                let element = Element::new(
                    ElementType::Wedge6,
                    vec![n1, n2, n3, n4, n5, n6], // Proper wedge connectivity
                    mat_id,
                )?;
                mesh.add_element(element)?;
            } else {
                // Create quad-based elements for other cases
                for i in 0..n_radial {
                    let i_next = (i + 1) % n_radial;

                    let n1 = node_layers[k][i];
                    let n2 = node_layers[k][i_next];
                    let n3 = node_layers[k + 1][i_next];
                    let n4 = node_layers[k + 1][i];
                    // Create quad-based prism (can be split into hex)
                    // For simplicity, we'll create a degenerate hex
                    let _center_bottom = mesh.add_node(Node::new_3d(
                        0.0,
                        0.0,
                        node_layers[k][0].0 as f64 * height / n_axial as f64,
                    ));
                    let _center_top = mesh.add_node(Node::new_3d(
                        0.0,
                        0.0,
                        node_layers[k + 1][0].0 as f64 * height / n_axial as f64,
                    ));

                    let element = Element::new(
                        ElementType::Hex8,
                        vec![n1, n2, n2, n1, n4, n3, n3, n4],
                        mat_id,
                    )?;
                    mesh.add_element(element)?;
                }
            }
        }

        // Add material
        mesh.materials.insert(
            mat_id,
            MaterialProperties {
                name: "cylinder".to_string(),
                youngs_modulus: 200e9,
                poissons_ratio: 0.3,
                density: 7850.0,
                thermal_conductivity: None,
                specific_heat: None,
                thermal_expansion: None,
            },
        );

        Ok(mesh)
    }
}