import struct
import numpy as np
import pyvista as pv


def read_binary_stl(filename):
    with open(filename, 'rb') as f:
        _ = f.read(80)  # ヘッダー
        num_faces = struct.unpack('<I', f.read(4))[0]
        print(f"num_faces={num_faces}")

        vertices = []
        faces = []
        scalars = []
        for i in range(num_faces):
            _ = struct.unpack("<3f", f.read(12))  # normal
            vertices.extend([struct.unpack("<3f", f.read(12))
                            for _ in range(3)])
            base_index = 3 * i
            # 各頂点インデックスの前に追加している自然数値はこのセルが含む頂点数. ここでの3は三角形を意味する. pyvistaの仕様で表示時に必要.
            faces.append([3, base_index, base_index + 1, base_index + 2])
            _ = struct.unpack("<H", f.read(2))[0]  # パディング

            # 表示時に着色するため、各頂点に対してスカラー値を設定.
            scalars.extend([base_index, base_index + 1, base_index + 2])
    return {
        "vertices": np.array(vertices),
        "faces": np.hstack(faces),
        "num_faces": num_faces,
        "scalars": scalars
    }


filename = 'data/data.stl'
triangles = read_binary_stl(filename)
mesh = pv.PolyData(triangles["vertices"],
                   triangles["faces"], triangles["num_faces"])
mesh.point_data['Scalars'] = triangles["scalars"]
mesh.plot(show_edges=True, color='w', smooth_shading=True, scalars='Scalars')