from typing import Union
import torch
import random


def _generate_binary_cut_mix_map(height: int, width: int,
                                 device: Union[str, torch.device] = "cpu") -> torch.Tensor:
    """
    This function generate a random binary map with two areas for cut mix augmentation of consistency regularization.
    :param height: (int) Height of the map
    :param width: (int) Width of the map
    :param device: (Union[str, torch.device]) Device to utilize
    :return: (torch.Tensor) Binary map
    """
    # Make target
    binary_map = torch.zeros(1, 1, 1, height, width, dtype=torch.float, device=device)
    # Generate a random coordinates to perform cut
    cut_coordinates_height = torch.randint(int(0.1 * height), int(0.9 * height), size=(1,))
    cut_coordinates_width = torch.randint(int(0.1 * width), int(0.9 * width), size=(1,))
    # Apply coordinates
    if random.random() > 0.5:
        binary_map[..., cut_coordinates_height:, cut_coordinates_width:] = 1.0
    else:
        binary_map[..., :cut_coordinates_height, :cut_coordinates_width] = 1.0
    # Invert target randomly
    if random.random() > 0.5:
        binary_map = - binary_map + 1.
    return binary_map