import math
from typing import Tuple

import torch
import torch.nn.functional as F
from jaxtyping import Bool, Float, Integer, Int, Num
from torch import Tensor

def tri_winding(tri: Float[Tensor, "*B 3 2"]) -> Float[Tensor, "*B 3 3"]:
    # One pad for determinant
    tri_sq = F.pad(tri, (0, 1), "constant", 1.0)
    det_tri = torch.det(tri_sq)
    tri_rev = torch.cat(
        (tri_sq[..., 0:1, :], tri_sq[..., 2:3, :], tri_sq[..., 1:2, :]), -2
    )
    tri_sq[det_tri < 0] = tri_rev[det_tri < 0]
    return tri_sq

def triangle_intersection_2d(
    t1: Float[Tensor, "*B 3 2"],
    t2: Float[Tensor, "*B 3 2"],
    eps=1e-12,
) -> Float[Tensor, "*B"]:  # noqa: F821
    """Returns True if triangles collide, False otherwise"""

    def chk_edge(x: Float[Tensor, "*B 3 3"]) -> Bool[Tensor, "*B"]:  # noqa: F821
        logdetx = torch.logdet(x.double())
        if eps is None:
            return ~torch.isfinite(logdetx)
        return ~(torch.isfinite(logdetx) & (logdetx > math.log(eps)))

    t1s = tri_winding(t1)
    t2s = tri_winding(t2)

    # Assume the triangles do not collide in the begging
    ret = torch.zeros(t1.shape[0], dtype=torch.bool, device=t1.device)
    for i in range(3):
        edge = torch.roll(t1s, i, dims=1)[:, :2, :]
        # Check if all points of triangle 2 lay on the external side of edge E.
        # If this is the case the triangle do not collide
        upd = (
            chk_edge(torch.cat((edge, t2s[:, 0:1]), 1))
            & chk_edge(torch.cat((edge, t2s[:, 1:2]), 1))
            & chk_edge(torch.cat((edge, t2s[:, 2:3]), 1))
        )
        # Here no collision is still True due to inversion
        ret = ret | upd

    for i in range(3):
        edge = torch.roll(t2s, i, dims=1)[:, :2, :]

        upd = (
            chk_edge(torch.cat((edge, t1s[:, 0:1]), 1))
            & chk_edge(torch.cat((edge, t1s[:, 1:2]), 1))
            & chk_edge(torch.cat((edge, t1s[:, 2:3]), 1))
        )
        # Here no collision is still True due to inversion
        ret = ret | upd

    return ~ret  # Do the inversion

def dot(x, y, dim=-1):
    return torch.sum(x * y, dim, keepdim=True)

def compute_vertex_normal(v_pos, t_pos_idx):
    i0 = t_pos_idx[:, 0]
    i1 = t_pos_idx[:, 1]
    i2 = t_pos_idx[:, 2]
    v0 = v_pos[i0, :]
    v1 = v_pos[i1, :]
    v2 = v_pos[i2, :]
    face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1)
    # Splat face normals to vertices
    v_nrm = torch.zeros_like(v_pos)
    v_nrm.scatter_add_(0, i0[:, None].repeat(1, 3), face_normals)
    v_nrm.scatter_add_(0, i1[:, None].repeat(1, 3), face_normals)
    v_nrm.scatter_add_(0, i2[:, None].repeat(1, 3), face_normals)
    # Normalize, replace zero (degenerated) normals with some default value
    v_nrm = torch.where(
        dot(v_nrm, v_nrm) > 1e-20, v_nrm, torch.as_tensor([0.0, 0.0, 1.0]).to(v_nrm)
    )
    v_nrm = F.normalize(v_nrm, dim=1)
    if torch.is_anomaly_enabled():
        assert torch.all(torch.isfinite(v_nrm))
    return v_nrm

def _box_assign_vertex_to_cube_face(
    vertex_positions: Float[Tensor, "Nv 3"],
    vertex_normals: Float[Tensor, "Nv 3"],
    triangle_idxs: Integer[Tensor, "Nf 3"],
    bbox: Float[Tensor, "2 3"],
) -> Tuple[Float[Tensor, "Nf 3 2"], Integer[Tensor, "Nf 3"]]:
    # Test to not have a scaled model to fit the space better
    # bbox_min = bbox[:1].mean(-1, keepdim=True)
    # bbox_max = bbox[1:].mean(-1, keepdim=True)
    # v_pos_normalized = (vertex_positions - bbox_min) / (bbox_max - bbox_min)

    # Create a [0, 1] normalized vertex position
    v_pos_normalized = (vertex_positions - bbox[:1]) / (bbox[1:] - bbox[:1])
    # And to [-1, 1]
    v_pos_normalized = 2.0 * v_pos_normalized - 1.0

    # Get all vertex positions for each triangle
    # Now how do we define to which face the triangle belongs? Mean face pos? Max vertex pos?
    v0 = v_pos_normalized[triangle_idxs[:, 0]]
    v1 = v_pos_normalized[triangle_idxs[:, 1]]
    v2 = v_pos_normalized[triangle_idxs[:, 2]]
    tri_stack = torch.stack([v0, v1, v2], dim=1)

    vn0 = vertex_normals[triangle_idxs[:, 0]]
    vn1 = vertex_normals[triangle_idxs[:, 1]]
    vn2 = vertex_normals[triangle_idxs[:, 2]]
    tri_stack_nrm = torch.stack([vn0, vn1, vn2], dim=1)

    # Just average the normals per face
    face_normal = F.normalize(torch.sum(tri_stack_nrm, 1), eps=1e-6, dim=-1)

    # Now decide based on the face normal in which box map we project
    # abs_x, abs_y, abs_z = tri_stack_nrm.abs().unbind(-1)
    abs_x, abs_y, abs_z = tri_stack.abs().unbind(-1)

    axis = torch.tensor(
        [
            [1, 0, 0],  # 0
            [-1, 0, 0],  # 1
            [0, 1, 0],  # 2
            [0, -1, 0],  # 3
            [0, 0, 1],  # 4
            [0, 0, -1],  # 5
        ],
        device=face_normal.device,
        dtype=face_normal.dtype,
    )
    face_normal_axis = (face_normal[:, None] * axis[None]).sum(-1)
    index = face_normal_axis.argmax(-1)

    max_axis, uc, vc = (
        torch.ones_like(abs_x),
        torch.zeros_like(tri_stack[..., :1]),
        torch.zeros_like(tri_stack[..., :1]),
    )
    mask_pos_x = index == 0
    max_axis[mask_pos_x] = abs_x[mask_pos_x]
    uc[mask_pos_x] = tri_stack[mask_pos_x][..., 1:2]
    vc[mask_pos_x] = -tri_stack[mask_pos_x][..., -1:]

    mask_neg_x = index == 1
    max_axis[mask_neg_x] = abs_x[mask_neg_x]
    uc[mask_neg_x] = tri_stack[mask_neg_x][..., 1:2]
    vc[mask_neg_x] = -tri_stack[mask_neg_x][..., -1:]

    mask_pos_y = index == 2
    max_axis[mask_pos_y] = abs_y[mask_pos_y]
    uc[mask_pos_y] = tri_stack[mask_pos_y][..., 0:1]
    vc[mask_pos_y] = -tri_stack[mask_pos_y][..., -1:]

    mask_neg_y = index == 3
    max_axis[mask_neg_y] = abs_y[mask_neg_y]
    uc[mask_neg_y] = tri_stack[mask_neg_y][..., 0:1]
    vc[mask_neg_y] = -tri_stack[mask_neg_y][..., -1:]

    mask_pos_z = index == 4
    max_axis[mask_pos_z] = abs_z[mask_pos_z]
    uc[mask_pos_z] = tri_stack[mask_pos_z][..., 0:1]
    vc[mask_pos_z] = tri_stack[mask_pos_z][..., 1:2]

    mask_neg_z = index == 5
    max_axis[mask_neg_z] = abs_z[mask_neg_z]
    uc[mask_neg_z] = tri_stack[mask_neg_z][..., 0:1]
    vc[mask_neg_z] = -tri_stack[mask_neg_z][..., 1:2]

    # UC from [-1, 1] to [0, 1]
    max_dim_div = max_axis.max(dim=0, keepdims=True).values
    uc = ((uc[..., 0] / max_dim_div + 1.0) * 0.5).clip(0, 1)
    vc = ((vc[..., 0] / max_dim_div + 1.0) * 0.5).clip(0, 1)

    uv = torch.stack([uc, vc], dim=-1)

    return uv, index


def _assign_faces_uv_to_atlas_index(
    vertex_positions: Float[Tensor, "Nv 3"],
    triangle_idxs: Integer[Tensor, "Nf 3"],
    face_uv: Float[Tensor, "Nf 3 2"],
    face_index: Integer[Tensor, "Nf 3"],
) -> Integer[Tensor, "Nf"]:  # noqa: F821
    triangle_pos = vertex_positions[triangle_idxs]
    # We need to do perform 3 overlap checks.
    # The first set is placed in the upper two thirds of the UV atlas.
    # Conceptually, this is the direct visible surfaces from the each cube side
    # The second set is placed in the lower thirds and the left half of the UV atlas.
    # This is the first set of occluded surfaces. They will also be saved in the projected fashion
    # The third pass finds all non assigned faces. They will be placed in the bottom right half of
    # the UV atlas in scattered fashion.
    assign_idx = face_index.clone()
    for overlap_step in range(3):
        overlapping_indicator = torch.zeros_like(assign_idx, dtype=torch.bool)
        for i in range(overlap_step * 6, (overlap_step + 1) * 6):
            mask = assign_idx == i
            if not mask.any():
                continue
            # Get all elements belonging to the projection face
            uv_triangle = face_uv[mask]
            cur_triangle_pos = triangle_pos[mask]
            # Find the center of the uv coordinates
            center_uv = uv_triangle.mean(dim=1, keepdim=True)
            # And also the radius of the triangle
            uv_triangle_radius = (uv_triangle - center_uv).norm(dim=-1).max(-1).values

            potentially_overlapping_mask = (
                # Find all close triangles
                (center_uv[None, ...] - center_uv[:, None]).norm(dim=-1)
                # Do not select the same element by offseting with an large valued identity matrix
                + torch.eye(
                    uv_triangle.shape[0],
                    device=uv_triangle.device,
                    dtype=uv_triangle.dtype,
                ).unsqueeze(-1)
                * 1000
            )
            # Mark all potentially overlapping triangles to reduce the number of triangle intersection tests
            potentially_overlapping_mask = (
                potentially_overlapping_mask
                <= (uv_triangle_radius.view(-1, 1, 1) * 3.0)
            ).squeeze(-1)
            overlap_coords = torch.stack(torch.where(potentially_overlapping_mask), -1)

            # Only unique triangles (A|B and B|A should be the same)
            f = torch.min(overlap_coords, dim=-1).values
            s = torch.max(overlap_coords, dim=-1).values
            overlap_coords = torch.unique(torch.stack([f, s], dim=1), dim=0)
            first, second = overlap_coords.unbind(-1)

            # Get the triangles
            tri_1 = uv_triangle[first]
            tri_2 = uv_triangle[second]

            # Perform the actual set with the reduced number of potentially overlapping triangles
            its = triangle_intersection_2d(tri_1, tri_2, eps=1e-6)

            # So we now need to detect which triangles are the occluded ones.
            # We always assume the first to be the visible one (the others should move)
            # In the previous step we use a lexigraphical sort to get the unique pairs
            # In this we use a sort based on the orthographic projection
            ax = 0 if i < 2 else 1 if i < 4 else 2
            use_max = i % 2 == 1

            tri1_c = cur_triangle_pos[first].mean(dim=1)
            tri2_c = cur_triangle_pos[second].mean(dim=1)

            mark_first = (
                (tri1_c[..., ax] > tri2_c[..., ax])
                if use_max
                else (tri1_c[..., ax] < tri2_c[..., ax])
            )
            first[mark_first] = second[mark_first]

            # Lastly the same index can be tested multiple times.
            # If one marks it as overlapping we keep it marked as such.
            # We do this by testing if it has been marked at least once.
            unique_idx, rev_idx = torch.unique(first, return_inverse=True)

            add = torch.zeros_like(unique_idx, dtype=torch.float32)
            add.index_add_(0, rev_idx, its.float())
            its_mask = add > 0

            # And fill it in the overlapping indicator
            idx = torch.where(mask)[0][unique_idx]
            overlapping_indicator[idx] = its_mask

        # Move the index to the overlap regions (shift by 6)
        assign_idx[overlapping_indicator] += 6

    # We do not care about the correct face placement after the first 2 slices
    max_idx = 6 * 2
    return assign_idx.clamp(0, max_idx)


def _find_slice_offset_and_scale(
    index: Integer[Tensor, "Nf"],  # noqa: F821
) -> Tuple[
    Float[Tensor, "Nf"], Float[Tensor, "Nf"], Float[Tensor, "Nf"], Float[Tensor, "Nf"]  # noqa: F821
]:  # noqa: F821
    # 6 due to the 6 cube faces
    off = 1 / 3
    dupl_off = 1 / 6

    # Here, we need to decide how to pack the textures in the case of overlap
    def x_offset_calc(x, i):
        offset_calc = i // 6
        # Initial coordinates - just 3x2 grid
        if offset_calc == 0:
            return off * x
        else:
            # Smaller 3x2 grid plus eventual shift to right for
            # second overlap
            return dupl_off * x + min(offset_calc - 1, 1) * 0.5

    def y_offset_calc(x, i):
        offset_calc = i // 6
        # Initial coordinates - just a 3x2 grid
        if offset_calc == 0:
            return off * x
        else:
            # Smaller coordinates in the lowest row
            return dupl_off * x + off * 2

    offset_x = torch.zeros_like(index, dtype=torch.float32)
    offset_y = torch.zeros_like(index, dtype=torch.float32)
    offset_x_vals = [0, 1, 2, 0, 1, 2]
    offset_y_vals = [0, 0, 0, 1, 1, 1]
    for i in range(index.max().item() + 1):
        mask = index == i
        if not mask.any():
            continue
        offset_x[mask] = x_offset_calc(offset_x_vals[i % 6], i)
        offset_y[mask] = y_offset_calc(offset_y_vals[i % 6], i)

    div_x = torch.full_like(index, 6 // 2, dtype=torch.float32)
    # All overlap elements are saved in half scale
    div_x[index >= 6] = 6
    div_y = div_x.clone()  # Same for y
    # Except for the random overlaps
    div_x[index >= 12] = 2
    # But the random overlaps are saved in a large block in the lower thirds
    div_y[index >= 12] = 3

    return offset_x, offset_y, div_x, div_y


def rotation_flip_matrix_2d(
    rad: float, flip_x: bool = False, flip_y: bool = False
) -> Float[Tensor, "2 2"]:
    cos = math.cos(rad)
    sin = math.sin(rad)
    rot_mat = torch.tensor([[cos, -sin], [sin, cos]], dtype=torch.float32)
    flip_mat = torch.tensor(
        [
            [-1 if flip_x else 1, 0],
            [0, -1 if flip_y else 1],
        ],
        dtype=torch.float32,
    )

    return flip_mat @ rot_mat


def calculate_tangents(
    vertex_positions: Float[Tensor, "Nv 3"],
    vertex_normals: Float[Tensor, "Nv 3"],
    triangle_idxs: Integer[Tensor, "Nf 3"],
    face_uv: Float[Tensor, "Nf 3 2"],
) -> Float[Tensor, "Nf 3 4"]:  # noqa: F821
    vn_idx = [None] * 3
    pos = [None] * 3
    tex = face_uv.unbind(1)
    for i in range(0, 3):
        pos[i] = vertex_positions[triangle_idxs[:, i]]
        # t_nrm_idx is always the same as t_pos_idx
        vn_idx[i] = triangle_idxs[:, i]

    tangents = torch.zeros_like(vertex_normals)
    tansum = torch.zeros_like(vertex_normals)

    # Compute tangent space for each triangle
    duv1 = tex[1] - tex[0]
    duv2 = tex[2] - tex[0]
    dpos1 = pos[1] - pos[0]
    dpos2 = pos[2] - pos[0]

    tng_nom = dpos1 * duv2[..., 1:2] - dpos2 * duv1[..., 1:2]

    denom = duv1[..., 0:1] * duv2[..., 1:2] - duv1[..., 1:2] * duv2[..., 0:1]

    # Avoid division by zero for degenerated texture coordinates
    denom_safe = denom.clip(1e-6)
    tang = tng_nom / denom_safe

    # Update all 3 vertices
    for i in range(0, 3):
        idx = vn_idx[i][:, None].repeat(1, 3)
        tangents.scatter_add_(0, idx, tang)  # tangents[n_i] = tangents[n_i] + tang
        tansum.scatter_add_(
            0, idx, torch.ones_like(tang)
        )  # tansum[n_i] = tansum[n_i] + 1
    # Also normalize it. Here we do not normalize the individual triangles first so larger area
    # triangles influence the tangent space more
    tangents = tangents / tansum

    # Normalize and make sure tangent is perpendicular to normal
    tangents = F.normalize(tangents, dim=1)
    tangents = F.normalize(tangents - dot(tangents, vertex_normals) * vertex_normals)

    return tangents


def _rotate_uv_slices_consistent_space(
    vertex_positions: Float[Tensor, "Nv 3"],
    vertex_normals: Float[Tensor, "Nv 3"],
    triangle_idxs: Integer[Tensor, "Nf 3"],
    uv: Float[Tensor, "Nf 3 2"],
    index: Integer[Tensor, "Nf"],  # noqa: F821
):
    tangents = calculate_tangents(vertex_positions, vertex_normals, triangle_idxs, uv)
    pos_stack = torch.stack(
        [
            -vertex_positions[..., 1],
            vertex_positions[..., 0],
            torch.zeros_like(vertex_positions[..., 0]),
        ],
        dim=-1,
    )
    expected_tangents = F.normalize(
        torch.linalg.cross(
            vertex_normals, torch.linalg.cross(pos_stack, vertex_normals)
        ),
        -1,
    )

    actual_tangents = tangents[triangle_idxs]
    expected_tangents = expected_tangents[triangle_idxs]

    def rotation_matrix_2d(theta):
        c, s = torch.cos(theta), torch.sin(theta)
        return torch.tensor([[c, -s], [s, c]])

    # Now find the rotation
    index_mod = index % 6  # Shouldn't happen. Just for safety
    for i in range(6):
        mask = index_mod == i
        if not mask.any():
            continue

        actual_mean_tangent = actual_tangents[mask].mean(dim=(0, 1))
        expected_mean_tangent = expected_tangents[mask].mean(dim=(0, 1))

        dot_product = torch.dot(actual_mean_tangent, expected_mean_tangent)
        cross_product = (
            actual_mean_tangent[0] * expected_mean_tangent[1]
            - actual_mean_tangent[1] * expected_mean_tangent[0]
        )
        angle = torch.atan2(cross_product, dot_product)

        rot_matrix = rotation_matrix_2d(angle).to(mask.device)
        # Center the uv coordinate to be in the range of -1 to 1 and 0 centered
        uv_cur = uv[mask] * 2 - 1  # Center it first
        # Rotate it
        uv[mask] = torch.einsum("ij,nfj->nfi", rot_matrix, uv_cur)

        # Rescale uv[mask] to be within the 0-1 range
        uv[mask] = (uv[mask] - uv[mask].min()) / (uv[mask].max() - uv[mask].min())

    return uv


def _handle_slice_uvs(
    uv: Float[Tensor, "Nf 3 2"],
    index: Integer[Tensor, "Nf"],  # noqa: F821
    island_padding: float,
    max_index: int = 6 * 2,
) -> Float[Tensor, "Nf 3 2"]:  # noqa: F821
    uc, vc = uv.unbind(-1)

    # Get the second slice (The first overlap)
    index_filter = [index == i for i in range(6, max_index)]

    # Normalize them to always fully fill the atlas patch
    for i, fi in enumerate(index_filter):
        if fi.sum() > 0:
            # Scale the slice but only up to a factor of 2
            # This keeps the texture resolution with the first slice in line (Half space in UV)
            uc[fi] = (uc[fi] - uc[fi].min()) / (uc[fi].max() - uc[fi].min()).clip(0.5)
            vc[fi] = (vc[fi] - vc[fi].min()) / (vc[fi].max() - vc[fi].min()).clip(0.5)

    uc_padded = (uc * (1 - 2 * island_padding) + island_padding).clip(0, 1)
    vc_padded = (vc * (1 - 2 * island_padding) + island_padding).clip(0, 1)

    return torch.stack([uc_padded, vc_padded], dim=-1)


def _handle_remaining_uvs(
    uv: Float[Tensor, "Nf 3 2"],
    index: Integer[Tensor, "Nf"],  # noqa: F821
    island_padding: float,
) -> Float[Tensor, "Nf 3 2"]:
    uc, vc = uv.unbind(-1)
    # Get all remaining elements
    remaining_filter = index >= 6 * 2
    squares_left = remaining_filter.sum()

    if squares_left == 0:
        return uv

    uc = uc[remaining_filter]
    vc = vc[remaining_filter]

    # Or remaining triangles are distributed in a rectangle
    # The rectangle takes 0.5 of the entire uv space in width and 1/3 in height
    ratio = 0.5 * (1 / 3)  # 1.5
    # sqrt(744/(0.5*(1/3)))

    mult = math.sqrt(squares_left / ratio)
    num_square_width = int(math.ceil(0.5 * mult))
    num_square_height = int(math.ceil(squares_left / num_square_width))

    width = 1 / num_square_width
    height = 1 / num_square_height

    # The idea is again to keep the texture resolution consistent with the first slice
    # This only occupys half the region in the texture chart but the scaling on the squares
    # assumes full coverage.
    clip_val = min(width, height) * 1.5
    # Now normalize the UVs with taking into account the maximum scaling
    uc = (uc - uc.min(dim=1, keepdim=True).values) / (
        uc.amax(dim=1, keepdim=True) - uc.amin(dim=1, keepdim=True)
    ).clip(clip_val)
    vc = (vc - vc.min(dim=1, keepdim=True).values) / (
        vc.amax(dim=1, keepdim=True) - vc.amin(dim=1, keepdim=True)
    ).clip(clip_val)
    # Add a small padding
    uc = (
        uc * (1 - island_padding * num_square_width * 0.5)
        + island_padding * num_square_width * 0.25
    ).clip(0, 1)
    vc = (
        vc * (1 - island_padding * num_square_height * 0.5)
        + island_padding * num_square_height * 0.25
    ).clip(0, 1)

    uc = uc * width
    vc = vc * height

    # And calculate offsets for each element
    idx = torch.arange(uc.shape[0], device=uc.device, dtype=torch.int32)
    x_idx = idx % num_square_width
    y_idx = idx // num_square_width
    # And move each triangle to its own spot
    uc = uc + x_idx[:, None] * width
    vc = vc + y_idx[:, None] * height

    uc = (uc * (1 - 2 * island_padding * 0.5) + island_padding * 0.5).clip(0, 1)
    vc = (vc * (1 - 2 * island_padding * 0.5) + island_padding * 0.5).clip(0, 1)

    uv[remaining_filter] = torch.stack([uc, vc], dim=-1)

    return uv


def _distribute_individual_uvs_in_atlas(
    face_uv: Float[Tensor, "Nf 3 2"],
    assigned_faces: Integer[Tensor, "Nf"],  # noqa: F821
    offset_x: Float[Tensor, "Nf"],  # noqa: F821
    offset_y: Float[Tensor, "Nf"],  # noqa: F821
    div_x: Float[Tensor, "Nf"],  # noqa: F821
    div_y: Float[Tensor, "Nf"],  # noqa: F821
    island_padding: float,
):
    # Place the slice first
    placed_uv = _handle_slice_uvs(face_uv, assigned_faces, island_padding)
    # Then handle the remaining overlap elements
    placed_uv = _handle_remaining_uvs(placed_uv, assigned_faces, island_padding)

    uc, vc = placed_uv.unbind(-1)
    uc = uc / div_x[:, None] + offset_x[:, None]
    vc = vc / div_y[:, None] + offset_y[:, None]

    uv = torch.stack([uc, vc], dim=-1).view(-1, 2)

    return uv


def _get_unique_face_uv(
    uv: Float[Tensor, "Nf 3 2"],
) -> Tuple[Float[Tensor, "Utex 3"], Integer[Tensor, "Nf"]]:  # noqa: F821
    unique_uv, unique_idx = torch.unique(uv, return_inverse=True, dim=0)
    # And add the face to uv index mapping
    vtex_idx = unique_idx.view(-1, 3)

    return unique_uv, vtex_idx


def _align_mesh_with_main_axis(
    vertex_positions: Float[Tensor, "Nv 3"], vertex_normals: Float[Tensor, "Nv 3"]
) -> Tuple[Float[Tensor, "Nv 3"], Float[Tensor, "Nv 3"]]:
    # Use pca to find the 2 main axis (third is derived by cross product)
    # Set the random seed so it's repeatable
    torch.manual_seed(0)
    _, _, v = torch.pca_lowrank(vertex_positions, q=2)
    main_axis, seconday_axis = v[:, 0], v[:, 1]

    main_axis: Float[Tensor, "3"] = F.normalize(main_axis, eps=1e-6, dim=-1)
    # Orthogonalize the second axis
    seconday_axis: Float[Tensor, "3"] = F.normalize(
        seconday_axis - dot(seconday_axis, main_axis) * main_axis, eps=1e-6, dim=-1
    )
    # Create perpendicular third axis
    third_axis: Float[Tensor, "3"] = F.normalize(
        torch.cross(main_axis, seconday_axis), dim=-1, eps=1e-6
    )

    # Check to which canonical axis each aligns
    main_axis_max_idx = main_axis.abs().argmax().item()
    seconday_axis_max_idx = seconday_axis.abs().argmax().item()
    third_axis_max_idx = third_axis.abs().argmax().item()

    # Now sort the axes based on the argmax so they align with thecanonoical axes
    # If two axes have the same argmax move one of them
    all_possible_axis = {0, 1, 2}
    cur_index = 1
    while len(set([main_axis_max_idx, seconday_axis_max_idx, third_axis_max_idx])) != 3:
        # Find missing axis
        missing_axis = all_possible_axis - set(
            [main_axis_max_idx, seconday_axis_max_idx, third_axis_max_idx]
        )
        missing_axis = missing_axis.pop()
        # Just assign it to third axis as it had the smallest contribution to the
        # overall shape
        if cur_index == 1:
            third_axis_max_idx = missing_axis
        elif cur_index == 2:
            seconday_axis_max_idx = missing_axis
        else:
            raise ValueError("Could not find 3 unique axis")
        cur_index += 1

    if len({main_axis_max_idx, seconday_axis_max_idx, third_axis_max_idx}) != 3:
        raise ValueError("Could not find 3 unique axis")

    axes = [None] * 3
    axes[main_axis_max_idx] = main_axis
    axes[seconday_axis_max_idx] = seconday_axis
    axes[third_axis_max_idx] = third_axis
    # Create rotation matrix from the individual axes
    rot_mat = torch.stack(axes, dim=1).T

    # Now rotate the vertex positions and vertex normals so the mesh aligns with the main axis
    vertex_positions = torch.einsum("ij,nj->ni", rot_mat, vertex_positions)
    vertex_normals = torch.einsum("ij,nj->ni", rot_mat, vertex_normals)

    return vertex_positions, vertex_normals


def box_projection_uv_unwrap(
    vertex_positions: Float[Tensor, "Nv 3"],
    vertex_normals: Float[Tensor, "Nv 3"],
    triangle_idxs: Integer[Tensor, "Nf 3"],
    island_padding: float,
) -> Tuple[Float[Tensor, "Utex 3"], Integer[Tensor, "Nf"]]:  # noqa: F821
    # Align the mesh with main axis directions first
    # vertex_positions, vertex_normals = _align_mesh_with_main_axis(
    #     vertex_positions, vertex_normals
    # )

    bbox: Float[Tensor, "2 3"] = torch.stack(
        [vertex_positions.min(dim=0).values, vertex_positions.max(dim=0).values], dim=0
    )
    # First decide in which cube face the triangle is placed
    face_uv, face_index = _box_assign_vertex_to_cube_face(
        vertex_positions, vertex_normals, triangle_idxs, bbox
    )

    # Rotate the UV islands in a way that they align with the radial z tangent space
    face_uv = _rotate_uv_slices_consistent_space(
        vertex_positions, vertex_normals, triangle_idxs, face_uv, face_index
    )

    # Then find where where the face is placed in the atlas.
    # This has to detect potential overlaps
    assigned_atlas_index = _assign_faces_uv_to_atlas_index(
        vertex_positions, triangle_idxs, face_uv, face_index
    )

    # Then figure out the final place in the atlas based on the assignment
    offset_x, offset_y, div_x, div_y = _find_slice_offset_and_scale(
        assigned_atlas_index
    )

    # Next distribute the faces in the uv atlas
    placed_uv = _distribute_individual_uvs_in_atlas(
        face_uv, assigned_atlas_index, offset_x, offset_y, div_x, div_y, island_padding
    )

    # And get the unique per-triangle UV coordinates
    return _get_unique_face_uv(placed_uv)