automesh 0.5.0

use super::{
    ErrorWrapper,
    io::{extension, invalid_input, read_segmentation, write_mesh},
    metrics::write_metrics,
    remesh::apply_remesh_subcommand,
    smooth::{MeshSmoothCommands, apply_smoothing_method},
};
use clap::Subcommand;
use conspire::{
    geometry::{
        Coordinate, Coordinates,
        grid::{Gradient, MarchingCubes, Voxels},
        mesh::{
            Class, Connectivity, Finish, Fitting, Freedom, Marching, Mesh, Placement, Tessellation,
        },
        ntree::{Balance, Balancing, CurvatureSizing, Dualization, Octree, Pairing},
        segmentation::Segmentation,
    },
    math::{Tensor, TensorVec},
    units::Length,
};
use std::{collections::HashSet, path::Path, time::Instant};

#[derive(Subcommand)]
pub enum MeshSubcommand {
    /// Creates an all-hexahedral mesh from a segmentation or tessellation
    Hex(MeshArgs),
    /// Creates a hex-dominant mesh from a tessellation, polyhedral at the boundary
    Hexdom(MeshArgs),
    /// Creates a polyhedral mesh from a tessellation
    Poly(MeshArgs),
    /// Creates an all-tetrahedral mesh from a tessellation
    Tet(MeshArgs),
    /// Creates all-triangular isosurface(s) from a segmentation
    Tri(MeshArgs),
}

#[derive(clap::Args)]
pub struct MeshArgs {
    #[command(subcommand)]
    pub smoothing: Option<MeshSmoothCommands>,

    /// Segmentation (npy | spn) or tessellation (stl) input file
    #[arg(long, short, value_name = "FILE")]
    pub input: String,

    /// Mesh output file (exo | inp | mesh | off | stl | vtu)
    #[arg(long, short, value_name = "FILE")]
    pub output: String,

    /// Defeature clusters with less than NUM voxels
    #[arg(long, short, value_name = "NUM")]
    pub defeature: Option<usize>,

    /// Number of voxels in the x-direction (spn)
    #[arg(long, short = 'x', value_name = "NEL")]
    pub nelx: Option<usize>,

    /// Number of voxels in the y-direction (spn)
    #[arg(long, short = 'y', value_name = "NEL")]
    pub nely: Option<usize>,

    /// Number of voxels in the z-direction (spn)
    #[arg(long, short = 'z', value_name = "NEL")]
    pub nelz: Option<usize>,

    /// Voxel IDs to remove from the mesh (npy | spn)
    #[arg(long, num_args = 1.., short, value_delimiter = ' ', value_name = "ID")]
    pub remove: Option<Vec<usize>>,

    /// Scaling (> 0.0) in the x-direction, applied before translation
    #[arg(default_value_t = 1.0, long, value_name = "SCALE")]
    pub xscale: f64,

    /// Scaling (> 0.0) in the y-direction, applied before translation
    #[arg(default_value_t = 1.0, long, value_name = "SCALE")]
    pub yscale: f64,

    /// Scaling (> 0.0) in the z-direction, applied before translation
    #[arg(default_value_t = 1.0, long, value_name = "SCALE")]
    pub zscale: f64,

    /// Translation in the x-direction
    #[arg(
        long,
        default_value_t = 0.0,
        allow_negative_numbers = true,
        value_name = "VAL"
    )]
    pub xtranslate: f64,

    /// Translation in the y-direction
    #[arg(
        long,
        default_value_t = 0.0,
        allow_negative_numbers = true,
        value_name = "VAL"
    )]
    pub ytranslate: f64,

    /// Translation in the z-direction
    #[arg(
        long,
        default_value_t = 0.0,
        allow_negative_numbers = true,
        value_name = "VAL"
    )]
    pub ztranslate: f64,

    /// Octree refinement scale for dualizing a tessellation (stl) input
    #[arg(long, default_value_t = 5.0, short = 's', value_name = "SCALE")]
    pub scale: f64,

    /// Uniform mesh of the given element size instead of an octree (stl)
    #[arg(long, short = 'u', value_name = "SPACING")]
    pub uniform: Option<f64>,

    /// Chord-error tolerance for curvature-driven refinement [default: disabled]
    #[arg(long, short = 't', value_name = "TOL")]
    pub tolerance: Option<f64>,

    /// Uses strong balancing instead of the default weak balancing
    #[arg(action, long)]
    pub strong: bool,

    /// Cuts the uniform lattice with marching cubes and inflates it onto the surface (stl)
    #[arg(action, long, requires = "uniform")]
    pub inflate: bool,

    /// Cuts the uniform lattice with marching cubes, leaving the boundary as cut (stl)
    #[arg(action, conflicts_with = "inflate", long, requires = "uniform")]
    pub marching: bool,

    /// Snaps the buffer layer onto the surface instead of a soft fit
    #[arg(action, long)]
    pub snap: bool,

    /// Replaces shell hexahedra below this minimum scaled Jacobian along
    /// features with pyramids and refits, snapping both fits under --snap
    /// [default: disabled]
    #[arg(conflicts_with_all = ["inflate", "marching"], long, value_name = "MSJ")]
    pub pyramids: Option<f64>,

    /// Level difference allowed between neighboring octree cells (poly)
    #[arg(long, default_value_t = 1, short = 'l', value_name = "NUM")]
    pub levels: usize,

    /// Quality metrics output file (csv | npy)
    #[arg(long, value_name = "FILE")]
    pub metrics: Option<String>,

    /// Isosurface method for triangles from a segmentation (tri)
    #[arg(long, value_enum, default_value_t = Cubes::Cuberille, value_name = "METHOD")]
    pub cubes: Cubes,
}

#[derive(Clone, Copy, clap::ValueEnum)]
pub enum Cubes {
    /// Faces of the voxels, shared between materials
    Cuberille,
    /// Smooth isosurface of each material
    Marching,
}

pub enum Element {
    Hexahedra,
    HexDominant,
    Polyhedra,
    Tetrahedra,
    Triangles,
}

fn read_voxels(args: &MeshArgs, quiet: bool) -> Result<Voxels<u8>, ErrorWrapper> {
    match extension(&args.input) {
        Some("npy") | Some("spn") => {
            let mut voxels =
                read_segmentation(&args.input, args.nelx, args.nely, args.nelz, quiet)?;
            if let Some(min) = args.defeature {
                let time = Instant::now();
                crate::echo!(
                    quiet,
                    " \x1b[1;96mDefeaturing\x1b[0m clusters of {min} voxels or less"
                );
                voxels = voxels.defeature(min)?;
                crate::echo!(quiet, "        \x1b[1;92mDone\x1b[0m {:?}", time.elapsed());
            }
            Ok(voxels)
        }
        extension => Err(invalid_input(&args.input, extension)),
    }
}

fn finish(mut mesh: Mesh<3>, args: MeshArgs, quiet: bool) -> Result<(), ErrorWrapper> {
    if let Some(MeshSmoothCommands::Smooth {
        remeshing,
        iterations,
        method,
        pass_band,
        scale,
        hierarchical,
    }) = args.smoothing
    {
        apply_smoothing_method(
            &mut mesh,
            iterations,
            method,
            pass_band,
            scale,
            hierarchical,
            quiet,
        )?;
        if let Some(subcommand) = remeshing {
            mesh = apply_remesh_subcommand(mesh, subcommand, quiet)?;
        }
    }
    if let Some(file) = &args.metrics {
        write_metrics(&mesh, file, quiet)?;
    }
    write_mesh(&args.output, mesh, quiet)
}

pub fn mesh(element: Element, args: MeshArgs, quiet: bool) -> Result<(), ErrorWrapper> {
    match (&element, extension(&args.input)) {
        (Element::Hexahedra, Some("stl")) => return hexahedralize(args, quiet),
        (Element::Tetrahedra, Some("stl")) => return tetrahedralize(args, quiet),
        (element @ (Element::HexDominant | Element::Polyhedra), Some("stl")) => {
            return cut(args, element, quiet);
        }
        (Element::HexDominant | Element::Polyhedra | Element::Tetrahedra, extension) => {
            return Err(invalid_input(&args.input, extension));
        }
        _ => {}
    }
    if args.uniform.is_some() {
        return Err(ErrorWrapper::from(
            "Uniform lattice meshing applies to tessellation (stl) inputs only",
        ));
    }
    let voxels = read_voxels(&args, quiet)?;
    let time = Instant::now();
    let mesh = match element {
        Element::Hexahedra => {
            crate::echo!(quiet, "     \x1b[1;96mMeshing\x1b[0m voxels into hexahedra");
            let remove: Option<Vec<u8>> = args
                .remove
                .as_ref()
                .map(|ids| ids.iter().map(|&id| id as u8).collect());
            let scale = Coordinate::from([args.xscale, args.yscale, args.zscale]);
            let translate = Coordinate::from([args.xtranslate, args.ytranslate, args.ztranslate]);
            let segmentation = Segmentation::new(voxels, scale, translate);
            Mesh::from_segmentation(segmentation, remove.as_deref())
        }
        Element::HexDominant | Element::Polyhedra => {
            unreachable!("cutting requires a tessellation input")
        }
        Element::Tetrahedra => unreachable!("tetrahedralizing requires a tessellation input"),
        Element::Triangles => {
            crate::echo!(quiet, "     \x1b[1;96mMeshing\x1b[0m voxels into triangles");
            let voxels = remove_materials(voxels, args.remove.as_deref());
            let mesh = match args.cubes {
                Cubes::Cuberille => Mesh::from(Tessellation::from(voxels)),
                Cubes::Marching => marching(&voxels)?,
            };
            scaled(
                mesh,
                [args.xscale, args.yscale, args.zscale],
                [args.xtranslate, args.ytranslate, args.ztranslate],
            )
        }
    };
    crate::echo!(
        quiet,
        "        \x1b[1;92mDone\x1b[0m {:?} \x1b[2m[{} elements, {} nodes]\x1b[0m",
        time.elapsed(),
        mesh.number_of_elements(),
        mesh.number_of_nodes()
    );
    finish(mesh, args, quiet)
}

/// Meshes a tessellation (stl) input into an all-hexahedral mesh.
///
/// A background mesh of the enclosed volume is built first — the dual of an
/// octree fitted to the surface, or a uniform lattice under `--uniform` — and
/// trimmed to the surface. A buffer layer is then fitted onto the surface.
/// Buffering is timed on its own because it dominates the total by far, while
/// the steps building the background are lumped together as one.
fn hexahedralize(args: MeshArgs, quiet: bool) -> Result<(), ErrorWrapper> {
    crate::echo!(quiet, "     \x1b[1;96mReading\x1b[0m {}", args.input);
    let mut time = Instant::now();
    let tessellation = Tessellation::try_from(Path::new(&args.input))?;
    crate::echo!(quiet, "        \x1b[1;92mDone\x1b[0m {:?}", time.elapsed());
    if args.inflate || args.marching {
        return marching_hex(args, tessellation, quiet);
    }
    let fitting = if args.snap {
        Fitting::Snap
    } else {
        Fitting::Soft
    };

    crate::echo!(
        quiet,
        "     \x1b[1;96mMeshing\x1b[0m hexahedra {}",
        if args.uniform.is_some() {
            "uniformly"
        } else {
            "adaptively"
        }
    );
    time = Instant::now();
    let mut mesh = if let Some(spacing) = args.uniform {
        tessellation.lattice_background(Length::meters(spacing))?.0
    } else {
        let balancing = if args.strong {
            Balancing::Strong(1)
        } else {
            Balancing::Weak(1)
        };
        let mut octree = Octree::<u16, usize>::from_features(
            &tessellation,
            args.scale,
            CurvatureSizing {
                tolerance: args.tolerance.map(Length::meters),
                ..Default::default()
            },
            0,
        )?;
        octree.equilibrate(balancing, Pairing::Regular)?;
        octree.dualize()
    };
    tessellation.trim(&mut mesh)?;
    crate::echo!(
        quiet,
        "        \x1b[1;92mDone\x1b[0m {:?} \x1b[2m[{} elements, {} nodes]\x1b[0m",
        time.elapsed(),
        mesh.number_of_elements(),
        mesh.number_of_nodes()
    );

    crate::echo!(
        quiet,
        "   \x1b[1;96mBuffering\x1b[0m hexahedra onto geometry"
    );
    time = Instant::now();
    let mesh = if let Some(threshold) = args.pyramids {
        mesh.buffer_targeted(&tessellation, fitting, threshold)?
    } else {
        mesh.buffer(&tessellation, fitting)?
    };
    let mesh = scaled(
        mesh,
        [args.xscale, args.yscale, args.zscale],
        [args.xtranslate, args.ytranslate, args.ztranslate],
    );
    crate::echo!(
        quiet,
        "        \x1b[1;92mDone\x1b[0m {:?} \x1b[2m[{} elements, {} nodes]\x1b[0m",
        time.elapsed(),
        mesh.number_of_elements(),
        mesh.number_of_nodes()
    );
    finish(mesh, args, quiet)
}

/// Cuts a uniform lattice with marching cubes along a tessellation.
///
/// The lattice is cut along the surface and split into hexahedra. Under
/// `--inflate` every node is then fitted onto the surface, and the boundary is
/// snapped onto it under `--snap`. Under `--marching` the boundary is left as
/// cut. No buffer layer is inserted, so the element count is set by the
/// lattice alone.
///
/// Each lattice cell is split into eight hexahedra, so the lattice is laid at
/// twice the requested element size.
fn marching_hex(
    args: MeshArgs,
    tessellation: Tessellation,
    quiet: bool,
) -> Result<(), ErrorWrapper> {
    if args.marching && args.snap {
        return Err(ErrorWrapper::from(
            "Snapping applies to an inflated mesh, not one left as cut",
        ));
    }
    let size = args
        .uniform
        .expect("marching cubes requires a uniform element size");
    crate::echo!(quiet, "     \x1b[1;96mMeshing\x1b[0m hexahedra uniformly");
    let mut time = Instant::now();
    let mut mesh = tessellation.marching_hex(
        Length::meters(size * 2.0),
        Marching {
            placement: Placement::Crossing(0.2),
            finish: Finish::Cut,
        },
    )?;
    crate::echo!(
        quiet,
        "        \x1b[1;92mDone\x1b[0m {:?} \x1b[2m[{} elements, {} nodes]\x1b[0m",
        time.elapsed(),
        mesh.number_of_elements(),
        mesh.number_of_nodes()
    );
    if args.inflate {
        crate::echo!(
            quiet,
            "   \x1b[1;96mInflating\x1b[0m hexahedra onto geometry"
        );
        time = Instant::now();
        mesh.inflate(
            &tessellation,
            Freedom::Whole,
            if args.snap {
                Fitting::Snap
            } else {
                Fitting::Soft
            },
        )?;
        crate::echo!(
            quiet,
            "        \x1b[1;92mDone\x1b[0m {:?} \x1b[2m[{} elements, {} nodes]\x1b[0m",
            time.elapsed(),
            mesh.number_of_elements(),
            mesh.number_of_nodes()
        );
    }
    let mesh = scaled(
        mesh,
        [args.xscale, args.yscale, args.zscale],
        [args.xtranslate, args.ytranslate, args.ztranslate],
    );
    finish(mesh, args, quiet)
}

/// Meshes a tessellation (stl) input into an all-tetrahedral mesh.
///
/// The tetrahedral counterpart of [`hexahedralize`]. A background tetrahedral
/// mesh of the enclosed volume is built — six tetrahedra per cell of an octree
/// fitted to the surface, or of a uniform lattice under `--uniform` — and
/// trimmed to the surface. A buffer layer of tetrahedra is then fitted onto
/// the surface. As in the hexahedral path, buffering is timed on its own.
fn tetrahedralize(args: MeshArgs, quiet: bool) -> Result<(), ErrorWrapper> {
    crate::echo!(quiet, "     \x1b[1;96mReading\x1b[0m {}", args.input);
    let mut time = Instant::now();
    let tessellation = Tessellation::try_from(Path::new(&args.input))?;
    crate::echo!(quiet, "        \x1b[1;92mDone\x1b[0m {:?}", time.elapsed());
    let fitting = if args.snap {
        Fitting::Snap
    } else {
        Fitting::Soft
    };

    crate::echo!(
        quiet,
        "     \x1b[1;96mMeshing\x1b[0m tetrahedra {}",
        if args.uniform.is_some() {
            "uniformly"
        } else {
            "adaptively"
        }
    );
    time = Instant::now();
    let (mut mesh, _classes) = if let Some(spacing) = args.uniform {
        tessellation.lattice_tet_background(Length::meters(spacing))?
    } else {
        tessellation.octree_tet_background(
            Balancing::Strong(1),
            Pairing::None,
            args.scale,
            args.tolerance.map(Length::meters),
        )?
    };
    tessellation.trim(&mut mesh)?;
    crate::echo!(
        quiet,
        "        \x1b[1;92mDone\x1b[0m {:?} \x1b[2m[{} elements, {} nodes]\x1b[0m",
        time.elapsed(),
        mesh.number_of_elements(),
        mesh.number_of_nodes()
    );

    crate::echo!(
        quiet,
        "   \x1b[1;96mBuffering\x1b[0m tetrahedra onto geometry"
    );
    time = Instant::now();
    let mesh = mesh.buffer_tets(&tessellation, fitting)?;
    let mesh = scaled(
        mesh,
        [args.xscale, args.yscale, args.zscale],
        [args.xtranslate, args.ytranslate, args.ztranslate],
    );
    crate::echo!(
        quiet,
        "        \x1b[1;92mDone\x1b[0m {:?} \x1b[2m[{} elements, {} nodes]\x1b[0m",
        time.elapsed(),
        mesh.number_of_elements(),
        mesh.number_of_nodes()
    );
    finish(mesh, args, quiet)
}

/// Cuts an octree fitted to a tessellation (stl) input to the surface.
///
/// [`Element::Polyhedra`] cuts the octree itself, while [`Element::HexDominant`]
/// cuts its dual, leaving hexahedra everywhere but at the boundary.
fn cut(args: MeshArgs, element: &Element, quiet: bool) -> Result<(), ErrorWrapper> {
    crate::echo!(quiet, "     \x1b[1;96mReading\x1b[0m {}", args.input);
    let mut time = Instant::now();
    let tessellation = Tessellation::try_from(Path::new(&args.input))?;
    crate::echo!(quiet, "        \x1b[1;92mDone\x1b[0m {:?}", time.elapsed());
    let polyhedral = matches!(element, Element::Polyhedra);
    if polyhedral && args.uniform.is_some() {
        return Err(ErrorWrapper::from(
            "Uniform lattice meshing applies to mesh hex and mesh hexdom only",
        ));
    }
    if !polyhedral && args.levels != 1 {
        return Err(ErrorWrapper::from(
            "Dualization requires 2:1 balancing, so levels applies to mesh poly only",
        ));
    }

    crate::echo!(
        quiet,
        "     \x1b[1;96mMeshing\x1b[0m {} {}",
        if polyhedral { "polyhedra" } else { "hexahedra" },
        if args.uniform.is_some() {
            "uniformly"
        } else {
            "adaptively"
        }
    );
    time = Instant::now();
    let (background, classes) = if let Some(spacing) = args.uniform {
        tessellation.lattice_background(Length::meters(spacing))
    } else {
        let balancing = if args.strong {
            Balancing::Strong(args.levels)
        } else {
            Balancing::Weak(args.levels)
        };
        if polyhedral {
            tessellation.octree_background(balancing, args.scale)
        } else {
            tessellation.dual_background(balancing, args.scale)
        }
    }?;
    let (elements, nodes) = retained(&background, &classes);
    crate::echo!(
        quiet,
        "        \x1b[1;92mDone\x1b[0m {:?} \x1b[2m[{elements} elements, {nodes} nodes]\x1b[0m",
        time.elapsed()
    );

    crate::echo!(
        quiet,
        "   \x1b[1;96mBuffering\x1b[0m polyhedra onto geometry"
    );
    time = Instant::now();
    let mesh = if polyhedral {
        tessellation.cut_polyhedral(background, &classes)
    } else {
        tessellation.cut(background, &classes)
    }?;
    let mesh = scaled(
        mesh,
        [args.xscale, args.yscale, args.zscale],
        [args.xtranslate, args.ytranslate, args.ztranslate],
    );
    crate::echo!(
        quiet,
        "        \x1b[1;92mDone\x1b[0m {:?} \x1b[2m[{} elements, {} nodes]\x1b[0m",
        time.elapsed(),
        mesh.number_of_elements(),
        mesh.number_of_nodes()
    );
    finish(mesh, args, quiet)
}

/// Counts the background cells the cut will keep, and the nodes they use.
///
/// A background mesh spans the whole octree or lattice, so most of its cells
/// can lie outside the tessellation and be discarded by the cut. Counting only
/// what survives keeps this step comparable to the one that follows it, and to
/// the hexahedral path, where trimming drops those cells before they are
/// counted at all.
fn retained(background: &Mesh<3>, classes: &[Class]) -> (usize, usize) {
    let mut nodes = HashSet::new();
    let elements = background
        .connectivities()
        .iter()
        .flatten()
        .zip(classes)
        .filter(|(_, class)| !matches!(class, Class::Outside))
        .inspect(|(element, _)| nodes.extend(element.iter().copied()))
        .count();
    (elements, nodes.len())
}

/// Zeroes out (treats as void) any voxels whose material is in `remove`.
fn remove_materials(voxels: Voxels<u8>, remove: Option<&[usize]>) -> Voxels<u8> {
    match remove {
        Some(remove) if !remove.is_empty() => {
            let nel = *voxels.nel();
            let data = voxels
                .data()
                .iter()
                .map(|&block| {
                    if remove.contains(&(block as usize)) {
                        0
                    } else {
                        block
                    }
                })
                .collect();
            Voxels::new(data, nel)
        }
        _ => voxels,
    }
}

/// Extracts a closed marching cubes isosurface of each nonzero material.
///
/// Each material is padded with void so its surface closes at the grid
/// boundary. Interfaces are represented once per adjacent material.
fn marching(voxels: &Voxels<u8>) -> Result<Mesh<3>, ErrorWrapper> {
    let [nx, ny, nz] = *voxels.nel();
    let padded = [nz + 2, ny + 2, nx + 2];
    let mut materials: Vec<u8> = voxels.data().iter().copied().filter(|&m| m != 0).collect();
    materials.sort_unstable();
    materials.dedup();
    let extractor = MarchingCubes {
        gradient: Gradient::Ascent,
        degenerate: false,
        ..Default::default()
    };
    let mut coordinates = Coordinates::new();
    let mut connectivities = Vec::with_capacity(materials.len());
    for material in materials {
        let mut data = vec![0.0; padded.iter().product()];
        for k in 0..nz {
            for j in 0..ny {
                for i in 0..nx {
                    if voxels.data()[voxels.flat([i, j, k])] == material {
                        data[((k + 1) * padded[1] + j + 1) * padded[2] + i + 1] = 1.0;
                    }
                }
            }
        }
        let surface = extractor
            .extract(&Voxels::new_row_major(data, padded), None)
            .map_err(ErrorWrapper::from)?;
        let offset = coordinates.len();
        for vertex in surface.vertices.iter() {
            coordinates.push(Coordinate::from([
                vertex[0] + Length::meters(-0.5),
                vertex[1] + Length::meters(-0.5),
                vertex[2] + Length::meters(-0.5),
            ]));
        }
        let faces: Vec<[usize; 3]> = surface
            .faces
            .iter()
            .map(|face| face.map(|node| node + offset))
            .collect();
        connectivities.push(Connectivity::Triangular(faces.into()));
    }
    Ok(Mesh::from((connectivities, coordinates)))
}

/// Applies per-axis scaling (before translation) to the mesh coordinates.
fn scaled(mesh: Mesh<3>, scale: [f64; 3], translate: [f64; 3]) -> Mesh<3> {
    if scale == [1.0, 1.0, 1.0] && translate == [0.0, 0.0, 0.0] {
        return mesh;
    }
    let (connectivities, coordinates) = mesh.into();
    let coordinates: Coordinates<3> = coordinates
        .iter()
        .map(|coordinate| {
            Coordinate::from([
                coordinate[0] * scale[0] + Length::meters(translate[0]),
                coordinate[1] * scale[1] + Length::meters(translate[1]),
                coordinate[2] * scale[2] + Length::meters(translate[2]),
            ])
        })
        .collect();
    Mesh::from((connectivities.into_members(), coordinates))
}