/*
 * File: transformer.h
 * Project: transformer
 * Author: koth (Koth Chen)
 * -----
 * Last Modified: 2019-09-14 10:08:03
 * Modified By: koth (nobody@verycool.com)
 * -----
 * Copyright 2020 - 2019
 */
#pragma once

#include <cstddef>
#include <vector>

#include "torch/nn/cloneable.h"
#include "torch/nn/module.h"
#include "torch/nn/modules/linear.h"
#include "torch/nn/pimpl.h"
#include "torch/types.h"

#include "radish/transformer/transformer_decoder.h"
#include "radish/transformer/transformer_encoder.h"

namespace radish {
using Tensor = torch::Tensor;
/// Options for the `Transformer` module.
struct TORCH_API TransformerOptions {
  TransformerOptions(int64_t n_src_vocab, int64_t n_tgt_vocab,
                     int64_t len_max_seq, int64_t d_word_vec = 200,
                     int64_t d_model = 200, int64_t d_inner = 800,
                     int64_t n_layers = 6, int64_t n_head = 8, int64_t d_k = 25,
                     int64_t d_v = 25, double dropout = 0.1,
                     bool tgt_emb_prj_weight_sharing = true,
                     bool emb_src_tgt_weight_sharing = true);
  TORCH_ARG(int64_t, n_src_vocab);
  TORCH_ARG(int64_t, n_tgt_vocab);
  TORCH_ARG(int64_t, len_max_seq);
  TORCH_ARG(int64_t, d_word_vec) = 200;
  TORCH_ARG(int64_t, d_model) = 200;
  TORCH_ARG(int64_t, d_inner) = 800;
  TORCH_ARG(int64_t, n_layers) = 6;
  TORCH_ARG(int64_t, n_head) = 8;
  TORCH_ARG(int64_t, d_k) = 25;
  TORCH_ARG(int64_t, d_v) = 25;
  TORCH_ARG(double, dropout) = 0.1;
  TORCH_ARG(bool, tgt_emb_prj_weight_sharing) = true;
  TORCH_ARG(bool, emb_src_tgt_weight_sharing) = true;
};

class TORCH_API TransformerImpl
    : public ::torch::nn::Cloneable<TransformerImpl> {
 public:
  TransformerImpl(int64_t n_src_vocab, int64_t n_tgt_vocab, int64_t len_max_seq)
      : TransformerImpl(
            TransformerOptions(n_src_vocab, n_tgt_vocab, len_max_seq)) {}
  explicit TransformerImpl(TransformerOptions options);

  void reset() override;

  void pretty_print(std::ostream& stream) const override;

  Tensor forward(const Tensor& src_seq, const Tensor& src_pos,
                 const Tensor& tgt_seq, const Tensor& tgt_pos);

  /// The options used to configure this module.
  TransformerOptions options;
  torch::nn::Linear tgt_word_prj = nullptr;
  TransformerEncoder encoder = nullptr;
  TransformerDecoder decoder = nullptr;
  double x_logit_scale = 1.0;
};

/// A `ModuleHolder` subclass for `TransformerImpl`.
/// See the documentation for `TransformerImpl` class to learn what
/// methods it provides, or the documentation for `ModuleHolder` to learn about
/// PyTorch's module storage semantics.
TORCH_MODULE(Transformer);

}  // namespace radish
