# Copyright (c) Facebook, Inc. and its affiliates.
import numpy as np
from typing import List, Optional, Tuple
import torch

from detectron2.data.detection_utils import read_image

from ..structures import DensePoseChartResult
from .base import Boxes, Image
from .densepose_results import DensePoseResultsVisualizer


def get_texture_atlas(path: Optional[str]) -> Optional[np.ndarray]:
    if path is None:
        return None

    # Reading images like that downsamples 16-bit images to 8-bit
    # If 16-bit images are needed, we can replace that by cv2.imread with the
    # cv2.IMREAD_UNCHANGED flag (with cv2 we also need it to keep alpha channels)
    # The rest of the pipeline would need to be adapted to 16-bit images too
    bgr_image = read_image(path)
    rgb_image = np.copy(bgr_image)  # Convert BGR -> RGB
    rgb_image[:, :, :3] = rgb_image[:, :, 2::-1]  # Works with alpha channel
    return rgb_image


class DensePoseResultsVisualizerWithTexture(DensePoseResultsVisualizer):
    """
    texture_atlas: An image, size 6N * 4N, with N * N squares for each of the 24 body parts.
            It must follow the grid found at https://github.com/facebookresearch/DensePose/blob/master/DensePoseData/demo_data/texture_atlas_200.png  # noqa
            For each body part, U is proportional to the x coordinate, and (1 - V) to y
    """

    def __init__(self, texture_atlas, **kwargs):
        self.texture_atlas = texture_atlas
        self.body_part_size = texture_atlas.shape[0] // 6
        assert self.body_part_size == texture_atlas.shape[1] // 4

    def visualize(
        self,
        image_bgr: Image,
        results_and_boxes_xywh: Tuple[Optional[List[DensePoseChartResult]], Optional[Boxes]],
    ) -> Image:
        densepose_result, boxes_xywh = results_and_boxes_xywh
        if densepose_result is None or boxes_xywh is None:
            return image_bgr

        boxes_xywh = boxes_xywh.int().cpu().numpy()
        texture_image, alpha = self.get_texture()
        for i, result in enumerate(densepose_result):
            iuv_array = torch.cat((result.labels[None], result.uv.clamp(0, 1)))
            x, y, w, h = boxes_xywh[i]
            bbox_image = image_bgr[y : y + h, x : x + w]
            image_bgr[y : y + h, x : x + w] = self.generate_image_with_texture(
                texture_image, alpha, bbox_image, iuv_array.cpu().numpy()
            )
        return image_bgr

    def get_texture(self):
        N = self.body_part_size
        texture_image = np.zeros([24, N, N, self.texture_atlas.shape[-1]])
        for i in range(4):
            for j in range(6):
                texture_image[(6 * i + j), :, :, :] = self.texture_atlas[
                    N * j : N * (j + 1), N * i : N * (i + 1), :
                ]

        if texture_image.shape[-1] == 4:  # Image with alpha channel
            alpha = texture_image[:, :, :, -1] / 255.0
            texture_image = texture_image[:, :, :, :3]
        else:
            alpha = texture_image.sum(axis=-1) > 0

        return texture_image, alpha

    def generate_image_with_texture(self, texture_image, alpha, bbox_image_bgr, iuv_array):

        I, U, V = iuv_array
        generated_image_bgr = bbox_image_bgr.copy()

        for PartInd in range(1, 25):
            x, y = np.where(I == PartInd)
            x_index = (U[x, y] * (self.body_part_size - 1)).astype(int)
            y_index = ((1 - V[x, y]) * (self.body_part_size - 1)).astype(int)
            part_alpha = np.expand_dims(alpha[PartInd - 1, y_index, x_index], -1)
            generated_image_bgr[I == PartInd] = (
                generated_image_bgr[I == PartInd] * (1 - part_alpha)
                + texture_image[PartInd - 1, y_index, x_index] * part_alpha
            )

        return generated_image_bgr.astype(np.uint8)