import torch


def mask_channels(mask_type, in_channels, out_channels, data_channels=3):
    """
    Creates an autoregressive channel mask.

    Input:
        mask_type: str
            Either 'A' or 'B'. 'A' for first layer of network, 'B' for all others.
        in_channels: int
            Number of input channels to layer.
        out_channels: int
            Number of output channels of layer.
        data_channels: int
            Number of channels in the input data, e.g. 3 for RGB images. (default = 3).
    Output:
        mask: torch.FloatTensor
            Shape (out_channels, in_channels).
            A mask with 0 in places for masked elements.
    """
    in_factor = in_channels // data_channels + 1
    out_factor = out_channels // data_channels + 1

    base_mask = torch.ones([data_channels,data_channels])
    if mask_type == 'A':
        base_mask = base_mask.tril(-1)
    else:
        base_mask = base_mask.tril(0)

    mask_p1 = torch.cat([base_mask]*in_factor, dim=1)
    mask_p2 = torch.cat([mask_p1]*out_factor, dim=0)

    mask = mask_p2[0:out_channels,0:in_channels]
    return mask