# -*- coding: utf-8 -*-

# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2023 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: mica@tue.mpg.de


import pytorch3d
import torch
import torch.nn as nn
from pytorch3d.io import load_obj
from pytorch3d.renderer import (
    FoVPerspectiveCameras, look_at_view_transform,
    RasterizationSettings, MeshRenderer, MeshRasterizer, SoftPhongShader, TexturesVertex
)


class MeshShapeRenderer(nn.Module):
    def __init__(self, obj_filename):
        super().__init__()

        verts, faces, aux = load_obj(obj_filename)
        faces = faces.verts_idx[None, ...].cuda()
        self.register_buffer('faces', faces)

        R, T = look_at_view_transform(2.7, 10.0, 10.0)
        self.cameras = FoVPerspectiveCameras(device='cuda:0', R=R, T=T, fov=6)
        raster_settings = RasterizationSettings(
            image_size=512,
            blur_radius=0.0,
            faces_per_pixel=1,
            perspective_correct=True
        )

        lights = pytorch3d.renderer.DirectionalLights(
            device='cuda:0',
            direction=((0, 0, 1),),
            ambient_color=((0.4, 0.4, 0.4),),
            diffuse_color=((0.35, 0.35, 0.35),),
            specular_color=((0.05, 0.05, 0.05),))

        self.renderer = MeshRenderer(
            rasterizer=MeshRasterizer(cameras=self.cameras, raster_settings=raster_settings),
            shader=SoftPhongShader(device='cuda:0', cameras=self.cameras, lights=lights)
        )

    def render_mesh(self, vertices, faces=None, verts_rgb=None):
        B, N, V = vertices.shape
        if faces is None:
            faces = self.faces.repeat(B, 1, 1)
        else:
            faces = faces.repeat(B, 1, 1)

        if verts_rgb is None:
            verts_rgb = torch.ones_like(vertices)
        textures = TexturesVertex(verts_features=verts_rgb.cuda())
        meshes = pytorch3d.structures.Meshes(verts=vertices, faces=faces, textures=textures)

        rendering = self.renderer(meshes).permute(0, 3, 1, 2)
        color = rendering[:, 0:3, ...]

        return color