import laspy
import numpy as np

def read_logits_from_las(las_path: str, n_classes: int) -> tuple[np.ndarray, np.ndarray]:
    las = laspy.read(las_path)
    coords = np.vstack((las.x, las.y, las.z)).T
    logits = np.vstack([getattr(las, f"sof_log{i}") for i in range(n_classes)]).T.astype(np.float32) / 255.0
    return coords, logits
    
def read_gt_from_las(las_path: str) -> tuple[np.ndarray, np.ndarray]:
    las = laspy.read(las_path)
    coords = np.vstack((las.x, las.y, las.z)).T
    gt = np.array(las.ground_truth, dtype=np.int32)
    return coords, gt

def get_las_length(path):
    las = laspy.open(path)
    return las.header.point_count