import numpy as np
import logging
import torch
import re


def random_per_point_translation_in_place(pcd_data) -> None:
    """
    Jittering the point cloud data by a random value between -0.02 and 0.02

    Args:
        pcd_data: point cloud data in the form x, y, z

    """
    translations = (
        np.random.rand(pcd_data.shape[0], 3) - 0.5
    ) * 0.04  # Random values between -0.02 and 0.02
    pcd_data[:, -3:] += translations


def compute_max_extent_and_centroid(pcd_data, epsilon=1e-4) -> tuple[np.ndarray, np.ndarray]:
    """_summary_

    Args:
        pcd_data : point cloud data in the form x, y, z
        epsilon (float, optional): buffer for the max_extent. Defaults to 1e-4.

    Returns:
        max_extent: maximum extent of the point cloud data in terms of the largest dimension
        centroid: centroid of the point cloud data
    """
    min_vals = pcd_data.min(axis=0)
    max_vals = pcd_data.max(axis=0)
    centroid = (min_vals + max_vals) / 2
    max_extent = np.max(max_vals - min_vals) + epsilon
    return max_extent, centroid


def unit_cube_normalization_in_place(
    pcd_data,
    max_extent,
    centroid,
):
    """
    Normalized data point in a unit cube between 0 and 1 for each x, y, z in-place

    Args:
        pcd_data: point cloud data in the form x, y, z

    """
    # translate the centroid to the origin
    pcd_data -= centroid

    # scale the data to fit within [-0.5, 0.5]
    pcd_data /= max_extent

    # translate it back to within [0, 1]
    pcd_data += 0.5


def point_to_index(point, grid_size):
    """
    Maps a point in the unit cube to a unique index based on the grid size.

    Args:
        point (tuple): a tuple of (x, y, z) coordinates of the point. Each coordinate should be in [0, 1].
        grid_size (int): the number of divisions along each axis.

    Returns:
        int: a unique index for the point.
    """

    xi = int(point[0] * grid_size)
    yi = int(point[1] * grid_size)
    zi = int(point[2] * grid_size)

    # Ensure that the point is inside the unit cube
    if not (0 <= xi < grid_size) or not (0 <= yi < grid_size) or not (0 <= zi < grid_size):
        logging.warning(
            f"The point is outside the unit cube: point: {point}, grid_index: ({xi}, {yi}, {zi})"
        )

    # Clamp the point to be inside the unit cube
    xi = min(max(xi, 0), grid_size - 1)
    yi = min(max(yi, 0), grid_size - 1)
    zi = min(max(zi, 0), grid_size - 1)

    # Compute the unique voxel ID, row-major order
    voxel_id = xi + yi * grid_size + zi * grid_size * grid_size

    return voxel_id


def scale_bbox(bbox_str, max_extent, centroid):
    """
    Scale the bounding box to be within a unit cube and output numerically tokenized bounding box.

    Args:
        bbox_str (str): A string representing a bounding box, in the format "<x_min,y_min,z_min,x_max,y_max,z_max>".
        max_extent (float): The maximum extent of the bounding box.
        centroid (np.array): The centroid of the bounding box.

    Returns:
        str: A string representing the scaled bounding box, in the same format as the input.
    """
    # Remove < and > from the bounding box string
    bbox_str = bbox_str.strip("<>")

    bbox_values = bbox_str.split(",")
    # Convert each string to a float and store in a list
    bbox_floats = [float(value) for value in bbox_values]
    # Convert the list to a numpy array
    bbox_array = np.array(bbox_floats)
    bbox_array[:3] -= centroid
    bbox_array[3:] -= centroid
    bbox_array /= max_extent
    bbox_array += 0.5
    x_min, y_min, z_min, x_max, y_max, z_max = bbox_array
    x_min, y_min, z_min, x_max, y_max, z_max = (
        x_min.item(),
        y_min.item(),
        z_min.item(),
        x_max.item(),
        y_max.item(),
        z_max.item(),
    )
    x_min, y_min, z_min, x_max, y_max, z_max = (
        round(x_min, 3),
        round(y_min, 3),
        round(z_min, 3),
        round(x_max, 3),
        round(y_max, 3),
        round(z_max, 3),
    )
    new_bbox_str = f"< {x_min}, {y_min}, {z_min}, {x_max}, {y_max}, {z_max}>"  # adding space after < because tokenizer will not merge < and first digit or negative sign
    return new_bbox_str


def voxelize_points(
    xyz_to_be_voxelized: np.array,
    scene_min_xyz: np.array,
    scene_max_xyz: np.array,
    num_voxels_per_axis: int,
):
    """Convert points to voxel indexes

    Args:
        xyz_to_be_voxelized (np.array): shape (num_points, 3)
        scene_min_xyz (np.array): shape (3,)
        scene_max_xyz (np.array): shape (3,)
        num_voxels_per_axis (int): number of voxels per axis

    Returns:
        voxel_id (np.array): shape (num_points,)
    """
    voxel_index = np.floor(
        (xyz_to_be_voxelized - scene_min_xyz)
        / (scene_max_xyz - scene_min_xyz)
        * num_voxels_per_axis
    ).astype(
        int
    )  # range after this overations: [0, num_voxels_per_axis]
    voxel_index = np.clip(
        voxel_index, 0, num_voxels_per_axis - 1
    )  # clamp range to [0, num_voxels_per_axis - 1]
    # calculate index using row-major order
    voxel_id = (
        voxel_index[:, 0]
        + voxel_index[:, 1] * num_voxels_per_axis
        + voxel_index[:, 2] * num_voxels_per_axis * num_voxels_per_axis
    )  # range after this operation: [0, num_voxels_per_axis ** 3 - 1]
    return voxel_id


def process_one_bbox_minkowski_loc_token(
    bbox_str, scene_min_xyz, scene_max_xyz, num_voxels_per_axis
):
    # Remove < and > from the bounding box string
    bbox_str = bbox_str.strip("<>")
    bbox_values = bbox_str.split(",")
    # Convert each string to a float and store in a list
    bbox_floats = [float(value) for value in bbox_values]
    # Convert the list to a numpy array
    bbox_array = np.array(bbox_floats)  # shape: (6,)
    bbox_array = bbox_array.reshape(2, 3)  # shape: (2, 3)

    voxel_indices = voxelize_points(
        bbox_array, scene_min_xyz, scene_max_xyz, num_voxels_per_axis
    )  # shape: (2,)

    new_bbox_str = f"<loc_{voxel_indices[0]}><loc_{voxel_indices[1]}>"
    return new_bbox_str


def scale_bbox_special_token(bbox_str, max_extent, centroid, num_grid_cells):
    """
    Special token for the bbox. The bbox is scaled to the unit cube and then converted
    to a unique index based on the grid size.

    Args:
        bbox_str (str): bbox string in the form "<x_min, y_min, z_min, x_max, y_max, z_max>"
        max_extent (float): max extent of the point cloud data in terms of the largest dimension
        centroid (np.array): centroid of the point cloud data
        num_grid_cells (int): number of grids along each axis

    Returns:
        two unique special tokens for the bbox as string
    """
    # Remove < and > from the bounding box string
    bbox_str = bbox_str.strip("<>")

    bbox_values = bbox_str.split(",")
    # Convert each string to a float and store in a list
    bbox_floats = [float(value) for value in bbox_values]
    # Convert the list to a numpy array
    bbox_floats = np.array(bbox_floats)
    bbox_floats[:3] -= centroid
    bbox_floats[3:] -= centroid
    bbox_floats /= max_extent
    bbox_floats += 0.5
    min_point = bbox_floats[:3]
    max_point = bbox_floats[3:]
    index_min = point_to_index(min_point, num_grid_cells)
    index_max = point_to_index(max_point, num_grid_cells)

    new_bbox_str = f"<loc_{index_min}><loc_{index_max}>"
    return new_bbox_str


def rotate_point_cloud_90_degrees(pcd_data):
    """
    Rotate the point cloud data by 90 degrees in the x-y plane

    Args:
        pcd_data: point cloud data in the form x, y, z

    Returns:
        pcd_data: rotated point cloud data in the form x, y, z
    """
    # Randomly select among no change, clockwise, and counterclockwise
    rotation_choices = ["no change", "clockwise", "counterclockwise"]
    direction = np.random.choice(rotation_choices)

    if direction == "clockwise":
        rotation_matrix = torch.tensor([[0, 1], [-1, 0]])
        # Apply rotation on x-y plane
        pcd_data[:, -3:-1] = torch.matmul(pcd_data[:, -3:-1], rotation_matrix)
    elif direction == "counterclockwise":
        rotation_matrix = torch.tensor([[0, -1], [1, 0]])
        # Apply rotation on x-y plane
        pcd_data[:, -3:-1] = torch.matmul(pcd_data[:, -3:-1], rotation_matrix)

    return pcd_data, direction


def adjust_bbox_after_rotation(bbox_str, direction):
    """_summary_

    Args:
        bbox_str (_type_): _description_
        direction (_type_): _description_

    Returns:
        _type_: _description_
    """

    if direction == "no change":
        return bbox_str

    values = list(map(float, re.findall(r"[-+]?\d*\.\d+|\d+", bbox_str)))
    x_min, y_min, z_min, x_max, y_max, z_max = values

    if direction == "clockwise":
        # adding space after < because tokenizer will not merge < and first digit or negative sign
        new_bbox_str = f"< {y_min}, {x_min}, {z_min}, {y_max}, {x_max}, {z_max}>"
    else:  # counterclockwise
        # adding space after < because tokenizer will not merge < and first digit or negative sign
        new_bbox_str = f"< {x_max}, {y_min}, {z_min}, {x_min}, {y_max}, {z_max}>"

    return new_bbox_str