import numpy as np
import torch
import torch.nn.functional as F

from typing import Tuple


class SoftErosion(torch.nn.Module):
    def __init__(
        self, kernel_size: int = 15, threshold: float = 0.6, iterations: int = 1
    ):
        super(SoftErosion, self).__init__()
        r = kernel_size // 2
        self.padding = r
        self.iterations = iterations
        self.threshold = threshold

        # Create kernel
        y_indices, x_indices = torch.meshgrid(
            torch.arange(0.0, kernel_size), torch.arange(0.0, kernel_size)
        )
        dist = torch.sqrt((x_indices - r) ** 2 + (y_indices - r) ** 2)
        kernel = dist.max() - dist
        kernel /= kernel.sum()
        kernel = kernel.view(1, 1, *kernel.shape)
        self.register_buffer("weight", kernel)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        for i in range(self.iterations - 1):
            x = torch.min(
                x,
                F.conv2d(
                    x, weight=self.weight, groups=x.shape[1], padding=self.padding
                ),
            )
        x = F.conv2d(x, weight=self.weight, groups=x.shape[1], padding=self.padding)

        mask = x >= self.threshold

        x[mask] = 1.0
        # add small epsilon to avoid Nans
        x[~mask] /= (x[~mask].max() + 1e-7)

        return x, mask


def encode_segmentation_rgb(
    segmentation: np.ndarray, no_neck: bool = True
) -> np.ndarray:
    parse = segmentation
    # https://github.com/zllrunning/face-parsing.PyTorch/blob/master/prepropess_data.py
    face_part_ids = (
        [1, 2, 3, 4, 5, 6, 10, 12, 13]
        if no_neck
        else [1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 13, 14]
    )
    mouth_id = 11
    # hair_id = 17
    face_map = np.zeros([parse.shape[0], parse.shape[1]])
    mouth_map = np.zeros([parse.shape[0], parse.shape[1]])
    # hair_map = np.zeros([parse.shape[0], parse.shape[1]])

    for valid_id in face_part_ids:
        valid_index = np.where(parse == valid_id)
        face_map[valid_index] = 255
    valid_index = np.where(parse == mouth_id)
    mouth_map[valid_index] = 255
    # valid_index = np.where(parse==hair_id)
    # hair_map[valid_index] = 255
    # return np.stack([face_map, mouth_map,hair_map], axis=2)
    return np.stack([face_map, mouth_map], axis=2)


def encode_segmentation_rgb_batch(
    segmentation: torch.Tensor, no_neck: bool = True
) -> torch.Tensor:
    # https://github.com/zllrunning/face-parsing.PyTorch/blob/master/prepropess_data.py
    face_part_ids = (
        [1, 2, 3, 4, 5, 6, 10, 12, 13]
        if no_neck
        else [1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 13, 14]
    )
    mouth_id = 11
    # hair_id = 17
    segmentation = segmentation.int()
    face_map = torch.zeros_like(segmentation)
    mouth_map = torch.zeros_like(segmentation)
    # hair_map = np.zeros([parse.shape[0], parse.shape[1]])

    white_tensor = face_map + 255
    for valid_id in face_part_ids:
        face_map = torch.where(segmentation == valid_id, white_tensor, face_map)
    mouth_map = torch.where(segmentation == mouth_id, white_tensor, mouth_map)

    return torch.cat([face_map, mouth_map], dim=1)


def postprocess(
    swapped_face: np.ndarray,
    target: np.ndarray,
    target_mask: np.ndarray,
    smooth_mask: torch.nn.Module,
) -> np.ndarray:
    # target_mask = cv2.resize(target_mask, (self.size,  self.size))

    mask_tensor = (
        torch.from_numpy(target_mask.copy().transpose((2, 0, 1)))
        .float()
        .mul_(1 / 255.0)
        .cuda()
    )
    face_mask_tensor = mask_tensor[0] + mask_tensor[1]

    soft_face_mask_tensor, _ = smooth_mask(face_mask_tensor.unsqueeze_(0).unsqueeze_(0))
    soft_face_mask_tensor.squeeze_()

    soft_face_mask = soft_face_mask_tensor.cpu().numpy()
    soft_face_mask = soft_face_mask[:, :, np.newaxis]

    result = swapped_face * soft_face_mask + target * (1 - soft_face_mask)
    result = result[:, :, ::-1]  # .astype(np.uint8)
    return result