from typing import Literal, Union

import numpy as np
import torch
import warp as wp


def generate_dense_grid_points(
    bbox_min: np.ndarray,
    bbox_max: np.ndarray,
    resolution_base: float,
    indexing: Literal["xy", "ij"] = "ij",
) -> tuple[np.ndarray, list[int], np.ndarray]:
    """
    Generate a dense grid of points within a bounding box.

    Parameters:
        bbox_min (np.ndarray): The minimum coordinates of the bounding box (3D).
        bbox_max (np.ndarray): The maximum coordinates of the bounding box (3D).
        resolution_base (float): The base resolution for the grid. The number of cells along each axis will be 2^resolution_base.
        indexing (Literal["xy", "ij"], optional): The indexing convention for the grid. "xy" for Cartesian indexing, "ij" for matrix indexing. Default is "ij".
    Returns:
        tuple: A tuple containing:
            - xyz (np.ndarray): A 2D array of shape (N, 3) where N is the total number of grid points. Each row represents the (x, y, z) coordinates of a grid point.
            - grid_size (list): A list of three integers representing the number of grid points along each axis.
            - length (np.ndarray): The length of the bounding box along each axis.
    """
    length = bbox_max - bbox_min
    num_cells = np.exp2(resolution_base)
    x = np.linspace(bbox_min[0], bbox_max[0], int(num_cells) + 1, dtype=np.float32)
    y = np.linspace(bbox_min[1], bbox_max[1], int(num_cells) + 1, dtype=np.float32)
    z = np.linspace(bbox_min[2], bbox_max[2], int(num_cells) + 1, dtype=np.float32)
    [xs, ys, zs] = np.meshgrid(x, y, z, indexing=indexing)
    xyz = np.stack((xs, ys, zs), axis=-1)
    xyz = xyz.reshape(-1, 3)
    grid_size = [int(num_cells) + 1, int(num_cells) + 1, int(num_cells) + 1]

    return xyz, grid_size, length


def marching_cubes_with_warp(
    grid_logits: torch.Tensor,
    level: float,
    device: Union[str, torch.device] = "cuda",
    max_verts: int = 3_000_000,
    max_tris: int = 3_000_000,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Perform the marching cubes algorithm on a 3D grid with warp support.
    Args:
        grid_logits (torch.Tensor): A 3D tensor containing the grid logits.
        level (float): The threshold level for the isosurface.
        device (Union[str, torch.device], optional): The device to perform the computation on. Defaults to "cuda".
        max_verts (int, optional): The maximum number of vertices. Defaults to 3,000,000.
        max_tris (int, optional): The maximum number of triangles. Defaults to 3,000,000.
    Returns:
        Tuple[np.ndarray, np.ndarray]: A tuple containing the vertices and faces of the isosurface.
    """
    if isinstance(device, torch.device):
        device = str(device)

    assert grid_logits.ndim == 3
    if "cuda" in device:
        assert wp.is_cuda_available()
    else:
        raise ValueError(
            f"Device {device} is not supported for marching_cubes_with_warp"
        )

    dim = grid_logits.shape[0]
    field = wp.from_torch(grid_logits)

    iso = wp.MarchingCubes(
        nx=dim,
        ny=dim,
        nz=dim,
        max_verts=int(max_verts),
        max_tris=int(max_tris),
        device=device,
    )
    iso.surface(field=field, threshold=level)
    vertices = iso.verts.numpy()
    faces = iso.indices.numpy().reshape(-1, 3)
    return vertices, faces