# -*- coding: utf-8 -*-
"""
simplify-stl.py —— 简化过大的 STL 网格 (窗口内 3D 预览用, 纯 numpy 顶点聚类, 零额外依赖)

用途: 屏幕预览不需要几十 MB 的精细网格; 本脚本把 mesh\\ 里 > 阈值的 STL
     按网格边长 (GRID_MM) 做顶点聚类, 体积缩小 3~10 倍, 外形基本不变。

用法: python simplify-stl.py   (处理 mesh\\ 下所有 > 5MB 的 .stl, 就地覆盖)
"""
import os
import re
import struct

import numpy as np

BASE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
MESH_DIR = os.path.join(BASE, 'catalog', 'mesh')
SIZE_MIN = 2 * 1024 * 1024    # 只处理 2MB 以上的网格 (控制窗口内渲染内存)
GRID_MM = 0.75                # 聚类网格边长 (mm): 越大越简 (0.5~1.0 预览合适)


def read_stl(path):
    """自动识别二进制/ASCII STL, 返回 (N,3,3) 三角顶点数组"""
    with open(path, 'rb') as f:
        head = f.read(5)
    if head[:5] == b'solid':
        # ASCII STL: 正则抽出所有 vertex 行
        with open(path, 'r', errors='ignore') as f:
            txt = f.read()
        arr = np.array(re.findall(r'vertex\s+([-\deE+.]+)\s+([-\deE+.]+)\s+([-\deE+.]+)', txt),
                       dtype=np.float32)
        return arr.reshape(-1, 3, 3)
    with open(path, 'rb') as f:
        f.read(80)
        (n,) = struct.unpack('<I', f.read(4))
        raw = np.frombuffer(f.read(n * 50), dtype=np.uint8).reshape(n, 50)
        # 每条记录 50 字节: 12 个 float (48B) + 2B 属性; 截取前 48B 再按 float 视图
        f32 = np.ascontiguousarray(raw[:, :48]).view('<f4').reshape(n, 12)
        return f32[:, 3:].reshape(n, 3, 3)


def write_binary_stl(path, tris):
    n = len(tris)
    with open(path, 'wb') as f:
        f.write(b'ONEBOT simplified mesh'.ljust(80, b'\0'))
        f.write(struct.pack('<I', n))
        for tri in tris:
            p0, p1, p2 = tri[0], tri[1], tri[2]
            nrm = np.cross(p1 - p0, p2 - p0)
            norm = np.linalg.norm(nrm)
            nrm = nrm / norm if norm > 1e-12 else np.array([0.0, 0.0, 1.0])
            f.write(struct.pack('<12f', nrm[0], nrm[1], nrm[2],
                                p0[0], p0[1], p0[2], p1[0], p1[1], p1[2], p2[0], p2[1], p2[2]))
            f.write(struct.pack('<H', 0))


def simplify(path):
    tris = read_stl(path)
    n_in = len(tris)
    verts = tris.reshape(-1, 3)
    # 顶点聚类: 落到同一网格单元的顶点合并为均值点
    cell = np.floor(verts / GRID_MM).astype(np.int64)
    _, inv, counts = np.unique(cell, axis=0, return_inverse=True, return_counts=True)
    sums = np.zeros((counts.shape[0], 3))
    np.add.at(sums, inv, verts)
    new_verts = sums / counts[:, None]
    # 重映射三角形, 丢弃退化 (两顶点落入同一单元)
    idx = inv.reshape(-1, 3)
    keep = (idx[:, 0] != idx[:, 1]) & (idx[:, 1] != idx[:, 2]) & (idx[:, 0] != idx[:, 2])
    out = new_verts[idx[keep]]
    write_binary_stl(path, out)
    return n_in, len(out)


def main():
    if not os.path.isdir(MESH_DIR):
        print('mesh 目录不存在')
        return
    for name in sorted(os.listdir(MESH_DIR)):
        if not name.lower().endswith('.stl'):
            continue
        p = os.path.join(MESH_DIR, name)
        size_in = os.path.getsize(p)
        if size_in < SIZE_MIN:
            continue
        n_in, n_out = simplify(p)
        print('%s: %d → %d 三角面 (%.1fMB → %.1fMB)' % (
            name, n_in, n_out, size_in / 1048576.0, os.path.getsize(p) / 1048576.0))
    print('完成')


if __name__ == '__main__':
    main()