import numpy as np
import torch
from collections import defaultdict
from torch_scatter import scatter_add


class Data(object):
    def __init__(self, x, y, adj, train_mask, val_mask, test_mask):
        self.x = x
        self.y = y
        self.num_nodes = x.shape[0]
        self.adj = adj
        self.train_mask = train_mask
        self.val_mask = val_mask
        self.test_mask = test_mask

    def to(self, device):
        self.x = self.x.to(device)
        self.y = self.y.to(device)
        self.adj = self.adj.to(device)


# sample partial edges
def sample_edge(edge_index, ratio, seed=123):
    # sample from half side of the symmetric adjacency matrix
    half_edge_index = []
    for i in range(edge_index.shape[1]):
        if edge_index[0, i] < edge_index[1, i]:
            half_edge_index.append(edge_index[:, i].view(2,-1))
    half_edge_index = torch.cat(half_edge_index, dim=1)
    np.random.seed(seed)
    num_edge = half_edge_index.shape[1]
    samples = np.random.choice(num_edge, size=int(ratio*num_edge), replace=False)
    sampled_edge_index = half_edge_index[:, samples]
    sampled_edge_index = torch.cat([sampled_edge_index, sampled_edge_index[torch.LongTensor([1,0])]], dim=1)
    return sampled_edge_index


def random_split(data, data_seed, args):
    if data_seed == -1:
        print("fixed split")
        return 0
    if args.dataset in ["Cora", "CiteSeer", "PubMed"]:
        if not args.keep_train_num:
            num_nodes = data.train_mask.shape[0]
            train_num = torch.sum(data.train_mask)
            val_num = torch.sum(data.val_mask)
            test_num = torch.sum(data.test_mask)
            print(train_num, val_num, test_num)
            data.train_mask = torch.zeros(num_nodes).type(torch.uint8)
            data.val_mask = torch.zeros(num_nodes).type(torch.uint8)
            data.test_mask = torch.zeros(num_nodes).type(torch.uint8)
            inds = np.arange(num_nodes)
            np.random.seed(data_seed)
            np.random.shuffle(inds)
            inds = torch.tensor(inds)
            data.train_mask[inds[:train_num]] = 1
            data.val_mask[inds[train_num:train_num + val_num]] = 1
            data.test_mask[inds[train_num + val_num:train_num + val_num + test_num]] = 1
        else:
            index = data.y
            num_nodes = data.y.shape[0]
            val_num = torch.sum(data.val_mask)
            test_num = torch.sum(data.test_mask)
            input = torch.ones(num_nodes)
            out = scatter_add(input, index)
            train_node_per_class = args.train_num
            data.train_mask = torch.zeros(num_nodes).type(torch.uint8)
            data.val_mask = torch.zeros(num_nodes).type(torch.uint8)
            data.test_mask = torch.zeros(num_nodes).type(torch.uint8)
            inds = np.arange(num_nodes)
            np.random.seed(data_seed)
            np.random.shuffle(inds)
            inds = torch.tensor(inds)
            data_y_shuffle = data.y[inds]
            for i in range(out.shape[0]):
                inds_i = inds[data_y_shuffle==i]
                data.train_mask[inds_i[:train_node_per_class]] = 1
            val_test_inds = np.arange(num_nodes)[data.train_mask.numpy() == 0]
            np.random.shuffle(val_test_inds)
            val_test_inds = torch.tensor(val_test_inds)
            data.val_mask[val_test_inds[:val_num]] = 1
            data.test_mask[val_test_inds[val_num:val_num + test_num]] = 1
    elif args.dataset in ["CoraFull", "Computers", "Photo", "CS"]:
        index = data.y
        num_nodes = data.y.shape[0]
        input = torch.ones(num_nodes)
        out = scatter_add(input, index)
        train_node_per_class = args.train_num
        val_node_per_class = 30
        data.train_mask = torch.zeros(num_nodes).type(torch.uint8)
        data.val_mask = torch.zeros(num_nodes).type(torch.uint8)
        data.test_mask = torch.zeros(num_nodes).type(torch.uint8)
        inds = np.arange(num_nodes)
        np.random.seed(data_seed)
        np.random.shuffle(inds)
        inds = torch.tensor(inds)
        data_y_shuffle = data.y[inds]
        for i in range(out.shape[0]):
            if out[i] <= train_node_per_class + val_node_per_class:
                continue
            inds_i = inds[data_y_shuffle==i]
            data.train_mask[inds_i[:train_node_per_class]] = 1
            data.val_mask[inds_i[train_node_per_class:train_node_per_class+val_node_per_class]] = 1
            data.test_mask[inds_i[train_node_per_class+val_node_per_class:]] = 1


def dataloader(data, Adj, batch_size, sparse=False, shuffle=False):
    inds = np.arange(data.num_nodes)
    if shuffle:
        np.random.shuffle(inds)
    inds = torch.tensor(inds)
    batch_data_list = []
    start = 0
    while start < data.num_nodes:
        end = min(start + batch_size, data.num_nodes)
        batch_inds = inds[start:end]
        if not sparse:
            batch_data = Data(x=data.x[batch_inds], y=data.y[batch_inds],
                    adj=Adj[batch_inds][:, batch_inds],
                    train_mask=data.train_mask[batch_inds],
                    val_mask=data.val_mask[batch_inds],
                    test_mask=data.test_mask[batch_inds])
        else:
            if len(batch_inds) == data.num_nodes:
                inds_value = torch.zeros(data.num_nodes).type(torch.LongTensor)
                inds_value[batch_inds] = torch.arange(len(batch_inds))
                indices = inds_value[Adj]
                values = torch.ones(indices.shape[1])
                sample_adj = torch.sparse.FloatTensor(indices, values, [len(batch_inds), len(batch_inds)])
                batch_data = Data(x=data.x[batch_inds], y=data.y[batch_inds],
                        adj=sample_adj.coalesce(),
                        train_mask=data.train_mask[batch_inds],
                        val_mask=data.val_mask[batch_inds],
                        test_mask=data.test_mask[batch_inds])
            else:
                inds_value = (torch.zeros(data.num_nodes) - data.num_nodes).type(torch.LongTensor)
                inds_value[batch_inds] = torch.arange(len(batch_inds))
                Adj_inds = inds_value[Adj]
                indices = Adj_inds[:, torch.sum(Adj_inds, dim=0) >= 0]
                values = torch.ones(indices.shape[1])
                sample_adj = torch.sparse.FloatTensor(indices, values, [len(batch_inds), len(batch_inds)])
                batch_data = Data(x=data.x[batch_inds], y=data.y[batch_inds],
                        adj=sample_adj.coalesce(),
                        train_mask=data.train_mask[batch_inds],
                        val_mask=data.val_mask[batch_inds],
                        test_mask=data.test_mask[batch_inds])

        batch_data_list.append(batch_data)
        start += batch_size
    return batch_data_list
