import cv2
import numpy as np


def merge_masks(
    masks: np.ndarray,
) -> np.ndarray:
    """Merge masks of split tiles back into a single image.
    
    Args:
        masks: A List of tiles masks in a form of numpy arrays.

    Returns:
        Merged mask in a form of numpy array.
    """
    up_left, up_right, down_left, down_right = masks

    upper = np.concatenate([up_left, up_right], axis=1)
    lower = np.concatenate([down_left, down_right], axis=1)

    mask = np.concatenate([upper, lower], axis=0)

    return mask


def combine_masks(
    mask1: np.ndarray,
    mask2: np.ndarray,
    beta1: float = 1,
    beta2: float = 1,
) -> np.ndarray:
    """Combine full image and tiled masks.

    Args:
        mask1: The whole image mask.
        mask2: The tiled mask.
        beta1: The whole image mask's weight.
        beta2: The tiled mask's weight.

    Returns:
        Combined and normalized to 0-1 mask in a form of numpy array.
    """
    mask_stack = np.stack([mask1 * beta1, mask2 * beta2])
    combined_mask = np.sum(mask_stack, axis=0)

    normalized_mask = combined_mask / (beta1 + beta2)

    return normalized_mask


def create_binary_mask(
    mask: np.ndarray,
    threshold: float = 0.5,
    close_kernel_size: int = 5,
    open_kernel_size: int = 3,
) -> np.ndarray:
    """Generate binary mask.

    The binary mask is generated by thresholding the probability mask and applying opening, then closing operations.

    Args:
        mask: The probability mask.
        threshold: The threshold to mark buildings on the binary mask.
        close_kernel_size: Size of the kernel for the closing operation.
        open_kernel_size: Size of the kernel for the opening operation.

    Returns:
        Binary mask in a form of numpy array.
    """
    binary_mask = np.where(mask >= threshold, 1, 0).astype('uint8')

    if open_kernel_size is not None:
        open_kernel = np.ones((open_kernel_size, open_kernel_size))
        binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_OPEN, open_kernel)

    if close_kernel_size is not None:
        close_kernel = np.ones((close_kernel_size, close_kernel_size))
        binary_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_CLOSE, close_kernel)

    return binary_mask
