#include <gtest/gtest.h>

#include <torch/torch.h>

#include <test/cpp/api/support.h>

#include <cmath>
#include <cstdlib>
#include <random>

using namespace torch::nn;
using namespace torch::test;

const double kPi = 3.1415926535898;

class CartPole {
  // Translated from openai/gym's cartpole.py
 public:
  double gravity = 9.8;
  double masscart = 1.0;
  double masspole = 0.1;
  double total_mass = (masspole + masscart);
  double length = 0.5; // actually half the pole's length;
  double polemass_length = (masspole * length);
  double force_mag = 10.0;
  double tau = 0.02; // seconds between state updates;

  // Angle at which to fail the episode
  double theta_threshold_radians = 12 * 2 * kPi / 360;
  double x_threshold = 2.4;
  int steps_beyond_done = -1;

  torch::Tensor state;
  double reward;
  bool done;
  int step_ = 0;

  torch::Tensor getState() {
    return state;
  }

  double getReward() {
    return reward;
  }

  double isDone() {
    return done;
  }

  void reset() {
    state = torch::empty({4}).uniform_(-0.05, 0.05);
    steps_beyond_done = -1;
    step_ = 0;
  }

  CartPole() {
    reset();
  }

  void step(int action) {
    auto x = state[0].item<float>();
    auto x_dot = state[1].item<float>();
    auto theta = state[2].item<float>();
    auto theta_dot = state[3].item<float>();

    auto force = (action == 1) ? force_mag : -force_mag;
    auto costheta = std::cos(theta);
    auto sintheta = std::sin(theta);
    auto temp = (force + polemass_length * theta_dot * theta_dot * sintheta) /
        total_mass;
    auto thetaacc = (gravity * sintheta - costheta * temp) /
        (length * (4.0 / 3.0 - masspole * costheta * costheta / total_mass));
    auto xacc = temp - polemass_length * thetaacc * costheta / total_mass;

    x = x + tau * x_dot;
    x_dot = x_dot + tau * xacc;
    theta = theta + tau * theta_dot;
    theta_dot = theta_dot + tau * thetaacc;
    state = torch::tensor({x, x_dot, theta, theta_dot});

    done = x < -x_threshold || x > x_threshold ||
        theta < -theta_threshold_radians || theta > theta_threshold_radians ||
        step_ > 200;

    if (!done) {
      reward = 1.0;
    } else if (steps_beyond_done == -1) {
      // Pole just fell!
      steps_beyond_done = 0;
      reward = 0;
    } else {
      if (steps_beyond_done == 0) {
        AT_ASSERT(false); // Can't do this
      }
    }
    step_++;
  }
};

template <typename M, typename F, typename O>
bool test_mnist(
    size_t batch_size,
    size_t number_of_epochs,
    bool with_cuda,
    M&& model,
    F&& forward_op,
    O&& optimizer) {
  std::string mnist_path = "mnist";
  if (const char* user_mnist_path = getenv("TORCH_CPP_TEST_MNIST_PATH")) {
    mnist_path = user_mnist_path;
  }

  auto train_dataset =
      torch::data::datasets::MNIST(
          mnist_path, torch::data::datasets::MNIST::Mode::kTrain)
          .map(torch::data::transforms::Stack<>());

  auto data_loader =
      torch::data::make_data_loader(std::move(train_dataset), batch_size);

  torch::Device device(with_cuda ? torch::kCUDA : torch::kCPU);
  model->to(device);

  for (size_t epoch = 0; epoch < number_of_epochs; epoch++) {
    for (torch::data::Example<> batch : *data_loader) {
      auto data = batch.data.to(device), targets = batch.target.to(device);
      torch::Tensor prediction = forward_op(std::move(data));
      torch::Tensor loss = torch::nll_loss(prediction, std::move(targets));
      AT_ASSERT(!torch::isnan(loss).any().item<int64_t>());
      optimizer.zero_grad();
      loss.backward();
      optimizer.step();
    }
  }

  torch::NoGradGuard guard;
  torch::data::datasets::MNIST test_dataset(
      mnist_path, torch::data::datasets::MNIST::Mode::kTest);
  auto images = test_dataset.images().to(device),
       targets = test_dataset.targets().to(device);

  auto result = std::get<1>(forward_op(images).max(/*dim=*/1));
  torch::Tensor correct = (result == targets).to(torch::kFloat32);
  return correct.sum().item<float>() > (test_dataset.size().value() * 0.8);
}

struct IntegrationTest : torch::test::SeedingFixture {};

TEST_F(IntegrationTest, CartPole) {
  torch::manual_seed(0);
  auto model = std::make_shared<SimpleContainer>();
  auto linear = model->add(Linear(4, 128), "linear");
  auto policyHead = model->add(Linear(128, 2), "policy");
  auto valueHead = model->add(Linear(128, 1), "action");
  auto optimizer = torch::optim::Adam(model->parameters(), 1e-3);

  std::vector<torch::Tensor> saved_log_probs;
  std::vector<torch::Tensor> saved_values;
  std::vector<float> rewards;

  auto forward = [&](torch::Tensor inp) {
    auto x = linear->forward(inp).clamp_min(0);
    torch::Tensor actions = policyHead->forward(x);
    torch::Tensor value = valueHead->forward(x);
    return std::make_tuple(torch::softmax(actions, -1), value);
  };

  auto selectAction = [&](torch::Tensor state) {
    // Only work on single state right now, change index to gather for batch
    auto out = forward(state);
    auto probs = torch::Tensor(std::get<0>(out));
    auto value = torch::Tensor(std::get<1>(out));
    auto action = probs.multinomial(1)[0].item<int32_t>();
    // Compute the log prob of a multinomial distribution.
    // This should probably be actually implemented in autogradpp...
    auto p = probs / probs.sum(-1, true);
    auto log_prob = p[action].log();
    saved_log_probs.emplace_back(log_prob);
    saved_values.push_back(value);
    return action;
  };

  auto finishEpisode = [&] {
    auto R = 0.;
    for (int i = rewards.size() - 1; i >= 0; i--) {
      R = rewards[i] + 0.99 * R;
      rewards[i] = R;
    }
    auto r_t = torch::from_blob(
        rewards.data(), {static_cast<int64_t>(rewards.size())});
    r_t = (r_t - r_t.mean()) / (r_t.std() + 1e-5);

    std::vector<torch::Tensor> policy_loss;
    std::vector<torch::Tensor> value_loss;
    for (auto i = 0U; i < saved_log_probs.size(); i++) {
      auto advantage = r_t[i] - saved_values[i].item<float>();
      policy_loss.push_back(-advantage * saved_log_probs[i]);
      value_loss.push_back(
          torch::smooth_l1_loss(saved_values[i], torch::ones(1) * r_t[i]));
    }

    auto loss =
        torch::stack(policy_loss).sum() + torch::stack(value_loss).sum();

    optimizer.zero_grad();
    loss.backward();
    optimizer.step();

    rewards.clear();
    saved_log_probs.clear();
    saved_values.clear();
  };

  auto env = CartPole();
  double running_reward = 10.0;
  for (size_t episode = 0;; episode++) {
    env.reset();
    auto state = env.getState();
    int t = 0;
    for (; t < 10000; t++) {
      auto action = selectAction(state);
      env.step(action);
      state = env.getState();
      auto reward = env.getReward();
      auto done = env.isDone();

      rewards.push_back(reward);
      if (done)
        break;
    }

    running_reward = running_reward * 0.99 + t * 0.01;
    finishEpisode();
    /*
    if (episode % 10 == 0) {
      printf("Episode %i\tLast length: %5d\tAverage length: %.2f\n",
              episode, t, running_reward);
    }
    */
    if (running_reward > 150) {
      break;
    }
    ASSERT_LT(episode, 3000);
  }
}

TEST_F(IntegrationTest, MNIST_CUDA) {
  torch::manual_seed(0);
  auto model = std::make_shared<SimpleContainer>();
  auto conv1 = model->add(Conv2d(1, 10, 5), "conv1");
  auto conv2 = model->add(Conv2d(10, 20, 5), "conv2");
  auto drop = Dropout(0.3);
  auto drop2d = Dropout2d(0.3);
  auto linear1 = model->add(Linear(320, 50), "linear1");
  auto linear2 = model->add(Linear(50, 10), "linear2");

  auto forward = [&](torch::Tensor x) {
    x = torch::max_pool2d(conv1->forward(x), {2, 2}).relu();
    x = conv2->forward(x);
    x = drop2d->forward(x);
    x = torch::max_pool2d(x, {2, 2}).relu();

    x = x.view({-1, 320});
    x = linear1->forward(x).clamp_min(0);
    x = drop->forward(x);
    x = linear2->forward(x);
    x = torch::log_softmax(x, 1);
    return x;
  };

  auto optimizer = torch::optim::SGD(
      model->parameters(), torch::optim::SGDOptions(1e-2).momentum(0.5));

  ASSERT_TRUE(test_mnist(
      32, // batch_size
      3, // number_of_epochs
      true, // with_cuda
      model,
      forward,
      optimizer));
}

TEST_F(IntegrationTest, MNISTBatchNorm_CUDA) {
  torch::manual_seed(0);
  auto model = std::make_shared<SimpleContainer>();
  auto conv1 = model->add(Conv2d(1, 10, 5), "conv1");
  auto batchnorm2d = model->add(BatchNorm2d(10), "batchnorm2d");
  auto conv2 = model->add(Conv2d(10, 20, 5), "conv2");
  auto linear1 = model->add(Linear(320, 50), "linear1");
  auto batchnorm1 = model->add(BatchNorm1d(50), "batchnorm1");
  auto linear2 = model->add(Linear(50, 10), "linear2");

  auto forward = [&](torch::Tensor x) {
    x = torch::max_pool2d(conv1->forward(x), {2, 2}).relu();
    x = batchnorm2d->forward(x);
    x = conv2->forward(x);
    x = torch::max_pool2d(x, {2, 2}).relu();

    x = x.view({-1, 320});
    x = linear1->forward(x).clamp_min(0);
    x = batchnorm1->forward(x);
    x = linear2->forward(x);
    x = torch::log_softmax(x, 1);
    return x;
  };

  auto optimizer = torch::optim::SGD(
      model->parameters(), torch::optim::SGDOptions(1e-2).momentum(0.5));

  ASSERT_TRUE(test_mnist(
      32, // batch_size
      3, // number_of_epochs
      true, // with_cuda
      model,
      forward,
      optimizer));
}