from argparse import Namespace
import os.path as osp
from pathlib import Path
import numpy as np

import torch
from torch_geometric.datasets import PPI, Amazon, Planetoid, WebKB, HeterophilousGraphDataset
from torch_geometric.data import (InMemoryDataset, Data, DataLoader)
import torch_geometric.transforms as T
import torch.nn.functional as F
from torch.autograd import Variable
import pandas as pd
import random

def dataloader(config):
    device = list(range(torch.cuda.device_count()))
    device_id = f'cuda:{device[0]}' if isinstance(device, list) else device

    if config.data.data in ['Roman-empire', 'Amazon-ratings']:
        data = HeterophilousGraphDataset(root = 'data', name = config.data.data)[0]
        nlabel = config.data.nlabel
        y = F.one_hot(data.y, nlabel).float() 

        return (data.x.to(device_id), y.to(device_id), data.edge_index.to(device_id), 
                data.train_mask[:,config.seed].to(device_id), data.val_mask[:,config.seed].to(device_id), data.test_mask[:,config.seed].to(device_id))
