#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
fill_gap_v3.py — 基于射线投射的眼镜-脸空隙填充
"""

import numpy as np
from scipy.spatial import cKDTree
from scipy.ndimage import binary_dilation, distance_transform_edt
from skimage import measure
import trimesh
from trimesh.smoothing import filter_taubin
import argparse
import time


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--obj", required=True)
    parser.add_argument("--index", required=True)
    parser.add_argument("--out", default="fill_mesh.obj")
    parser.add_argument("--res", type=int, default=60, help="分辨率（先跑低分辨率验证）")
    parser.add_argument("--gap", type=float, default=1.5, help="填充间隙宽度 mm")
    args = parser.parse_args()

    t0 = time.time()
    print("=== 眼镜-脸部空隙填充 v3 ===\n")

    # 1. 加载
    mesh = trimesh.load(args.obj, force="mesh")
    V = mesh.vertices
    F = mesh.faces
    idx = np.loadtxt(args.index, dtype=np.int64) - 1
    print(f"[load] 顶点 {len(V)}, 面 {len(F)}, 眼镜顶点 {len(idx)}")

    # 2. 分类
    is_glass = np.zeros(len(V), dtype=bool)
    is_glass[idx] = True
    glass_verts = V[idx]
    face_verts = V[~is_glass]
    print(f"[classify] 眼镜顶点 {len(glass_verts)}, 脸部顶点 {len(face_verts)}")

    # 3. 对每个眼镜顶点，找最近的脸部顶点
    face_tree = cKDTree(face_verts)
    dists, nearest_idx = face_tree.query(glass_verts, k=1)
    nearest_face_verts = face_verts[nearest_idx]
    
    # 方向向量：从脸部指向眼镜
    directions = glass_verts - nearest_face_verts
    dir_lengths = np.linalg.norm(directions, axis=1)
    directions = directions / dir_lengths[:, None]
    
    # 统计平均方向
    avg_dir = directions.mean(axis=0)
    avg_dir = avg_dir / np.linalg.norm(avg_dir)
    print(f"[dir] 平均方向（脸部→眼镜）: {avg_dir.round(3)}")

    # 4. 确定填充范围
    # 眼镜中心
    glass_center = glass_verts.mean(axis=0)
    # 沿方向投影
    proj = (glass_verts - glass_center) @ avg_dir
    front_pos = proj.max()  # 眼镜前表面（远离脸部）
    back_pos = proj.min()   # 眼镜后表面（靠近脸部）
    print(f"[surf] 前 {front_pos:.2f}, 后 {back_pos:.2f}")

    # 5. 包围盒（只用眼镜区域）
    margin = 4.0
    lo = glass_verts.min(axis=0) - margin
    hi = glass_verts.max(axis=0) + margin
    print(f"[bbox] {lo.round(2)} -> {hi.round(2)}")

    # 6. 体素网格（低分辨率测试）
    shape = np.array([args.res] * 3)
    pitch = (hi - lo).max() / args.res
    xs = np.linspace(lo[0], hi[0], shape[0])
    ys = np.linspace(lo[1], hi[1], shape[1])
    zs = np.linspace(lo[2], hi[2], shape[2])
    X, Y, Z = np.meshgrid(xs, ys, zs, indexing="ij")
    points = np.stack([X.ravel(), Y.ravel(), Z.ravel()], axis=1)
    print(f"[grid] {shape} = {points.shape[0]} 体素, 大小 {pitch:.3f}mm")

    # 7. KD-Tree
    glass_tree = cKDTree(glass_verts)
    face_tree = cKDTree(face_verts)
    print("[tree] KD-Tree 构建完成")

    # 8. 计算距离
    print("[dist] 计算距离场...")
    dist_to_glass, _ = glass_tree.query(points, k=1)
    dist_to_face, _ = face_tree.query(points, k=1)

    # 9. 计算每个体素在"方向"上的投影
    proj_val = (points - glass_center) @ avg_dir

    # 10. 填充规则
    # 在眼镜后表面附近（靠近脸部的那侧）
    near_glass = dist_to_glass < args.gap * 2
    
    # 不在脸内部
    outside_face = dist_to_face > 0.3
    
    # 在后表面附近（back_pos 到 back_pos + gap 之间）
    in_gap = (proj_val > back_pos - 0.5) & (proj_val < back_pos + args.gap)
    
    # 组合
    fill_grid = near_glass & outside_face & in_gap
    fill_grid = fill_grid.reshape(shape)

    n_fill = fill_grid.sum()
    print(f"[fill] 填充体素: {n_fill}")

    if n_fill < 30:
        print("[warn] 太少，尝试仅用近眼镜条件...")
        fill_grid = near_glass.reshape(shape)
        n_fill = fill_grid.sum()
        print(f"[fill] 仅近眼镜: {n_fill}")

    if n_fill < 10:
        print("[error] 填充为空")
        return

    # 11. 膨胀
    fill_grid = binary_dilation(fill_grid, iterations=3)

    # 12. MC
    print("[mc] Marching Cubes...")
    sdf = distance_transform_edt(~fill_grid).astype(np.float32)
    sdf[fill_grid] = -distance_transform_edt(fill_grid)[fill_grid]

    verts, faces, _, _ = measure.marching_cubes(
        sdf, level=0.0,
        spacing=(pitch, pitch, pitch)
    )
    verts += lo

    result = trimesh.Trimesh(vertices=verts, faces=faces, process=True)

    # 取最大分量
    if len(result.split()) > 1:
        comps = result.split(only_watertight=False)
        result = max(comps, key=lambda m: m.volume)
        print(f"[clean] 取最大分量")

    filter_taubin(result, lamb=0.5, nu=0.5, iterations=5)

    result.export(args.out)
    print(f"\n[output] {args.out}")
    print(f"  顶点: {len(result.vertices)}, 面: {len(result.faces)}")
    print(f"  体积: {result.volume:.3f}, 水密: {result.is_watertight}")
    print(f"  耗时: {time.time()-t0:.1f}s")


if __name__ == "__main__":
    main()