import torch
import numpy as np


def paint_points(
    points_lidar: np.ndarray,    # N × 3  (x, y, z) LiDAR 좌표계
    seg_map: np.ndarray,          # H × W × C  소프트맥스 점수
    T_cam_lidar: np.ndarray,      # 4×4 외부 매개변수 (lidar → camera)
    K: np.ndarray,                # 3×4 내부 매개변수 / 투영 행렬
) -> np.ndarray:
    # 각 LiDAR 포인트에 픽셀별 의미론적 점수를 추가합니다.
    # 반환: painted_points: N × (3 + C)
    H, W, C = seg_map.shape
    N = points_lidar.shape[0]

    # 동차 LiDAR 포인트: 4 × N
    ones = np.ones((N, 1))
    pts_h = np.hstack([points_lidar, ones]).T       # 4 × N

    # 카메라 프레임에 투영
    pts_cam = T_cam_lidar @ pts_h                   # 4 × N  (Rt 적용)
    pts_2d  = K @ pts_cam                           # 3 × N  (K · [X;Y;Z;1])

    # 픽셀 좌표로 정규화
    depth = pts_2d[2, :]                            # 카메라 프레임의 Z
    u = pts_2d[0, :] / (depth + 1e-6)
    v = pts_2d[1, :] / (depth + 1e-6)
    u_px = np.round(u).astype(int)
    v_px = np.round(v).astype(int)

    # 유효성 마스크: 범위 내이고 양의 깊이
    valid = (depth > 0) & (u_px >= 0) & (u_px < W) & (v_px >= 0) & (v_px < H)

    # 의미론적 점수 조회
    scores = np.zeros((N, C), dtype=np.float32)
    scores[valid] = seg_map[v_px[valid], u_px[valid], :]  # N_valid × C

    painted = np.hstack([points_lidar, scores])     # N × (3 + C)
    return painted


# --- 배치 GPU 추론을 위한 PyTorch 버전 ---
def paint_points_torch(
    points: torch.Tensor,   # N × 3
    seg_map: torch.Tensor,  # H × W × C
    T: torch.Tensor,        # 4 × 4
    K: torch.Tensor,        # 3 × 4
) -> torch.Tensor:
    N = points.shape[0]
    H, W, C = seg_map.shape
    ones = torch.ones(N, 1, device=points.device)
    pts_h = torch.cat([points, ones], dim=1).T      # 4 × N
    pts_cam = T @ pts_h                             # 4 × N
    pts_2d  = K @ pts_cam                           # 3 × N
    depth = pts_2d[2]
    u = (pts_2d[0] / depth.clamp(min=1e-6)).long()
    v = (pts_2d[1] / depth.clamp(min=1e-6)).long()
    valid = (depth > 0) & (u >= 0) & (u < W) & (v >= 0) & (v < H)
    scores = torch.zeros(N, C, device=points.device)
    scores[valid] = seg_map[v[valid], u[valid]]
    return torch.cat([points, scores], dim=1)       # N × (3+C)