# Given edges and targets as .csv files and features in a .json file,
# this is a data loader that outputs labels as an array, edges in an adjacency matrix A and features as a matrix X.
# and saves all of them in a single file

import pandas as pd
import json
import numpy as np


def snap_data_loader(edge_filename, label_filename, feature_filename, filename, header_flag=1):

    if header_flag:
        edges = pd.read_csv(edge_filename)
        labels = pd.read_csv(label_filename)
    else:
        edges = pd.read_csv(edge_filename, header=None)
        labels = pd.read_csv(label_filename, header=None)

    with open(feature_filename) as f:
        data = json.load(f)
    num_features = int(max([max(data[key]) for key in data]))
    num_nodes = int(max(data, key=int))

    features = np.zeros((num_nodes+1, num_features+1))
    A = np.zeros((num_nodes+1, num_nodes+1))

    edges = edges.to_numpy()
    labels = labels.to_numpy() # populate the labels
    N = num_nodes+1
    Ne = np.size(edges,0)

    A[edges[:,0], edges[:,1]] = 1 # populate the adjacency matrix

    A=(A+A.transpose())

    for key in data: # populate the features matrix
        features[int(key),data[key]] = 1

    D = SNAPdata(A, features, labels, filename,N,Ne)
    return D


class SNAPdata():
    def __init__(self, A, X, L, filename, N, Ne):
        self.A = A
        self.X = X
        self.L = L
        self.dataset = filename
        self.N = N
        self.Ne = Ne








