#pragma once
/* Copyright Ramakrishnan Kannan 2018 */

#include "distntf/distauntf.hpp"

namespace planc {

class DistNTFHALS : public DistAUNTF {
 protected:
   /**
   * This is HALS update function.
   * Given the MTTKRP and the hadamard of all the grams, we
   * determine the factor matrix to be updated.
   * @param[in] Mode of the factor to be updated
   * @returns The new updated factor
   */
  MAT update(const int mode) {
    MAT H(this->m_local_ncp_factors.factor(mode));
    // iterate over all columns of H
    for (int i = 0; i < this->m_local_ncp_factors.rank(); i++) {
      arma::vec updHi;
      if (m_nls_sizes[mode] > 0) {
        updHi = H.col(i) + ((this->ncp_local_mttkrp_t[mode].row(i)).t() -
                            H * this->global_gram.col(i));
      } else {
        updHi = H.col(i);
        updHi.zeros();
      }

      fixNumericalError<arma::vec>(&updHi, EPSILON_1EMINUS16, EPSILON_1EMINUS16);
      double normHi = arma::norm(updHi, 2);
      normHi *= normHi;
      double globalnormHi = normHi;
      MPI_Allreduce(&normHi, &globalnormHi, 1, MPI_DOUBLE, MPI_SUM,
                    MPI_COMM_WORLD);
      if (globalnormHi > 0) {
        H.col(i) = updHi;
      }
    }
    return H.t();
  }

 public:
  DistNTFHALS(const Tensor &i_tensor, const int i_k, algotype i_algo,
              const arma::uvec &i_global_dims, const arma::uvec &i_local_dims,
              const arma::uvec &i_nls_sizes, const arma::uvec &i_nls_idxs,
              const NTFMPICommunicator &i_mpicomm)
      : DistAUNTF(i_tensor, i_k, i_algo, i_global_dims, i_local_dims,
                  i_nls_sizes, i_nls_idxs, i_mpicomm) {}
};  // class DistNTFHALS

}  // namespace planc
