#!/usr/bin/env python3

import argparse
import meshio
import pyvista as pv
import numpy as np
import matplotlib.pyplot as plt
import warnings

def analyze_mesh_quality(mesh_filepath):
    """
    Reads a mesh file, computes various quality metrics using the current
    PyVista API (cell_quality), and prints statistics.
    """
    print(f"--- Analyzing Mesh: {mesh_filepath} ---")

    try:
        mesh_mshio = meshio.read(mesh_filepath)
    except FileNotFoundError:
        print(f"ERROR: Mesh file '{mesh_filepath}' not found.")
        return
    except Exception as e:
        print(f"ERROR: Could not read mesh file '{mesh_filepath}' with meshio: {e}")
        return

    points = mesh_mshio.points
    cells_list = []
    cell_types_list = []

    supported_3d_types = {
        "tetra": pv.CellType.TETRA,
        "tetra10": pv.CellType.QUADRATIC_TETRA
    }

    found_3d_cells = False
    for cells_block in mesh_mshio.cells:
        if cells_block.type in supported_3d_types:
            print(f"Found cell block of type '{cells_block.type}' with {len(cells_block.data)} cells.")
            num_points_per_cell = cells_block.data.shape[1]
            pv_cell_block = np.hstack((
                np.full((len(cells_block.data), 1), num_points_per_cell, dtype=cells_block.data.dtype),
                cells_block.data
            )).ravel()
            cells_list.append(pv_cell_block)
            cell_types_list.extend([supported_3d_types[cells_block.type]] * len(cells_block.data))
            found_3d_cells = True

    if not found_3d_cells:
        print("ERROR: No supported 3D cell types (e.g., tetra, tetra10, hexahedron) found in the mesh.")
        print("Meshio cell types found in file:", [block.type for block in mesh_mshio.cells])
        return

    if len(cells_list) > 1:
        print("Warning: Multiple 3D cell blocks found and concatenated. Ensure these are compatible for a single UnstructuredGrid.")
        flat_cells = np.concatenate(cells_list)
    elif len(cells_list) == 1:
        flat_cells = cells_list[0]
    else: # Should have been caught by "not found_3d_cells"
        return

    try:
        mesh_pv = pv.UnstructuredGrid(flat_cells, cell_types_list, points)
        print(f"\nSuccessfully converted to PyVista mesh with {mesh_pv.n_points} points and {mesh_pv.n_cells} cells.")
    except Exception as e:
        print(f"ERROR: Could not create PyVista UnstructuredGrid: {e}")
        return

    quality_measures_to_try = [
        ('aspect_ratio', "Aspect Ratio"),
        ('condition', "Condition Number of Element Jacobian"),
        ('distortion', "Distortion (0=perfect, 1=degenerated)"),
        ('jacobian', "Determinant of Element Jacobian (should be > 0)"),
        ('radius_ratio', "Radius Ratio (for tets: inscribed/circumscribed, 1 is best)"),
        ('scaled_jacobian', "Scaled Jacobian (1 is ideal, <0 bad)"),
        ('skew', "Skewness (0 is ideal, higher is worse)"),
        ('volume', "Volume (should be > 0)"),
        ('min_angle', "Minimum Dihedral/Face Angle (larger is better)"),
        ('max_angle', "Maximum Dihedral/Face Angle (smaller is better, less 'flat')")
    ]

    # Suppress common warnings for cleaner output, adjust as needed
    warnings.filterwarnings("ignore", category=UserWarning, message="`quality_measure`.*not implemented for this cell type.")
    warnings.filterwarnings("ignore", category=UserWarning, message="Cannot compute quality of a point cell.") # If points/lines exist
    warnings.filterwarnings("ignore", category=UserWarning, message="Mesh contains non-tetrahedral cells. Radius ratio is only defined for tetrahedra.")
    warnings.filterwarnings("ignore", category=RuntimeWarning) # For potential NaN in stats from empty arrays etc.


    for measure, description in quality_measures_to_try:
        if not mesh_pv.n_cells > 0:
            continue
        # Use the recommended cell_quality method
        quality_array = mesh_pv.cell_quality(quality_measure=measure)
        range = quality_array.get_data_range(measure)
        print(f"Range for {measure} is: {range[0]:0.3E} → {range[1]:0.3E}")



    warnings.resetwarnings()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Analyze mesh quality using PyVista.")
    parser.add_argument("mesh_filepath", type=str, help="Path to the mesh file (e.g., .msh, .vtk).")

    args = parser.parse_args()

    analyze_mesh_quality(args.mesh_filepath)