rustytorch/crates/specialized/rtx-cfd/src/mesh/structured.rs

// TDD: GREEN phase - Implement structured mesh

use super::{Mesh, MeshBounds, MeshStatistics};
use crate::error::{CfdError, CfdResult};
use crate::mesh::entities::{Cell, Face, Node};
use crate::traits::MeshEntity;
use indexmap::IndexMap;
use nalgebra::Vector3;

/// Structured (regular) mesh for rectangular domains
#[derive(Debug, Clone)]
pub struct StructuredMesh {
    /// Number of nodes in x direction
    nx: usize,
    /// Number of nodes in y direction
    ny: usize,
    /// Number of nodes in z direction
    nz: usize,
    /// Domain width (x direction)
    width: f64,
    /// Domain height (y direction)
    height: f64,
    /// Domain depth (z direction)
    depth: f64,
    /// Grid spacing in x direction
    dx: f64,
    /// Grid spacing in y direction
    dy: f64,
    /// Grid spacing in z direction
    dz: f64,
    /// Nodes storage
    nodes: IndexMap<usize, Node>,
    /// Cells storage
    cells: IndexMap<usize, Cell>,
    /// Faces storage
    faces: IndexMap<usize, Face>,
    /// Whether this is a 2D mesh
    is_2d: bool,
}

impl StructuredMesh {
    /// Create a new 2D structured mesh
    pub fn new(nx: usize, ny: usize, width: f64, height: f64) -> CfdResult<Self> {
        if nx < 2 || ny < 2 {
            return Err(CfdError::mesh("Mesh dimensions must be at least 2x2"));
        }
        if width <= 0.0 || height <= 0.0 {
            return Err(CfdError::mesh("Mesh dimensions must be positive"));
        }

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

        let mut mesh = Self {
            nx,
            ny,
            nz: 1,
            width,
            height,
            depth: 0.0,
            dx,
            dy,
            dz: 0.0,
            nodes: IndexMap::new(),
            cells: IndexMap::new(),
            faces: IndexMap::new(),
            is_2d: true,
        };

        mesh.generate_nodes()?;
        mesh.generate_cells()?;
        mesh.generate_faces()?;

        Ok(mesh)
    }

    /// Create a new 3D structured mesh
    pub fn new_3d(
        nx: usize,
        ny: usize,
        nz: usize,
        width: f64,
        height: f64,
        depth: f64,
    ) -> CfdResult<Self> {
        if nx < 2 || ny < 2 || nz < 2 {
            return Err(CfdError::mesh("Mesh dimensions must be at least 2x2x2"));
        }
        if width <= 0.0 || height <= 0.0 || depth <= 0.0 {
            return Err(CfdError::mesh("Mesh dimensions must be positive"));
        }

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

        let mut mesh = Self {
            nx,
            ny,
            nz,
            width,
            height,
            depth,
            dx,
            dy,
            dz,
            nodes: IndexMap::new(),
            cells: IndexMap::new(),
            faces: IndexMap::new(),
            is_2d: false,
        };

        mesh.generate_nodes()?;
        mesh.generate_cells()?;
        mesh.generate_faces()?;

        Ok(mesh)
    }

    /// Get number of nodes in x direction
    #[must_use]
    pub fn nx(&self) -> usize {
        self.nx
    }

    /// Get number of nodes in y direction
    #[must_use]
    pub fn ny(&self) -> usize {
        self.ny
    }

    /// Get number of nodes in z direction
    #[must_use]
    pub fn nz(&self) -> usize {
        self.nz
    }

    /// Get domain width
    #[must_use]
    pub fn width(&self) -> f64 {
        self.width
    }

    /// Get domain height
    #[must_use]
    pub fn height(&self) -> f64 {
        self.height
    }

    /// Get domain depth
    #[must_use]
    pub fn depth(&self) -> f64 {
        self.depth
    }

    /// Get grid spacing in x direction
    #[must_use]
    pub fn dx(&self) -> f64 {
        self.dx
    }

    /// Get grid spacing in y direction
    #[must_use]
    pub fn dy(&self) -> f64 {
        self.dy
    }

    /// Get grid spacing in z direction
    #[must_use]
    pub fn dz(&self) -> f64 {
        self.dz
    }

    /// Convert (i,j,k) indices to linear node index
    fn node_index(&self, i: usize, j: usize, k: usize) -> usize {
        k * self.nx * self.ny + j * self.nx + i
    }

    /// Convert (i,j,k) indices to linear cell index
    fn cell_index(&self, i: usize, j: usize, k: usize) -> usize {
        k * (self.nx - 1) * (self.ny - 1) + j * (self.nx - 1) + i
    }

    /// Get node at (i,j,k) coordinates
    pub fn get_node(&self, i: usize, j: usize, k: usize) -> CfdResult<&Node> {
        if i >= self.nx || j >= self.ny || k >= self.nz {
            return Err(CfdError::mesh("Node indices out of bounds"));
        }

        let index = self.node_index(i, j, k);
        self.nodes
            .get(&index)
            .ok_or_else(|| CfdError::mesh("Node not found"))
    }

    /// Get cell at (i,j,k) coordinates
    pub fn get_cell(&self, i: usize, j: usize, k: usize) -> CfdResult<&Cell> {
        let nz_cells = if self.is_2d { 1 } else { self.nz - 1 };
        if i >= self.nx - 1 || j >= self.ny - 1 || k >= nz_cells {
            return Err(CfdError::mesh("Cell indices out of bounds"));
        }

        let index = self.cell_index(i, j, k);
        self.cells
            .get(&index)
            .ok_or_else(|| CfdError::mesh("Cell not found"))
    }

    /// Generate all nodes
    fn generate_nodes(&mut self) -> CfdResult<()> {
        for k in 0..self.nz {
            for j in 0..self.ny {
                for i in 0..self.nx {
                    let x = i as f64 * self.dx;
                    let y = j as f64 * self.dy;
                    let z = if self.is_2d { 0.0 } else { k as f64 * self.dz };

                    let position = Vector3::new(x, y, z);
                    let index = self.node_index(i, j, k);
                    let node = Node::new(index, position);

                    self.nodes.insert(index, node);
                }
            }
        }
        Ok(())
    }

    /// Generate all cells
    fn generate_cells(&mut self) -> CfdResult<()> {
        let nz_cells = if self.is_2d { 1 } else { self.nz - 1 };
        for k in 0..nz_cells {
            for j in 0..self.ny - 1 {
                for i in 0..self.nx - 1 {
                    let cell_id = self.cell_index(i, j, k);

                    if self.is_2d {
                        // 2D quadrilateral cell
                        let vertices = vec![
                            self.node_index(i, j, 0),
                            self.node_index(i + 1, j, 0),
                            self.node_index(i + 1, j + 1, 0),
                            self.node_index(i, j + 1, 0),
                        ];

                        let centroid = Vector3::new(
                            (i as f64 + 0.5) * self.dx,
                            (j as f64 + 0.5) * self.dy,
                            0.0,
                        );

                        let volume = self.dx * self.dy; // Area for 2D
                        let mut cell = Cell::new(cell_id, vertices, centroid, volume);

                        // Check if cell is on boundary
                        if i == 0 || i == self.nx - 2 || j == 0 || j == self.ny - 2 {
                            cell.set_boundary(true);
                        }

                        self.cells.insert(cell_id, cell);
                    } else {
                        // 3D hexahedral cell
                        let vertices = vec![
                            self.node_index(i, j, k),
                            self.node_index(i + 1, j, k),
                            self.node_index(i + 1, j + 1, k),
                            self.node_index(i, j + 1, k),
                            self.node_index(i, j, k + 1),
                            self.node_index(i + 1, j, k + 1),
                            self.node_index(i + 1, j + 1, k + 1),
                            self.node_index(i, j + 1, k + 1),
                        ];

                        let centroid = Vector3::new(
                            (i as f64 + 0.5) * self.dx,
                            (j as f64 + 0.5) * self.dy,
                            (k as f64 + 0.5) * self.dz,
                        );

                        let volume = self.dx * self.dy * self.dz;
                        let mut cell = Cell::new(cell_id, vertices, centroid, volume);

                        // Check if cell is on boundary
                        if i == 0
                            || i == self.nx - 2
                            || j == 0
                            || j == self.ny - 2
                            || k == 0
                            || k == self.nz - 2
                        {
                            cell.set_boundary(true);
                        }

                        self.cells.insert(cell_id, cell);
                    }
                }
            }
        }
        Ok(())
    }

    /// Generate all faces
    fn generate_faces(&mut self) -> CfdResult<()> {
        let mut face_id = 0;

        if self.is_2d {
            // Generate 2D faces (edges)
            // Horizontal faces
            for j in 0..self.ny {
                for i in 0..self.nx - 1 {
                    let vertices = vec![self.node_index(i, j, 0), self.node_index(i + 1, j, 0)];

                    let centroid =
                        Vector3::new((i as f64 + 0.5) * self.dx, j as f64 * self.dy, 0.0);
                    let area = self.dx;
                    let normal = Vector3::new(
                        0.0,
                        if j == 0 {
                            -1.0
                        } else if j == self.ny - 1 {
                            1.0
                        } else {
                            0.0
                        },
                        0.0,
                    );

                    let is_boundary = j == 0 || j == self.ny - 1;
                    let face = if is_boundary {
                        Face::new_boundary(face_id, vertices, centroid, area, normal)
                    } else {
                        Face::new(face_id, vertices, centroid, area)
                    };

                    self.faces.insert(face_id, face);
                    face_id += 1;
                }
            }

            // Vertical faces
            for j in 0..self.ny - 1 {
                for i in 0..self.nx {
                    let vertices = vec![self.node_index(i, j, 0), self.node_index(i, j + 1, 0)];

                    let centroid =
                        Vector3::new(i as f64 * self.dx, (j as f64 + 0.5) * self.dy, 0.0);
                    let area = self.dy;
                    let normal = Vector3::new(
                        if i == 0 {
                            -1.0
                        } else if i == self.nx - 1 {
                            1.0
                        } else {
                            0.0
                        },
                        0.0,
                        0.0,
                    );

                    let is_boundary = i == 0 || i == self.nx - 1;
                    let face = if is_boundary {
                        Face::new_boundary(face_id, vertices, centroid, area, normal)
                    } else {
                        Face::new(face_id, vertices, centroid, area)
                    };

                    self.faces.insert(face_id, face);
                    face_id += 1;
                }
            }
        } else {
            // 3D face generation - create faces for all 6 directions
            self.generate_3d_faces(&mut face_id)?;
        }

        Ok(())
    }

    /// Generate 3D faces for structured mesh
    fn generate_3d_faces(&mut self, face_id: &mut usize) -> CfdResult<()> {
        // X-direction faces (YZ planes)
        for k in 0..self.nz - 1 {
            for j in 0..self.ny - 1 {
                for i in 0..self.nx {
                    let vertices = vec![
                        self.node_index(i, j, k),
                        self.node_index(i, j + 1, k),
                        self.node_index(i, j + 1, k + 1),
                        self.node_index(i, j, k + 1),
                    ];

                    let centroid = Vector3::new(
                        i as f64 * self.dx,
                        (j as f64 + 0.5) * self.dy,
                        (k as f64 + 0.5) * self.dz,
                    );

                    let area = self.dy * self.dz;
                    let is_boundary = i == 0 || i == self.nx - 1;
                    let normal = Vector3::new(
                        if i == 0 {
                            -1.0
                        } else if i == self.nx - 1 {
                            1.0
                        } else {
                            0.0
                        },
                        0.0,
                        0.0,
                    );

                    let face = if is_boundary {
                        Face::new_boundary(*face_id, vertices, centroid, area, normal)
                    } else {
                        Face::new(*face_id, vertices, centroid, area)
                    };

                    self.faces.insert(*face_id, face);
                    *face_id += 1;
                }
            }
        }

        // Y-direction faces (XZ planes)
        for k in 0..self.nz - 1 {
            for j in 0..self.ny {
                for i in 0..self.nx - 1 {
                    let vertices = vec![
                        self.node_index(i, j, k),
                        self.node_index(i, j, k + 1),
                        self.node_index(i + 1, j, k + 1),
                        self.node_index(i + 1, j, k),
                    ];

                    let centroid = Vector3::new(
                        (i as f64 + 0.5) * self.dx,
                        j as f64 * self.dy,
                        (k as f64 + 0.5) * self.dz,
                    );

                    let area = self.dx * self.dz;
                    let is_boundary = j == 0 || j == self.ny - 1;
                    let normal = Vector3::new(
                        0.0,
                        if j == 0 {
                            -1.0
                        } else if j == self.ny - 1 {
                            1.0
                        } else {
                            0.0
                        },
                        0.0,
                    );

                    let face = if is_boundary {
                        Face::new_boundary(*face_id, vertices, centroid, area, normal)
                    } else {
                        Face::new(*face_id, vertices, centroid, area)
                    };

                    self.faces.insert(*face_id, face);
                    *face_id += 1;
                }
            }
        }

        // Z-direction faces (XY planes)
        for k in 0..self.nz {
            for j in 0..self.ny - 1 {
                for i in 0..self.nx - 1 {
                    let vertices = vec![
                        self.node_index(i, j, k),
                        self.node_index(i + 1, j, k),
                        self.node_index(i + 1, j + 1, k),
                        self.node_index(i, j + 1, k),
                    ];

                    let centroid = Vector3::new(
                        (i as f64 + 0.5) * self.dx,
                        (j as f64 + 0.5) * self.dy,
                        k as f64 * self.dz,
                    );

                    let area = self.dx * self.dy;
                    let is_boundary = k == 0 || k == self.nz - 1;
                    let normal = Vector3::new(
                        0.0,
                        0.0,
                        if k == 0 {
                            -1.0
                        } else if k == self.nz - 1 {
                            1.0
                        } else {
                            0.0
                        },
                    );

                    let face = if is_boundary {
                        Face::new_boundary(*face_id, vertices, centroid, area, normal)
                    } else {
                        Face::new(*face_id, vertices, centroid, area)
                    };

                    self.faces.insert(*face_id, face);
                    *face_id += 1;
                }
            }
        }

        Ok(())
    }

    /// Get all faces
    #[must_use]
    pub fn get_faces(&self) -> Vec<&Face> {
        self.faces.values().collect()
    }

    /// Check if a node exists
    #[must_use]
    pub fn has_node(&self, node_id: usize) -> bool {
        self.nodes.contains_key(&node_id)
    }
}

impl Mesh for StructuredMesh {
    fn cell_count(&self) -> usize {
        self.cells.len()
    }

    fn node_count(&self) -> usize {
        self.nodes.len()
    }

    fn bounds(&self) -> MeshBounds {
        MeshBounds::new(
            Vector3::new(0.0, 0.0, 0.0),
            Vector3::new(self.width, self.height, self.depth),
        )
    }

    fn validate(&self) -> CfdResult<()> {
        // Check that all cells have valid vertices
        for (_, cell) in &self.cells {
            for &vertex_id in cell.vertex_indices() {
                if !self.nodes.contains_key(&vertex_id) {
                    return Err(CfdError::mesh("Cell references non-existent vertex"));
                }
            }
        }

        // Check mesh connectivity
        if self.cells.is_empty() {
            return Err(CfdError::mesh("Mesh has no cells"));
        }

        if self.nodes.is_empty() {
            return Err(CfdError::mesh("Mesh has no nodes"));
        }

        Ok(())
    }

    fn statistics(&self) -> MeshStatistics {
        let mut stats = MeshStatistics::new();
        stats.total_cells = self.cells.len();
        stats.total_nodes = self.nodes.len();
        stats.total_faces = self.faces.len();

        // Calculate volume statistics
        if !self.cells.is_empty() {
            let volumes: Vec<f64> = self
                .cells
                .values()
                .map(super::super::traits::MeshEntity::volume)
                .collect();
            stats.min_cell_volume = volumes.iter().fold(f64::INFINITY, |a, &b| a.min(b));
            stats.max_cell_volume = volumes.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b));
            stats.average_cell_volume = volumes.iter().sum::<f64>() / volumes.len() as f64;
        }

        // Count boundary faces
        stats.boundary_faces = self
            .faces
            .values()
            .filter(|face| face.is_boundary())
            .count();

        // Cell aspect ratio: the longest cell edge over the shortest.
        //
        // This is the quantity that describes mesh quality — a 4 x 2 domain
        // resolved by 8 x 4 cells is perfectly regular, and reporting its
        // *domain* aspect ratio of 2 would say otherwise.
        //
        // The previous form took the bounding-box extents and guarded the
        // z-extent with `.max(1e-10)`. On a 2-D mesh the depth is exactly
        // zero, so that guard became the minimum and the reported aspect
        // ratio was the domain width divided by 1e-10 — around 2e10 for a
        // unit square.
        let spacings = if self.is_2d {
            vec![self.dx, self.dy]
        } else {
            vec![self.dx, self.dy, self.dz]
        };
        let max_spacing = spacings.iter().copied().fold(0.0_f64, f64::max);
        let min_spacing = spacings.iter().copied().fold(f64::INFINITY, f64::min);
        stats.aspect_ratio = if min_spacing > 1e-12 {
            max_spacing / min_spacing
        } else {
            f64::INFINITY
        };

        stats
    }

    fn is_boundary_cell(&self, cell_id: usize) -> bool {
        self.cells
            .get(&cell_id)
            .is_some_and(super::super::traits::MeshEntity::is_boundary)
    }

    fn get_cell_neighbors(&self, cell_id: usize) -> CfdResult<Vec<usize>> {
        // For structured mesh, we can calculate neighbors directly
        let mut neighbors = Vec::new();

        // Find the (i,j,k) coordinates of this cell
        let nz_cells = if self.is_2d { 1 } else { self.nz - 1 };
        for k in 0..nz_cells {
            for j in 0..self.ny - 1 {
                for i in 0..self.nx - 1 {
                    if self.cell_index(i, j, k) == cell_id {
                        // Add neighboring cells
                        if i > 0 {
                            neighbors.push(self.cell_index(i - 1, j, k));
                        }
                        if i < self.nx - 2 {
                            neighbors.push(self.cell_index(i + 1, j, k));
                        }
                        if j > 0 {
                            neighbors.push(self.cell_index(i, j - 1, k));
                        }
                        if j < self.ny - 2 {
                            neighbors.push(self.cell_index(i, j + 1, k));
                        }
                        if !self.is_2d {
                            if k > 0 {
                                neighbors.push(self.cell_index(i, j, k - 1));
                            }
                            if k < nz_cells - 1 {
                                neighbors.push(self.cell_index(i, j, k + 1));
                            }
                        }
                        return Ok(neighbors);
                    }
                }
            }
        }

        Err(CfdError::mesh("Cell not found"))
    }

    fn refine(&mut self) -> CfdResult<()> {
        // Simple uniform refinement - double the resolution
        let new_nx = (self.nx - 1) * 2 + 1;
        let new_ny = (self.ny - 1) * 2 + 1;

        if self.is_2d {
            *self = Self::new(new_nx, new_ny, self.width, self.height)?;
        } else {
            let new_nz = (self.nz - 1) * 2 + 1;
            *self = Self::new_3d(new_nx, new_ny, new_nz, self.width, self.height, self.depth)?;
        }

        Ok(())
    }

    fn coarsen(&mut self) -> CfdResult<()> {
        // Simple uniform coarsening - halve the resolution
        // Ensure we have enough cells to coarsen
        if self.nx < 3 || self.ny < 3 {
            return Err(CfdError::mesh("Mesh too small to coarsen"));
        }

        let new_nx = (self.nx - 1) / 2 + 1;
        let new_ny = (self.ny - 1) / 2 + 1;

        if new_nx < 2 || new_ny < 2 {
            return Err(CfdError::mesh("Coarsening would result in invalid mesh"));
        }

        if self.is_2d {
            *self = Self::new(new_nx, new_ny, self.width, self.height)?;
        } else {
            if self.nz < 3 {
                return Err(CfdError::mesh("3D mesh too small to coarsen"));
            }
            let new_nz = (self.nz - 1) / 2 + 1;
            if new_nz < 2 {
                return Err(CfdError::mesh("Coarsening would result in invalid 3D mesh"));
            }
            *self = Self::new_3d(new_nx, new_ny, new_nz, self.width, self.height, self.depth)?;
        }
        Ok(())
    }
}

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

    #[test]
    fn test_structured_mesh_2d() {
        let mesh = StructuredMesh::new(3, 3, 1.0, 1.0).unwrap();
        assert_eq!(mesh.nx(), 3);
        assert_eq!(mesh.ny(), 3);
        assert_eq!(mesh.cell_count(), 4); // (3-1) * (3-1)
        assert_eq!(mesh.node_count(), 9); // 3 * 3
    }

    #[test]
    fn test_structured_mesh_3d() {
        let mesh = StructuredMesh::new_3d(2, 2, 2, 1.0, 1.0, 1.0).unwrap();
        assert_eq!(mesh.nx(), 2);
        assert_eq!(mesh.ny(), 2);
        assert_eq!(mesh.nz(), 2);
        assert_eq!(mesh.cell_count(), 1); // (2-1) * (2-1) * (2-1)
        assert_eq!(mesh.node_count(), 8); // 2 * 2 * 2
    }

    #[test]
    fn test_invalid_dimensions() {
        assert!(StructuredMesh::new(0, 5, 1.0, 1.0).is_err());
        assert!(StructuredMesh::new(5, 0, 1.0, 1.0).is_err());
        assert!(StructuredMesh::new(5, 5, 0.0, 1.0).is_err());
        assert!(StructuredMesh::new(5, 5, 1.0, 0.0).is_err());
    }

    #[test]
    fn test_node_access() {
        let mesh = StructuredMesh::new(3, 3, 1.0, 1.0).unwrap();
        let node = mesh.get_node(0, 0, 0).unwrap();
        assert_eq!(node.position(), Vector3::new(0.0, 0.0, 0.0));

        let node = mesh.get_node(2, 2, 0).unwrap();
        assert_eq!(node.position(), Vector3::new(1.0, 1.0, 0.0));
    }

    #[test]
    fn test_cell_access() {
        let mesh = StructuredMesh::new(3, 3, 1.0, 1.0).unwrap();
        let cell = mesh.get_cell(0, 0, 0).unwrap();
        assert_eq!(cell.vertex_count(), 4);
        assert!((cell.volume() - 0.25).abs() < 1e-10);
    }

    #[test]
    fn test_mesh_bounds() {
        let mesh = StructuredMesh::new(4, 5, 2.0, 3.0).unwrap();
        let bounds = mesh.bounds();
        assert_eq!(bounds.min, Vector3::new(0.0, 0.0, 0.0));
        assert_eq!(bounds.max, Vector3::new(2.0, 3.0, 0.0));
    }

    #[test]
    fn test_mesh_refinement() {
        let mut mesh = StructuredMesh::new(3, 3, 1.0, 1.0).unwrap();
        let original_cells = mesh.cell_count();

        mesh.refine().unwrap();
        assert_eq!(mesh.cell_count(), original_cells * 4);
    }
}