# Path: smore/common/embedding/embed_rw.py
# class DummyJob(object):
# class EmbeddingReadOnly(object):
# class EmbeddingRW(EmbeddingReadOnly):
#     def sync(self):
#     def __init__(self, embed, gpu_id=-1):
#     def get_or_create_buf(self, name, buf_size, buf_dict, gpu_buf=True):
#     def read(self, indices, name=None):
#     def async_read(self, indices, name):
#     def __init__(self, embed, gpu_id=-1):
#     def write(self, indices, values, name=None, additive=False):
#     def async_write(self, indices, values, name, additive=False):
#     def add(self, indices, values, name="default"):


import torch
import sys
import numpy as np
import extlib_cuda as extlib
import time
import random
from smore.common.embedding.embed_rw import AsyncEmbeddingRW


def test_print():
    print(indices)
    print(embed)
    rw = AsyncEmbeddingRW(embed, buf_size=n_embed, gpu_id=gpu_id)
    submat = rw.read(indices)
    submat.job_handle.sync()
    print(submat)
    
    new_embed = torch.randn(n_embed, m).cuda(gpu_id)
    job = rw.add(indices, new_embed)
    print(new_embed)
    job.sync()
    print(embed)

    sys.exit()

if __name__ == '__main__':
    n = 5
    m = 2
    gpu_id = 0
    embed = torch.randn(n, m)
    n_embed = 3
    ids = list(range(n))
    random.shuffle(ids)
    ids = ids[:n_embed]
    indices = torch.LongTensor(ids)
    test_print()


# Path: smore/common/embedding/embed_rw.py
# class DummyJob(object):
# class EmbeddingReadOnly(object):
# class EmbeddingRW(EmbeddingReadOnly):
#     def sync(self):
#     def __init__(self, embed, gpu_id=-1):
#     def get_or_create_buf(self, name, buf_size, buf_dict, gpu_buf=True):
#     def read(self, indices, name=None):
#     def async_read(self, indices, name):
#     def __init__(self, embed, gpu_id=-1):
#     def write(self, indices, values, name=None, additive=False):
#     def async_write(self, indices, values, name, additive=False):
#     def add(self, indices, values, name="default"):


# Path: smore/common/embedding/test_async.py
import torch
import sys
import numpy as np
import extlib_cuda as extlib
import time
import random
from smore.common.embedding.embed_rw import AsyncEmbeddingRW



def test_print():
    print(indices)
    print(embed)
    rw = AsyncEmbeddingRW(embed, buf_size=n_embed, gpu_id=gpu_id)
    submat = rw.read(indices)
    submat.job_handle.sync()
    print(submat)
    
    new_embed = torch.randn(n_embed, m).cuda(gpu_id)
    job = rw.add(indices, new_embed)
    print(new_embed)
    job.sync()
    print(embed)

    sys.exit()

if __name__ == '__main__':
    n = 5
    m = 2
    gpu_id = 0
    embed = torch.randn(n, m)
    n_embed = 3
    ids = list(range(n))
    random.shuffle(ids)
    ids = ids[:n_embed]
    indices = torch.LongTensor(ids)
    test_print()


    rw = AsyncEmbeddingRW(embed, buf_size=n_embed, gpu_id=gpu_id) 