#include "OneWordNetwork.hpp"

OneWordNetworkImpl::OneWordNetworkImpl(int nbOutputs, int focusedIndex)
{
  constexpr int embeddingsSize = 30;

  wordEmbeddings = register_module("word_embeddings", torch::nn::Embedding(torch::nn::EmbeddingOptions(200000, embeddingsSize).sparse(true)));
  auto params = wordEmbeddings->parameters();
  _sparseParameters.insert(_sparseParameters.end(), params.begin(), params.end());

  linear = register_module("linear", torch::nn::Linear(embeddingsSize, nbOutputs));
  params = linear->parameters();
  _denseParameters.insert(_denseParameters.end(), params.begin(), params.end());

  int leftBorder = 0;
  int rightBorder = 0;
  if (focusedIndex < 0)
    leftBorder = -focusedIndex;
  if (focusedIndex > 0)
    rightBorder = focusedIndex;

  this->focusedIndex = focusedIndex <= 0 ? 0 : focusedIndex;

  setLeftBorder(leftBorder);
  setRightBorder(rightBorder);
  setNbStackElements(0);
}

std::vector<torch::Tensor> & OneWordNetworkImpl::denseParameters()
{
  return _denseParameters;
}

std::vector<torch::Tensor> & OneWordNetworkImpl::sparseParameters()
{
  return _sparseParameters;
}

torch::Tensor OneWordNetworkImpl::forward(torch::Tensor input)
{
  // input dim = {batch, sequence, embeddings}
  auto wordsAsEmb = wordEmbeddings(input);
  auto reshaped = wordsAsEmb;
  // reshaped dim = {sequence, batch, embeddings}
  if (reshaped.dim() == 3)
    reshaped = wordsAsEmb.permute({1,0,2});

  auto res = torch::softmax(linear(reshaped[focusedIndex]), reshaped.dim() == 3 ? 1 : 0);

  return res;
}