import torch


def generate_train_cards_matrix(train_cards, max_count, v):
    """
    Generate a binary matrix of shape max_count * V * V, where
    row c in the i_th channel represents the there are at least 
    i cards of color c in train cards
    """
    matrix = torch.zeros(max_count, v, v)
    for i in range(max_count):
        for j in range(v):
            if j < len(train_cards) and train_cards[j] > i:
                for k in range(v):
                    matrix[i][j][k] = 1
    return matrix