import numpy as np
import torch

from anydata.utils.types import is_seq

def points_to_camera_frame(points, extrinsics):
    """
    Transform 3D points from world/base frame to camera frame.

    Args:
        points: (N, 3) numpy array of 3D points in world/base frame
        extrinsics: (4, 4) or (N, 4, 4) numpy array camera extrinsics (camera-to-world)

    Returns:
        (N, 3) numpy array of 3D points in camera frame
    """
    N, _ = points.shape
    points_homo = np.concatenate([points, np.ones((N, 1))], axis=1)

    if extrinsics.ndim == 2:  # Single transform (4, 4)
        points_cam = np.linalg.inv(extrinsics) @ points_homo.T
        return points_cam.T[:, :3]
    else:  # Batch of transforms (N, 4, 4)
        extrinsics_inv = np.linalg.inv(extrinsics) # Invert to get world-to-camera
        points_cam = np.einsum('nij,nj->ni', extrinsics_inv, points_homo)
        return points_cam[:, :3]


def project_to_image(points_cam, intrinsics):
    """
    Project 3D points in camera frame to 2D pixel coordinates.

    Args:
        points_cam: (N, 3) numpy array of 3D points in camera frame
        intrinsics: (3, 3) or (3, 4) numpy array camera intrinsics

    Returns:
        (N, 3) numpy array of pixel coordinates (u, v, 1)
    """
    # Handle case where intrinsics has extra batch dimension
    if intrinsics.ndim == 3 and intrinsics.shape[0] == 1:
        intrinsics = intrinsics.squeeze(0)

    # Convert 3x3 intrinsics to 3x4 if needed
    if intrinsics.shape == (3, 3):
        intrinsics = np.hstack([intrinsics, np.zeros((3, 1))])

    N, _ = points_cam.shape
    points_homo = np.concatenate([points_cam, np.ones((N, 1))], axis=1)

    px_val = intrinsics @ points_homo.T

    # Safety check: avoid division by very small Z values
    z_vals = px_val[2, :]
    eps = 1e-6
    safe_z = np.where(np.abs(z_vals) < eps, eps * np.sign(z_vals), z_vals)
    safe_z = np.where(safe_z == 0, eps, safe_z)

    px_val = px_val / safe_z

    return px_val.T


def multiply_extrinsics(data, extrinsics):
    """Transform list of points given extrinsics"""
    return extrinsics[:, :3, :3].bmm(data) + extrinsics[:, :3, -1].unsqueeze(-1)


def invert_intrinsics(K):
    """Invert camera intrinsics"""
    Kinv = K.clone()
    Kinv[:, 0, 0] = 1. / K[:, 0, 0]
    Kinv[:, 1, 1] = 1. / K[:, 1, 1]
    Kinv[:, 0, 2] = -1. * K[:, 0, 2] / K[:, 0, 0]
    Kinv[:, 1, 2] = -1. * K[:, 1, 2] / K[:, 1, 1]
    return Kinv


def invert_extrinsics(T, mode='trans'):
    """Inverts a [B,4,4] torch.tensor pose"""
    if mode == 'trans':
        Tinv = torch.eye(4, device=T.device, dtype=T.dtype).repeat([len(T), 1, 1])
        Tinv[:, :3, :3] = torch.transpose(T[:, :3, :3], -2, -1)
        Tinv[:, :3, -1] = torch.bmm(-1. * Tinv[:, :3, :3], T[:, :3, -1].unsqueeze(-1)).squeeze(-1)
        return Tinv
    elif mode == 'linalg':
        return torch.linalg.inv(T)
    else:
        raise ValueError('Invalid mode {}'.format(mode))


def scale_intrinsics(K, ratio):
    """Scale intrinsics given x_scale and y_scale factors"""

    if is_seq(ratio):
        ratio_h, ratio_w = ratio
    else:
        ratio_h = ratio_w = ratio

    K = K.clone()

    K[..., 0, 0] *= ratio_w
    K[..., 1, 1] *= ratio_h

    # if align_corners():
    K[..., 0, 2] = K[..., 0, 2] * ratio_w
    K[..., 1, 2] = K[..., 1, 2] * ratio_h
    # else:
    #     K[..., 0, 2] = (K[..., 0, 2] - 0.5) * ratio_w + 0.5
    #     K[..., 1, 2] = (K[..., 1, 2] - 0.5) * ratio_h + 0.5

    return K