import torch

def target_transform(labels):
    count_labels = torch.zeros(60)
    for shape_index in range(len(labels)):
        count_labels[int(shape_index * 10 + labels[shape_index])] = 1
    return count_labels

def loss_function(output, target):
    a = torch.tensor(range(10)).repeat(6 * target.shape[0]).reshape(target.shape)
    errors = a - target

    return torch.sum(output * errors ** 2) / output.shape[0]

def count_correct(output, target):
    output = output.reshape(output.shape[0], 6, 10)
    target = target.reshape(target.shape[0], 6, 10)
    target = torch.argmax(target, dim=2)
    predictions = torch.argmax(output, dim=2)
    correct = torch.min(predictions.reshape(target.shape) == target, dim=1).values
    return torch.sum(correct).tolist()

def finalizer(x):
    x = x.reshape(x.shape[0], 6, 10)
    x = torch.softmax(x, 1)
    return x.reshape(x.shape[0], 60)

outputs = 60