import pyvista as pv
import json
import numpy as np
from stpyvista import stpyvista
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors


    # Load the .obj file
mesh = pv.read('file.obj')

    # Load the JSON file
with open('dental-labels4.json', 'r') as file:
    labels_data = json.load(file)

    # Assuming labels_data['labels'] is a list of labels
labels = labels_data['labels']


    # Make sure the number of labels matches the number of vertices or faces
assert len(labels) == mesh.n_points or len(labels) == mesh.n_cells

# If labels correspond to vertices
if len(labels) == mesh.n_points:
    mesh.point_data['Labels'] = labels
# If labels correspond to faces
elif len(labels) == mesh.n_cells:
    mesh.cell_data['Labels'] = labels
    
# Create a pyvista plotter
plotter = pv.Plotter()
cmap = plt.cm.get_cmap('jet', 27)  # Using a colormap with sufficient distinct colors

colors = cmap(np.linspace(0, 1, 27))  # Generate colors

# Convert colors to a format acceptable by PyVista
colormap = mcolors.ListedColormap(colors)

# Add the mesh to the plotter with labels as a scalar field
#plotter.add_mesh(mesh, scalars='Labels', show_scalar_bar=True, cmap='jet')
plotter.add_mesh(mesh, scalars='Labels', show_scalar_bar=True, cmap=colormap, clim=[0, 27])

# Show the plot
#plotter.show()
## Send to streamlit
stpyvista(plotter)