import torch
from ogb.nodeproppred import PygNodePropPredDataset
from torch_geometric import datasets
import time
from .Utils import GraphData

def load_ogb_dataset(name, data_path):
    dataset = PygNodePropPredDataset(name=name, root=data_path)
    split_idx = dataset.get_idx_split()
    g = dataset[0]
    n_node = g.num_nodes
    node_data={}
    node_data['train_mask'] = torch.zeros(n_node, dtype=torch.bool)
    node_data['val_mask'] = torch.zeros(n_node, dtype=torch.bool)
    node_data['test_mask'] = torch.zeros(n_node, dtype=torch.bool)
    node_data['train_mask'][split_idx["train"]] = True
    node_data['val_mask'][split_idx["valid"]] = True
    node_data['test_mask'][split_idx["test"]] = True
    return g, node_data

def test():
    g, node_data = load_ogb_dataset('ogbn-products', "/home/zlj/hzq/code/gnn/dataset/")
    print(g)
    # for worker in [1,2,3,4,5,6,7,8,9,10,20,30]:
    g_data = GraphData(id=1, edge_index=g.edge_index, data=g, partptr=torch.tensor([0, g.num_nodes//4, g.num_nodes//4*2, g.num_nodes//4*3, g.num_nodes]))

    row, col = g.edge_index

    # edge_weight = torch.ones(g.num_edges).float()
    # indices = [x for x in range(0, g.num_edges, 5)]
    # edge_weight[indices] = 2.0

    # g_data.eid = None
    edge_weight = None
    timestamp = None

    from .neighbor_sampler import NeighborSampler, SampleType
    from .neighbor_sampler import get_neighbors

    update_edge_row =     row
    update_edge_col =     col
    update_edge_w =  torch.DoubleTensor([i for i in range(g.num_edges)])
    # print('begin update')
    # pre = time.time()
    # # update_edge_weight(tnb, update_edge_row.contiguous(), update_edge_col.contiguous(), update_edge_w.contiguous())
    # end = time.time()
    # print("update time:", end-pre)

    print('begin tnb')
    pre = time.time()
    tnb = get_neighbors("a",
                        row.contiguous(), 
                        col.contiguous(), 
                        g.num_nodes, 0, 
                        g_data.eid, 
                        edge_weight, 
                        timestamp)
    end = time.time()
    print("init tnb time:", end-pre)
    # torch.save(tnb, "/home/zlj/hzq/code/gnn/my_sampler/MergeSample/tnb_static.my")


    # print('begin load')
    # pre = time.time()
    # tnb = torch.load("/home/zlj/hzq/code/gnn/my_sampler/MergeSample/tnb_static.my")
    # end = time.time()
    # print("load time:", end-pre)

    print('begin init')
    pre = time.time()
    sampler = NeighborSampler(g.num_nodes, 
                            tnb = tnb, 
                            num_layers=2, 
                            fanout=[100,100], 
                            graph_data=g_data, 
                            workers=10, 
                            graph_name='a',
                            policy="uniform")
    end = time.time()
    print("init time:", end-pre)

    # from torch_geometric.sampler import NeighborSampler, NumNeighbors, NodeSamplerInput, SamplerOutput
    # pre = time.time()
    # num_nei = NumNeighbors([100, 100])
    # node_idx = NodeSamplerInput(input_id=None, node=torch.tensor(range(g.num_nodes//4, g.num_nodes//4+600000)))# (input_id=None, node=torch.masked_select(torch.arange(g.num_nodes),node_data['train_mask']))
    # sampler = NeighborSampler(g, num_nei)
    # end = time.time()
    # print("init time:", end-pre)

    pre = time.time()

    out = sampler.sample_from_nodes(torch.tensor(range(g.num_nodes//4, g.num_nodes//4+600000)))# sampler.sample_from_nodes(torch.masked_select(torch.arange(g.num_nodes),node_data['train_mask']))

    # out = sampler.sample_from_nodes(node_idx)
    # node = out.node
    # edge = [out.row, out.col]

    end = time.time()
    print('node1:\t', out[0].sample_nodes())
    print('eid1:\t', out[0].eid())
    print('edge1:\t', g.edge_index[:, out[0].eid()])
    print('node2:\t', out[1].sample_nodes())
    print('eid2:\t', out[1].eid())
    print('edge2:\t', g.edge_index[:, out[1].eid()])
    print("sample time", end-pre)
    
    
if __name__ == "__main__":
    test()