// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "cc/model/batching_model.h"

#include <utility>

#include "absl/memory/memory.h"
#include "absl/time/clock.h"
#include "cc/logging.h"
#include "cc/model/buffered_model.h"
#include "cc/model/loader.h"
#include "wtf/macros.h"

namespace minigo {

namespace internal {

ModelBatcher::ModelBatcher(std::unique_ptr<Model> model_impl, int buffer_count)
    : model_impl_(std::move(model_impl)), stats_(buffer_count) {}

ModelBatcher::~ModelBatcher() {
  MG_LOG(INFO) << "Ran " << num_batches_ << " batches with an average size of "
               << static_cast<float>(num_inferences_) / num_batches_;
}

void ModelBatcher::StartGame() {
  absl::MutexLock lock(&mutex_);
  num_active_clients_ += 1;
}

void ModelBatcher::EndGame() {
  absl::MutexLock lock(&mutex_);
  num_active_clients_ -= 1;
  MaybeRunBatchesLocked();
}

void ModelBatcher::RunMany(ModelBatcher* other_batcher,
                           const std::vector<const ModelInput*>& inputs,
                           std::vector<ModelOutput*>* outputs,
                           std::string* model_name) {
  WTF_SCOPE("ModelBatcher::RunMany", size_t)(inputs.size());

  absl::Notification notification;

  {
    absl::MutexLock lock(&mutex_);
    queue_.push({other_batcher, &inputs, outputs, model_name, &notification});
    if (other_batcher != nullptr) {
      other_batcher->num_waiting_ += 1;
    }
    MaybeRunBatchesLocked();
  }

  if (other_batcher != nullptr) {
    absl::MutexLock lock(&other_batcher->mutex_);
    other_batcher->MaybeRunBatchesLocked();
  }

  notification.WaitForNotification();
}

BatchingModelStats ModelBatcher::FlushStats() {
  mutex_.Lock();
  auto result = stats_;
  stats_ = BatchingModelStats(stats_.buffer_count);
  mutex_.Unlock();
  return result;
}

size_t ModelBatcher::GetBatchSize() const {
  return std::max<size_t>(1, num_active_clients_ / stats_.buffer_count);
}

void ModelBatcher::MaybeRunBatchesLocked() {
  while (!queue_.empty()) {
    auto queue_size = queue_.size();
    if (queue_size < GetBatchSize()) {
      // The queue doesn't have enough requests to fill a batch: see if we
      // can run a smaller batch instead.
      //
      // We run a small batch if all clients of this model have either
      // submitted inference requests, or are in a two player game and
      // waiting for the other player's inference.
      //
      // Additionally... when starting a bunch of games in parallel, we
      // will initially submit several smaller batches until all the
      // clients have been created. This has a ripple effect across all
      // subsequent reads, making the batching irregular. To counteract
      // this, we additionally enforce the constraint that a small batch
      // can't be run until at least half of the clients have submitted
      // inference requests. This has the effect of forcing those clients
      // to run their batches in lock-step.
      bool can_run_small_batch =
          queue_size >= num_active_clients_ / stats_.buffer_count &&
          queue_size + num_waiting_ >= num_active_clients_;
      if (!can_run_small_batch) {
        break;
      }
    }

    RunBatch();
  }
}

void ModelBatcher::RunBatch() {
  WTF_SCOPE0("ModelBatcher::RunBatch");
  auto run_batch_start_time = absl::Now();

  auto batch_size = GetBatchSize();

  // TODO(tommadams): reserve GetBatchSize() * virtual_losses elements.
  std::vector<const ModelInput*> inputs;
  std::vector<ModelOutput*> outputs;
  std::vector<InferenceRequest> inferences;

  while (!queue_.empty() && inferences.size() < batch_size) {
    auto& inference = queue_.front();
    size_t num_features = inference.inputs->size();

    std::copy_n(inference.inputs->begin(), num_features,
                std::back_inserter(inputs));
    std::copy_n(inference.outputs->begin(), num_features,
                std::back_inserter(outputs));
    inferences.push_back(inference);

    queue_.pop();
  }

  num_batches_ += 1;
  num_inferences_ += inputs.size();
  auto num_inferences_in_batch = inputs.size();

  // Unlock the mutex while running inference. This allows more inferences
  // to be enqueued while inference is running.
  mutex_.Unlock();

  std::string model_name;
  auto run_many_start_time = absl::Now();

  MG_CHECK(inputs.size() == outputs.size());
  model_impl_->RunMany(inputs, &outputs, &model_name);
  auto run_many_time = absl::Now() - run_many_start_time;

  for (auto& inference : inferences) {
    if (inference.model_name != nullptr) {
      *inference.model_name = model_name;
    }
    // For all two player games, tell the batcher of the opponent model that
    // it isn't blocked on this inference any more.
    if (inference.other_batcher != nullptr) {
      inference.other_batcher->num_waiting_ -= 1;
    }
  }

  // All the required work is done, unblock all the waiting clients.
  for (auto& inference : inferences) {
    inference.notification->Notify();
  }

  // Lock the mutex again.
  mutex_.Lock();

  stats_.run_batch_time +=
      (absl::Now() - run_batch_start_time) / stats_.buffer_count;
  stats_.run_many_time += (run_many_time) / stats_.buffer_count;
  stats_.num_inferences += num_inferences_in_batch;
}

}  // namespace internal

BatchingModel::BatchingModel(std::shared_ptr<internal::ModelBatcher> batcher)
    : Model(batcher->name(), batcher->feature_descriptor()),
      batcher_(std::move(batcher)) {}

void BatchingModel::RunMany(const std::vector<const ModelInput*>& inputs,
                            std::vector<ModelOutput*>* outputs,
                            std::string* model) {
  batcher_->RunMany(other_batcher_.get(), inputs, outputs, model);
}

void BatchingModel::StartGame() { batcher_->StartGame(); }

void BatchingModel::EndGame() { batcher_->EndGame(); }

void BatchingModel::SetOther(BatchingModel* other) {
  if (other == nullptr) {
    MG_CHECK(other_batcher_ != nullptr);
    other_batcher_ = nullptr;
  } else {
    MG_CHECK(other_batcher_ == nullptr);
    other_batcher_ = other->batcher_;
  }
}

BatchingModelFactory::BatchingModelFactory(std::string device, int buffer_count)
    : device_(std::move(device)), buffer_count_(buffer_count) {}

std::unique_ptr<BatchingModel> BatchingModelFactory::NewModel(
    const std::string& path) {
  absl::MutexLock lock(&mutex_);

  // Find or create a service for the requested model.
  auto it = batchers_.find(path);
  if (it == batchers_.end()) {
    std::vector<std::unique_ptr<Model>> models;
    auto def = LoadModelDefinition(path);
    auto* factory = GetModelFactory(def, device_);
    for (int i = 0; i < buffer_count_; ++i) {
      models.push_back(factory->NewModel(def));
    }
    auto batcher = std::make_shared<internal::ModelBatcher>(
        absl::make_unique<BufferedModel>(std::move(models)), buffer_count_);
    it = batchers_.emplace(path, std::move(batcher)).first;
  }

  auto model = absl::make_unique<BatchingModel>(it->second);

  // Take this opportunity to prune any services that have no clients.
  it = batchers_.begin();
  while (it != batchers_.end()) {
    // If the factory is the only one left with a reference to the batcher,
    // delete it.
    if (it->second.use_count() == 1) {
      batchers_.erase(it++);
    } else {
      ++it;
    }
  }

  return model;
}

void BatchingModelFactory::StartGame(Model* black, Model* white) {
  // TODO(tommadams): figure out if we can refactor the code somehow to take
  // BatchingModel pointers and avoid these dynamic_casts.
  auto* b = dynamic_cast<BatchingModel*>(black);
  auto* w = dynamic_cast<BatchingModel*>(white);
  MG_CHECK(b != nullptr && w != nullptr);

  if (b != w) {
    // This is a two player game, inform each client who the other one is.
    b->SetOther(w);
    w->SetOther(b);
  }

  b->StartGame();
  if (b != w) {
    w->StartGame();
  }
}

void BatchingModelFactory::EndGame(Model* black, Model* white) {
  // TODO(tommadams): figure out if we can refactor the code somehow to take
  // BatchingModel pointers and avoid these dynamic_casts.
  auto* b = dynamic_cast<BatchingModel*>(black);
  auto* w = dynamic_cast<BatchingModel*>(white);
  MG_CHECK(b != nullptr && w != nullptr);

  if (b != w) {
    b->SetOther(nullptr);
    w->SetOther(nullptr);
  }

  b->EndGame();
  if (b != w) {
    w->EndGame();
  }
}

std::vector<std::pair<std::string, BatchingModelStats>>
BatchingModelFactory::FlushStats() {
  absl::MutexLock lock(&mutex_);
  std::vector<std::pair<std::string, BatchingModelStats>> result;
  for (const auto& kv : batchers_) {
    result.emplace_back(kv.first, kv.second->FlushStats());
  }
  return result;
}

}  // namespace minigo
