text: // Copyright 2018 The Simons Foundation, Inc. - All Rights Reserved.
  //
  // 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
  //
  // Unless required by applicable law or agreed to in writing, software
  // distributed under the License is distributed on an "AS IS" BASIS,
  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  // See the License for the specific language governing permissions and
  // limitations under the License.
  //
  // by S. Efthymiou, October 2018
  
  #include "Machine/mps_periodic.hpp"
  
  #include <set>
  
  #include "Utils/json_utils.hpp"
  #include "Utils/messages.hpp"
  
  namespace netket {
  
  MPSPeriodic::MPSPeriodic(std::shared_ptr<const AbstractHilbert> hilbert,
                           int bond_dim, bool diag, int symperiod)
      : AbstractMachine(hilbert),
        N_{hilbert->Size()},
        d_{hilbert->LocalSize()},
        D_{bond_dim},
        symperiod_{symperiod},
        is_diag_{diag} {
    if (symperiod_ == -1) {
      symperiod_ = N_;
    }
    Init();
  }
  
  MPSPeriodic::MatrixType MPSPeriodic::prod(const MatrixType &m1,
                                            const MatrixType &m2) const {
    if (is_diag_) {
      return m1.cwiseProduct(m2);
    }
    return m1 * m2;
  }
  
  Complex MPSPeriodic::trace(const MatrixType &m) const {
    if (is_diag_) {
      return m.sum();
    }
    return m.trace();
  }
  
  void MPSPeriodic::setparamsident(MatrixType &m, VectorConstRefType pars) const {
    if (is_diag_) {
      for (int i = 0; i < D_; i++) {
        m(i, 0) = Complex(1, 0) + pars(i);
      }
    } else {
      for (int i = 0; i < D_; i++) {
        for (int j = 0; j < D_; j++) {
          m(i, j) = pars(i * D_ + j);
          if (i == j) {
            m(i, j) += Complex(1, 0);
          }
        }
      }
    }
  }
  
  void MPSPeriodic::Init() {
    // Initialize parameters
    std::vector<MatrixType> pushback_vec;
    if (is_diag_) {
      Dsec_ = 1;
      identity_mat_ = MatrixType::Ones(D_, 1);
    } else {
      Dsec_ = D_;
      identity_mat_ = MatrixType::Identity(D_, D_);
    }
    MatrixType init_mat = MatrixType::Zero(D_, Dsec_);
    Dsq_ = D_ * Dsec_;
    npar_ = symperiod_ * d_ * Dsq_;
  
    for (int site = 0; site < symperiod_; site++) {
      W_.push_back(pushback_vec);
      for (int spin = 0; spin < d_; spin++) {
        W_[site].push_back(init_mat);
      }
    }
  
    // Initialize tree parameters
    InitTree();
  
    // Machine creation messages
    if (is_diag_) {
      InfoMessage() << "Periodic diagonal MPS machine with " << N_
                    << " sites created" << std::endl;
    } else {
      InfoMessage() << "Periodic MPS machine with " << N_ << " sites created"
                    << std::endl;
    }
    InfoMessage() << "Physical dimension d = " << d_
                  << " and bond dimension D = " << D_ << std::endl;
    if (symperiod_ < N_) {
      InfoMessage() << "Translation invariance is used. Number of "
                       "variational parameters is "
                    << npar_ << " instead of " << npar_ * N_ / symperiod_
                    << std::endl;
    } else {
      InfoMessage() << "Number of variational parameters is " << npar_
                    << std::endl;
    }
    // Initialize map from Hilbert space states to MPS indices
    auto localstates = GetHilbert().LocalStates();
    for (int i = 0; i < d_; i++) {
      confindex_[localstates[i]] = i;
    }
  }
  
  void MPSPeriodic::InitTree() {
    // Initializes vectors used for tree look up tables
    // leaves_of_site_ and leaf_contractions_
  
    std::vector<int> two_vector(2), empty_vector, level_start;
    std::vector<std::vector<int>> above;
    int level = 1, available_ind = 0, available_level = 0;
    bool available = false;
  
    level_start.push_back(0);
    level_start.push_back(N_);
  
    int length = level_start[level] - level_start[level - 1];
    while (length > 1 or available) {
      above.push_back(empty_vector);
      // Iterate level-1
      for (int i = 0; i < length - 1; i += 2) {
        // Construct above for level-1
        above.back().push_back(level_start[level] + i / 2);
        above.back().push_back(level_start[level] + i / 2);
        // Construct leaf_contractions_ for level
        two_vector[0] = level_start[level - 1] + i;
        two_vector[1] = level_start[level - 1] + i + 1;
        leaf_contractions_.push_back(two_vector);
      }
      if (length % 2 == 1) {
        if (available) {
          // Connect the two odd leaves
          above.back().push_back(level_start[level] + (length - 1) / 2);
          above[available_level].push_back(level_start[level] + (length - 1) / 2);
  
          two_vector[0] = level_start[level] - 1;
          two_vector[1] = available_ind;
          leaf_contractions_.push_back(two_vector);
  
          level_start.push_back(level_start.back() + length / 2 + 1);
          available = false;
        } else {
          available = true;
          available_ind = level_start[level] - 1;
          available_level = level - 1;
          level_start.push_back(level_start.back() + length / 2);
        }
      } else {
        level_start.push_back(level_start.back() + length / 2);
      }
      level++;
      length = level_start[level] - level_start[level - 1];
    }
    Nleaves_ = level_start.back();
  
    // Flatten above vector
    std::vector<int> flat_above;
    for (std::size_t l = 0; l < above.size(); l++) {
      for (std::size_t k = 0; k < above[l].size(); k++) {
        flat_above.push_back(above[l][k]);
      }
    }
  
    // Create leaves_of_site_ from above vector
    for (int i = 0; i < N_; i++) {
      std::vector<int> leaves;
      leaves.push_back(flat_above[i]);
      while (flat_above[leaves.back()] < Nleaves_ - 1) {
        leaves.push_back(flat_above[leaves.back()]);
      }
      leaves.push_back(Nleaves_ - 1);
      leaves_of_site_.push_back(leaves);
      leaves.clear();
    }
    Nleaves_ += -N_;
  }
  
  int MPSPeriodic::Npar() const { return npar_; }
  int MPSPeriodic::Nvisible() const { return N_; }
  
  MPSPeriodic::VectorType MPSPeriodic::GetParameters() {
    int k = 0;
    VectorType pars(npar_);
  
    for (int site = 0; site < symperiod_; site++) {
      for (int spin = 0; spin < d_; spin++) {
        for (int i = 0; i < D_; i++) {
          for (int j = 0; j < Dsec_; j++) {
            pars(k) = W_[site][spin](i, j);
            k++;
          }
        }
      }
    }
    return pars;
  }
  
  void MPSPeriodic::SetParameters(VectorConstRefType pars) {
    int k = 0;
  
    for (int site = 0; site < symperiod_; site++) {
      for (int spin = 0; spin < d_; spin++) {
        for (int i = 0; i < D_; i++) {
          for (int j = 0; j < Dsec_; j++) {
            W_[site][spin](i, j) = pars(k);
            k++;
          }
        }
      }
    }
  }
  
  // Auxiliary function used for setting initial random parameters and adding
  // identities in every matrix
  void MPSPeriodic::SetParametersIdentity(VectorConstRefType pars) {
    int k = 0;
    for (int site = 0; site < symperiod_; site++) {
      for (int spin = 0; spin < d_; spin++) {
        setparamsident(W_[site][spin], pars.segment(k, Dsq_));
        k += Dsq_;
      }
    }
  }
  
  any MPSPeriodic::InitLookup(VisibleConstType v) {
    LookupType lt;
    for (Index k = 0; k < Nleaves_; k++) {
      _InitLookup_check(lt, k);
      if (leaf_contractions_[k][0] < N_) {
        if (leaf_contractions_[k][1] < N_) {
          lt[k] = prod(W_[leaf_contractions_[k][0] % symperiod_]
                         [confindex_[v(leaf_contractions_[k][0])]],
                       W_[leaf_contractions_[k][1] % symperiod_]
                         [confindex_[v(leaf_contractions_[k][1])]]);
        } else {
          lt[k] = prod(W_[leaf_contractions_[k][0] % symperiod_]
                         [confindex_[v(leaf_contractions_[k][0])]],
                       lt[leaf_contractions_[k][1] - N_]);
        }
  
      } else {
        if (leaf_contractions_[k][1] < N_) {
          lt[k] = prod(lt[leaf_contractions_[k][0] - N_],
                       W_[leaf_contractions_[k][1] % symperiod_]
                         [confindex_[v(leaf_contractions_[k][1])]]);
        } else {
          lt[k] = prod(lt[leaf_contractions_[k][0] - N_],
                       lt[leaf_contractions_[k][1] - N_]);
        }
      }
    }
    return any{std::move(lt)};
  }
  
  // Auxiliary function
  void MPSPeriodic::_InitLookup_check(LookupType &lt, Index i) {
    if (static_cast<Index>(lt.size()) == i) {
      lt.push_back(MatrixType(D_, Dsec_));
    } else {
      lt[i].resize(D_, Dsec_);
    }
  }
  
  // Auxiliary function for sorting indeces
  // (copied from stackexchange - original answer by Lukasz Wiklendt)
  std::vector<std::size_t> MPSPeriodic::sort_indeces(const std::vector<int> &v) {
    // initialize original index locations
    std::vector<std::size_t> idx(v.size());
    std::iota(idx.begin(), idx.end(), 0);
    // sort indexes based on comparing values in v
    std::sort(idx.begin(), idx.end(),
              [&v](std::size_t i1, std::size_t i2) { return v[i1] < v[i2]; });
    return idx;
  }
  
  void MPSPeriodic::UpdateLookup(VisibleConstType v,
                                 const std::vector<int> &tochange,
                                 const std::vector<double> &newconf,
                                 any &lookup) {
    std::size_t nchange = tochange.size();
    if (nchange <= 0) {
      return;
    }
  
    auto &lt = any_cast_ref<LookupType>(lookup);
    MatrixType empty_matrix = MatrixType::Zero(D_, Dsec_);
    std::vector<std::size_t> sorted_ind = sort_indeces(tochange);
  
    std::set<int> leaves2update;
    std::size_t set_len = 0;
    std::map<int, MatrixType *> base_mat;
  
    for (std::size_t k = 0; k < nchange; k++) {
      int site = tochange[sorted_ind[k]];
      // Add changed visible matrices to map
      base_mat[site] =
          &(W_[site % symperiod_][confindex_[newconf[sorted_ind[k]]]]);
      // Add the rest of matrices that affects
      for (std::size_t l = 0; l < leaves_of_site_[site].size(); l++) {
        int leaf = leaves_of_site_[site][l];
        leaves2update.insert(leaf);
        if (leaves2update.size() > set_len) {
          set_len++;
          for (int i = 0; i < 2; i++) {
            int lc = leaf_contractions_[leaf - N_][i];
            if (lc < N_) {
              if (base_mat.count(lc) == 0) {
                base_mat[lc] = &(W_[lc % symperiod_][confindex_[v(lc)]]);
              }
            } else {
              base_mat[lc] = &(lt[lc - N_]);
            }
          }
        }
      }
    }
    for (auto leaf : leaves2update) {
      lt[leaf - N_] = prod(*(base_mat[leaf_contractions_[leaf - N_][0]]),
                           *(base_mat[leaf_contractions_[leaf - N_][1]]));
    }
  }
  
  // Auxiliary function that calculates contractions from site1 to site2
  MPSPeriodic::MatrixType MPSPeriodic::mps_contraction(VisibleConstType v,
                                                       const int &site1,
                                                       const int &site2) {
    MatrixType c = identity_mat_;
    for (int site = site1; site < site2; site++) {
      c = prod(c, W_[site % symperiod_][confindex_[v(site)]]);
    }
    return c;
  }
  
  Complex MPSPeriodic::LogValSingle(VisibleConstType v, const any &lt) {
    if (lt.empty()) return std::log(trace(mps_contraction(v, 0, N_)));
    return std::log(trace(any_cast_ref<LookupType>(lt)[Nleaves_ - 1]));
  }
  
  void MPSPeriodic::LogValDiff(VisibleConstType v,
                               const std::vector<std::vector<int>> &tochange,
                               const std::vector<std::vector<double>> &newconf,
                               Eigen::Ref<Eigen::VectorXcd> logvaldiffs) {
    const std::size_t nconn = tochange.size();
  
    std::vector<std::size_t> sorted_ind;
    logvaldiffs = VectorType::Zero(nconn);
    Complex current_psi = trace(mps_contraction(v, 0, N_));
    MatrixType new_prods(D_, Dsec_);
  
    for (std::size_t k = 0; k < nconn; k++) {
      std::size_t nchange = tochange[k].size();
      if (nchange > 0) {
        sorted_ind = sort_indeces(tochange[k]);
        int site = tochange[k][sorted_ind[0]];
  
        if (site == 0) {
          new_prods = W_[0][confindex_[newconf[k][sorted_ind[0]]]];
        } else {
          new_prods =
              prod(mps_contraction(v, 0, site),
                   W_[site % symperiod_][confindex_[newconf[k][sorted_ind[0]]]]);
        }
  
        for (std::size_t i = 1; i < nchange; i++) {
          site = tochange[k][sorted_ind[i]];
          new_prods = prod(
              new_prods,
              prod(mps_contraction(v, tochange[k][sorted_ind[i - 1]] + 1, site),
                   W_[site % symperiod_][confindex_[newconf[k][sorted_ind[i]]]]));
        }
        site = tochange[k][sorted_ind[nchange - 1]];
        if (site < N_ - 1) {
          new_prods = prod(new_prods, mps_contraction(v, site + 1, N_));
        }
        logvaldiffs(k) = std::log(trace(new_prods) / current_psi);
      }
    }
  }
  
  // Derivative with full calculation
  MPSPeriodic::VectorType MPSPeriodic::DerLogSingle(VisibleConstType v,
                                                    const any & /*unused*/) {
    MatrixType temp_product(D_, Dsec_);
    std::vector<MatrixType> left_prods, right_prods;
    VectorType der = VectorType::Zero(npar_);
  
    // Calculate products
    left_prods.push_back(W_[0][confindex_[v(0)]]);
    right_prods.push_back(W_[(N_ - 1) % symperiod_][confindex_[v(N_ - 1)]]);
    for (int site = 1; site < N_ - 1; site++) {
      left_prods.push_back(
          prod(left_prods[site - 1], W_[site % symperiod_][confindex_[v(site)]]));
      right_prods.push_back(
          prod(W_[(N_ - 1 - site) % symperiod_][confindex_[v(N_ - 1 - site)]],
               right_prods[site - 1]));
    }
    left_prods.push_back(prod(left_prods[N_ - 2],
                              W_[(N_ - 1) % symperiod_][confindex_[v(N_ - 1)]]));
    right_prods.push_back(prod(W_[0][confindex_[v(0)]], right_prods[N_ - 2]));
  
    der.segment(confindex_[v(0)] * Dsq_, Dsq_) +=
        Eigen::Map<VectorType>(right_prods[N_ - 2].transpose().data(), Dsq_);
    for (int site = 1; site < N_ - 1; site++) {
      temp_product = prod(right_prods[N_ - site - 2], left_prods[site - 1]);
      der.segment((d_ * (site % symperiod_) + confindex_[v(site)]) * Dsq_,
                  Dsq_) +=
          Eigen::Map<VectorType>(temp_product.transpose().data(), Dsq_);
    }
    der.segment((d_ * ((N_ - 1) % symperiod_) + confindex_[v(N_ - 1)]) * Dsq_,
                Dsq_) +=
        Eigen::Map<VectorType>(left_prods[N_ - 2].transpose().data(), Dsq_);
  
    return der / trace(left_prods[N_ - 1]);
  }
  
  // Json functions
  
  void MPSPeriodic::Save(const std::string &filename) const {
    json state;
    state["Name"] = "MPSperiodic";
    state["Length"] = N_;
    state["BondDim"] = D_;
    state["PhysDim"] = d_;
    state["Diagonal"] = is_diag_;
    state["SymmetryPeriod"] = symperiod_;
    for (int i = 0; i < symperiod_; i++) {
      for (int k = 0; k < d_; k++) {
        state["W" + std::to_string(d_ * i + k)] = W_[i][k];
      }
    }
    WriteJsonToFile(state, filename);
  }
  
  void MPSPeriodic::Load(const std::string &filename) {
    auto const pars = ReadJsonFromFile(filename);
    if (pars.at("Name") != "MPSperiodic") {
      throw InvalidInputError("Error while constructing MPS from Json input");
    }
  
    if (FieldExists(pars, "Length")) {
      N_ = pars["Length"];
    }
    if (N_ != GetHilbert().Size()) {
      throw InvalidInputError(
          "Number of spins is incompatible with given Hilbert space");
    }
  
    if (FieldExists(pars, "PhysDim")) {
      d_ = pars["PhysDim"];
    }
    if (d_ != GetHilbert().LocalSize()) {
      throw InvalidInputError(
          "Number of spins is incompatible with given Hilbert space");
    }
  
    if (FieldExists(pars, "BondDim")) {
      D_ = pars["BondDim"];
    } else {
      throw InvalidInputError("Unspecified bond dimension");
    }
  
    if (FieldExists(pars, "Diagonal")) {
      is_diag_ = pars["Diagonal"];
    } else {
      throw InvalidInputError("Unspecified whether the MPS is diagonal");
    }
  
    if (FieldExists(pars, "SymmetryPeriod")) {
      symperiod_ = pars["SymmetryPeriod"];
    } else {
      // Default is symperiod = N, resp. no translational symmetry
      symperiod_ = N_;
    }
  
    Init();
  
    // Loading parameters, if defined in the input
    for (int i = 0; i < symperiod_; i++) {
      for (int k = 0; k < d_; k++) {
        if (FieldExists(pars, "W" + std::to_string(d_ * i + k))) {
          W_[i][k] = pars["W" + std::to_string(d_ * i + k)];
        }
      }
    }
  }
  
  bool MPSPeriodic::IsHolomorphic() const noexcept { return true; }
  
  }  // namespace netket
  