import pyvista as pv
import numpy as np
import ansys.fluent.core as pyfluent
from ansys.fluent.core.services.field_data import CellElementType


def convert_fluent_element_type_to_pv_cell_type(element_type):
    if element_type == CellElementType.TETRAHEDRON:
        return pv.CellType.TETRA
    if element_type == CellElementType.HEXAHEDRON:
        return pv.CellType.HEXAHEDRON
    # Add other element types
    raise ValueError(f"Unknown element type {element_type}")


solver = pyfluent.launch_fluent(processor_count=4)
solver.settings.file.read_case_data(file_name="d:/work/elbow.cas.h5")
mesh = solver.fields.field_data.get_mesh(zone_id=2)
points = np.array([(node.x, node.y, node.z) for node in mesh.nodes])
cells = np.array([[len(element.node_indices)] + element.node_indices for element in mesh.elements]).ravel()
cell_types = [convert_fluent_element_type_to_pv_cell_type(element.element_type) for element in mesh.elements]
grid = pv.UnstructuredGrid(cells, cell_types, points)
pressure_data = solver.fields.solution_variable_data.get_data(solution_variable_name="SV_P", zone_names=["fluid"], domain_name="mixture")["fluid"]
grid.cell_data["pressure"] = pressure_data
grid.plot(scalars="pressure")