from typing import List, Tuple

import numpy as np


def connected_component(r: np.ndarray, c: np.ndarray) -> List[List[int]]:
    """Find connected components in the given row and column indices.

    Args:
    ----
        r (np.ndarray): Row indices.
        c (np.ndarray): Column indices.

    Yields:
    ------
        List[int]: Indices of connected components.

    """
    indices = [0]
    for i in range(1, r.size):
        if r[i] == r[indices[-1]] and c[i] == c[indices[-1]] + 1:
            indices.append(i)
        else:
            yield indices
            indices = [i]
    yield indices


def nms_horizontal(ratio: np.ndarray, threshold: float) -> np.ndarray:
    """Apply Non-Maximum Suppression (NMS) horizontally on the given ratio matrix.

    Args:
    ----
        ratio (np.ndarray): Input ratio matrix.
        threshold (float): Threshold for NMS.

    Returns:
    -------
        np.ndarray: Binary mask after applying NMS.

    """
    mask = np.zeros_like(ratio, dtype=bool)
    r, c = np.nonzero(ratio > threshold)
    if len(r) == 0:
        return mask
    for ids in connected_component(r, c):
        values = [ratio[r[i], c[i]] for i in ids]
        mi = np.argmax(values)
        mask[r[ids[mi]], c[ids[mi]]] = True
    return mask


def nms_vertical(ratio: np.ndarray, threshold: float) -> np.ndarray:
    """Apply Non-Maximum Suppression (NMS) vertically on the given ratio matrix.

    Args:
    ----
        ratio (np.ndarray): Input ratio matrix.
        threshold (float): Threshold for NMS.

    Returns:
    -------
        np.ndarray: Binary mask after applying NMS.

    """
    return np.transpose(nms_horizontal(np.transpose(ratio), threshold))


def fgbg_depth(
    d: np.ndarray, t: float
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Find foreground-background relations between neighboring pixels.

    Args:
    ----
        d (np.ndarray): Depth matrix.
        t (float): Threshold for comparison.

    Returns:
    -------
        Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: Four matrices indicating
        left, top, right, and bottom foreground-background relations.

    """
    right_is_big_enough = (d[..., :, 1:] / d[..., :, :-1]) > t
    left_is_big_enough = (d[..., :, :-1] / d[..., :, 1:]) > t
    bottom_is_big_enough = (d[..., 1:, :] / d[..., :-1, :]) > t
    top_is_big_enough = (d[..., :-1, :] / d[..., 1:, :]) > t
    return (
        left_is_big_enough,
        top_is_big_enough,
        right_is_big_enough,
        bottom_is_big_enough,
    )


def fgbg_depth_thinned(
    d: np.ndarray, t: float
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Find foreground-background relations between neighboring pixels with Non-Maximum Suppression.

    Args:
    ----
        d (np.ndarray): Depth matrix.
        t (float): Threshold for NMS.

    Returns:
    -------
        Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: Four matrices indicating
        left, top, right, and bottom foreground-background relations with NMS applied.

    """
    right_is_big_enough = nms_horizontal(d[..., :, 1:] / d[..., :, :-1], t)
    left_is_big_enough = nms_horizontal(d[..., :, :-1] / d[..., :, 1:], t)
    bottom_is_big_enough = nms_vertical(d[..., 1:, :] / d[..., :-1, :], t)
    top_is_big_enough = nms_vertical(d[..., :-1, :] / d[..., 1:, :], t)
    return (
        left_is_big_enough,
        top_is_big_enough,
        right_is_big_enough,
        bottom_is_big_enough,
    )


def fgbg_binary_mask(
    d: np.ndarray,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Find foreground-background relations between neighboring pixels in binary masks.

    Args:
    ----
        d (np.ndarray): Binary depth matrix.

    Returns:
    -------
        Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: Four matrices indicating
        left, top, right, and bottom foreground-background relations in binary masks.

    """
    assert d.dtype == bool
    right_is_big_enough = d[..., :, 1:] & ~d[..., :, :-1]
    left_is_big_enough = d[..., :, :-1] & ~d[..., :, 1:]
    bottom_is_big_enough = d[..., 1:, :] & ~d[..., :-1, :]
    top_is_big_enough = d[..., :-1, :] & ~d[..., 1:, :]
    return (
        left_is_big_enough,
        top_is_big_enough,
        right_is_big_enough,
        bottom_is_big_enough,
    )


def edge_recall_matting(pr: np.ndarray, gt: np.ndarray, t: float) -> float:
    """Calculate edge recall for image matting.

    Args:
    ----
        pr (np.ndarray): Predicted depth matrix.
        gt (np.ndarray): Ground truth binary mask.
        t (float): Threshold for NMS.

    Returns:
    -------
        float: Edge recall value.

    """
    assert gt.dtype == bool
    ap, bp, cp, dp = fgbg_depth_thinned(pr, t)
    ag, bg, cg, dg = fgbg_binary_mask(gt)
    return 0.25 * (
        np.count_nonzero(ap & ag) / max(np.count_nonzero(ag), 1)
        + np.count_nonzero(bp & bg) / max(np.count_nonzero(bg), 1)
        + np.count_nonzero(cp & cg) / max(np.count_nonzero(cg), 1)
        + np.count_nonzero(dp & dg) / max(np.count_nonzero(dg), 1)
    )


def boundary_f1(
    pr: np.ndarray,
    gt: np.ndarray,
    t: float,
    return_p: bool = False,
    return_r: bool = False,
) -> float:
    """Calculate Boundary F1 score.

    Args:
    ----
        pr (np.ndarray): Predicted depth matrix.
        gt (np.ndarray): Ground truth depth matrix.
        t (float): Threshold for comparison.
        return_p (bool, optional): If True, return precision. Defaults to False.
        return_r (bool, optional): If True, return recall. Defaults to False.

    Returns:
    -------
        float: Boundary F1 score, or precision, or recall depending on the flags.

    """
    ap, bp, cp, dp = fgbg_depth(pr, t)
    ag, bg, cg, dg = fgbg_depth(gt, t)

    r = 0.25 * (
        np.count_nonzero(ap & ag) / max(np.count_nonzero(ag), 1)
        + np.count_nonzero(bp & bg) / max(np.count_nonzero(bg), 1)
        + np.count_nonzero(cp & cg) / max(np.count_nonzero(cg), 1)
        + np.count_nonzero(dp & dg) / max(np.count_nonzero(dg), 1)
    )
    p = 0.25 * (
        np.count_nonzero(ap & ag) / max(np.count_nonzero(ap), 1)
        + np.count_nonzero(bp & bg) / max(np.count_nonzero(bp), 1)
        + np.count_nonzero(cp & cg) / max(np.count_nonzero(cp), 1)
        + np.count_nonzero(dp & dg) / max(np.count_nonzero(dp), 1)
    )
    if r + p == 0:
        return 0.0
    if return_p:
        return p
    if return_r:
        return r
    return 2 * (r * p) / (r + p)


def get_thresholds_and_weights(
    t_min: float, t_max: float, N: int
) -> Tuple[np.ndarray, np.ndarray]:
    """Generate thresholds and weights for the given range.

    Args:
    ----
        t_min (float): Minimum threshold.
        t_max (float): Maximum threshold.
        N (int): Number of thresholds.

    Returns:
    -------
        Tuple[np.ndarray, np.ndarray]: Array of thresholds and corresponding weights.

    """
    thresholds = np.linspace(t_min, t_max, N)
    weights = thresholds / thresholds.sum()
    return thresholds, weights


def invert_depth(depth: np.ndarray, eps: float = 1e-6) -> np.ndarray:
    """Inverts a depth map with numerical stability.

    Args:
    ----
        depth (np.ndarray): Depth map to be inverted.
        eps (float): Minimum value to avoid division by zero (default is 1e-6).

    Returns:
    -------
    np.ndarray: Inverted depth map.

    """
    inverse_depth = 1.0 / depth.clip(min=eps)
    return inverse_depth


def SI_boundary_F1(
    predicted_depth: np.ndarray,
    target_depth: np.ndarray,
    t_min: float = 1.05,
    t_max: float = 1.25,
    N: int = 10,
) -> float:
    """Calculate Scale-Invariant Boundary F1 Score for depth-based ground-truth.

    Args:
    ----
        predicted_depth (np.ndarray): Predicted depth matrix.
        target_depth (np.ndarray): Ground truth depth matrix.
        t_min (float, optional): Minimum threshold. Defaults to 1.05.
        t_max (float, optional): Maximum threshold. Defaults to 1.25.
        N (int, optional): Number of thresholds. Defaults to 10.

    Returns:
    -------
        float: Scale-Invariant Boundary F1 Score.

    """
    assert predicted_depth.ndim == target_depth.ndim == 2
    thresholds, weights = get_thresholds_and_weights(t_min, t_max, N)
    f1_scores = np.array(
        [
            boundary_f1(invert_depth(predicted_depth), invert_depth(target_depth), t)
            for t in thresholds
        ]
    )
    return np.sum(f1_scores * weights)


def SI_boundary_Recall(
    predicted_depth: np.ndarray,
    target_mask: np.ndarray,
    t_min: float = 1.05,
    t_max: float = 1.25,
    N: int = 10,
    alpha_threshold: float = 0.1,
) -> float:
    """Calculate Scale-Invariant Boundary Recall Score for mask-based ground-truth.

    Args:
    ----
        predicted_depth (np.ndarray): Predicted depth matrix.
        target_mask (np.ndarray): Ground truth binary mask.
        t_min (float, optional): Minimum threshold. Defaults to 1.05.
        t_max (float, optional): Maximum threshold. Defaults to 1.25.
        N (int, optional): Number of thresholds. Defaults to 10.
        alpha_threshold (float, optional): Threshold for alpha masking. Defaults to 0.1.

    Returns:
    -------
        float: Scale-Invariant Boundary Recall Score.

    """
    assert predicted_depth.ndim == target_mask.ndim == 2
    thresholds, weights = get_thresholds_and_weights(t_min, t_max, N)
    thresholded_target = target_mask > alpha_threshold

    recall_scores = np.array(
        [
            edge_recall_matting(
                invert_depth(predicted_depth), thresholded_target, t=float(t)
            )
            for t in thresholds
        ]
    )
    weighted_recall = np.sum(recall_scores * weights)
    return weighted_recall