""" 
Based on https://pypi.org/project/pyEdgeEval/0.2.6
Cite: http://arxiv.org/abs/2304.09427
"""

from typing import List

import cv2
import numpy as np


def mask2bdry(mask: np.ndarray, ignore_mask: np.ndarray, thickness: int, quality: int = 5) -> np.ndarray:
    """
    Convert binary mask to boundaries.

    Args:
        mask (np.ndarray): 2D binary mask
        ignore_mask (np.ndarray): 2D binary mask
        thickness (int): boundary thickness
        quality (int): distance transform quality

    Returns:
        bdry (np.ndarray): 2D binary boundary mask
    """
    inner = cv2.distanceTransform(((mask + ignore_mask) > 0).astype(np.uint8), cv2.DIST_L2, quality)
    outer = cv2.distanceTransform(((1.0 - mask) > 0).astype(np.uint8), cv2.DIST_L2, quality)
    dist = outer + inner

    dist[dist > thickness] = 0
    bdry = (dist > 0).astype(np.uint8)
    return bdry


def mask2sbd(mask: np.ndarray, ignore_indices: List, thickness: int = 4, quality: int = 5) -> np.ndarray:
    """
    Convert Segmentation Mask to Semantic Boundaries.

    Args:
        mask (np.ndarray): segmentation mask
        ignore_indicies (List[int]): list of indices to ignore
        thickness (int): boundary thickness
        quality (int): distance transform quality

    Returns:
        bdrys (np.ndarray): 3D array containing boundaries
    """
    assert mask.ndim == 3
    num_labels, h, w = mask.shape

    # make ignore mask
    ignore_mask = np.zeros((h, w), dtype=np.uint8)
    for i in ignore_indices:
        ignore_mask += mask[i]

    bdrys = np.zeros_like(mask)
    for label in range(num_labels):
        m = mask[label]

        if label in ignore_indices:
            continue

        # if there are no class labels in the mask
        if not np.count_nonzero(m):
            continue

        edge = mask2bdry(
            mask=m,
            ignore_mask=ignore_mask,
            thickness=thickness,
            quality=quality,
        )

        bdrys[label] = edge

    return bdrys