import neural_renderer as nr
import numpy as np
import torch
from core import path_config
from models import SMPL


class PartRenderer():
    """Renderer used to render segmentation masks and part segmentations.
    Internally it uses the Neural 3D Mesh Renderer
    """
    def __init__(self, focal_length=5000., render_res=224):
        # Parameters for rendering
        self.focal_length = focal_length
        self.render_res = render_res
        # We use Neural 3D mesh renderer for rendering masks and part segmentations
        self.neural_renderer = nr.Renderer(
            dist_coeffs=None,
            orig_size=self.render_res,
            image_size=render_res,
            light_intensity_ambient=1,
            light_intensity_directional=0,
            anti_aliasing=False
        )
        self.faces = torch.from_numpy(SMPL(path_config.SMPL_MODEL_DIR).faces.astype(np.int32)
                                     ).cuda()
        textures = np.load(path_config.VERTEX_TEXTURE_FILE)
        self.textures = torch.from_numpy(textures).cuda().float()
        self.cube_parts = torch.cuda.FloatTensor(np.load(path_config.CUBE_PARTS_FILE))

    def get_parts(self, parts, mask):
        """Process renderer part image to get body part indices."""
        bn, c, h, w = parts.shape
        mask = mask.view(-1, 1)
        parts_index = torch.floor(100 * parts.permute(0, 2, 3, 1).contiguous().view(-1, 3)).long()
        parts = self.cube_parts[parts_index[:, 0], parts_index[:, 1], parts_index[:, 2], None]
        parts *= mask
        parts = parts.view(bn, h, w).long()
        return parts

    def __call__(self, vertices, camera):
        """Wrapper function for rendering process."""
        # Estimate camera parameters given a fixed focal length
        cam_t = torch.stack([
            camera[:, 1], camera[:, 2], 2 * self.focal_length /
            (self.render_res * camera[:, 0] + 1e-9)
        ],
                            dim=-1)
        batch_size = vertices.shape[0]
        K = torch.eye(3, device=vertices.device)
        K[0, 0] = self.focal_length
        K[1, 1] = self.focal_length
        K[2, 2] = 1
        K[0, 2] = self.render_res / 2.
        K[1, 2] = self.render_res / 2.
        K = K[None, :, :].expand(batch_size, -1, -1)
        R = torch.eye(3, device=vertices.device)[None, :, :].expand(batch_size, -1, -1)
        faces = self.faces[None, :, :].expand(batch_size, -1, -1)
        parts, _, mask = self.neural_renderer(
            vertices,
            faces,
            textures=self.textures.expand(batch_size, -1, -1, -1, -1, -1),
            K=K,
            R=R,
            t=cam_t.unsqueeze(1)
        )
        parts = self.get_parts(parts, mask)
        return mask, parts