# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

from itertools import product

import torch
from fvcore.common.benchmark import benchmark
from pytorch3d.renderer.cameras import FoVPerspectiveCameras, look_at_view_transform
from pytorch3d.renderer.mesh.rasterizer import (
    Fragments,
    MeshRasterizer,
    RasterizationSettings,
)
from pytorch3d.renderer.mesh.utils import (
    _clip_barycentric_coordinates,
    _interpolate_zbuf,
)
from pytorch3d.utils.ico_sphere import ico_sphere


def baryclip_cuda(
    num_meshes: int = 8,
    ico_level: int = 5,
    image_size: int = 64,
    faces_per_pixel: int = 50,
    device="cuda",
):
    # Init meshes
    sphere_meshes = ico_sphere(ico_level, device).extend(num_meshes)
    # Init transform
    R, T = look_at_view_transform(1.0, 0.0, 0.0)
    cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
    # Init rasterizer
    raster_settings = RasterizationSettings(
        image_size=image_size,
        blur_radius=1e-4,
        faces_per_pixel=faces_per_pixel,
        clip_barycentric_coords=True,
    )
    rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings)

    torch.cuda.synchronize()

    def raster_fn():
        rasterizer(sphere_meshes)
        torch.cuda.synchronize()

    return raster_fn


def baryclip_pytorch(
    num_meshes: int = 8,
    ico_level: int = 5,
    image_size: int = 64,
    faces_per_pixel: int = 50,
    device="cuda",
):
    # Init meshes
    sphere_meshes = ico_sphere(ico_level, device).extend(num_meshes)
    # Init transform
    R, T = look_at_view_transform(1.0, 0.0, 0.0)
    cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
    # Init rasterizer
    raster_settings = RasterizationSettings(
        image_size=image_size,
        blur_radius=1e-4,
        faces_per_pixel=faces_per_pixel,
        clip_barycentric_coords=False,
    )
    rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings)

    torch.cuda.synchronize()

    def raster_fn():
        fragments = rasterizer(sphere_meshes)

        # Clip bary and reinterpolate
        clipped_bary_coords = _clip_barycentric_coordinates(fragments.bary_coords)
        clipped_zbuf = _interpolate_zbuf(
            fragments.pix_to_face, clipped_bary_coords, sphere_meshes
        )
        fragments = Fragments(
            bary_coords=clipped_bary_coords,
            zbuf=clipped_zbuf,
            dists=fragments.dists,
            pix_to_face=fragments.pix_to_face,
        )
        torch.cuda.synchronize()

    return raster_fn


def bm_barycentric_clip() -> None:
    if torch.cuda.is_available():
        kwargs_list = []
        num_meshes = [1, 8]
        ico_level = [0, 4]
        image_size = [64, 128, 256]
        faces_per_pixel = [10, 75, 100]
        test_cases = product(num_meshes, ico_level, image_size, faces_per_pixel)
        for case in test_cases:
            n, ic, im, nf = case
            kwargs_list.append(
                {
                    "num_meshes": n,
                    "ico_level": ic,
                    "image_size": im,
                    "faces_per_pixel": nf,
                }
            )

        benchmark(baryclip_cuda, "BARY_CLIP_CUDA", kwargs_list, warmup_iters=1)
        benchmark(baryclip_pytorch, "BARY_CLIP_PYTORCH", kwargs_list, warmup_iters=1)


if __name__ == "__main__":
    bm_barycentric_clip()