#include "catch_utils.hpp"

#include <torch/nn/module.h>
#include <torch/nn/modules/batchnorm.h>
#include <torch/nn/modules/conv.h>
#include <torch/nn/modules/dropout.h>
#include <torch/nn/modules/embedding.h>
#include <torch/nn/modules/functional.h>
#include <torch/nn/modules/linear.h>
#include <torch/tensor.h>
#include <torch/utils.h>

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

using Catch::StartsWith;

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

class TestModel : public torch::nn::Module {
 public:
  TestModel()
      : l1(register_module("l1", Linear(10, 3))),
        l2(register_module("l2", Linear(3, 5))),
        l3(register_module("l3", Linear(5, 100))) {}

  Linear l1, l2, l3;
};

class NestedModel : public torch::nn::Module {
 public:
  NestedModel()
      : param_(register_parameter("param", torch::empty({3, 2, 21}))),
        l1(register_module("l1", Linear(5, 20))),
        t(register_module("test", std::make_shared<TestModel>())) {}

  torch::Tensor param_;
  Linear l1;
  std::shared_ptr<TestModel> t;
};

CATCH_TEST_CASE("modules") {
  torch::manual_seed(0);
  CATCH_SECTION("conv") {
    CATCH_SECTION("1d") {
      Conv1d model(Conv1dOptions(3, 2, 3).stride(2));
      auto x = torch::randn({2, 3, 5}, torch::requires_grad());
      auto y = model->forward(x);
      torch::Tensor s = y.sum();

      s.backward();
      CATCH_REQUIRE(y.ndimension() == 3);
      CATCH_REQUIRE(s.ndimension() == 0);
      for (auto i = 0; i < 3; i++) {
        CATCH_REQUIRE(y.size(i) == 2);
      }

      CATCH_REQUIRE(model->parameters()["weight"].grad().numel() == 3 * 2 * 3);
    }
    CATCH_SECTION("2d") {
      CATCH_SECTION("even") {
        Conv2d model(Conv2dOptions(3, 2, 3).stride(2));
        auto x = torch::randn({2, 3, 5, 5}, torch::requires_grad());
        auto y = model->forward(x);
        torch::Tensor s = y.sum();

        s.backward();
        CATCH_REQUIRE(y.ndimension() == 4);
        CATCH_REQUIRE(s.ndimension() == 0);
        for (auto i = 0; i < 4; i++) {
          CATCH_REQUIRE(y.size(i) == 2);
        }

        CATCH_REQUIRE(model->parameters()["weight"].grad().numel() == 3 * 2 * 3 * 3);
      }

      CATCH_SECTION("uneven") {
        Conv2d model(Conv2dOptions(3, 2, {3, 2}).stride({2, 2}));
        auto x = torch::randn({2, 3, 5, 4}, torch::requires_grad());
        auto y = model->forward(x);
        torch::Tensor s = y.sum();

        s.backward();
        CATCH_REQUIRE(y.ndimension() == 4);
        CATCH_REQUIRE(s.ndimension() == 0);
        for (auto i = 0; i < 4; i++) {
          CATCH_REQUIRE(y.size(i) == 2);
        }

        CATCH_REQUIRE(model->parameters()["weight"].grad().numel() == 3 * 2 * 3 * 2);
      }
    }
    CATCH_SECTION("3d") {
      Conv3d model(Conv3dOptions(3, 2, 3).stride(2));
      auto x = torch::randn({2, 3, 5, 5, 5}, torch::requires_grad());
      auto y = model->forward(x);
      torch::Tensor s = y.sum();

      s.backward();
      CATCH_REQUIRE(y.ndimension() == 5);
      CATCH_REQUIRE(s.ndimension() == 0);
      for (auto i = 0; i < 5; i++) {
        CATCH_REQUIRE(y.size(i) == 2);
      }

      CATCH_REQUIRE(
          model->parameters()["weight"].grad().numel() == 3 * 2 * 3 * 3 * 3);
    }
  }
  CATCH_SECTION("linear") {
    CATCH_SECTION("basic1") {
      Linear model(5, 2);
      auto x = torch::randn({10, 5}, torch::requires_grad());
      auto y = model->forward(x);
      torch::Tensor s = y.sum();

      s.backward();
      CATCH_REQUIRE(y.ndimension() == 2);
      CATCH_REQUIRE(s.ndimension() == 0);
      CATCH_REQUIRE(y.size(0) == 10);
      CATCH_REQUIRE(y.size(1) == 2);

      CATCH_REQUIRE(model->parameters()["weight"].grad().numel() == 2 * 5);
    }
  }

  CATCH_SECTION("simple") {
    auto model = std::make_shared<SimpleContainer>();
    auto l1 = model->add(Linear(10, 3), "l1");
    auto l2 = model->add(Linear(3, 5), "l2");
    auto l3 = model->add(Linear(5, 100), "l3");

    auto x = torch::randn({1000, 10}, torch::requires_grad());
    x = l1->forward(x).clamp_min(0);
    x = l2->forward(x).clamp_min(0);
    x = l3->forward(x).clamp_min(0);

    x.backward();
    CATCH_REQUIRE(x.ndimension() == 2);
    CATCH_REQUIRE(x.size(0) == 1000);
    CATCH_REQUIRE(x.size(1) == 100);
    CATCH_REQUIRE(x.min().toCFloat() == 0);
  }

  CATCH_SECTION("embedding") {
    CATCH_SECTION("basic") {
      const int64_t dict_size = 10;
      Embedding model(dict_size, 2);
      CATCH_REQUIRE(model->parameters().contains("weight"));
      CATCH_REQUIRE(model->weight.ndimension() == 2);
      CATCH_REQUIRE(model->weight.size(0) == dict_size);
      CATCH_REQUIRE(model->weight.size(1) == 2);

      // Cannot get gradients to change indices (input) - only for embedding
      // params
      auto x = torch::full({10}, dict_size - 1, torch::kInt64);
      auto y = model->forward(x);
      torch::Tensor s = y.sum();

      s.backward();
      CATCH_REQUIRE(y.ndimension() == 2);
      CATCH_REQUIRE(s.ndimension() == 0);
      CATCH_REQUIRE(y.size(0) == 10);
      CATCH_REQUIRE(y.size(1) == 2);

      CATCH_REQUIRE(model->parameters()["weight"].grad().numel() == 2 * dict_size);
    }

    CATCH_SECTION("list") {
      Embedding model(6, 4);
      auto x = torch::full({2, 3}, 5, torch::kInt64);
      auto y = model->forward(x);
      torch::Tensor s = y.sum();

      s.backward();
      CATCH_REQUIRE(y.ndimension() == 3);
      CATCH_REQUIRE(y.size(0) == 2);
      CATCH_REQUIRE(y.size(1) == 3);
      CATCH_REQUIRE(y.size(2) == 4);
    }
  }

  CATCH_SECTION("dropout") {
    Dropout dropout(0.5);
    torch::Tensor x = torch::ones(100, torch::requires_grad());
    torch::Tensor y = dropout->forward(x);

    y.backward();
    CATCH_REQUIRE(y.ndimension() == 1);
    CATCH_REQUIRE(y.size(0) == 100);
    CATCH_REQUIRE(y.sum().toCFloat() < 130); // Probably
    CATCH_REQUIRE(y.sum().toCFloat() > 70); // Probably

    dropout->eval();
    y = dropout->forward(x);
    CATCH_REQUIRE(y.sum().toCFloat() == 100);
  }

  CATCH_SECTION("param") {
    auto model = std::make_shared<NestedModel>();
    auto parameters = model->parameters();
    CATCH_REQUIRE(parameters["param"].size(0) == 3);
    CATCH_REQUIRE(parameters["param"].size(1) == 2);
    CATCH_REQUIRE(parameters["param"].size(2) == 21);
    CATCH_REQUIRE(parameters["l1.bias"].size(0) == 20);
    CATCH_REQUIRE(parameters["l1.weight"].size(0) == 20);
    CATCH_REQUIRE(parameters["l1.weight"].size(1) == 5);
    CATCH_REQUIRE(parameters["test.l1.bias"].size(0) == 3);
    CATCH_REQUIRE(parameters["test.l1.weight"].size(0) == 3);
    CATCH_REQUIRE(parameters["test.l1.weight"].size(1) == 10);
    CATCH_REQUIRE(parameters["test.l2.bias"].size(0) == 5);
    CATCH_REQUIRE(parameters["test.l2.weight"].size(0) == 5);
    CATCH_REQUIRE(parameters["test.l2.weight"].size(1) == 3);
    CATCH_REQUIRE(parameters["test.l3.bias"].size(0) == 100);
    CATCH_REQUIRE(parameters["test.l3.weight"].size(0) == 100);
    CATCH_REQUIRE(parameters["test.l3.weight"].size(1) == 5);
  }

  CATCH_SECTION("functional") {
    {
      bool was_called = false;
      auto functional = Functional([&was_called](torch::Tensor input) {
        was_called = true;
        return input;
      });
      auto output = functional->forward(torch::ones(5, torch::requires_grad()));
      CATCH_REQUIRE(was_called);
      CATCH_REQUIRE(output.equal(torch::ones(5, torch::requires_grad())));

      was_called = false;
      // Use the call operator overload here.
      output = functional(torch::ones(5, torch::requires_grad()));
      CATCH_REQUIRE(was_called);
      CATCH_REQUIRE(output.equal(torch::ones(5, torch::requires_grad())));
    }
    {
      auto functional = Functional(torch::relu);
      CATCH_REQUIRE(functional(torch::ones({})).toCFloat() == 1);
      CATCH_REQUIRE(functional(torch::ones({})).toCFloat() == 1);
      CATCH_REQUIRE(functional(torch::ones({}) * -1).toCFloat() == 0);
    }
    {
      auto functional =
          Functional(torch::elu, /*alpha=*/1, /*scale=*/0, /*input_scale=*/1);
      CATCH_REQUIRE(functional(torch::ones({})).toCFloat() == 0);
    }
  }

  CATCH_SECTION("batchnorm") {
    {
      BatchNorm bn(5);

      // Is stateful by default.
      CATCH_REQUIRE(bn->options.stateful());

      CATCH_REQUIRE(bn->running_mean.defined());
      CATCH_REQUIRE(bn->running_mean.dim() == 1);
      CATCH_REQUIRE(bn->running_mean.size(0) == 5);

      CATCH_REQUIRE(bn->running_variance.defined());
      CATCH_REQUIRE(bn->running_variance.dim() == 1);
      CATCH_REQUIRE(bn->running_variance.size(0) == 5);

      // Is affine by default.
      CATCH_REQUIRE(bn->options.affine());

      CATCH_REQUIRE(bn->weight.defined());
      CATCH_REQUIRE(bn->weight.dim() == 1);
      CATCH_REQUIRE(bn->weight.size(0) == 5);

      CATCH_REQUIRE(bn->bias.defined());
      CATCH_REQUIRE(bn->bias.dim() == 1);
      CATCH_REQUIRE(bn->bias.size(0) == 5);
    }
    {
      BatchNorm bn(BatchNormOptions(5).stateful(false).affine(false));

      CATCH_REQUIRE(!bn->running_mean.defined());
      CATCH_REQUIRE(!bn->running_variance.defined());
      CATCH_REQUIRE(!bn->weight.defined());
      CATCH_REQUIRE(!bn->bias.defined());

      CATCH_REQUIRE_THROWS_WITH(
          bn->forward(torch::ones({2, 5})),
          StartsWith("Calling BatchNorm::forward is only permitted "
                     "when the 'stateful' option is true (was false). "
                     "Use BatchNorm::pure_forward instead."));
    }
    {
      BatchNorm bn(BatchNormOptions(5).affine(false));
      bn->eval();

      // Want to make sure we use the supplied values in `pure_forward` even if
      // we are stateful.
      auto input = torch::randn({2, 5});
      auto mean = torch::randn(5);
      auto variance = torch::rand(5);
      auto output = bn->pure_forward(input, mean, variance);
      auto expected =
          (input - mean) / torch::sqrt(variance + bn->options.eps());
      CATCH_REQUIRE(output.allclose(expected));
    }
  }
}

CATCH_TEST_CASE("modules_cuda", "[cuda]") {
  torch::manual_seed(0);
  CATCH_SECTION("1") {
    Linear model(5, 2);
    model->to(torch::kCUDA);
    auto x =
        torch::randn({10, 5}, torch::device(torch::kCUDA).requires_grad(true));
    auto y = model->forward(x);
    torch::Tensor s = y.sum();

    s.backward();
    CATCH_REQUIRE(y.ndimension() == 2);
    CATCH_REQUIRE(s.ndimension() == 0);
    CATCH_REQUIRE(y.size(0) == 10);
    CATCH_REQUIRE(y.size(1) == 2);

    CATCH_REQUIRE(model->parameters()["weight"].grad().numel() == 2 * 5);
  }

  CATCH_SECTION("2") {
    Linear model(5, 2);
    model->to(torch::kCUDA);
    model->to(torch::kCPU);
    auto x = torch::randn({10, 5}, torch::requires_grad());
    auto y = model->forward(x);
    torch::Tensor s = y.sum();

    s.backward();
    CATCH_REQUIRE(y.ndimension() == 2);
    CATCH_REQUIRE(s.ndimension() == 0);
    CATCH_REQUIRE(y.size(0) == 10);
    CATCH_REQUIRE(y.size(1) == 2);

    CATCH_REQUIRE(model->parameters()["weight"].grad().numel() == 2 * 5);
  }
}