import numpy as np
from scipy.spatial.transform import Rotation as R


def pad_0001(Ts):
    Ts = np.asarray(Ts)
    if Ts.shape[-2:] == (4, 4):
        return Ts
    if Ts.shape[-2:] != (3, 4):
        raise ValueError("Ts must have shape (..., 3, 4) or (..., 4, 4)")
    pad = np.zeros((*Ts.shape[:-2], 1, 4), dtype=Ts.dtype)
    pad[..., 0, 3] = 1
    return np.concatenate([Ts, pad], axis=-2)


def T_to_C(T):
    T = np.asarray(T)
    if T.shape != (4, 4):
        raise ValueError("T must be shape (4,4)")
    Rm = T[:3, :3]
    t = T[:3, 3]
    return -Rm.T @ t


def im_distance_to_im_depth(im_dist, K):
    im_dist = np.asarray(im_dist)
    H, W = im_dist.shape[:2]
    ys, xs = np.indices((H, W), dtype=np.float32)
    fx, fy = K[0, 0], K[1, 1]
    cx, cy = K[0, 2], K[1, 2]
    x = (xs - cx) / max(fx, 1e-8)
    y = (ys - cy) / max(fy, 1e-8)
    ray_norm = np.sqrt(x * x + y * y + 1.0)
    return im_dist / np.clip(ray_norm, 1e-8, None)


def im_depth_to_point_cloud(im_depth, K, T, to_image=False, ignore_invalid=False):
    im_depth = np.asarray(im_depth)
    H, W = im_depth.shape[:2]
    ys, xs = np.indices((H, W), dtype=np.float32)

    fx, fy = K[0, 0], K[1, 1]
    cx, cy = K[0, 2], K[1, 2]

    z = im_depth.reshape(-1)
    x = (xs.reshape(-1) - cx) / max(fx, 1e-8) * z
    y = (ys.reshape(-1) - cy) / max(fy, 1e-8) * z
    pts_cam = np.stack([x, y, z], axis=-1)

    T = np.asarray(T)
    if T.shape != (4, 4):
        raise ValueError("T must be shape (4,4)")
    Rm = T[:3, :3]
    t = T[:3, 3]

    # Assume T is world->camera; invert to camera->world for point transform.
    pts_world = (Rm.T @ (pts_cam - t).T).T

    if ignore_invalid:
        valid = np.isfinite(pts_world).all(axis=1) & (z > 0)
        pts_world = pts_world[valid]

    if to_image:
        return pts_world.reshape(H, W, 3)
    return pts_world

def rotx(x, theta=90):
    """
    Rotate x by theta degrees around the x-axis
    """
    theta = np.deg2rad(theta)
    rot_matrix = np.array(
        [
            [1, 0, 0, 0],
            [0, np.cos(theta), -np.sin(theta), 0],
            [0, np.sin(theta), np.cos(theta), 0],
            [0, 0, 0, 1],
        ]
    )
    return rot_matrix@ x


def Coord2zup(points, extrinsics, normals = None):
    """
    Convert the dust3r coordinate system to the z-up coordinate system
    """
    points = np.concatenate([points, np.ones([points.shape[0], 1])], axis=1).T
    points = rotx(points, -90)[:3].T
    if normals is not None:
        normals = np.concatenate([normals, np.ones([normals.shape[0], 1])], axis=1).T
        normals = rotx(normals, -90)[:3].T
        normals = normals / np.linalg.norm(normals, axis=1, keepdims=True)
    t = np.min(points,axis=0)
    points -= t
    extrinsics = rotx(extrinsics, -90)
    extrinsics[:, :3, 3] -= t.T
    return points, extrinsics, normals

def _ransac_plane(points, distance_threshold=0.01, ransac_n=3, num_iterations=1000):
    if points.shape[0] < ransac_n:
        raise ValueError("Not enough points for plane fitting.")

    best_inliers = None
    best_plane = None
    rng = np.random.default_rng(42)

    for _ in range(num_iterations):
        sample_idx = rng.choice(points.shape[0], size=ransac_n, replace=False)
        p0, p1, p2 = points[sample_idx]
        normal = np.cross(p1 - p0, p2 - p0)
        norm = np.linalg.norm(normal)
        if norm < 1e-8:
            continue
        normal = normal / norm
        d = -np.dot(normal, p0)

        dist = np.abs(points @ normal + d)
        inliers = np.where(dist < distance_threshold)[0]
        if best_inliers is None or len(inliers) > len(best_inliers):
            best_inliers = inliers
            best_plane = np.array([normal[0], normal[1], normal[2], d], dtype=np.float64)

    if best_inliers is None or best_plane is None:
        raise ValueError("Failed to fit plane with RANSAC.")

    return best_plane, best_inliers


def extract_and_align_ground_plane(points,
                                   colors=None,
                                   normals=None,
                                   height_percentile=20, 
                                   ransac_distance_threshold=0.01, 
                                   ransac_n=3, 
                                   ransac_iterations=1000,
                                   max_angle_degree=40,
                                   max_trials=6):
    points = np.asarray(points)
    if points.ndim != 2 or points.shape[1] != 3:
        raise ValueError("points must be shaped (N, 3)")

    aligned_colors = np.asarray(colors) if colors is not None else None
    aligned_normals = np.asarray(normals) if normals is not None else None

    z_vals = points[:, 2]
    z_thresh = np.percentile(z_vals, height_percentile)
    low_indices = np.where(z_vals <= z_thresh)[0]

    remaining_indices = low_indices.copy()

    for trial in range(max_trials):
        if len(remaining_indices) < ransac_n:
            raise ValueError("Not enough points left to fit a plane.")

        candidate_points = points[remaining_indices]
        plane_model, inliers = _ransac_plane(
            candidate_points,
            distance_threshold=ransac_distance_threshold,
            ransac_n=ransac_n,
            num_iterations=ransac_iterations,
        )
        a, b, c, d = plane_model
        normal = np.array([a, b, c])
        normal /= np.linalg.norm(normal)

        angle = np.arccos(np.clip(np.dot(normal, [0, 0, 1]), -1.0, 1.0)) * 180 / np.pi
        if angle <= max_angle_degree:
            inliers_global = remaining_indices[inliers]

            target = np.array([0, 0, 1])
            axis = np.cross(normal, target)
            axis_norm = np.linalg.norm(axis)

            if axis_norm < 1e-6:
                rotation_matrix = np.eye(3)
            else:
                axis /= axis_norm
                rot_angle = np.arccos(np.clip(np.dot(normal, target), -1.0, 1.0))
                rotation = R.from_rotvec(axis * rot_angle)
                rotation_matrix = rotation.as_matrix()

            rotated_points = points @ rotation_matrix.T
            ground_points_z = rotated_points[inliers_global, 2]
            offset = np.mean(ground_points_z)
            rotated_points[:, 2] -= offset

            if aligned_normals is not None and len(aligned_normals) == len(points):
                aligned_normals = aligned_normals @ rotation_matrix.T
                aligned_normals = aligned_normals / np.clip(np.linalg.norm(aligned_normals, axis=-1, keepdims=True), 1e-8, None)

            return rotated_points, aligned_colors, aligned_normals, inliers_global, rotation_matrix, offset

        else:
            rejected_indices = remaining_indices[inliers]
            remaining_indices = np.setdiff1d(remaining_indices, rejected_indices)

    raise ValueError("Failed to find a valid ground plane within max trials.")