// nnet/nnet-nnet.cc
  
  // 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.
  
  #include "nnet/nnet-nnet.h"
  #include "nnet/nnet-component.h"
  #include "nnet/nnet-parallel-component.h"
  #include "nnet/nnet-multibasis-component.h"
  #include "nnet/nnet-activation.h"
  #include "nnet/nnet-affine-transform.h"
  #include "nnet/nnet-various.h"
  
  namespace kaldi {
  namespace nnet1 {
  
  Nnet::Nnet() {
  }
  
  Nnet::~Nnet() {
    Destroy();
  }
  
  Nnet::Nnet(const Nnet& other) {
    // copy the components
    for (int32 i = 0; i < other.NumComponents(); i++) {
      components_.push_back(other.GetComponent(i).Copy());
    }
    // create empty buffers
    propagate_buf_.resize(NumComponents()+1);
    backpropagate_buf_.resize(NumComponents()+1);
    // copy train opts
    SetTrainOptions(other.opts_);
    Check();
  }
  
  Nnet& Nnet::operator= (const Nnet& other) {
    Destroy();
    // copy the components
    for (int32 i = 0; i < other.NumComponents(); i++) {
      components_.push_back(other.GetComponent(i).Copy());
    }
    // create empty buffers
    propagate_buf_.resize(NumComponents()+1);
    backpropagate_buf_.resize(NumComponents()+1);
    // copy train opts
    SetTrainOptions(other.opts_);
    Check();
    return *this;
  }
  
  /**
   * Forward propagation through the network,
   * (from first component to last).
   */
  void Nnet::Propagate(const CuMatrixBase<BaseFloat> &in,
                       CuMatrix<BaseFloat> *out) {
    // In case of empty network copy input to output,
    if (NumComponents() == 0) {
      (*out) = in;  // copy,
      return;
    }
    // We need C+1 buffers,
    if (propagate_buf_.size() != NumComponents()+1) {
      propagate_buf_.resize(NumComponents()+1);
    }
    // Copy input to first buffer,
    propagate_buf_[0] = in;
    // Propagate through all the components,
    for (int32 i = 0; i < static_cast<int32>(components_.size()); i++) {
      components_[i]->Propagate(propagate_buf_[i], &propagate_buf_[i+1]);
    }
    // Copy the output from the last buffer,
    (*out) = propagate_buf_[NumComponents()];
  }
  
  
  /**
   * Error back-propagation through the network,
   * (from last component to first).
   */
  void Nnet::Backpropagate(const CuMatrixBase<BaseFloat> &out_diff,
                           CuMatrix<BaseFloat> *in_diff) {
    // Copy the derivative in case of empty network,
    if (NumComponents() == 0) {
      (*in_diff) = out_diff;  // copy,
      return;
    }
    // We need C+1 buffers,
    KALDI_ASSERT(static_cast<int32>(propagate_buf_.size()) == NumComponents()+1);
    if (backpropagate_buf_.size() != NumComponents()+1) {
      backpropagate_buf_.resize(NumComponents()+1);
    }
    // Copy 'out_diff' to last buffer,
    backpropagate_buf_[NumComponents()] = out_diff;
    // Loop from last Component to the first,
    for (int32 i = NumComponents()-1; i >= 0; i--) {
      // Backpropagate through 'Component',
      components_[i]->Backpropagate(propagate_buf_[i],
                                    propagate_buf_[i+1],
                                    backpropagate_buf_[i+1],
                                    &backpropagate_buf_[i]);
      // Update 'Component' (if applicable),
      if (components_[i]->IsUpdatable()) {
        UpdatableComponent* uc =
          dynamic_cast<UpdatableComponent*>(components_[i]);
        uc->Update(propagate_buf_[i], backpropagate_buf_[i+1]);
      }
    }
    // Export the derivative (if applicable),
    if (NULL != in_diff) {
      (*in_diff) = backpropagate_buf_[0];
    }
  }
  
  
  void Nnet::Feedforward(const CuMatrixBase<BaseFloat> &in,
                         CuMatrix<BaseFloat> *out) {
    KALDI_ASSERT(NULL != out);
    (*out) = in;  // works even with 0 components,
    CuMatrix<BaseFloat> tmp_in;
    for (int32 i = 0; i < NumComponents(); i++) {
      out->Swap(&tmp_in);
      components_[i]->Propagate(tmp_in, out);
    }
  }
  
  
  int32 Nnet::OutputDim() const {
    KALDI_ASSERT(!components_.empty());
    return components_.back()->OutputDim();
  }
  
  int32 Nnet::InputDim() const {
    KALDI_ASSERT(!components_.empty());
    return components_.front()->InputDim();
  }
  
  const Component& Nnet::GetComponent(int32 c) const {
    return *(components_.at(c));
  }
  
  Component& Nnet::GetComponent(int32 c) {
    return *(components_.at(c));
  }
  
  const Component& Nnet::GetLastComponent() const {
    return *(components_.at(NumComponents()-1));
  }
  
  Component& Nnet::GetLastComponent() {
    return *(components_.at(NumComponents()-1));
  }
  
  void Nnet::ReplaceComponent(int32 c, const Component& comp) {
    delete components_.at(c);
    components_.at(c) = comp.Copy();  // deep copy,
    Check();
  }
  
  void Nnet::SwapComponent(int32 c, Component** comp) {
    Component* tmp = components_.at(c);
    components_.at(c) = *comp;
    (*comp) = tmp;
    Check();
  }
  
  void Nnet::AppendComponent(const Component& comp) {
    components_.push_back(comp.Copy());  // append,
    Check();
  }
  
  void Nnet::AppendComponentPointer(Component* dynamically_allocated_comp) {
    components_.push_back(dynamically_allocated_comp);  // append,
    Check();
  }
  
  void Nnet::AppendNnet(const Nnet& other) {
    for (int32 i = 0; i < other.NumComponents(); i++) {
      AppendComponent(other.GetComponent(i));
    }
    Check();
  }
  
  void Nnet::RemoveComponent(int32 c) {
    Component* ptr = components_.at(c);
    components_.erase(components_.begin()+c);
    delete ptr;
    Check();
  }
  
  void Nnet::RemoveLastComponent() {
    RemoveComponent(NumComponents()-1);
  }
  
  int32 Nnet::NumParams() const {
    int32 n_params = 0;
    for (int32 n = 0; n < components_.size(); n++) {
      if (components_[n]->IsUpdatable()) {
        n_params +=
          dynamic_cast<UpdatableComponent*>(components_[n])->NumParams();
      }
    }
    return n_params;
  }
  
  void Nnet::GetGradient(Vector<BaseFloat>* gradient) const {
    gradient->Resize(NumParams());
    int32 pos = 0;
    // loop over Components,
    for (int32 i = 0; i < components_.size(); i++) {
      if (components_[i]->IsUpdatable()) {
        UpdatableComponent& c =
          dynamic_cast<UpdatableComponent&>(*components_[i]);
        SubVector<BaseFloat> grad_range(gradient->Range(pos, c.NumParams()));
        c.GetGradient(&grad_range);  // getting gradient,
        pos += c.NumParams();
      }
    }
    KALDI_ASSERT(pos == NumParams());
  }
  
  void Nnet::GetParams(Vector<BaseFloat>* params) const {
    params->Resize(NumParams());
    int32 pos = 0;
    // loop over Components,
    for (int32 i = 0; i < components_.size(); i++) {
      if (components_[i]->IsUpdatable()) {
        UpdatableComponent& c =
          dynamic_cast<UpdatableComponent&>(*components_[i]);
        SubVector<BaseFloat> params_range(params->Range(pos, c.NumParams()));
        c.GetParams(&params_range);  // getting params,
        pos += c.NumParams();
      }
    }
    KALDI_ASSERT(pos == NumParams());
  }
  
  void Nnet::SetParams(const VectorBase<BaseFloat>& params) {
    KALDI_ASSERT(params.Dim() == NumParams());
    int32 pos = 0;
    // loop over Components,
    for (int32 i = 0; i < components_.size(); i++) {
      if (components_[i]->IsUpdatable()) {
        UpdatableComponent& c =
          dynamic_cast<UpdatableComponent&>(*components_[i]);
        c.SetParams(params.Range(pos, c.NumParams()));  // setting params,
        pos += c.NumParams();
      }
    }
    KALDI_ASSERT(pos == NumParams());
  }
  
  void Nnet::SetDropoutRate(BaseFloat r)  {
    for (int32 c = 0; c < NumComponents(); c++) {
      if (GetComponent(c).GetType() == Component::kDropout) {
        Dropout& comp = dynamic_cast<Dropout&>(GetComponent(c));
        BaseFloat r_old = comp.GetDropoutRate();
        comp.SetDropoutRate(r);
        KALDI_LOG << "Setting dropout-rate in component " << c
                  << " from " << r_old << " to " << r;
      }
    }
  }
  
  
  void Nnet::ResetStreams(const std::vector<int32> &stream_reset_flag) {
    for (int32 c = 0; c < NumComponents(); c++) {
      if (GetComponent(c).IsMultistream()) {
        MultistreamComponent& comp =
          dynamic_cast<MultistreamComponent&>(GetComponent(c));
        comp.ResetStreams(stream_reset_flag);
      }
    }
  }
  
  void Nnet::SetSeqLengths(const std::vector<int32> &sequence_lengths) {
    for (int32 c = 0; c < NumComponents(); c++) {
      if (GetComponent(c).IsMultistream()) {
        MultistreamComponent& comp =
          dynamic_cast<MultistreamComponent&>(GetComponent(c));
        comp.SetSeqLengths(sequence_lengths);
      }
    }
  }
  
  void Nnet::Init(const std::string &proto_file) {
    Input in(proto_file);
    std::istream &is = in.Stream();
    std::string proto_line, token;
  
    // Initialize from the prototype, where each line
    // contains the description for one component.
    while (is >> std::ws, !is.eof()) {
      KALDI_ASSERT(is.good());
  
      // get a line from the proto file,
      std::getline(is, proto_line);
      if (proto_line == "") continue;
      KALDI_VLOG(1) << proto_line;
  
      // get the 1st token from the line,
      std::istringstream(proto_line) >> std::ws >> token;
      // ignore these tokens:
      if (token == "<NnetProto>" || token == "</NnetProto>") continue;
  
      // create new component, append to Nnet,
      this->AppendComponentPointer(Component::Init(proto_line+"
  "));
    }
    // cleanup
    in.Close();
    Check();
  }
  
  
  /**
   * I/O wrapper for converting 'rxfilename' to 'istream',
   */
  void Nnet::Read(const std::string &rxfilename) {
    bool binary;
    Input in(rxfilename, &binary);
    Read(in.Stream(), binary);
    in.Close();
    // Warn if the NN is empty
    if (NumComponents() == 0) {
      KALDI_WARN << "The network '" << rxfilename << "' is empty.";
    }
  }
  
  
  void Nnet::Read(std::istream &is, bool binary) {
    // Read the Components through the 'factory' Component::Read(...),
    Component* comp(NULL);
    while (comp = Component::Read(is, binary), comp != NULL) {
      // Check dims,
      if (NumComponents() > 0) {
        if (components_.back()->OutputDim() != comp->InputDim()) {
          KALDI_ERR << "Dimensionality mismatch!"
                    << " Previous layer output:" << components_.back()->OutputDim()
                    << " Current layer input:" << comp->InputDim();
        }
      }
      // Append to 'this' Nnet,
      AppendComponentPointer(comp);
    }
    Check();
  }
  
  
  /**
   * I/O wrapper for converting 'wxfilename' to 'ostream',
   */
  void Nnet::Write(const std::string &wxfilename, bool binary) const {
    Output out(wxfilename, binary, true);
    Write(out.Stream(), binary);
    out.Close();
  }
  
  
  void Nnet::Write(std::ostream &os, bool binary) const {
    Check();
    WriteToken(os, binary, "<Nnet>");
    if (binary == false) os << std::endl;
    for (int32 i = 0; i < NumComponents(); i++) {
      components_[i]->Write(os, binary);
    }
    WriteToken(os, binary, "</Nnet>");
    if (binary == false) os << std::endl;
  }
  
  
  std::string Nnet::Info() const {
    // global info
    std::ostringstream ostr;
    ostr << "num-components " << NumComponents() << std::endl;
    if (NumComponents() == 0)
      return ostr.str();
    ostr << "input-dim " << InputDim() << std::endl;
    ostr << "output-dim " << OutputDim() << std::endl;
    ostr << "number-of-parameters " << static_cast<float>(NumParams())/1e6
         << " millions" << std::endl;
    // topology & weight stats
    for (int32 i = 0; i < NumComponents(); i++) {
      ostr << "component " << i+1 << " : "
           << Component::TypeToMarker(components_[i]->GetType())
           << ", input-dim " << components_[i]->InputDim()
           << ", output-dim " << components_[i]->OutputDim()
           << ", " << components_[i]->Info() << std::endl;
    }
    return ostr.str();
  }
  
  std::string Nnet::InfoGradient(bool header) const {
    std::ostringstream ostr;
    // gradient stats
    if (header) ostr << "
  ### GRADIENT STATS :
  ";
    for (int32 i = 0; i < NumComponents(); i++) {
      ostr << "Component " << i+1 << " : "
           << Component::TypeToMarker(components_[i]->GetType())
           << ", " << components_[i]->InfoGradient() << std::endl;
    }
    if (header) ostr << "### END GRADIENT
  ";
    return ostr.str();
  }
  
  std::string Nnet::InfoPropagate(bool header) const {
    std::ostringstream ostr;
    // forward-pass buffer stats
    if (header) ostr << "
  ### FORWARD PROPAGATION BUFFER CONTENT :
  ";
    ostr << "[0] output of <Input> " << MomentStatistics(propagate_buf_[0])
         << std::endl;
    for (int32 i = 0; i < NumComponents(); i++) {
      ostr << "[" << 1+i << "] output of "
           << Component::TypeToMarker(components_[i]->GetType())
           << MomentStatistics(propagate_buf_[i+1]) << std::endl;
      // nested networks too...
      if (Component::kParallelComponent == components_[i]->GetType()) {
        ostr <<
          dynamic_cast<ParallelComponent*>(components_[i])->InfoPropagate();
      }
      if (Component::kMultiBasisComponent == components_[i]->GetType()) {
        ostr << dynamic_cast<MultiBasisComponent*>(components_[i])->InfoPropagate();
      }
    }
    if (header) ostr << "### END FORWARD
  ";
    return ostr.str();
  }
  
  std::string Nnet::InfoBackPropagate(bool header) const {
    std::ostringstream ostr;
    // forward-pass buffer stats
    if (header) ostr << "
  ### BACKWARD PROPAGATION BUFFER CONTENT :
  ";
    ostr << "[0] diff of <Input> " << MomentStatistics(backpropagate_buf_[0])
         << std::endl;
    for (int32 i = 0; i < NumComponents(); i++) {
      ostr << "["<<1+i<< "] diff-output of "
           << Component::TypeToMarker(components_[i]->GetType())
           << MomentStatistics(backpropagate_buf_[i+1]) << std::endl;
      // nested networks too...
      if (Component::kParallelComponent == components_[i]->GetType()) {
        ostr <<
          dynamic_cast<ParallelComponent*>(components_[i])->InfoBackPropagate();
      }
      if (Component::kMultiBasisComponent == components_[i]->GetType()) {
        ostr << dynamic_cast<MultiBasisComponent*>(components_[i])->InfoBackPropagate();
      }
    }
    if (header) ostr << "### END BACKWARD
  
  ";
    return ostr.str();
  }
  
  
  void Nnet::Check() const {
    // check dims,
    for (size_t i = 0; i + 1 < components_.size(); i++) {
      KALDI_ASSERT(components_[i] != NULL);
      int32 output_dim = components_[i]->OutputDim(),
        next_input_dim = components_[i+1]->InputDim();
      // show error message,
      if (output_dim != next_input_dim) {
        KALDI_ERR << "Component dimension mismatch!"
                  << " Output dim of [" << i << "] "
                  << Component::TypeToMarker(components_[i]->GetType())
                  << " is " << output_dim << ". "
                  << "Input dim of next [" << i+1 << "] "
                  << Component::TypeToMarker(components_[i+1]->GetType())
                  << " is " << next_input_dim << ".";
      }
    }
    // check for nan/inf in network weights,
    Vector<BaseFloat> weights;
    GetParams(&weights);
    BaseFloat sum = weights.Sum();
    if (KALDI_ISINF(sum)) {
      KALDI_ERR << "'inf' in network parameters "
                << "(weight explosion, need lower learning rate?)";
    }
    if (KALDI_ISNAN(sum)) {
      KALDI_ERR << "'nan' in network parameters (need lower learning rate?)";
    }
  }
  
  
  void Nnet::Destroy() {
    for (int32 i = 0; i < NumComponents(); i++) {
      delete components_[i];
    }
    components_.resize(0);
    propagate_buf_.resize(0);
    backpropagate_buf_.resize(0);
  }
  
  
  void Nnet::SetTrainOptions(const NnetTrainOptions& opts) {
    opts_ = opts;
    // set values to individual components,
    for (int32 l = 0; l < NumComponents(); l++) {
      if (GetComponent(l).IsUpdatable()) {
        dynamic_cast<UpdatableComponent&>(GetComponent(l)).SetTrainOptions(opts_);
      }
    }
  }
  
  
  }  // namespace nnet1
  }  // namespace kaldi