#include "tde/id_transformer.h"
#include "tde/details/move_only_function.h"
namespace tde {

IDTransformer::IDTransformer(int64_t num_embedding, nlohmann::json json)
    : transformer_(num_embedding, std::move(json)),
      time_(-1),
      last_save_time_(-1) {}

c10::intrusive_ptr<TransformResult> IDTransformer::Transform(
    c10::intrusive_ptr<TensorList> global_id_list,
    c10::intrusive_ptr<TensorList> cache_id_list,
    int64_t time) {
  std::lock_guard<std::mutex> lock(mu_);
  torch::NoGradGuard no_grad;
  TORCH_CHECK(time >= 0);
  TORCH_CHECK(time >= time_, "Time cannot go backward");
  time_ = time;
  TORCH_CHECK(global_id_list->size() == cache_id_list->size());
  transformer_.strategy_.UpdateTime(static_cast<uint32_t>(time));
  {
    int64_t total_num_embeddings = std::accumulate(
        global_id_list->begin(),
        global_id_list->end(),
        int64_t(0),
        [](int64_t v, auto&& tensor) -> int64_t { return v + tensor.numel(); });
    try {
      ids_to_fetch_.resize(2 * total_num_embeddings);
    } catch (std::bad_alloc& ex) {
      TORCH_CHECK(
          false,
          "bad allocate ",
          ex.what(),
          " the total_num_embeddings=",
          total_num_embeddings);
    }
  }

  std::atomic<int64_t> next_fetch_offset{0};
  bool ok = true;
  for (int64_t i = 0; i < global_id_list->size(); ++i) {
    auto& global_ids = (*global_id_list)[i];
    auto& cache_ids = (*cache_id_list)[i];
    ok = transformer_.Transform(
        tcb::span{
            global_ids.data_ptr<int64_t>(),
            static_cast<size_t>(global_ids.numel())},
        tcb::span{
            cache_ids.data_ptr<int64_t>(),
            static_cast<size_t>(cache_ids.numel())},
        [&](int64_t global_id, int64_t cache_id) {
          int64_t offset = next_fetch_offset.fetch_add(1);
          ids_to_fetch_[2 * offset] = global_id;
          ids_to_fetch_[2 * offset + 1] = cache_id;
        });
    if (!ok) {
      break;
    }
  }

  return c10::make_intrusive<TransformResult>(
      ok,
      at::from_blob(
          ids_to_fetch_.data(),
          {next_fetch_offset.load(), 2},
          torch::TensorOptions().dtype(c10::kLong).device(c10::kCPU)));
}

torch::Tensor IDTransformer::Evict(int64_t num_to_evict) {
  std::lock_guard<std::mutex> lock(mu_);
  torch::NoGradGuard no_grad;
  std::vector<int64_t> ids_to_evict = transformer_.Evict(num_to_evict);
  int64_t num_ids_to_evict = ids_to_evict.size() / 2;
  return torch::tensor(ids_to_evict, torch::dtype(torch::kLong))
      .reshape({num_ids_to_evict, 2});
}

torch::Tensor IDTransformer::Save() {
  std::lock_guard<std::mutex> lock(mu_);
  torch::NoGradGuard no_grad;
  std::vector<int64_t> ids = transformer_.Save(last_save_time_);
  last_save_time_ = time_;
  int64_t num_ids = ids.size() / 2;
  return torch::tensor(ids, torch::dtype(torch::kLong)).reshape({num_ids, 2});
}

} // namespace tde
