#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
fill_all_empty_faces_v2_harmonic_atlas.py

目标：替换原来的“一个缺色孤岛 -> 一个均值色块”的逻辑。

核心逻辑：
1. 正常 face 保持原有 UV，只因为 texture 高度变大而 remap v 坐标。
2. missing face 不再按 region 写一个颜色点，也不再整组 face 指向同一个 vt。
3. 在 3D mesh 拓扑上，对 missing region 的“顶点颜色”做 harmonic/Laplace 扩散：
   - 边界顶点颜色来自相邻正常 face 的原始 texture 采样；
   - 内部顶点颜色由图拉普拉斯扩散得到；
   - 所以颜色从边缘向内部连续变化。
4. 每个 missing face 在底部新增 atlas 中分配一个小三角 tile。
   - tile 内不是纯色，而是按三个顶点颜色做 barycentric 插值；
   - 相邻 missing face 即使在 atlas 里不相邻，只要共享 3D 顶点，边上的插值颜色也连续；
   - tile 周围做 padding，避免纹理过滤时采到白底/其他 tile。

输入输出参数保持原脚本一致：
    --input_obj
    --input_texture
    --missing_faces
    --output_obj
    --output_texture

额外参数都是可选。
"""

import argparse
import os
import time
from collections import defaultdict, deque
from typing import Dict, List, Tuple, Set

import cv2
import numpy as np
import tqdm


# ============================================================
# OBJ / texture IO
# ============================================================

def _parse_obj_index(s: str, current_len: int) -> int:
    """OBJ index: positive is 1-based, negative is relative to current list end."""
    idx = int(s)
    if idx > 0:
        return idx - 1
    if idx < 0:
        return current_len + idx
    raise ValueError("OBJ index 不能为 0")


def read_obj_basic(obj_path: str):
    vertices: List[List[float]] = []
    uvs: List[List[float]] = []
    vertex_indices: List[List[int]] = []
    uv_indices: List[List[int]] = []

    mtllib_lines: List[str] = []
    usemtl_name = "material_0"

    with open(obj_path, "r", encoding="utf-8", errors="ignore") as f:
        for line in f:
            if line.startswith("mtllib "):
                mtllib_lines.append(line.strip())
            elif line.startswith("usemtl "):
                parts = line.strip().split(maxsplit=1)
                if len(parts) == 2:
                    usemtl_name = parts[1]
            elif line.startswith("v "):
                parts = line.strip().split()
                if len(parts) < 4:
                    continue
                vertices.append([float(parts[1]), float(parts[2]), float(parts[3])])
            elif line.startswith("vt "):
                parts = line.strip().split()
                if len(parts) < 3:
                    continue
                uvs.append([float(parts[1]), float(parts[2])])
            elif line.startswith("f "):
                parts = line.strip().split()[1:]
                if len(parts) != 3:
                    raise ValueError(f"当前脚本只支持三角面，发现非三角 face: {line.strip()}")

                v_row = []
                vt_row = []
                for token in parts:
                    sub = token.split("/")
                    if len(sub) < 2 or sub[1] == "":
                        raise ValueError(f"face 缺少 vt，无法处理: {line.strip()}")
                    v_row.append(_parse_obj_index(sub[0], len(vertices)))
                    vt_row.append(_parse_obj_index(sub[1], len(uvs)))

                vertex_indices.append(v_row)
                uv_indices.append(vt_row)

    vertices_np = np.asarray(vertices, dtype=np.float32)
    uvs_np = np.asarray(uvs, dtype=np.float32)
    vertex_indices_np = np.asarray(vertex_indices, dtype=np.int64)
    uv_indices_np = np.asarray(uv_indices, dtype=np.int64)

    if len(vertices_np) == 0 or len(uvs_np) == 0 or len(vertex_indices_np) == 0:
        raise ValueError("OBJ 内没有读到完整的 v / vt / f 数据。")

    return vertices_np, uvs_np, vertex_indices_np, uv_indices_np, mtllib_lines, usemtl_name


def read_missing_faces(path: str, index_base: int = 0) -> np.ndarray:
    arr = []
    with open(path, "r", encoding="utf-8", errors="ignore") as f:
        for line in f:
            s = line.strip()
            if not s:
                continue
            arr.append(int(s) - index_base)
    return np.asarray(arr, dtype=np.int64)


def read_texture_rgb(path: str) -> np.ndarray:
    bgr = cv2.imread(path, cv2.IMREAD_COLOR)
    if bgr is None:
        raise FileNotFoundError(f"无法读取贴图: {path}")
    return cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)


def write_texture_rgb(path: str, rgb: np.ndarray):
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
    ok = cv2.imwrite(path, bgr, [cv2.IMWRITE_PNG_COMPRESSION, 3])
    if not ok:
        raise IOError(f"写入贴图失败: {path}")


def write_obj_with_uv_coordinates(
    filename: str,
    vertices: np.ndarray,
    uvs: np.ndarray,
    vertex_indices: np.ndarray,
    uv_indices: np.ndarray,
    mtllib_lines: List[str] = None,
    usemtl_name: str = "material_0",
):
    os.makedirs(os.path.dirname(os.path.abspath(filename)), exist_ok=True)

    if not mtllib_lines:
        mtllib_lines = ["mtllib mesh.mtl"]

    estimated_size = len(vertices) * 48 + len(uvs) * 32 + len(vertex_indices) * 64
    buffer_size = min(max(int(estimated_size * 1.2), 64 * 1024 * 1024), 1024 * 1024 * 1024)

    lines: List[str] = []
    lines.extend(mtllib_lines)

    for v in vertices:
        lines.append("v %.6f %.6f %.6f" % (v[0], v[1], v[2]))

    lines.append("")
    for uv in uvs:
        lines.append("vt %.8f %.8f" % (uv[0], uv[1]))

    lines.append("")
    lines.append(f"usemtl {usemtl_name}")

    for v_idx, vt_idx in zip(vertex_indices, uv_indices):
        lines.append(
            "f %d/%d %d/%d %d/%d" % (
                int(v_idx[0]) + 1, int(vt_idx[0]) + 1,
                int(v_idx[1]) + 1, int(vt_idx[1]) + 1,
                int(v_idx[2]) + 1, int(vt_idx[2]) + 1,
            )
        )

    with open(filename, "w", buffering=buffer_size, encoding="utf-8") as f:
        f.write("\n".join(lines))


# ============================================================
# Mesh topology
# ============================================================

def build_face_adjacency_and_vertex_faces(faces: np.ndarray, num_vertices: int):
    """
    基于共享边构建 face adjacency，同时构建 vertex -> faces。
    返回：face_adjacency, vertex_faces
    """
    faces = np.asarray(faces, dtype=np.int64)
    num_faces = len(faces)

    # vertex -> faces
    vertex_faces: List[List[int]] = [[] for _ in range(num_vertices)]
    for fi, row in enumerate(faces):
        vertex_faces[int(row[0])].append(fi)
        vertex_faces[int(row[1])].append(fi)
        vertex_faces[int(row[2])].append(fi)

    # edge -> faces，向量化排序
    edges = np.stack([
        np.column_stack((faces[:, 0], faces[:, 1])),
        np.column_stack((faces[:, 1], faces[:, 2])),
        np.column_stack((faces[:, 2], faces[:, 0])),
    ], axis=1).reshape(-1, 2)
    edges.sort(axis=1)
    edge_faces = np.repeat(np.arange(num_faces, dtype=np.int64), 3)

    key = edges[:, 0] * num_vertices + edges[:, 1]
    order = np.argsort(key)
    key = key[order]
    edge_faces = edge_faces[order]

    adj = defaultdict(list)
    same = key[1:] == key[:-1]
    pos = np.where(same)[0]
    for p in pos:
        f0 = int(edge_faces[p])
        f1 = int(edge_faces[p + 1])
        if f0 != f1:
            adj[f0].append(f1)
            adj[f1].append(f0)

    return dict(adj), vertex_faces


def find_missing_regions(face_adjacency: Dict[int, List[int]], missing_faces: np.ndarray, total_faces: int):
    missing_all = set(int(x) for x in missing_faces if 0 <= int(x) < total_faces)
    unused = set(missing_all)
    regions: List[Tuple[Set[int], Set[int]]] = []

    pbar = tqdm.tqdm(total=len(unused), desc="Finding missing islands")
    done = 0

    while unused:
        start = unused.pop()
        group = {start}
        boundary = set()
        stack = [start]

        while stack:
            f = stack.pop()
            for nb in face_adjacency.get(f, []):
                if nb in unused:
                    unused.remove(nb)
                    group.add(nb)
                    stack.append(nb)
                elif nb not in missing_all:
                    boundary.add(nb)

        regions.append((group, boundary))
        new_done = len(missing_all) - len(unused)
        pbar.update(new_done - done)
        done = new_done

    pbar.close()

    if regions:
        sizes = [len(x[0]) for x in regions]
        print(f"Total missing islands: {len(regions)}")
        print(f"Island size: min={min(sizes)}, max={max(sizes)}, avg={sum(sizes)/len(sizes):.2f}")
    else:
        print("No missing islands found.")

    return regions


# ============================================================
# Texture sampling
# ============================================================

def sample_texture_bilinear_rgb(texture_rgb: np.ndarray, uv: np.ndarray) -> np.ndarray:
    """uv: [2], v is OBJ convention, image y is flipped."""
    H, W = texture_rgb.shape[:2]
    u = float(np.clip(uv[0], 0.0, 1.0))
    v = float(np.clip(uv[1], 0.0, 1.0))

    x = u * (W - 1)
    y = (1.0 - v) * (H - 1)

    x0 = int(np.floor(x))
    y0 = int(np.floor(y))
    x1 = min(x0 + 1, W - 1)
    y1 = min(y0 + 1, H - 1)

    dx = x - x0
    dy = y - y0

    c00 = texture_rgb[y0, x0].astype(np.float32)
    c10 = texture_rgb[y0, x1].astype(np.float32)
    c01 = texture_rgb[y1, x0].astype(np.float32)
    c11 = texture_rgb[y1, x1].astype(np.float32)

    c0 = c00 * (1.0 - dx) + c10 * dx
    c1 = c01 * (1.0 - dx) + c11 * dx
    return c0 * (1.0 - dy) + c1 * dy


def sample_face_corner_color(
    face_id: int,
    corner: int,
    texture_rgb: np.ndarray,
    uvs: np.ndarray,
    face_uv_indices: np.ndarray,
    corner_inset: float = 0.90,
) -> np.ndarray:
    """
    在正常 face 的某个顶点附近采样。
    不直接采顶点本身，而是往三角形内部缩一点，避免 seam 边界采样不稳。
    corner_inset 越接近 1，越贴近该顶点。
    """
    tri_uv = uvs[face_uv_indices[face_id]]  # [3,2]
    w = np.full(3, (1.0 - corner_inset) / 2.0, dtype=np.float32)
    w[corner] = corner_inset
    uv = w @ tri_uv
    return sample_texture_bilinear_rgb(texture_rgb, uv)


def median_color(colors: List[np.ndarray]) -> np.ndarray:
    if len(colors) == 0:
        return np.asarray([128, 128, 128], dtype=np.float32)
    return np.median(np.stack(colors, axis=0).astype(np.float32), axis=0)


# ============================================================
# Harmonic vertex color diffusion
# ============================================================

def build_region_vertex_graph(region_faces: Set[int], faces: np.ndarray):
    """
    对一个 missing island 构建 local vertex graph。
    返回：unique_vertices, vertex_to_local, unique_edges, local_face_ids
    """
    face_ids = np.asarray(sorted(region_faces), dtype=np.int64)
    tri_v = faces[face_ids]  # [F,3]
    unique_vertices = np.unique(tri_v.reshape(-1))
    v2local = {int(v): i for i, v in enumerate(unique_vertices)}

    local_faces = np.vectorize(lambda x: v2local[int(x)], otypes=[np.int64])(tri_v)

    edges = np.concatenate([
        local_faces[:, [0, 1]],
        local_faces[:, [1, 2]],
        local_faces[:, [2, 0]],
    ], axis=0)
    edges.sort(axis=1)
    edges = np.unique(edges, axis=0)

    return face_ids, unique_vertices, v2local, edges, local_faces


def compute_boundary_vertex_constraints(
    unique_vertices: np.ndarray,
    vertex_faces: List[List[int]],
    faces: np.ndarray,
    face_uv_indices: np.ndarray,
    uvs: np.ndarray,
    texture_rgb: np.ndarray,
    missing_mask: np.ndarray,
    corner_inset: float = 0.90,
):
    """
    对 region 内的每个顶点，如果它邻接任何非 missing face，就从这些正常 face 的对应 corner 采颜色。
    这比“对边界 face 求一个均值”更合理，因为边界不同位置的颜色会保留下来。
    """
    n = len(unique_vertices)
    constrained = np.zeros(n, dtype=bool)
    constraint_colors = np.zeros((n, 3), dtype=np.float32)

    for local_i, gv in enumerate(unique_vertices):
        samples = []
        for f in vertex_faces[int(gv)]:
            if missing_mask[f]:
                continue
            row = faces[f]
            corners = np.where(row == gv)[0]
            if len(corners) == 0:
                continue
            for c in corners:
                samples.append(sample_face_corner_color(
                    f,
                    int(c),
                    texture_rgb,
                    uvs,
                    face_uv_indices,
                    corner_inset=corner_inset,
                ))

        if samples:
            constrained[local_i] = True
            constraint_colors[local_i] = median_color(samples)

    return constrained, constraint_colors


def initialize_vertex_colors_by_bfs(edges: np.ndarray, constrained: np.ndarray, constraint_colors: np.ndarray) -> np.ndarray:
    n = len(constrained)
    colors = np.zeros((n, 3), dtype=np.float32)

    if constrained.any():
        colors[constrained] = constraint_colors[constrained]
        fallback = np.mean(constraint_colors[constrained], axis=0)
    else:
        fallback = np.asarray([128, 128, 128], dtype=np.float32)
        colors[:] = fallback
        return colors

    adj = [[] for _ in range(n)]
    for a, b in edges:
        adj[int(a)].append(int(b))
        adj[int(b)].append(int(a))

    q = deque(np.where(constrained)[0].tolist())
    visited = np.zeros(n, dtype=bool)
    visited[constrained] = True

    while q:
        cur = q.popleft()
        for nb in adj[cur]:
            if visited[nb]:
                continue
            colors[nb] = colors[cur]
            visited[nb] = True
            q.append(nb)

    colors[~visited] = fallback
    return colors


def solve_harmonic_vertex_colors(
    edges: np.ndarray,
    constrained: np.ndarray,
    constraint_colors: np.ndarray,
    smooth_iters: int = 80,
    self_weight: float = 0.0,
) -> np.ndarray:
    """
    对 region 顶点颜色做 Dirichlet boundary harmonic extension。
    constrained 顶点保持边界颜色；非 constrained 顶点反复取邻居平均。
    """
    n = len(constrained)
    colors = initialize_vertex_colors_by_bfs(edges, constrained, constraint_colors)

    if n == 0 or len(edges) == 0:
        return colors

    e0 = edges[:, 0].astype(np.int64)
    e1 = edges[:, 1].astype(np.int64)

    deg = np.zeros(n, dtype=np.float32)
    np.add.at(deg, e0, 1.0)
    np.add.at(deg, e1, 1.0)
    denom = deg[:, None] + float(self_weight)
    denom = np.maximum(denom, 1e-8)

    for _ in range(max(0, int(smooth_iters))):
        acc = np.zeros_like(colors)
        np.add.at(acc, e0, colors[e1])
        np.add.at(acc, e1, colors[e0])

        if self_weight > 0:
            acc += colors * float(self_weight)

        new_colors = acc / denom
        # Dirichlet boundary：边界顶点锁住，不允许被内部均化掉。
        new_colors[constrained] = constraint_colors[constrained]
        colors = new_colors

    return np.clip(colors, 0, 255).astype(np.float32)


# ============================================================
# Per-face triangular atlas baking
# ============================================================

def precompute_triangle_tile(tile_size: int, pad: int):
    """
    预计算一个标准三角 tile 的 barycentric weights 和 padding 最近邻映射。
    后续每个 missing face 只需要 weights @ 三个顶点颜色。
    """
    tile_size = int(tile_size)
    pad = int(pad)
    if tile_size < 4:
        raise ValueError("tile_size 至少要 >= 4")
    if pad < 1:
        pad = 1
    if pad * 2 + 2 >= tile_size:
        pad = max(1, tile_size // 4)

    # 标准右三角，三个 UV 角点都留 padding，避免双线性采样采到 tile 外。
    p0 = np.asarray([pad, pad], dtype=np.float32)
    p1 = np.asarray([tile_size - pad - 1, pad], dtype=np.float32)
    p2 = np.asarray([pad, tile_size - pad - 1], dtype=np.float32)

    yy, xx = np.meshgrid(np.arange(tile_size, dtype=np.float32),
                         np.arange(tile_size, dtype=np.float32), indexing="ij")
    pts = np.stack([xx + 0.5, yy + 0.5], axis=-1)  # [T,T,2], pixel centers

    # barycentric
    v0 = p1 - p0
    v1 = p2 - p0
    v2 = pts - p0
    d00 = float(np.dot(v0, v0))
    d01 = float(np.dot(v0, v1))
    d11 = float(np.dot(v1, v1))
    denom = d00 * d11 - d01 * d01
    if abs(denom) < 1e-8:
        raise ValueError("退化 tile triangle")

    d20 = v2[..., 0] * v0[0] + v2[..., 1] * v0[1]
    d21 = v2[..., 0] * v1[0] + v2[..., 1] * v1[1]
    w1 = (d11 * d20 - d01 * d21) / denom
    w2 = (d00 * d21 - d01 * d20) / denom
    w0 = 1.0 - w1 - w2
    weights = np.stack([w0, w1, w2], axis=-1).astype(np.float32)

    inside = (weights[..., 0] >= -1e-4) & (weights[..., 1] >= -1e-4) & (weights[..., 2] >= -1e-4)

    # padding：tile 外部/三角外部像素复制最近的三角内部像素颜色。
    inside_coords = np.argwhere(inside)
    nearest_y = np.zeros((tile_size, tile_size), dtype=np.int32)
    nearest_x = np.zeros((tile_size, tile_size), dtype=np.int32)
    for y in range(tile_size):
        for x in range(tile_size):
            if inside[y, x]:
                nearest_y[y, x] = y
                nearest_x[y, x] = x
            else:
                d2 = (inside_coords[:, 0] - y) ** 2 + (inside_coords[:, 1] - x) ** 2
                k = int(np.argmin(d2))
                nearest_y[y, x] = int(inside_coords[k, 0])
                nearest_x[y, x] = int(inside_coords[k, 1])

    padded_weights = weights[nearest_y, nearest_x]

    # UV 角点使用像素中心，更稳定。
    uv_points = np.stack([p0, p1, p2], axis=0).astype(np.float32)

    return padded_weights, uv_points


def build_harmonic_atlas_texture(
    texture_rgb: np.ndarray,
    original_uvs: np.ndarray,
    faces: np.ndarray,
    face_uv_indices: np.ndarray,
    missing_faces: np.ndarray,
    face_corner_colors: Dict[int, np.ndarray],
    tile_size: int = 12,
    tile_pad: int = 2,
):
    """
    给每个 missing face 分配一个小三角 tile。
    tile 内按三个顶点颜色插值，不是纯色块。
    """
    H, W = texture_rgb.shape[:2]
    missing_faces_sorted = np.asarray(sorted(set(int(x) for x in missing_faces)), dtype=np.int64)
    n_missing = len(missing_faces_sorted)

    cols = max(1, W // tile_size)
    rows = int(np.ceil(n_missing / cols))
    atlas_h = rows * tile_size
    new_H = H + atlas_h

    print(f"atlas tiles: {n_missing}, cols={cols}, rows={rows}, atlas_h={atlas_h}")

    new_texture = np.full((new_H, W, 3), 255, dtype=np.uint8)
    new_texture[:H, :, :] = texture_rgb

    # 原始 UV remap：保持原图像素位置不变，只是贴图高度变大。
    remapped_uvs = original_uvs.copy().astype(np.float32)
    old_y = (1.0 - remapped_uvs[:, 1]) * (H - 1)
    remapped_uvs[:, 1] = 1.0 - old_y / max(new_H - 1, 1)

    new_uvs: List[List[float]] = remapped_uvs.tolist()
    new_face_uv_indices = face_uv_indices.copy()

    weights_map, uv_points = precompute_triangle_tile(tile_size, tile_pad)
    weights_flat = weights_map.reshape(-1, 3)  # [T*T,3]

    for i, f in enumerate(tqdm.tqdm(missing_faces_sorted, desc="Baking missing face tiles")):
        col = i % cols
        row = i // cols
        x0 = col * tile_size
        y0 = H + row * tile_size

        corner_colors = face_corner_colors.get(int(f))
        if corner_colors is None:
            corner_colors = np.full((3, 3), 128, dtype=np.float32)
        corner_colors = np.asarray(corner_colors, dtype=np.float32)

        tile = weights_flat @ corner_colors  # [T*T,3]
        tile = np.clip(np.rint(tile), 0, 255).astype(np.uint8).reshape(tile_size, tile_size, 3)

        new_texture[y0:y0 + tile_size, x0:x0 + tile_size, :] = tile

        new_vt = []
        for p in uv_points:
            px = x0 + float(p[0])
            py = y0 + float(p[1])
            u = px / max(W - 1, 1)
            v = 1.0 - py / max(new_H - 1, 1)
            new_uvs.append([u, v])
            new_vt.append(len(new_uvs) - 1)

        new_face_uv_indices[int(f)] = np.asarray(new_vt, dtype=new_face_uv_indices.dtype)

    return new_texture, np.asarray(new_uvs, dtype=np.float32), new_face_uv_indices


# ============================================================
# Main process
# ============================================================

def process(
    input_obj_path: str,
    input_texture_path: str,
    missing_faces_path: str,
    output_obj_path: str,
    output_texture_path: str,
    missing_index_base: int = 0,
    tile_size: int = 12,
    tile_pad: int = 2,
    smooth_iters: int = 80,
    self_weight: float = 0.0,
    corner_inset: float = 0.90,
):
    start = time.time()

    print("Reading input files...")
    vertices, uvs, faces, face_uv_indices, mtllib_lines, usemtl_name = read_obj_basic(input_obj_path)
    texture_rgb = read_texture_rgb(input_texture_path)
    missing_faces = read_missing_faces(missing_faces_path, index_base=missing_index_base)

    total_faces = len(faces)
    missing_faces = missing_faces[(missing_faces >= 0) & (missing_faces < total_faces)]
    missing_faces = np.unique(missing_faces)
    missing_mask = np.zeros(total_faces, dtype=bool)
    missing_mask[missing_faces] = True

    print(f"vertices: {len(vertices)}")
    print(f"uvs: {len(uvs)}")
    print(f"faces: {len(faces)}")
    print(f"missing faces: {len(missing_faces)}")
    print(f"texture: {texture_rgb.shape[1]} x {texture_rgb.shape[0]}")

    if len(missing_faces) == 0:
        write_obj_with_uv_coordinates(output_obj_path, vertices, uvs, faces, face_uv_indices, mtllib_lines, usemtl_name)
        write_texture_rgb(output_texture_path, texture_rgb)
        return

    t0 = time.time()
    print("Building adjacency...")
    face_adjacency, vertex_faces = build_face_adjacency_and_vertex_faces(faces, len(vertices))
    print(f"adjacency using: {time.time() - t0:.2f}s")

    t0 = time.time()
    regions = find_missing_regions(face_adjacency, missing_faces, total_faces)
    print(f"regions using: {time.time() - t0:.2f}s")

    # face -> three corner colors
    face_corner_colors: Dict[int, np.ndarray] = {}

    for region_idx, (region_faces, _boundary_faces) in enumerate(tqdm.tqdm(regions, desc="Solving region colors")):
        face_ids, unique_vertices, v2local, edges, local_faces = build_region_vertex_graph(region_faces, faces)

        constrained, constraint_colors = compute_boundary_vertex_constraints(
            unique_vertices=unique_vertices,
            vertex_faces=vertex_faces,
            faces=faces,
            face_uv_indices=face_uv_indices,
            uvs=uvs,
            texture_rgb=texture_rgb,
            missing_mask=missing_mask,
            corner_inset=corner_inset,
        )

        if not constrained.any():
            # 极少数情况：region 没有任何正常面顶点约束。
            # 这种输入本身没有颜色来源，只能给灰色，避免崩。
            print(f"Warning: region {region_idx} has no boundary vertex color constraints, using gray.")
            vertex_colors = np.full((len(unique_vertices), 3), 128, dtype=np.float32)
        else:
            vertex_colors = solve_harmonic_vertex_colors(
                edges=edges,
                constrained=constrained,
                constraint_colors=constraint_colors,
                smooth_iters=smooth_iters,
                self_weight=self_weight,
            )

        # 写成每个 missing face 的三个角颜色。
        for local_fi, f in enumerate(face_ids):
            lf = local_faces[local_fi]
            face_corner_colors[int(f)] = vertex_colors[lf]

    print("Building harmonic atlas...")
    new_texture, new_uvs, new_face_uv_indices = build_harmonic_atlas_texture(
        texture_rgb=texture_rgb,
        original_uvs=uvs,
        faces=faces,
        face_uv_indices=face_uv_indices,
        missing_faces=missing_faces,
        face_corner_colors=face_corner_colors,
        tile_size=tile_size,
        tile_pad=tile_pad,
    )

    print("Writing outputs...")
    write_obj_with_uv_coordinates(
        output_obj_path,
        vertices,
        new_uvs,
        faces,
        new_face_uv_indices,
        mtllib_lines=mtllib_lines,
        usemtl_name=usemtl_name,
    )
    write_texture_rgb(output_texture_path, new_texture)

    print(f"output texture: {new_texture.shape[1]} x {new_texture.shape[0]}")
    print(f"new uv count: {len(new_uvs)}")
    print(f"Total using: {time.time() - start:.2f}s")


def main():
    parser = argparse.ArgumentParser(description="Fill missing color faces with harmonic vertex-color atlas baking.")

    # 原始接口保持一致
    parser.add_argument("--input_obj", type=str, required=True, help="Path to the input OBJ file")
    parser.add_argument("--input_texture", type=str, required=True, help="Path to the input texture file")
    parser.add_argument("--missing_faces", type=str, required=True, help="Path to missing face index file")
    parser.add_argument("--output_obj", type=str, required=True, help="Path to the output OBJ file")
    parser.add_argument("--output_texture", type=str, required=True, help="Path to the output texture file")

    # 可选参数
    parser.add_argument("--missing_index_base", type=int, default=0,
                        help="missing_faces.txt 的索引基准。原脚本默认 0-based；如果文件是 1-based，传 1。")
    parser.add_argument("--tile_size", type=int, default=12,
                        help="每个 missing face 在新增 atlas 中的 tile 尺寸。越大越细，贴图越大。推荐 8~14。")
    parser.add_argument("--tile_pad", type=int, default=2,
                        help="每个 tile 内三角形到边界的 padding。用于减少纹理过滤串色。推荐 2。")
    parser.add_argument("--smooth_iters", type=int, default=80,
                        help="顶点颜色 harmonic 扩散迭代次数。越大越平滑，稍慢。推荐 50~120。")
    parser.add_argument("--self_weight", type=float, default=0.0,
                        help="扩散时保留当前颜色的权重。一般用 0。")
    parser.add_argument("--corner_inset", type=float, default=0.90,
                        help="边界顶点采样时往正常 face 内部缩进的比例。0.85~0.95 比较稳。")

    args = parser.parse_args()

    process(
        input_obj_path=args.input_obj,
        input_texture_path=args.input_texture,
        missing_faces_path=args.missing_faces,
        output_obj_path=args.output_obj,
        output_texture_path=args.output_texture,
        missing_index_base=args.missing_index_base,
        tile_size=args.tile_size,
        tile_pad=args.tile_pad,
        smooth_iters=args.smooth_iters,
        self_weight=args.self_weight,
        corner_inset=args.corner_inset,
    )


if __name__ == "__main__":
    main()