"""
weld_seam.py — 블럭: STL→포인트클라우드 + 심검출 + 경로 생성
(welding_pipeline.py에서 2026-07-04 분할 — 로직 원문 그대로, 수정 없음)
출처: romi-lab/robotic-welding-demo (demo_all.py)
"""
import numpy as np
from scipy.interpolate import splprep, splev
from scipy.spatial.transform import Rotation as R
from scipy.spatial import KDTree
import open3d as o3d


# ─────────────────────────────────────────────
# 1. STL → 포인트 클라우드
# ─────────────────────────────────────────────

def stl_to_pointcloud(stl_path, sample_points=100000):
    """trimesh로 STL 샘플링 → open3d 포인트 클라우드"""
    import trimesh
    mesh = trimesh.load(stl_path)
    pts, face_idx = trimesh.sample.sample_surface(mesh, sample_points)
    normals = mesh.face_normals[face_idx]
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(pts)
    pcd.normals = o3d.utility.Vector3dVector(normals)
    return pcd


# ─────────────────────────────────────────────
# 2. 심검출 (romi-lab demo_all.py 알고리즘)
# ─────────────────────────────────────────────

def find_feature_value(feature, pcd, voxel_size):
    """법선벡터 비대칭도(asymmetry)로 홈 엣지 검출"""
    pc_number = len(pcd.points)
    n_list = np.asarray(pcd.normals)
    pcd_tree = o3d.geometry.KDTreeFlann(pcd)
    neighbor = 30
    radius = voxel_size * 5

    feature_value = np.zeros(pc_number)

    if feature == "asymmetry":
        for index in range(pc_number):
            [k, idx, _] = pcd_tree.search_radius_vector_3d(pcd.points[index], radius)
            if k < 3:
                continue
            vector = np.mean(n_list[list(idx), :], axis=0)
            dot = np.dot(vector, n_list[index, :])
            norm = np.linalg.norm(n_list[index, :])
            if norm < 1e-9:
                continue
            proj = dot / (norm * norm) * n_list[index, :]
            feature_value[index] = np.linalg.norm(vector - proj)

    return feature_value


def cluster_groove(pcd, feature_value, threshold_ratio=0.7, eps=0.01, min_points=10):
    """비대칭도 높은 포인트 → DBSCAN 클러스터링"""
    threshold = np.max(feature_value) * threshold_ratio
    high_idx = np.where(feature_value > threshold)[0]

    pts = np.asarray(pcd.points)[high_idx]
    groove_pcd = o3d.geometry.PointCloud()
    groove_pcd.points = o3d.utility.Vector3dVector(pts)

    labels = np.array(groove_pcd.cluster_dbscan(eps=eps, min_points=min_points))
    if len(labels) == 0 or labels.max() < 0:
        return groove_pcd  # 클러스터 없으면 전체 반환

    # 가장 큰 클러스터 선택
    best = np.argmax(np.bincount(labels[labels >= 0]))
    mask = labels == best
    result = o3d.geometry.PointCloud()
    result.points = o3d.utility.Vector3dVector(pts[mask])
    return result


def thin_line(points, thickness=0.5):
    """SVD로 포인트들을 회귀선 위로 투영 → 중심선 추출"""
    centroid = np.mean(points, axis=0)
    centered = points - centroid
    _, _, vh = np.linalg.svd(centered)
    direction = vh[0]  # 주축

    # 포인트를 직선 위 투영
    t_vals = centered.dot(direction)
    projected = np.outer(t_vals, direction) + centroid
    return projected, direction


def sort_points(points, step=0.005):
    """KDTree로 순차 정렬"""
    remaining = list(range(len(points)))
    sorted_pts = [points[0]]
    remaining.remove(0)
    tree = KDTree(points)

    while remaining:
        last = sorted_pts[-1]
        dists, idxs = tree.query(last, k=min(10, len(remaining) + 1))
        for idx in idxs:
            if idx in remaining:
                sorted_pts.append(points[idx])
                remaining.remove(idx)
                break

    return np.array(sorted_pts)


def generate_trajectory(sorted_points, smooth_factor=0.0, num_output=100):
    """B-스플라인 보간으로 부드러운 경로 생성"""
    pts = sorted_points.T  # (3, N)
    tck, u = splprep(pts, s=smooth_factor, k=min(3, len(sorted_points) - 1))
    u_new = np.linspace(0, 1, num_output)
    smooth = np.array(splev(u_new, tck)).T  # (num_output, 3)
    return smooth


def find_normal(pcd, method="ransac"):
    """RANSAC 평면 피팅으로 면법선 추출"""
    plane_model, _ = pcd.segment_plane(
        distance_threshold=0.01,
        ransac_n=3,
        num_iterations=1000
    )
    normal = np.array(plane_model[:3])
    normal /= np.linalg.norm(normal)
    return normal


def find_orientation(trajectory_pts, surface_normal):
    """경로 접선 + 면법선 → 6DOF 토치 방향 (romi-lab find_orientation)"""
    poses = []
    z_dir = surface_normal / np.linalg.norm(surface_normal)

    for i in range(len(trajectory_pts) - 1):
        pos_diff = trajectory_pts[i + 1] - trajectory_pts[i]
        proj = np.dot(pos_diff, z_dir) * z_dir
        x_dir = pos_diff - proj
        norm = np.linalg.norm(x_dir)
        if norm < 1e-9:
            x_dir = np.array([1.0, 0.0, 0.0])
        else:
            x_dir /= norm

        y_dir = np.cross(z_dir, x_dir)
        y_dir /= np.linalg.norm(y_dir)

        rot_matrix = np.column_stack([x_dir, y_dir, z_dir])
        r = R.from_matrix(rot_matrix)
        poses.append((trajectory_pts[i], r))

    # 마지막 포인트: 이전 방향 유지
    poses.append((trajectory_pts[-1], poses[-1][1]))
    return poses  # list of (position_xyz, Rotation)