#include "model.h"

#include <algorithm>
#include <cmath>
#include <vector>

namespace {

torch::nn::Sequential siren_trunk(int64_t dim_in, int width, int depth) {
  auto net = torch::nn::Sequential(SirenLayer(dim_in, width));
  for (int i = 0; i < depth; i++) net->push_back(SirenLayer(width, width));
  return net;
}

torch::nn::Sequential softplus_head(int width) {
  return torch::nn::Sequential(LinearLayer(width, 1), torch::nn::Softplus());
}

}  // namespace

torch::Tensor positional_encoding(const torch::Tensor &x, int num_freqs) {
  std::vector<torch::Tensor> parts{x};
  for (int i = 0; i < num_freqs; i++) {
    const float freq = std::pow(2.0f, static_cast<float>(i)) * static_cast<float>(M_PI);
    parts.push_back(torch::sin(freq * x));
    parts.push_back(torch::cos(freq * x));
  }
  return torch::cat(parts, -1);
}

LinearLayer::LinearLayer(int64_t dim_in, int64_t dim_out) {
  weight_ = register_parameter("weight", torch::empty({dim_out, dim_in}));
  bias_ = register_parameter("bias", torch::full({dim_out}, 0.1f));
  torch::nn::init::xavier_normal_(weight_);
}

torch::Tensor LinearLayer::forward(const torch::Tensor &x) {
  return torch::nn::functional::linear(x, weight_, bias_);
}

SirenLayer::SirenLayer(int64_t dim_in, int64_t dim_out) {
  const float bound = std::sqrt(6.0f / static_cast<float>(dim_in));
  weight_ = register_parameter("weight", torch::empty({dim_out, dim_in}));
  bias_ = register_parameter("bias", torch::empty({dim_out}));
  torch::nn::init::uniform_(weight_, -bound, bound);
  torch::nn::init::uniform_(bias_, -bound, bound);
}

torch::Tensor SirenLayer::forward(const torch::Tensor &x) {
  return torch::sin(torch::nn::functional::linear(x, weight_, bias_));
}

SirenNeRF::SirenNeRF(torch::Device device, int width, int depth) {
  constexpr int kPropWidth = 128, kPropDepth = 2;
  depth = std::max(depth, 1);

  trunk_ = register_module("nerf_net", siren_trunk(kPosDim, width, depth));
  sigma_head_ = register_module("sigma_head", softplus_head(width));
  // View direction enters only here, so density stays view-independent.
  auto rgb = torch::nn::Sequential(SirenLayer(width + kViewDim, width));
  rgb->push_back(SirenLayer(width, width));
  rgb->push_back(LinearLayer(width, 3));
  rgb->push_back(torch::nn::Sigmoid());
  rgb_head_ = register_module("rgb_head", rgb);

  prop_trunk_ = register_module("prop_net", siren_trunk(kPosDim, kPropWidth, kPropDepth));
  prop_head_ = register_module("prop_sigma_head", softplus_head(kPropWidth));

  to(device);
}

NeRFOutput SirenNeRF::forward(const torch::Tensor &pts, const torch::Tensor &view_dirs) {
  auto features = trunk_->forward(positional_encoding(pts, kPosFreqs));
  auto view_features = positional_encoding(view_dirs, kViewFreqs);
  return {rgb_head_->forward(torch::cat({features, view_features}, -1)),
          sigma_head_->forward(features)};
}

torch::Tensor SirenNeRF::proposal_sigma(const torch::Tensor &pts) {
  auto enc = positional_encoding(pts, kPosFreqs).detach();
  return prop_head_->forward(prop_trunk_->forward(enc)).squeeze(-1);
}