// nnet/nnet-component.h

// Copyright 2011-2016  Brno University of Technology (Author: Karel Vesely)

// See ../../COPYING for clarification regarding multiple authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//  http://www.apache.org/licenses/LICENSE-2.0
//
// THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
// WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
// MERCHANTABLITY OR NON-INFRINGEMENT.
// See the Apache 2 License for the specific language governing permissions and
// limitations under the License.



#ifndef KALDI_NNET_NNET_COMPONENT_H_
#define KALDI_NNET_NNET_COMPONENT_H_

#include <iostream>
#include <string>

#include "base/kaldi-common.h"
#include "matrix/matrix-lib.h"
#include "cudamatrix/cu-matrix.h"
#include "cudamatrix/cu-vector.h"
#include "nnet/nnet-trnopts.h"

namespace kaldi {
namespace nnet1 {

/**
 * Abstract class, building block of the network.
 * It is able to propagate (PropagateFnc: compute the output based on its input)
 * and backpropagate (BackpropagateFnc: i.e. transform loss derivative w.r.t. output to derivative w.r.t. the input)
 * the formulas are implemented in descendant classes (AffineTransform,Sigmoid,Softmax,...).
 */
class Component {
 /// Component type identification mechanism,
 public:
  /// Types of Components,
  typedef enum {
    kUnknown = 0x0,

    kUpdatableComponent = 0x0100,
    kAffineTransform,
    kLinearTransform,
    kConvolutionalComponent,
    kLstmProjected,
    kBlstmProjected,
    kRecurrentComponent,

    kActivationFunction = 0x0200,
    kSoftmax,
    kHiddenSoftmax,
    kBlockSoftmax,
    kSigmoid,
    kTanh,
    kParametricRelu,
    kDropout,
    kLengthNormComponent,

    kTranform = 0x0400,
    kRbm,
    kSplice,
    kCopy,
    kTranspose,
    kBlockLinearity,
    kAddShift,
    kRescale,

    kKlHmm = 0x0800,
    kSentenceAveragingComponent, /* deprecated */
    kSimpleSentenceAveragingComponent,
    kAveragePoolingComponent,
    kMaxPoolingComponent,
    kFramePoolingComponent,
    kParallelComponent,
    kMultiBasisComponent
  } ComponentType;

  /// A pair of type and marker,
  struct key_value {
    const Component::ComponentType key;
    const char *value;
  };

  /// The table with pairs of Component types and markers
  /// (defined in nnet-component.cc),
  static const struct key_value kMarkerMap[];

  /// Converts component type to marker,
  static const char* TypeToMarker(ComponentType t);

  /// Converts marker to component type (case insensitive),
  static ComponentType MarkerToType(const std::string &s);

 /// Generic interface of a component,
 public:
  Component(int32 input_dim, int32 output_dim):
    input_dim_(input_dim),
    output_dim_(output_dim)
  { }

  virtual ~Component()
  { }

  /// Copy component (deep copy),
  virtual Component* Copy() const = 0;

  /// Get Type Identification of the component,
  virtual ComponentType GetType() const = 0;

  /// Check if componeny has 'Updatable' interface (trainable components),
  virtual bool IsUpdatable() const {
    return false;
  }

  /// Check if component has 'Recurrent' interface (trainable and recurrent),
  virtual bool IsMultistream() const {
    return false;
  }

  /// Get the dimension of the input,
  int32 InputDim() const {
    return input_dim_;
  }

  /// Get the dimension of the output,
  int32 OutputDim() const {
    return output_dim_;
  }

  /// Perform forward-pass propagation 'in' -> 'out',
  void Propagate(const CuMatrixBase<BaseFloat> &in, CuMatrix<BaseFloat> *out);

  /// Perform backward-pass propagation 'out_diff' -> 'in_diff'.
  /// Note: 'in' and 'out' will be used only sometimes...
  void Backpropagate(const CuMatrixBase<BaseFloat> &in,
                     const CuMatrixBase<BaseFloat> &out,
                     const CuMatrixBase<BaseFloat> &out_diff,
                     CuMatrix<BaseFloat> *in_diff);

  /// Initialize component from a line in config file,
  static Component* Init(const std::string &conf_line);

  /// Read the component from a stream (static method),
  static Component* Read(std::istream &is, bool binary);

  /// Write the component to a stream,
  void Write(std::ostream &os, bool binary) const;

  /// Print some additional info (after <ComponentName> and the dims),
  virtual std::string Info() const { return ""; }

  /// Print some additional info about gradient (after <...> and dims),
  virtual std::string InfoGradient() const { return ""; }


 /// Abstract interface for propagation/backpropagation
 protected:
  /// Forward pass transformation (to be implemented by descending class...)
  virtual void PropagateFnc(const CuMatrixBase<BaseFloat> &in,
                            CuMatrixBase<BaseFloat> *out) = 0;

  /// Backward pass transformation (to be implemented by descending class...)
  virtual void BackpropagateFnc(const CuMatrixBase<BaseFloat> &in,
                                const CuMatrixBase<BaseFloat> &out,
                                const CuMatrixBase<BaseFloat> &out_diff,
                                CuMatrixBase<BaseFloat> *in_diff) = 0;

 /// Virtual interface for initialization and I/O,
 protected:
  /// Initialize internal data of a component
  virtual void InitData(std::istream &is) { }

  /// Reads the component content
  virtual void ReadData(std::istream &is, bool binary) { }

  /// Writes the component content
  virtual void WriteData(std::ostream &os, bool binary) const { }

 /// Data members,
 protected:
  int32 input_dim_;  ///< Dimension of the input of the Component,
  int32 output_dim_;  ///< Dimension of the output of the Component,

 /// Private members (descending classes cannot call this),
 private:
  /// Create a new intance of component,
  static Component* NewComponentOfType(
    ComponentType t, int32 input_dim, int32 output_dim
  );
};


/**
 * Class UpdatableComponent is a Component which has trainable parameters,
 * it contains SGD training hyper-parameters in NnetTrainOptions.
 * The constants 'learning_rate_coef_' and 'bias_learn_rate_coef_'
 * are separate, and should be stored by ::WriteData(...),
 */
class UpdatableComponent : public Component {
 public:
  UpdatableComponent(int32 input_dim, int32 output_dim):
    Component(input_dim, output_dim),
    learn_rate_coef_(1.0),
    bias_learn_rate_coef_(1.0)
  { }

  virtual ~UpdatableComponent()
  { }

  /// Check if contains trainable parameters,
  bool IsUpdatable() const {
    return true;
  }

  /// Number of trainable parameters,
  virtual int32 NumParams() const = 0;

  /// Get gradient reshaped as a vector,
  virtual void GetGradient(VectorBase<BaseFloat> *gradient) const = 0;

  /// Get the trainable parameters reshaped as a vector,
  virtual void GetParams(VectorBase<BaseFloat> *params) const = 0;

  /// Set the trainable parameters from, reshaped as a vector,
  virtual void SetParams(const VectorBase<BaseFloat> &params) = 0;

  /// Compute gradient and update parameters,
  virtual void Update(const CuMatrixBase<BaseFloat> &input,
                      const CuMatrixBase<BaseFloat> &diff) = 0;

  /// Set the training options to the component,
  virtual void SetTrainOptions(const NnetTrainOptions &opts) {
    opts_ = opts;
  }

  /// Get the training options from the component,
  const NnetTrainOptions& GetTrainOptions() const {
    return opts_;
  }

  /// Set the learn-rate coefficient,
  virtual void SetLearnRateCoef(BaseFloat val) {
    learn_rate_coef_ = val;
  }

  /// Set the learn-rate coefficient for bias,
  virtual void SetBiasLearnRateCoef(BaseFloat val) {
    bias_learn_rate_coef_ = val;
  }

  /// Initialize the content of the component by the 'line' from the prototype,
  virtual void InitData(std::istream &is) = 0;

 protected:
  /// Option-class with training hyper-parameters,
  NnetTrainOptions opts_;

  /// Scalar applied to learning rate for weight matrices
  /// (to be used in ::Update method),
  BaseFloat learn_rate_coef_;

  /// Scalar applied to learning rate for bias
  /// (to be used in ::Update method),
  BaseFloat bias_learn_rate_coef_;
};


/**
 * Class MultistreamComponent is an extension of UpdatableComponent
 * for recurrent networks, which are trained with parallel sequences.
 */
class MultistreamComponent : public UpdatableComponent {
 public:
  MultistreamComponent(int32 input_dim, int32 output_dim):
    UpdatableComponent(input_dim, output_dim)
  { }

  bool IsMultistream() const {
    return true;
  }

  virtual void SetSeqLengths(const std::vector<int32>& sequence_lengths) {
    sequence_lengths_ = sequence_lengths;
  }

  int32 NumStreams() const {
    return std::max<int32>(1, sequence_lengths_.size());
  }

  /// Optional function to reset the transfer of context (not used for BLSTMs
  virtual void ResetStreams(const std::vector<int32>& stream_reset_flag)
  { }

 protected:
  std::vector<int32> sequence_lengths_;
};


/*
 * Inline methods for ::Component,
 */
inline void Component::Propagate(const CuMatrixBase<BaseFloat> &in,
                                 CuMatrix<BaseFloat> *out) {
  // Check the dims
  if (input_dim_ != in.NumCols()) {
    KALDI_ERR << "Non-matching dims on the input of " << TypeToMarker(GetType())
              << " component. The input-dim is " << input_dim_
              << ", the data had " << in.NumCols() << " dims.";
  }
  // Allocate target buffer
  out->Resize(in.NumRows(), output_dim_, kSetZero);  // reset
  // Call the propagation implementation of the component
  PropagateFnc(in, out);
}

inline void Component::Backpropagate(const CuMatrixBase<BaseFloat> &in,
                                     const CuMatrixBase<BaseFloat> &out,
                                     const CuMatrixBase<BaseFloat> &out_diff,
                                     CuMatrix<BaseFloat> *in_diff) {
  // Check the dims,
  if (OutputDim() != out_diff.NumCols()) {
    KALDI_ERR << "Non-matching dims! Component output dim " << OutputDim()
              << ", the dim of output derivatives " << out_diff.NumCols();
  }

  int32 num_frames = out_diff.NumRows();
  KALDI_ASSERT(num_frames == in.NumRows());
  KALDI_ASSERT(num_frames == out.NumRows());

  KALDI_ASSERT(InputDim() == in.NumCols());
  KALDI_ASSERT(OutputDim() == out.NumCols());

  // Allocate target buffer,
  KALDI_ASSERT(in_diff != NULL);
  in_diff->Resize(num_frames, InputDim(), kSetZero);  // reset,

  // Call the 'virtual' backprop function,
  BackpropagateFnc(in, out, out_diff, in_diff);
}


}  // namespace nnet1
}  // namespace kaldi


#endif  // KALDI_NNET_NNET_COMPONENT_H_