# Path: smore/cpp_sampler/online_sampler.py
# def is_all_relation(query_structure):
# def has_negation(query_structure):
# def build_query_tree(query_structure, fn_qt_create):
#     def __init__(self, kg, query_names, negative_sample_size, 
#                  sample_mode, normalized_structure_prob, sampler_type='naive', 
#                  share_negative=False, same_in_batch=False,
#                  weighted_answer_sampling=False, weighted_negative_sampling=False,
#                  nprefetch=10, num_threads=8):
#     def print_queries(self):
#     def sample_entities(self, weighted, num):
#     def set_seed(self, seed):
#     def batch_generator(self, batch_size):
# def has_negation(st):
# def print_qt(qt, g, idx):
# class OnlineSampler(object):
# 
# Path: smore/cpp_sampler/sampler_clib.py
# class KGMem(object):
#     def __init__(self, dtype, mem=None):
#         self.mem = mem
#         self.dtype = dtype
# 
#     def load(self, fname):
#         fsize = os.stat(fname).st_size
#         if self.dtype == 'uint32':
#             assert fsize % 4 == 0
#             n_int = fsize // 4
#             self.mem = torch.IntTensor(n_int)
#         else:
#             assert fsize % 8 == 0
#             n_int = fsize // 8
#             self.mem = torch.LongTensor(n_int)
#         SCLIB.lib.load_binary_file(ctypes.c_char_p(fname.encode()),
#                                    ctypes.c_void_p(self.mem.data_ptr()),
#                                    ctypes.c_ulonglong(n_int),
#                                    ctypes.c_char_p(self.dtype.encode()))
# 
#     def create_kg(self):
#         assert self.mem is not None
#         kg = libsampler.KG32() if self.dtype == 'uint32' else libsampler.KG64()
#         SCLIB.lib.load_kg_from_binary(ctypes.c_void_p(kg.ptr()),
#                                       ctypes.c_void_p(self.mem.data_ptr()),
#                                       ctypes.c_ulonglong(self.mem.shape[0]),
#                                       ctypes.c_char_p(self.dtype.encode()))
#         return kg
# 
#     def share_memory(self):
#         assert self.mem is not None
#         self.mem.share_memory_()


import os
import sys
import numpy as np
import torch
import torch.multiprocessing as mp
import torch.distributed as dist
from smore.cpp_sampler.online_sampler import OnlineSampler, query_name_dict
from collections import defaultdict
from tqdm import tqdm
from smore.cpp_sampler.sampler_clib import KGMem
    query_structures = list(query_name_dict.keys())
    sampler_type = 'naive'
    negative_sample_size = 128
    sample_mode = ('sepa', 500, 'w', 'wstruct', 500000)
    for i in range(4):
        barrier.wait()
        if i == rank:
            sampler = OnlineSampler(kg, query_structures, negative_sample_size, sample_mode, [1.0 / len(query_structures)] * len(query_structures), 
                                    sampler_type=sampler_type, num_threads=2)
            print('num ent', kg.num_ent)
            print('num rel', kg.num_rel)
            print('num edges', kg.num_edges)
    data_gen = sampler.batch_generator(1024)
    for _ in tqdm(data_gen):
        pass

if __name__ == '__main__':
    kg_mem = KGMem()
    db_name = 'FB15k-237-betae'
    data_folder = os.path.join(os.path.expanduser('~'), 'data/knowledge_graphs/%s' % db_name)
    kg_mem.load(data_folder + '/train_bidir.bin')
    kg_mem.share_memory()
    mp.set_start_method('spawn')

    procs = []
    barrier = mp.Barrier(4)
    for i in range(4):
        proc = mp.Process(target=test_kg, args=(i, data_folder, kg_mem, barrier))
        proc.start()
        procs.append(proc)

# Path: smore/cpp_sampler/online_sampler.py
# def is_all_relation(query_structure):
# def has_negation(query_structure):
# def build_query_tree(query_structure, fn_qt_create):
#     def __init__(self, kg, query_names, negative_sample_size, 
#                  sample_mode, normalized_structure_prob, sampler_type='naive', 
#                  share_negative=False, same_in_batch=False,
#                  weighted_answer_sampling=False, weighted_negative_sampling=False,
#                  nprefetch=10, num_threads=8):
#     def print_queries(self):
#     def sample_entities(self, weighted, num):
#     def set_seed(self, seed):
#     def batch_generator(self, batch_size):
# def has_negation(st):
# def print_qt(qt, g, idx):
# class OnlineSampler(object):
# 
# Path: smore/cpp_sampler/sampler_clib.py
# class KGMem(object):
#     def __init__(self, dtype, mem=None):
#         self.mem = mem
#         self.dtype = dtype
# 
#     def load(self, fname):
#         fsize = os.stat(fname).st_size
#         if self.dtype == 'uint32':
#             assert fsize % 4 == 0
#             n_int = fsize // 4
#             self.mem = torch.IntTensor(n_int)
#         else:
#             assert fsize % 8 == 0
#             n_int = fsize // 8
#             self.mem = torch.LongTensor(n_int)
#         SCLIB.lib.load_binary_file(ctypes.c_char_p(fname.encode()),
#                                    ctypes.c_void_p(self.mem.data_ptr()),
#                                    ctypes.c_ulonglong(n_int),
#                                    ctypes.c_char_p(self.dtype.encode()))
# 
#     def create_kg(self):
#         assert self.mem is not None
#         kg = libsampler.KG32() if self.dtype == 'uint32' else libsampler.KG64()
#         SCLIB.lib.load_kg_from_binary(ctypes.c_void_p(kg.ptr()),
#                                       ctypes.c_void_p(self.mem.data_ptr()),
#                                       ctypes.c_ulonglong(self.mem.shape[0]),
#                                       ctypes.c_char_p(self.dtype.encode()))
#         return kg
# 
#     def share_memory(self):
#         assert self.mem is not None
#         self.mem.share_memory_()


# Path: smore/cpp_sampler/test_kg.py
import os
import sys
import numpy as np
import torch
import torch.multiprocessing as mp
import torch.distributed as dist
from smore.cpp_sampler.online_sampler import OnlineSampler, query_name_dict
from collections import defaultdict
from tqdm import tqdm
from smore.cpp_sampler.sampler_clib import KGMem

    query_structures = list(query_name_dict.keys())
    sampler_type = 'naive'
    negative_sample_size = 128
    sample_mode = ('sepa', 500, 'w', 'wstruct', 500000)
    for i in range(4):
        barrier.wait()
        if i == rank:
            sampler = OnlineSampler(kg, query_structures, negative_sample_size, sample_mode, [1.0 / len(query_structures)] * len(query_structures), 
                                    sampler_type=sampler_type, num_threads=2)
            print('num ent', kg.num_ent)
            print('num rel', kg.num_rel)
            print('num edges', kg.num_edges)
    data_gen = sampler.batch_generator(1024)
    for _ in tqdm(data_gen):
        pass

if __name__ == '__main__':
    kg_mem = KGMem()
    db_name = 'FB15k-237-betae'
    data_folder = os.path.join(os.path.expanduser('~'), 'data/knowledge_graphs/%s' % db_name)
    kg_mem.load(data_folder + '/train_bidir.bin')
    kg_mem.share_memory()
    mp.set_start_method('spawn')

    procs = []
    barrier = mp.Barrier(4)
    for i in range(4):
        proc = mp.Process(target=test_kg, args=(i, data_folder, kg_mem, barrier))
        proc.start()
        procs.append(proc)

    for proc in procs: