// 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.
//
// Standalone test code that runs the Minigo model on a Cloud TPU.
// Helpful for debugging issues.
//
// Example usage (you will need to supply your own values for tpu_name, model_a
// and model_b):
//   bazel build --define=tf=1 --define=tpu=1 cc:tpu_test
//   ./bazel-bin/cc/tpu_test \
//     --tpu_name=grpc://10.240.2.10:8470 \
//     --model_a=gs://tmadams-sandbox/tpu_cpp/000674-neptune.pb \
//     --model_b=gs://tmadams-sandbox/tpu_cpp/000001-bootstrap.pb

#include <algorithm>
#include <iostream>
#include <string>
#include <utility>
#include <vector>

#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "gflags/gflags.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/public/session.h"

DEFINE_string(
    tpu_name, "",
    "Cloud TPU name to run inference on, e.g. \"grpc://10.240.2.10:8470\"");
DEFINE_string(model_a, "", "Path to first model to load");
DEFINE_string(model_b, "", "Path to second model to load");

using tensorflow::DT_FLOAT;
using tensorflow::Env;
using tensorflow::GraphDef;
using tensorflow::NewSession;
using tensorflow::ReadBinaryProto;
using tensorflow::Session;
using tensorflow::SessionOptions;
using tensorflow::Tensor;
using tensorflow::TensorShape;

constexpr int kNumReplicas = 8;
constexpr int kN = 19;
constexpr int kNumMoves = kN * kN + 1;
constexpr int kNumStoneFeatures = 17;
constexpr int kNumBoardFeatures = kN * kN * kNumStoneFeatures;

using Features = std::array<float, kNumBoardFeatures>;
using Policy = std::array<float, kNumMoves>;

// Simple wrapper around the Minigo model.
class Model {
 public:
  explicit Model(const std::string& path) : path_(path) {
    // Load model.
    auto* env = Env::Default();
    GraphDef graph_def;
    TF_CHECK_OK(ReadBinaryProto(env, path, &graph_def));

    // Create a session.
    SessionOptions options;
    options.target = FLAGS_tpu_name;
    options.config.set_allow_soft_placement(true);
    options.config.set_log_device_placement(true);
    session_.reset(NewSession(options));
    TF_CHECK_OK(session_->Create(graph_def));

    // Initialize model inputs & outputs.
    for (int i = 0; i < kNumReplicas; ++i) {
      inputs_.emplace_back(
          absl::StrCat("pos_tensor_", i),
          Tensor(DT_FLOAT, TensorShape({1, kN, kN, kNumStoneFeatures})));
      output_names_.push_back(absl::StrCat("policy_output_", i));
      output_names_.push_back(absl::StrCat("value_output_", i));
    }
  }

  ~Model() {
    Log() << "Closing session" << std::endl;
    TF_CHECK_OK(session_->Close());
  }

  void InitializeTpu() {
    Log() << "Initializing TPU" << std::endl;
    TF_CHECK_OK(session_->Run({}, {}, {"ConfigureDistributedTPU"}, nullptr));
  }

  void ShutdownTpu() {
    Log() << "Shutting down TPU" << std::endl;
    TF_CHECK_OK(session_->Run({}, {}, {"ShutdownDistributedTPU"}, nullptr));
  }

  void Run(const Features& features) {
    Log() << "Running inference" << std::endl;

    // Copy features into all input tensors.
    for (int replica = 0; replica < kNumReplicas; ++replica) {
      auto* data = inputs_[replica].second.flat<float>().data();
      memcpy(data, features.data(), sizeof(features));
    }

    // Run inference.
    TF_CHECK_OK(session_->Run(inputs_, output_names_, {}, &outputs_));

    // Copy results out of the output tensors.
    Policy policy[kNumReplicas];
    float value[kNumReplicas];
    for (int replica = 0; replica < kNumReplicas; ++replica) {
      const auto* policy_data = outputs_[replica * 2].flat<float>().data();
      const auto* value_data = outputs_[replica * 2 + 1].flat<float>().data();
      memcpy(policy[replica].data(), policy_data, sizeof(Policy));
      value[replica] = *value_data;
    }

    // Check the outputs from all replicas are the same.
    for (int replica = 1; replica < kNumReplicas; ++replica) {
      for (int i = 0; i < kNumMoves; ++i) {
        if (policy[0][i] != policy[replica][i]) {
          Log() << absl::StreamFormat("policy[0][%d] == %f\n", i, policy[0][i]);
          Log() << absl::StreamFormat("policy[%d][%d] == %f\n", replica, i,
                                      policy[replica][i]);
          LOG(FATAL) << ":(";
        }
        if (value[0] != value[replica]) {
          Log() << absl::StreamFormat("value[0] == %f\n", value[0]);
          Log() << absl::StreamFormat("value[%d] == %f\n", replica,
                                      value[replica]);
          LOG(FATAL) << ":(";
        }
      }
    }

    // Log the output of the first replica (since the other replica outputs
    // all match the first).
    for (int i = 0; i < kNumMoves; ++i) {
      if (i != 0 && (i % kN) == 0) {
        std::cerr << "\n";
      }
      std::cerr << absl::StrFormat(" %0.2f", policy[0][i]);
    }
    std::cerr << "\n" << value[0] << "\n";
  }

 private:
  std::ostream& Log() { return std::cerr << "(" << path_ << ") "; }

  std::string path_;
  std::unique_ptr<Session> session_;
  std::vector<std::pair<std::string, tensorflow::Tensor>> inputs_;
  std::vector<tensorflow::Tensor> outputs_;
  std::vector<std::string> output_names_;
};

void SimpleTest() {
  // Initialize some features that represent an empty board.
  Features features;
  int i = 0;
  for (int pos = 0; pos < kN * kN; ++pos) {
    for (int f = 0; f < kNumStoneFeatures; ++f) {
      features[i++] = (f == kNumStoneFeatures - 1) ? 1 : 0;
    }
  }

  Model model_a(FLAGS_model_a);
  Model model_b(FLAGS_model_b);

  // -----------------------------------

  // This works as expected: the outputs of model_a.Run() and model_b.Run()
  // are different.
  // model_a.InitializeTpu();
  // model_a.Run(features);
  // model_a.ShutdownTpu();
  //
  // model_b.InitializeTpu();
  // model_b.Run(features);
  // model_b.ShutdownTpu();

  // -----------------------------------

  // This does not work.
  model_a.InitializeTpu();
  model_a.Run(features);

  model_b.InitializeTpu();
  model_b.Run(features);

  // This call produces the output from model_b.
  model_a.Run(features);

  model_a.ShutdownTpu();
  // Calling model_b.ShutdownTpu() here crashes because the TPU is already
  // shutdown.

  // -----------------------------------
}

int main(int argc, char* argv[]) {
  gflags::SetUsageMessage(argv[0]);
  gflags::ParseCommandLineFlags(&argc, &argv, true);

  SimpleTest();

  return 0;
}
