import torch
import math
import random


def prepare_data_for_node_classification(datalist, train_ratio=0.3, rnd_labeled_nodes=True):
    """For each graph split the nodes for training and testing. It creates a train_mask
    ehere elements equal to 1 are for training.
    rnd_labeled_nodes=True means that the nodes that are given labels for training are choosen at random (at 
    different epochs, the same graph can have different labeled nodes)."""
    for graph in datalist:
        num_nodes = graph.num_nodes
        num_classes = graph.node_y.size(1)
        nodes_per_class = {}
        for node, y in enumerate(graph.node_y.argmax(1)):
            y = y.item()
            if y not in nodes_per_class:
                nodes_per_class[y] = []
            nodes_per_class[y].append(node)
        
        train_mask = torch.zeros((num_nodes, 1))
        for y in nodes_per_class.keys():
            num_nodes_in_class = len(nodes_per_class[y])
            num_train_nodes = math.floor(num_nodes_in_class*train_ratio)
            if rnd_labeled_nodes:
                train_nodes = random.sample(nodes_per_class[y], num_train_nodes)
            else:
                train_nodes = nodes_per_class[y][:num_train_nodes]
            for node in train_nodes:
                train_mask[node] = 1
        graph.train_mask = train_mask
    return datalist