# Copyright (c) 2023 NVIDIA CORPORATION.  All rights reserved.
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

"""Helpers to author mesh geometries represented as OmniGraph bundles."""

from typing import Optional

import numpy as np
import omni.graph.core as og
import warp as wp

from omni.warp.nodes._impl.attributes import (
    attr_get,
    attr_get_array_on_gpu,
    attr_set,
)
from omni.warp.nodes._impl.bundles import (
    bundle_copy_attr_value,
    bundle_create_attr,
    bundle_create_child,
    bundle_create_metadata_attr,
    bundle_get_attr,
    bundle_set_prim_type,
    bundle_set_world_xform,
)
from omni.warp.nodes._impl.points import (
    points_get_display_color,
    points_get_points,
    points_get_velocities,
    points_get_local_extent,
    points_get_world_extent,
)


def mesh_create_bundle(
    dst_bundle: og.BundleContents,
    point_count: int,
    vertex_count: int,
    face_count: int,
    xform: Optional[np.ndarray] = None,
    create_display_color: bool = False,
    create_normals: bool = False,
    create_uvs: bool = False,
    child_idx: int = 0,
) -> None:
    """Creates and initializes mesh attributes within a bundle."""
    child_bundle = bundle_create_child(dst_bundle, child_idx)
    bundle_create_attr(
        child_bundle,
        "points",
        og.Type(
            og.BaseDataType.FLOAT,
            tuple_count=3,
            array_depth=1,
            role=og.AttributeRole.POSITION,
        ),
        size=point_count,
    )
    bundle_create_attr(
        child_bundle,
        "faceVertexCounts",
        og.Type(
            og.BaseDataType.INT,
            tuple_count=1,
            array_depth=1,
            role=og.AttributeRole.NONE,
        ),
        size=face_count,
    )
    bundle_create_attr(
        child_bundle,
        "faceVertexIndices",
        og.Type(
            og.BaseDataType.INT,
            tuple_count=1,
            array_depth=1,
            role=og.AttributeRole.NONE,
        ),
        size=vertex_count,
    )

    bundle_set_prim_type(dst_bundle, "Mesh", child_idx=child_idx)

    if xform is not None:
        bundle_set_world_xform(dst_bundle, xform, child_idx=child_idx)

    if create_display_color:
        bundle_create_attr(
            child_bundle,
            "primvars:displayColor",
            og.Type(
                og.BaseDataType.FLOAT,
                tuple_count=3,
                array_depth=1,
                role=og.AttributeRole.COLOR,
            ),
            size=point_count,
        )
        interp_attr = bundle_create_metadata_attr(
            child_bundle,
            "primvars:displayColor",
            "interpolation",
            og.Type(
                og.BaseDataType.TOKEN,
                tuple_count=1,
                array_depth=0,
                role=og.AttributeRole.NONE,
            ),
        )
        attr_set(interp_attr, "vertex")

    if create_normals:
        bundle_create_attr(
            child_bundle,
            "normals",
            og.Type(
                og.BaseDataType.FLOAT,
                tuple_count=3,
                array_depth=1,
                role=og.AttributeRole.NORMAL,
            ),
            size=vertex_count,
        )

    if create_uvs:
        bundle_create_attr(
            child_bundle,
            "primvars:st",
            og.Type(
                og.BaseDataType.FLOAT,
                tuple_count=2,
                array_depth=1,
                role=og.AttributeRole.TEXCOORD,
            ),
            size=vertex_count,
        )


def mesh_copy_bundle(
    dst_bundle: og.BundleContents,
    src_bundle: og.BundleContents,
    deep_copy: bool = False,
    child_idx: int = 0,
) -> None:
    """Creates and initializes mesh attributes from an existing bundle."""
    dst_child_bundle = bundle_create_child(dst_bundle, child_idx)
    src_child_bundle = src_bundle.bundle.get_child_bundle(child_idx)
    dst_child_bundle.copy_bundle(src_child_bundle)

    if deep_copy:
        bundle_copy_attr_value(dst_child_bundle, src_child_bundle, "points", wp.vec3)
        bundle_copy_attr_value(dst_child_bundle, src_child_bundle, "faceVertexCounts", int)
        bundle_copy_attr_value(dst_child_bundle, src_child_bundle, "faceVertexIndices", int)
        bundle_copy_attr_value(dst_child_bundle, src_child_bundle, "primvars:displayColor", wp.vec3)
        bundle_copy_attr_value(dst_child_bundle, src_child_bundle, "normals", wp.vec3)
        bundle_copy_attr_value(dst_child_bundle, src_child_bundle, "primvars:st", wp.vec2)


def mesh_get_point_count(
    bundle: og.BundleContents,
    child_idx: int = 0,
) -> int:
    """Retrieves the number of points."""
    return bundle_get_attr(bundle, "points", child_idx).size()


def mesh_get_vertex_count(
    bundle: og.BundleContents,
    child_idx: int = 0,
) -> int:
    """Retrieves the number of vertices."""
    attr = bundle_get_attr(bundle, "faceVertexCounts", child_idx)
    return int(np.sum(attr_get(attr)))


def mesh_get_face_count(
    bundle: og.BundleContents,
    child_idx: int = 0,
) -> int:
    """Retrieves the number of faces."""
    return bundle_get_attr(bundle, "faceVertexCounts", child_idx).size()


def mesh_get_points(
    bundle: og.BundleContents,
    child_idx: int = 0,
) -> wp.array(dtype=wp.vec3):
    """Retrieves the bundle points attribute as a Warp array."""
    return points_get_points(bundle, child_idx=child_idx)


def mesh_get_velocities(
    bundle: og.BundleContents,
    child_idx: int = 0,
) -> wp.array(dtype=wp.vec3):
    """Retrieves the bundle velocities attribute as a Warp array."""
    return points_get_velocities(bundle, child_idx=child_idx)


def mesh_get_normals(
    bundle: og.BundleContents,
    child_idx: int = 0,
) -> wp.array(dtype=wp.vec3):
    """Retrieves the bundle normals attribute as a Warp array."""
    attr = bundle_get_attr(bundle, "normals", child_idx)
    return attr_get_array_on_gpu(attr, wp.vec3, read_only=bundle.read_only)


def mesh_get_uvs(
    bundle: og.BundleContents,
    child_idx: int = 0,
) -> wp.array(dtype=wp.vec2):
    """Retrieves the bundle UVs attribute as a Warp array."""
    attr = bundle_get_attr(bundle, "primvars:st", child_idx)
    return attr_get_array_on_gpu(attr, wp.vec2, read_only=bundle.read_only)


def mesh_get_display_color(
    bundle: og.BundleContents,
    child_idx: int = 0,
) -> wp.array(dtype=wp.vec3):
    """Retrieves the bundle display color attribute as a Warp array."""
    return points_get_display_color(bundle, child_idx=child_idx)


def mesh_get_face_vertex_counts(
    bundle: og.BundleContents,
    child_idx: int = 0,
) -> wp.array(dtype=int):
    """Retrieves the bundle face vertex counts attribute as a Warp array."""
    attr = bundle_get_attr(bundle, "faceVertexCounts", child_idx)
    return attr_get_array_on_gpu(attr, int, read_only=bundle.read_only)


def mesh_get_face_vertex_indices(
    bundle: og.BundleContents,
    child_idx: int = 0,
) -> wp.array(dtype=int):
    """Retrieves the bundle face vertex indices attribute as a Warp array."""
    attr = bundle_get_attr(bundle, "faceVertexIndices", child_idx)
    return attr_get_array_on_gpu(attr, int, read_only=bundle.read_only)


def mesh_get_local_extent(
    bundle: og.BundleContents,
    child_idx: int = 0,
) -> np.ndarray:
    """Retrieves the local extent of the geometry mesh."""
    return points_get_local_extent(bundle, child_idx=child_idx)


def mesh_get_world_extent(
    bundle: og.BundleContents,
    axis_aligned: bool = False,
    child_idx: int = 0,
) -> np.ndarray:
    """Retrieves the world extent of the geometry mesh."""
    return points_get_world_extent(
        bundle,
        axis_aligned=axis_aligned,
        child_idx=child_idx,
    )


def mesh_triangulate(
    bundle: og.BundleContents,
    child_idx: int = 0,
) -> wp.array(dtype=int):
    """Computes a triangulated version of the face vertex indices."""
    counts = mesh_get_face_vertex_counts(bundle, child_idx=child_idx).numpy()
    if np.all(counts == 3):
        return mesh_get_face_vertex_indices(bundle, child_idx=child_idx)

    indices = mesh_get_face_vertex_indices(bundle, child_idx=child_idx).numpy()
    tri_face_count = np.sum(np.subtract(counts, 2))
    out = np.empty(tri_face_count * 3, dtype=int)

    dst_offset = 0
    src_offset = 0
    for count in counts:
        for i in range(count - 2):
            out[dst_offset] = indices[src_offset]
            out[dst_offset + 1] = indices[src_offset + i + 1]
            out[dst_offset + 2] = indices[src_offset + i + 2]
            dst_offset += 3

        src_offset += count

    return wp.array(out, dtype=int, copy=True)
