# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
import numpy as np
from tqdm import tqdm

from gem.utils.vis.renderer import Renderer


def simple_render_mesh(render_dict):
    """Render an camera-space mesh, blank background"""
    width, height, focal_length = render_dict["whf"]
    faces = render_dict["faces"]
    verts = render_dict["verts"]

    renderer = Renderer(width, height, focal_length, device="cuda", faces=faces)
    outputs = []
    for i in tqdm(range(len(verts)), desc="Rendering"):
        img = renderer.render_mesh(verts[i].cuda(), colors=[0.8, 0.8, 0.8])
        outputs.append(img)
    outputs = np.stack(outputs, axis=0)
    return outputs


def simple_render_mesh_background(render_dict, VI=50, colors=None):
    """Render an camera-space mesh, blank background"""
    if colors is None:
        colors = [0.8, 0.8, 0.8]
    K = render_dict["K"]
    faces = render_dict["faces"]
    verts = render_dict["verts"]
    background = render_dict["background"]
    N_frames = len(verts)
    if len(background.shape) == 3:
        background = [background] * N_frames
    height, width = background[0].shape[:2]

    renderer = Renderer(width, height, device="cuda", faces=faces, K=K)
    outputs = []
    for i in tqdm(range(len(verts)), desc="Rendering"):
        img = renderer.render_mesh(verts[i].cuda(), colors=colors, background=background[i], VI=VI)
        outputs.append(img)
    outputs = np.stack(outputs, axis=0)
    return outputs