/*!@file
 * @copyright This code is licensed under the 3-clause BSD license.
 *   Copyright ETH Zurich, Laboratory of Physical Chemistry, Reiher Group.
 *   See LICENSE.txt for details.
 */

#include "boost/program_options.hpp"

#include "Molassembler/Shapes/ContinuousMeasures.h"

#include "Molassembler/Shapes/Data.h"
#include "Molassembler/Shapes/PropertyCaching.h"

#include "boost/math/tools/minima.hpp"
#include "Molassembler/Temple/constexpr/Numeric.h"
#include "Molassembler/Temple/Adaptors/AllPairs.h"
#include "Molassembler/Temple/Adaptors/Iota.h"
#include "Molassembler/Temple/Functional.h"
#include "Molassembler/Temple/Stringify.h"
#include "Molassembler/Temple/Random.h"
#include "Molassembler/Temple/constexpr/Jsf.h"
#include "Molassembler/Temple/Permutations.h"
#include "Molassembler/Temple/Loops.h"

#include <Eigen/SparseCore>
#include <Eigen/Eigenvalues>
#include <fstream>
#include <iostream>
#include <iomanip>
#include <memory>
#include <chrono>

/* Summary
 *
 * This file contains a number of methods trying to cheapen finding the correct
 * vertex mapping to evaluate the continuous shape measure.
 *
 * Note that the problem does have some rotational redundancy depending on the
 * number of rotations that the shape has. But exploiting that gets you maybe
 * (N - 1)!. Crasser solutions are necessary.
 */

using namespace Scine;
using namespace Molassembler;
using namespace Shapes;
using namespace Continuous;

constexpr unsigned factorial(unsigned x) {
  if(x <= 1) {
    return 1;
  }

  return x * factorial(x - 1);
}

template<std::size_t ... Inds>
constexpr auto makeFactorials(std::index_sequence<Inds ...> /* inds */) {
  return std::array<unsigned, sizeof...(Inds)> {
    factorial(Inds + 1)...
  };
}

constexpr auto facs = makeFactorials(std::make_index_sequence<14> {});

Continuous::PositionCollection addOrigin(const Continuous::PositionCollection& vs) {
  const unsigned N = vs.cols();
  Continuous::PositionCollection positions(3, N + 1);
  for(unsigned i = 0; i < N; ++i) {
    positions.col(i) = vs.col(i);
  }

  // Add origin point explicitly to consideration
  positions.col(N) = Eigen::Vector3d::Zero(3);
  return positions;
}

void distort(Eigen::Ref<Continuous::PositionCollection> positions, const double distortionNorm = 0.01) {
  const unsigned N = positions.cols();
  for(unsigned i = 0; i < N; ++i) {
    positions.col(i) += distortionNorm * Eigen::Vector3d::Random().normalized();
  }
}

template<typename Derived>
bool centroidIsZero(const Eigen::MatrixBase<Derived>& a) {
  assert(a.rows() == 3);
  return (a.rowwise().sum() / a.cols()).squaredNorm() < 1e-8;
}

template<typename DerivedA, typename DerivedB>
Eigen::Matrix3d fitQuaternion(const Eigen::MatrixBase<DerivedA>& stator, const Eigen::MatrixBase<DerivedB>& rotor) {
  assert(centroidIsZero(stator));
  assert(centroidIsZero(rotor));

  Eigen::Matrix4d b = Eigen::Matrix4d::Zero();
  // generate decomposable matrix per atom and add them
  for (int i = 0; i < rotor.cols(); i++) {
    auto& rotorCol = rotor.col(i);
    auto& statorCol = stator.col(i);

    Eigen::Matrix4d a = Eigen::Matrix4d::Zero();
    a.block<1, 3>(0, 1) = (rotorCol - statorCol).transpose();
    a.block<3, 1>(1, 0) = statorCol - rotorCol;
    a.block<3, 3>(1, 1) = Eigen::Matrix3d::Identity().rowwise().cross(statorCol + rotorCol);
    b += a.transpose() * a;
  }

  // Decompose b
  Eigen::SelfAdjointEigenSolver<Eigen::Matrix4d> eigensolver(b);

  // Do not allow improper rotation
  const Eigen::Vector4d& q = eigensolver.eigenvectors().col(0);
  return Eigen::Quaterniond(q[0], q[1], q[2], q[3]).toRotationMatrix();
}

using Parameters = std::vector<Vertex>;

struct Energy {
  const PositionCollection& referencePositions;
  const PositionCollection& shapePositions;

  double operator() (const Parameters& p) const {
    const unsigned N = shapePositions.cols();

    static PositionCollection permuted(3, N);
    for(unsigned i = 0; i < N; ++i) {
      permuted.col(i) = shapePositions.col(p[i]);
    }

    auto rotationMatrix = fitQuaternion(referencePositions, permuted);
    permuted = rotationMatrix * permuted;

    return (referencePositions - permuted).colwise().squaredNorm().sum();

    // // Minimize over isotropic scaling factor
    // auto scalingMinimizationResult = boost::math::tools::brent_find_minima(
    //   [&](const double scaling) -> double {
    //     return (referencePositions - scaling * permuted).colwise().squaredNorm().sum();
    //   },
    //   0.5,
    //   1.1,
    //   std::numeric_limits<double>::digits
    // );

    // return scalingMinimizationResult.second;
  }
};

struct OldShapeAlgorithm {
  virtual ~OldShapeAlgorithm() = default;

  static PositionCollection shapeCoordinates(const Shape shape) {
    return normalize(addOrigin(coordinates(shape)));
  }

  virtual double shape(const PositionCollection& positions, Shape shape) = 0;
  virtual std::string name() const = 0;
};

/**
 * @brief Abstract base class for algorithms to conform to
 */
struct ShapeAlgorithm {
  virtual ~ShapeAlgorithm() = default;

  using Permutation = std::vector<Vertex>;
  using ResultPair = std::pair<double, Permutation>;

  static Permutation invert(const Permutation& p) {
    const unsigned P = p.size();
    Permutation inverse(P);
    for(unsigned i = 0; i < P; ++i) {
      inverse.at(p.at(i)) = i;
    }
    return inverse;
  }

  /* A forwards permutation maps from positions to shape indices,
   * a backwards permutation maps from shape indices to positions
   *
   * Only backwards permutations have the (added) centroid at the back, and
   * hence their [O, S] = [0, N - 1] = [0, N) ranges can be rotated.
   */
  static std::set<ShapeAlgorithm::Permutation> generateRotations(
    const Permutation& forwardsPermutation,
    const Shape shape
  ) {
    auto backwards = ShapeAlgorithm::invert(forwardsPermutation);
    Vertex centroid = backwards.back();
    backwards.pop_back();
    auto rotations = Properties::generateAllRotations(shape, backwards);

    std::set<ShapeAlgorithm::Permutation> forwardsPermutationRotations;
    for(auto rotation : rotations) {
      rotation.push_back(centroid);
      forwardsPermutationRotations.emplace(
        ShapeAlgorithm::invert(rotation)
      );
    }
    return forwardsPermutationRotations;
  }

  static PositionCollection shapeCoordinates(const Shape shape) {
    return normalize(addOrigin(coordinates(shape)));
  }

  virtual ResultPair shape(
    const PositionCollection& positions,
    Shape shape,
    const Permutation& correctPermutation
  ) = 0;
  virtual std::string name() const = 0;
};

/**
 * @brief Reference continuous shape measure calculation method
 *
 * Complexity: Theta(N!)
 */
struct AllPermutations {
  ShapeAlgorithm::ResultPair shape(const PositionCollection& positions, Shape shape) {
    Energy energy {positions, ShapeAlgorithm::shapeCoordinates(shape)};
    ShapeAlgorithm::Permutation permutation = Temple::iota<Vertex>(Vertex(positions.cols()));
    ShapeAlgorithm::ResultPair minimal {std::numeric_limits<double>::max(), {}};
    do {
      double permutationEnergy = energy(permutation);

      if(permutationEnergy < minimal.first) {
        minimal = {permutationEnergy, permutation};
      }
    } while(Temple::next_permutation(permutation));

    return minimal;
  }

  std::string name() const {
    return "Reference";
  }
};

struct Cooling {
  // This works best for annealing and tunneling
  static double linear(const unsigned step, const unsigned steps) {
    return 1.0 - (static_cast<double>(step) / steps);
  }

  // This doesn't work well at all for annealing or tunneling, but it's here
  static double exponential(const unsigned step, const unsigned steps) {
    const double alpha = std::exp(std::log(1e-2) / steps);
    return std::pow(alpha, step);
  }
};

/**
 * @brief Simulated annealing method
 *
 * With well tuned temperature, this works pretty well, but it's stochastic in
 * nature and its good results are due to tracking the lowest energy state
 * visited instead of using the final annealing state. The energy surface for
 * small distortion norms is super jagged and the "correct" state has no
 * low-energy neighbors, making this type of approach problematic. Behaves
 * better for high distortion norms as there state space can be better sampled.
 *
 * Complexity: Needs at least O((N - 2)!) steps to explore enough state space
 * to encounter the minimum.
 */
struct Anneal final : OldShapeAlgorithm {
  using Permutation = std::vector<Vertex>;

  Temple::JSF64 prng;
  std::ofstream trace {"anneal_trace.csv"};

  Anneal() {
    prng.seed(0);
  }

  static double acceptanceProbability(
    const double energy,
    const double prospectiveEnergy,
    const double temperature
  ) {
    if(prospectiveEnergy < energy) {
      return 1.0;
    }

    return std::exp(-(prospectiveEnergy - energy) / temperature);
  }

  static double temperature(const unsigned step, const unsigned steps) {
    return Cooling::linear(step, steps);
  }

  //! Assumes temperature in 0, 1
  Permutation generateMove(Permutation state, const double temperature) {
    const unsigned N = state.size();

    const unsigned nSwaps = 1 + std::round(temperature * Temple::Random::getSingle<unsigned>(0, N - 1, prng));
    for(unsigned s = 0; s < nSwaps; ++s) {
      const unsigned i = Temple::Random::getSingle<unsigned>(0, N - 2, prng);
      std::swap(state.at(i), state.at(i + 1));
    }

    return state;
  }

  double shape(const PositionCollection& positions, Shape shape) final {
    const unsigned steps = facs.at(positions.cols() - 2);
    Energy energy {positions, shapeCoordinates(shape)};

    auto parameters = Temple::iota<Vertex>(positions.cols());
    decltype(parameters) prospectiveParameters;

    double temperatureMultiplier = Temple::accumulate(
      Temple::Adaptors::range(10),
      0.0,
      [&](const double carry, unsigned /* i */) -> double {
        Temple::Random::shuffle(parameters, prng);
        return carry + energy(parameters);
      }
    ) / 10;

    double currentEnergy = energy(parameters);
    double minimalEnergy = currentEnergy;
    for(unsigned step = 0; step < steps; ++step) {
      const double currentTemperature = temperature(step, steps);
      prospectiveParameters = generateMove(parameters, currentTemperature);
      const double prospectiveEnergy = energy(prospectiveParameters);
      const double p = acceptanceProbability(currentEnergy, prospectiveEnergy, temperatureMultiplier * currentTemperature);

      trace << step << ", " << currentEnergy << ", " << temperatureMultiplier * currentTemperature << ", " << prospectiveEnergy << ", " << p << "\n";

      if(p >= Temple::Random::getSingle(0.0, 1.0, prng)) {
        std::swap(parameters, prospectiveParameters);
        currentEnergy = prospectiveEnergy;
      }

      minimalEnergy = std::min(minimalEnergy, prospectiveEnergy);
    }

    return minimalEnergy;
  }

  std::string name() const final {
    return "Annealing";
  }
};

/**
 * @brief Stochastic tunneling method
 *
 * This works by transforming the energy function to level out minima that have
 * higher function values than the lowest found. This then should help losing
 * time in unhelpful minima.
 *
 * Works pretty well even when the temperature is not well tuned. You have to
 * adjust the gamma parameter for the particular problem you have, and the
 * value we have now works well for CShMs.
 *
 * But now that simulated annealing is well tuned and since there aren't really
 * and multi-state minima to get stuck in in this energy surface for small
 * distortions, this doesn't perform any better.
 *
 * The move generator for this and for simulated annealing generate multi-swap
 * moves for high temperatures and increase locality as the temperature is
 * lowered, generating only single adjacent swap moves.
 *
 * Complexity: Needs at least O((N - 2)!) steps to explore enough state space
 * to encounter the minimum.
 */
struct Tunnel final : OldShapeAlgorithm {
  using Permutation = std::vector<Vertex>;

  Temple::JSF64 prng;
  std::ofstream trace {"tunnel_trace.csv"};
  double gamma = 0.1;

  Tunnel() {
    prng.seed(0);
  }

  static double acceptanceProbability(
    const double energy,
    const double prospectiveEnergy,
    const double temperature
  ) {
    if(prospectiveEnergy < energy) {
      return 1.0;
    }

    return std::exp(-(prospectiveEnergy - energy) / temperature);
  }

  //! Assumes parameters are untransformed!
  double energyTransformation(const double energy, const double lowestEnergy) const {
    return 1.0 - std::exp(-gamma * (energy - lowestEnergy));
  }

  static double temperature(const unsigned step, const unsigned steps) {
    return Cooling::linear(step, steps);
  }

  //! Assumes temperature in 0, 1
  Permutation generateMove(Permutation state, const double temperature) {
    const unsigned N = state.size();

    const unsigned nSwaps = 1 + std::round(temperature * Temple::Random::getSingle<unsigned>(0, N - 1, prng));
    for(unsigned s = 0; s < nSwaps; ++s) {
      const unsigned i = Temple::Random::getSingle<unsigned>(0, N - 2, prng);
      std::swap(state.at(i), state.at(i + 1));
    }

    return state;
  }

  double shape(const PositionCollection& positions, Shape shape) final {
    const unsigned steps = facs.at(positions.cols() - 2);
    Energy energy {positions, shapeCoordinates(shape)};

    auto parameters = Temple::iota<Vertex>(positions.cols());
    decltype(parameters) prospectiveParameters;

    double currentEnergy = energy(parameters);
    double lowestEnergy = currentEnergy;
    double transformedCurrentEnergy = energyTransformation(currentEnergy, lowestEnergy);

    for(unsigned step = 0; step < steps; ++step) {
      const double currentTemperature = temperature(step, steps);
      prospectiveParameters = generateMove(parameters, currentTemperature);
      const double prospectiveEnergy = energy(prospectiveParameters);
      const double transformedProspectiveEnergy = energyTransformation(prospectiveEnergy, lowestEnergy);
      const double p = acceptanceProbability(transformedCurrentEnergy, transformedProspectiveEnergy, 0.7 * currentTemperature);

      trace << step << ", " << currentEnergy << ", " << 0.7 * currentTemperature << ", " << prospectiveEnergy << ", " << p << "\n";

      if(p >= Temple::Random::getSingle(0.0, 1.0, prng)) {
        std::swap(parameters, prospectiveParameters);
        currentEnergy = prospectiveEnergy;
        transformedCurrentEnergy = transformedProspectiveEnergy;
      }

      if(prospectiveEnergy < lowestEnergy) {
        lowestEnergy = prospectiveEnergy;
        transformedCurrentEnergy = energyTransformation(currentEnergy, lowestEnergy);
      }
    }

    return lowestEnergy;
  }

  std::string name() const final {
    return "Tunneling";
  }
};

template<typename T, std::size_t N>
struct CircularBuffer {
  void insert(T value) {
    if(size_ < N) {
      buffer_[size_] = std::move(value);
      ++size_;
    } else {
      buffer_[start_] = std::move(value);
      start_ = (start_ + 1) % N;
    }
  }

  T min() const {
    assert(size_ > 0);
    return *std::min_element(
      std::begin(buffer_),
      std::end(buffer_)
    );
  }

  double average() const {
    return static_cast<double>(
      std::accumulate(
        std::begin(buffer_),
        std::end(buffer_),
        T {0},
        std::plus<>()
      )
    ) / size_;
  }

  double variance(const double average) const {
    return static_cast<double>(
      std::accumulate(
        std::begin(buffer_),
        std::end(buffer_),
        T {0},
        [&](const double carry, const double value) -> double {
          return carry + std::pow(value - average, 2);
        }
      )
    ) / size_;
  }

  void clear() {
    size_ = 0;
    start_ = 0;
  }

  std::size_t size() const {
    return size_;
  }

  std::array<T, N> buffer_;
  std::size_t start_ = 0;
  std::size_t size_ = 0;
};

/**
 * @brief An attempt at thermodynamic simulated annealing
 *
 * In thermodynamic simulated annealing, you try to exploit statistical
 * mechanics / thermodynamics to drive your cooling schedule optimally.
 *
 * However, I couldn't ever get this to work. The temperature adjustments are
 * always too large and no amount of messing about ever got them right.
 *
 * Complexity: ???
 */
struct ThermodynamicAnneal final : OldShapeAlgorithm {
  Temple::JSF64 prng;
  using Permutation = std::vector<Vertex>;

  std::unordered_map<unsigned, unsigned> stateIndexReduction;
  std::vector<double> energies;
  Eigen::SparseMatrix<unsigned> Q;
  CircularBuffer<double, 1000> lastEnergies;
  std::ofstream trace {"thermo_trace.csv"};

  ThermodynamicAnneal() {
    prng.seed(0);
  }

  static double acceptanceProbability(
    const double energy,
    const double prospectiveEnergy,
    const double temperature
  ) {
    if(prospectiveEnergy < energy) {
      return 1.0;
    }

    return std::exp(-(prospectiveEnergy - energy) / temperature);
  }

  //! Assumes temperature in 0, 100
  Permutation generateMove(Permutation state, const double temperature) {
    const unsigned N = state.size();

    const unsigned nSwaps = 1 + std::round((temperature / 10) * Temple::Random::getSingle<unsigned>(0, N - 1, prng));
    for(unsigned s = 0; s < nSwaps; ++s) {
      const unsigned i = Temple::Random::getSingle<unsigned>(0, N - 2, prng);
      std::swap(state.at(i), state.at(i + 1));
    }

    return state;
  }

  double partitionFunction(const double temperature) const {
    return Temple::accumulate(
      energies,
      0.0,
      [&](const double carry, const double energy) -> double {
        return carry + std::exp(-energy / temperature);
      }
    );
  }

  template<typename F>
  static double centralDifference(F&& f, const double x, const double h) {
    return (
      f(x + h) - f(x - h)
    ) / (2 * h);
  }

  double averageEnergy(const double temperature) const {
    // Estimate dln Z/dT by central finite difference
    return temperature * temperature * centralDifference(
      [&](const double x) { return std::log(partitionFunction(x)); },
      temperature,
      1e-4
    );
  }

  double heatCapacity(const double temperature) const {
    return centralDifference(
      [&](const double x) { return averageEnergy(x); },
      temperature,
      1e-4
    );
  }

  double relaxationTime(const double temperature) const {
    // Q is strictly lower triangular (i.e. always access i > j)
    Eigen::SparseMatrix<double> G(Q.rows(), Q.cols());
    for(Eigen::Index k = 0; k < Q.outerSize(); ++k) {
      // Sum up all values in Q's column
      double sum = 0;
      for(Eigen::SparseMatrix<unsigned>::InnerIterator it(Q,k); it; ++it) {
        sum += it.value();
      }

      for(Eigen::SparseMatrix<unsigned>::InnerIterator it(Q,k); it; ++it) {
        if(it.row() != it.col()) {
          const double energyDiff = energies.at(it.row()) - energies.at(it.col());
          if(energyDiff > 0) {
            G.coeffRef(it.row(), it.col()) = it.value() * std::exp(-energyDiff / temperature);
          } else {
            G.coeffRef(it.row(), it.col()) = it.value();
          }
        }
      }

      // Diagonal entries of G are 1 - other column's entries
      sum = 0;
      for(Eigen::SparseMatrix<double>::InnerIterator it(G,k); it; ++it) {
        sum += it.value();
      }
      G.coeffRef(k, k) = 1 - sum;
    }

    Eigen::SelfAdjointEigenSolver<Eigen::SparseMatrix<double>> solver(G);
    return - 1.0 / solver.eigenvalues()(1);
  }

  double updateTemperature(const double temperature, double minimalEnergy) const {
    const double hundredEnergiesAverage = lastEnergies.average();
    const double hundredEnergiesVariance = lastEnergies.variance(hundredEnergiesAverage);
    const double thermodynamicSpeed = (hundredEnergiesAverage - minimalEnergy) / hundredEnergiesVariance;

    const double epsilon = relaxationTime(temperature);
    const double C = heatCapacity(temperature);

    double Theta = 1 + temperature * centralDifference([&](double x) { return heatCapacity(x);}, temperature, 1e-4) / (2 * C);
    auto calculateDelta = [&](const double delta) {
      double correction = std::sqrt(1 + (Theta * epsilon * delta) / temperature);
      return - thermodynamicSpeed * temperature / (
        epsilon * std::sqrt(C) * correction
      );
    };

    // NOTE: this can diverge too :(
    double delta = -1.0;
    for(unsigned i = 0; i < 10; ++i) {
      delta = calculateDelta(delta);
    }

    // NOTE 1e-5 is abitrary, the updates are always too large!
    return std::max(0.0, temperature + 1e-5 * delta);
  }

  double shape(const PositionCollection& positions, Shape shape) final {
    const unsigned steps = 5e5;

    stateIndexReduction.clear();
    energies.clear();
    energies.reserve(100);
    Q.resize(0,0);
    Q.reserve(100);
    lastEnergies.clear();

    const unsigned N = positions.cols();
    Energy energy {positions, shapeCoordinates(shape)};

    auto permutation = Temple::iota<Vertex>(N);
    decltype(permutation) prospectiveMove;
    double currentEnergy = energy(permutation);
    double minimalEnergy = currentEnergy;

    unsigned iop = Temple::permutationIndex(permutation);
    stateIndexReduction.emplace(iop, 0);
    energies.push_back(currentEnergy);
    unsigned currentStateIndex = 0;

    double temperature = 10.0;

    /* First 1000 steps of annealing without changing temperature to collect
     * statistics we can use to update the temperature
     */
    unsigned step = 0;
    for(; step < 1000; ++step) {
      prospectiveMove = generateMove(permutation, temperature);
      const double prospectiveEnergy = energy(prospectiveMove);
      const double p = acceptanceProbability(currentEnergy, prospectiveEnergy, temperature);

      iop = Temple::permutationIndex(prospectiveMove);
      unsigned prospectiveStateIndex;
      auto stateIndexFindIter = stateIndexReduction.find(iop);
      if(stateIndexFindIter == std::end(stateIndexReduction)) {
        prospectiveStateIndex = stateIndexReduction.size();
        stateIndexReduction.emplace(iop, prospectiveStateIndex);
        energies.push_back(prospectiveEnergy);
        Q.conservativeResize(energies.size(), energies.size());
        Q.insert(prospectiveStateIndex, currentStateIndex) = 1;
      } else {
        prospectiveStateIndex = stateIndexFindIter->second;
        unsigned col = prospectiveStateIndex;
        unsigned row = currentStateIndex;
        if(col > row) {
          std::swap(col, row);
        }
        Q.coeffRef(row, col) = Q.coeff(row, col) + 1;
      }

      trace << step << ", " << currentEnergy << ", " << temperature << ", " << prospectiveEnergy << ", " << p << "\n";

      if(p >= Temple::Random::getSingle(0.0, 1.0, prng)) {
        std::swap(permutation, prospectiveMove);
        currentEnergy = prospectiveEnergy;
        currentStateIndex = prospectiveStateIndex;
      }

      minimalEnergy = std::min(minimalEnergy, prospectiveEnergy);
      lastEnergies.insert(currentEnergy);
    }

    /* Now let temperature freely vary according to update formula */
    while(temperature > 0.01 && step < steps) {
      temperature = updateTemperature(temperature, minimalEnergy);
      prospectiveMove = generateMove(permutation, temperature);
      const double prospectiveEnergy = energy(prospectiveMove);
      const double p = acceptanceProbability(currentEnergy, prospectiveEnergy, temperature);

      /* Update tracking state */
      iop = Temple::permutationIndex(prospectiveMove);
      unsigned prospectiveStateIndex;
      auto stateIndexFindIter = stateIndexReduction.find(iop);
      if(stateIndexFindIter == std::end(stateIndexReduction)) {
        prospectiveStateIndex = stateIndexReduction.size();
        stateIndexReduction.emplace(iop, prospectiveStateIndex);
        energies.push_back(prospectiveEnergy);
        Q.conservativeResize(energies.size(), energies.size());
        Q.insert(prospectiveStateIndex, currentStateIndex) = 1;
      } else {
        prospectiveStateIndex = stateIndexFindIter->second;
        unsigned col = prospectiveStateIndex;
        unsigned row = currentStateIndex;
        if(col > row) {
          std::swap(col, row);
        }
        Q.coeffRef(row, col) = Q.coeff(row, col) + 1;
      }

      trace << step << ", " << currentEnergy << ", " << temperature << ", " << prospectiveEnergy << ", " << p << "\n";

      // Conditionally accept the move
      if(p >= Temple::Random::getSingle(0.0, 1.0, prng)) {
        std::swap(permutation, prospectiveMove);
        currentEnergy = prospectiveEnergy;
        currentStateIndex = prospectiveStateIndex;
      }

      minimalEnergy = std::min(minimalEnergy, prospectiveEnergy);
      lastEnergies.insert(currentEnergy);
      ++step;
    }

    return minimalEnergy;
  }

  std::string name() const final {
    return "Thermodyn.";
  }
};

/**
 * @brief Fixed-step number greedy minimization with shuffling
 *
 * This just greedily minimizes the permutation and shuffles when it's at the
 * minimum. Unreliable.
 *
 * Complexity: Needs at least Theta(N - 2)! steps to discover enough state
 * space that it might find the minimum.
 */
struct Greedy final : OldShapeAlgorithm {
  Temple::JSF64 prng;

  double shape(const PositionCollection& positions, Shape shape) final {
    const unsigned N = positions.cols();
    const unsigned steps = 4 * facs.at(N - 2);

    Energy energy {positions, shapeCoordinates(shape)};
    auto permutation = Temple::iota<Vertex>(positions.cols());
    double lowestEnergy = std::numeric_limits<double>::max();

    bool foundBetterPermutation;
    for(unsigned m = 0; m < steps; ++m) {
      Temple::Random::shuffle(permutation, prng);
      double minimizationEnergy = energy(permutation);

      do {
        foundBetterPermutation = false;

        // Single swaps
        for(unsigned i = 0; i < N - 1 && m < steps; ++i) {
          for(unsigned j = i + 1; j < N && m < steps; ++j) {
            std::swap(permutation.at(i), permutation.at(j));

            double value = energy(permutation);
            ++m;
            if(value < minimizationEnergy) {
              minimizationEnergy = value;
              foundBetterPermutation = true;
              break;
            }

            // Swap back
            std::swap(permutation.at(i), permutation.at(j));
          }
        }

        // Propose double-swaps
        for(unsigned i = 0; i < N && !foundBetterPermutation; ++i) {
          for(unsigned j = i + 1; j < N && !foundBetterPermutation; ++j) {
            for(unsigned k = 0; k < N && k != i && !foundBetterPermutation; ++k) {
              for(unsigned l = k + 1; l < N && l != j; ++l) {
                std::swap(permutation.at(i), permutation.at(j));
                std::swap(permutation.at(k), permutation.at(l));

                // Calculate the value
                double value = energy(permutation);
                if(value < minimizationEnergy) {
                  minimizationEnergy = value;
                  foundBetterPermutation = true;
                  break;
                }
                ++m;
                if(m >= steps) {
                  return std::min(lowestEnergy, minimizationEnergy);
                }

                // Swap back
                std::swap(permutation.at(k), permutation.at(l));
                std::swap(permutation.at(i), permutation.at(j));
              }
            }
          }
        }

      } while(foundBetterPermutation);

      lowestEnergy = std::min(lowestEnergy, minimizationEnergy);
    }

    return lowestEnergy;
  }

  std::string name() const final {
    return "Greedy";
  }
};

/**
 * @brief Fixed steps steepest descent minimizer with shuffles
 *
 * Finds best swap to reduce permutation the most at each position and
 * minimizes. Shuffles when it's at a minimum. Unreliable.
 *
 * Complexity: Needs at least Theta(N - 2)! steps to discover enough state
 * space that it might find the minimum.
 */
struct SteepestDescent final : OldShapeAlgorithm {
  Temple::JSF64 prng;

  double shape(const PositionCollection& positions, Shape shape) final {
    const unsigned N = positions.cols();
    const unsigned steps = 4 * facs.at(N - 2);

    Energy energy {positions, shapeCoordinates(shape)};
    auto bestCandidate = Temple::iota<Vertex>(positions.cols());
    auto permutation = bestCandidate;
    double lowestEnergy = std::numeric_limits<double>::max();

    bool foundBetterPermutation;
    for(unsigned m = 0; m < steps; ++m) {
      Temple::Random::shuffle(bestCandidate, prng);
      double minimizationEnergy = energy(bestCandidate);
      lowestEnergy = std::min(lowestEnergy, minimizationEnergy);

      do {
        permutation = bestCandidate;
        foundBetterPermutation = false;

        // Propose double-swaps
        for(unsigned i = 0; i < N && !foundBetterPermutation; ++i) {
          for(unsigned j = i + 1; j < N && !foundBetterPermutation; ++j) {
            for(unsigned k = 0; k < N && !foundBetterPermutation; ++k) {
              for(unsigned l = k + 1; l < N; ++l) {
                std::swap(permutation.at(i), permutation.at(j));
                std::swap(permutation.at(k), permutation.at(l));

                // Calculate the value
                double value = energy(permutation);
                ++m;
                if(value < minimizationEnergy) {
                  minimizationEnergy = value;
                  foundBetterPermutation = true;
                  bestCandidate = permutation;
                }
                if(m >= steps) {
                  return std::min(lowestEnergy, minimizationEnergy);
                }

                // Swap back
                std::swap(permutation.at(k), permutation.at(l));
                std::swap(permutation.at(i), permutation.at(j));
              }
            }
          }
        }
      } while(foundBetterPermutation);

      lowestEnergy = std::min(lowestEnergy, minimizationEnergy);
    }

    return lowestEnergy;
  }

  std::string name() const final {
    return "Steepest";
  }
};

/**
 * @brief Aligns four atoms, then greedily zips the rest
 *
 * For each tuple of four atoms, aligns the positions, then greedily chooses
 * the best next sequence alignment.
 *
 * Really good for small distortions (<= 0.5). Suffers from its greediness
 * afterwards. Some branching might do it good.
 *
 * Complexity: Theta(N! / (N - 5)!) quaternion fits
 *
 * Potentially there is another optimization possible that could reduce the
 * scaling but there is a tradeoff involved that might counteract the formal
 * complexity decrease: Principally, spatial rotations of index mappings from
 * the underlying shape mappings need not be repeated. However, the question is
 * how to store them efficiently.
 */
struct AlignFive final : ShapeAlgorithm {
  using M = Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>;
  using IncrementalPermutation = std::unordered_map<unsigned, unsigned>;

  static Eigen::Matrix3d fitQuaternion(
    const PositionCollection& stator,
    const PositionCollection& rotor,
    const IncrementalPermutation& p
  ) {
    assert(centroidIsZero(stator));
    assert(centroidIsZero(rotor));

    Eigen::Matrix4d b = Eigen::Matrix4d::Zero();
    // generate decomposable matrix per atom and add them
    for(auto& iterPair : p) {
      auto& statorCol = stator.col(iterPair.first);
      auto& rotorCol = rotor.col(iterPair.second);

      Eigen::Matrix4d a = Eigen::Matrix4d::Zero();
      a.block<1, 3>(0, 1) = (rotorCol - statorCol).transpose();
      a.block<3, 1>(1, 0) = statorCol - rotorCol;
      a.block<3, 3>(1, 1) = Eigen::Matrix3d::Identity().rowwise().cross(statorCol + rotorCol);
      b += a.transpose() * a;
    }

    // Decompose b
    Eigen::SelfAdjointEigenSolver<Eigen::Matrix4d> eigensolver(b);

    // Do not allow improper rotation
    const Eigen::Vector4d& q = eigensolver.eigenvectors().col(0);
    return Eigen::Quaterniond(q[0], q[1], q[2], q[3]).toRotationMatrix();
  }

  static Permutation flatten(const IncrementalPermutation& p) {
    const unsigned P = p.size();
    Permutation flat(P);
    for(unsigned i = 0; i < P; ++i) {
      flat.at(i) = p.at(i);
    }
    return flat;
  }

  ResultPair narrow(
    const PositionCollection& stator,
    const PositionCollection& rotor,
    const std::set<Permutation>& correctRotations,
    bool inCorrectBranch,
    std::unordered_map<unsigned, unsigned> permutation,
    std::vector<unsigned> freeLeftVertices,
    std::vector<unsigned> freeRightVertices
  ) {
    const unsigned N = stator.cols();

    const unsigned V = freeLeftVertices.size();
    Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> costs (V, V);
    for(unsigned i = 0; i < V; ++i) {
      for(unsigned j = 0; j < V; ++j) {
        costs(i, j) = (
          stator.col(freeLeftVertices.at(i))
          - rotor.col(freeRightVertices.at(j))
        ).squaredNorm();
      }
    }

    // Maps freeLeftVertices onto freeRightVertices
    auto subPermutation = Temple::iota<unsigned>(V);

    auto isCorrectBranch = [&](const auto& s) {
      return Temple::any_of(
        correctRotations,
        [&](const auto& rot) -> bool {
          // Rotation has to match the current permutation
          if(!Temple::all_of(
            permutation,
            [&](const unsigned left, const unsigned right) -> bool {
              return rot.at(left) == right;
            }
          )) {
            return false;
          }

          // Rotation has to match the sub-permutation too
          for(unsigned i = 0; i < V; ++i) {
            unsigned left = freeLeftVertices.at(i);
            unsigned right = freeRightVertices.at(s.at(i));
            if(rot.at(left) != right) {
              return false;
            }
          }

          return true;
        }
      );
    };

    decltype(subPermutation) bestPermutation;
    double minimalCost = std::numeric_limits<double>::max();
    do {
      double cost = 0.0;
      for(unsigned i = 0; i < V; ++i) {
        cost += costs(i, subPermutation.at(i));
      }

      if(cost < minimalCost) {
        if(inCorrectBranch && !bestPermutation.empty() && isCorrectBranch(bestPermutation) && !isCorrectBranch(subPermutation)) {
          std::cerr << "Excluding correct branch, decreasing cost from " << minimalCost << " to " << cost << " with non-optimum sub-permutation\n";
        }

        minimalCost = cost;
        bestPermutation = subPermutation;
      } else if(inCorrectBranch && isCorrectBranch(subPermutation)) {
        std::cerr << "Excluding correct branch due to cost " << cost << " > " << minimalCost << "\n";
      }
    } while(Temple::next_permutation(subPermutation));

    // Fuse permutation and subpermutation
    for(unsigned i = 0; i < V; ++i) {
      permutation.emplace(freeLeftVertices.at(i), freeRightVertices.at(bestPermutation.at(i)));
    }

    assert(permutation.size() == N);

    auto R = fitQuaternion(stator, rotor, permutation);
    auto rotated = R * rotor;

    const double energy = Temple::accumulate(
      Temple::Adaptors::range(N),
      0.0,
      [&](const double carry, const unsigned i) -> double {
        return carry + (
          stator.col(i) - rotated.col(permutation.at(i))
        ).squaredNorm();
      }
    );

    return ResultPair {energy, flatten(permutation)};
  }

  ResultPair shape(
    const PositionCollection& positions,
    Shape shape,
    const Permutation& correctPermutation
  ) final {
    const unsigned N = positions.cols();
    if(N <= 4) {
      return AllPermutations {}.shape(positions, shape);
    }

    const auto shapeCoords = shapeCoordinates(shape);

    ResultPair minimal {std::numeric_limits<double>::max(), {}};
    std::unordered_map<unsigned, unsigned> permutation;

    const auto correctRotations = generateRotations(correctPermutation, shape);

    // i != j != k != l != m with {i, j, k, l, m} in [0, N)
    Temple::Loops::different(
      [&](const std::vector<unsigned>& indices) {
        permutation.clear();
        for(unsigned i = 0; i < 5; ++i) {
          permutation.emplace(i, indices[i]);
        }

        const unsigned i = indices[0];
        const unsigned j = indices[1];
        const unsigned k = indices[2];
        const unsigned l = indices[3];
        const unsigned m = indices[4];

        Eigen::Matrix3d R = fitQuaternion(positions, shapeCoords, permutation);
        auto rotatedShape = R * shapeCoords;

        double penalty = (
          (positions.col(0) - rotatedShape.col(i)).squaredNorm()
          + (positions.col(1) - rotatedShape.col(j)).squaredNorm()
          + (positions.col(2) - rotatedShape.col(k)).squaredNorm()
          + (positions.col(3) - rotatedShape.col(l)).squaredNorm()
          + (positions.col(4) - rotatedShape.col(m)).squaredNorm()
        );

        bool inCorrectBranch = Temple::any_of(
          correctRotations,
          [&](const auto& rot) {
            return rot[0] == i && rot[1] == j && rot[2] == k && rot[3] == l && rot[4] == m;
          }
        );
        if(inCorrectBranch) {
          std::cerr << "At correct permutation quintuple!\n";
        }

        if(penalty > minimal.first) {
          return;
        }

        std::vector<unsigned> freeLeftVertices;
        freeLeftVertices.reserve(N - 5);
        for(unsigned a = 5; a < N; ++a) {
          freeLeftVertices.push_back(a);
        }
        std::vector<unsigned> freeRightVertices;
        freeRightVertices.reserve(N - 5);
        for(unsigned a = 0; a < N; ++a) {
          if(a != i && a != j && a != k && a != l && a != m) {
            freeRightVertices.push_back(a);
          }
        }

        auto narrowed = narrow(
          positions,
          rotatedShape,
          correctRotations,
          inCorrectBranch,
          permutation,
          std::move(freeLeftVertices),
          std::move(freeRightVertices)
        );

        if(narrowed.first < minimal.first) {
          minimal = narrowed;
        }
      },
      5,
      N
    );

    return minimal;
  }

  std::string name() const final {
    return "AlignFive";
  }
};

void writeEnergyStatistics() {
  /* Write an R file with all of the energy values for a particular shape
   */
  std::ofstream rFile("energy_statistics.R");
  std::vector<Shape> shapes;
  unsigned shapeNumber = 1;
  for(const Shape shape : allShapes) {
    if(size(shape) <= 4) {
      continue;
    }

    if(size(shape) > 8) {
      break;
    }

    std::cout << "Number of permutations for " << name(shape) << " shape: " << facs.at(size(shape)) << "\n";

    shapes.push_back(shape);

    auto shapeCoordinates = normalize(addOrigin(coordinates(shape)));
    const unsigned N = shapeCoordinates.cols();

    auto distorted = shapeCoordinates;
    const double distortionNorm = 0.2;
    distort(distorted, distortionNorm);

    Energy energy {shapeCoordinates, distorted};

    rFile << "shape" << shapeNumber << " <- c(";
    auto permutation = Temple::iota<Vertex>(N);
    rFile << energy(permutation);
    while(Temple::next_permutation(permutation)) {
      rFile << ", " << energy(permutation);
    }
    rFile << ")\n";

    ++shapeNumber;
  }

  rFile << "shapeNames <- c(" << Temple::condense(Temple::map(shapes, [](auto s) { return "\"" + name(s) + "\""; })) << ")\n";
}

template<typename PRNG>
PositionCollection shuffle(const PositionCollection& positions, PRNG& prng) {
  const unsigned C = positions.cols();
  auto permutation = Temple::iota<Vertex>(C);
  Temple::Random::shuffle(permutation, prng);
  PositionCollection shuffled(3, C);
  for(unsigned i = 0; i < C; ++i) {
    shuffled.col(permutation.at(i)) = positions.col(i);
  }
  return shuffled;
}

namespace color {

std::ostream& boldMagenta(std::ostream& os) {
  os << "\033[1;35m";
  return os;
}

std::ostream& red(std::ostream& os) {
  os << "\033[31m";
  return os;
}

std::ostream& green(std::ostream& os) {
  os << "\033[32m";
  return os;
}

std::ostream& reset(std::ostream& os) {
  os << "\033[0m";
  return os;
}

} // namespace color

int main(int argc, char* argv[]) {
  boost::program_options::options_description options_description("Recognized options");
  options_description.add_options()
    ("help,h", "Produce help message")
    (
      "prng,p",
      boost::program_options::value<int>(),
      "Seed to initialize PRNG with."
    )
    (
      "shape,s",
      boost::program_options::value<unsigned>(),
      "Shape to run algorithms with"
    )
    (
      "distortion,d",
      boost::program_options::value<double>(),
      "Distortion norm to apply to shapes."
    )
    (
      "repeats,r",
      boost::program_options::value<unsigned>(),
      "Number of repetitions"
    )
  ;

  /* Parse */
  boost::program_options::variables_map options_variables_map;
  boost::program_options::store(
    boost::program_options::command_line_parser(argc, argv).
    options(options_description).
    style(
      boost::program_options::command_line_style::unix_style
      | boost::program_options::command_line_style::allow_long_disguise
    ).run(),
    options_variables_map
  );
  boost::program_options::notify(options_variables_map);

  if(options_variables_map.count("help") > 0) {
    std::cout << options_description << "\n";
    return 0;
  }

  using namespace std::chrono;

  Temple::JSF64 prng;
  if(options_variables_map.count("seed")) {
    const int seed = options_variables_map["seed"].as<int>();
    prng.seed(seed);
    std::cout << "PRNG seeded from parameters: " << seed << ".\n";
  } else {
    std::random_device randomDevice;
    const int seed = std::random_device {}();
    std::cout << "PRNG seeded from random_device: " << seed << ".\n";
    prng.seed(seed);
  }

  std::vector<std::unique_ptr<ShapeAlgorithm>> algorithmPtrs;
  // algorithmPtrs.emplace_back(std::make_unique<Anneal>());
  // algorithmPtrs.emplace_back(std::make_unique<Tunnel>());
  // algorithmPtrs.emplace_back(std::make_unique<ThermodynamicAnneal>());
  // algorithmPtrs.emplace_back(std::make_unique<Greedy>());
  // algorithmPtrs.emplace_back(std::make_unique<SteepestDescent>());
  // algorithmPtrs.emplace_back(std::make_unique<CentroidElimination>());
  algorithmPtrs.emplace_back(std::make_unique<AlignFive>());

  const unsigned nameColWidth = 12;
  const unsigned timeColWidth = 6;
  const unsigned repeats = (options_variables_map.count("repeats") == 1)
    ? options_variables_map["repeats"].as<unsigned>()
    : 10;
  const double distortionNorm = (options_variables_map.count("distortion") == 1)
    ? options_variables_map["distortion"].as<double>()
    : 0.1;
  Shape shape = Shape::CappedSquareAntiprism;
  if(options_variables_map.count("shape") == 1) {
    unsigned shapeIndex = options_variables_map["shape"].as<unsigned>();
    if(shapeIndex < nShapes) {
      shape = static_cast<Shape>(shapeIndex);
    }
  }

  std::cout << "For shape " << name(shape) << " (size " << size(shape) << ") with distortion norm " << distortionNorm << "\n\n";

  AllPermutations referenceAlgorithm;
  std::cout << std::setw(nameColWidth) << referenceAlgorithm.name() << std::setw(timeColWidth) << "msec";

  for(auto& algorithmPtr : algorithmPtrs) {
    std::cout << std::setw(nameColWidth) << algorithmPtr->name() << std::setw(timeColWidth) << "msec" << std::setw(timeColWidth) << "S";
  }
  std::cout << "\n";

  const unsigned A = algorithmPtrs.size();
  std::vector<
    std::vector<double>
  > errors (A);

  std::vector<
    std::vector<double>
  > latencies (A);

  std::vector<
    std::vector<double>
  > speedups (A);

  for(unsigned i = 0; i < repeats; ++i) {
    auto coordinates = ShapeAlgorithm::shapeCoordinates(shape);
    distort(coordinates, distortionNorm);
    coordinates = normalize(shuffle(coordinates, prng));

    time_point<steady_clock> start, end;
    start = steady_clock::now();
    const auto referencePair = referenceAlgorithm.shape(coordinates, shape);
    end = steady_clock::now();
    const unsigned referenceLatency = duration_cast<microseconds>(end - start).count();
    std::cout << std::setw(nameColWidth) << referencePair.first << std::setw(timeColWidth) << (referenceLatency / 1000);
    auto rotations = ShapeAlgorithm::generateRotations(referencePair.second, shape);

    for(unsigned j = 0; j < A; ++j) {
      auto& algorithmPtr = algorithmPtrs.at(j);
      start = steady_clock::now();
      const auto valuePair = algorithmPtr->shape(coordinates, shape, referencePair.second);
      end = steady_clock::now();

      if(valuePair.second != referencePair.second) {
        std::cerr << "\n" << std::setw(25) << "Reference permutation" << " " << Temple::condense(referencePair.second) << "\n"
          << std::setw(25) << algorithmPtr->name() << " " << Temple::condense(valuePair.second) << "\n";
        if(rotations.count(valuePair.second) == 1) {
          std::cerr << "Algorithm result is a shape rotation of the reference permutation\n";
        } else {
          std::cerr << "Algorithm result is not a shape rotation of the reference permutation:\n";
          for(const auto& rotation : rotations) {
            std::cerr << "Rotation " << Temple::condense(rotation) << "\n";
          }
        }
      }

      if(valuePair.first < referencePair.first - 1e-5) {
        std::cout << color::boldMagenta << std::setw(nameColWidth) << valuePair.first << color::reset;
      } else if(std::fabs(valuePair.first - referencePair.first) < 1e-5) {
        std::cout << color::green << std::setw(nameColWidth) << valuePair.first << color::reset;
      } else {
        std::cout << color::red << std::setw(nameColWidth) << valuePair.first << color::reset;
      }
      errors.at(j).push_back(std::fabs(referencePair.first - valuePair.first));

      unsigned latency = duration_cast<microseconds>(end - start).count();
      std::cout << std::setw(timeColWidth) << (latency / 1000);
      latencies.at(j).push_back(latency);

      if(latency == 0) {
        std::cout << std::setw(timeColWidth) << "inf";
      } else {
        std::cout << std::setw(timeColWidth) << std::round(static_cast<double>(referenceLatency) / latency);
      }
      speedups.at(j).push_back(static_cast<double>(referenceLatency) / latency);
    }
    std::cout << "\n";
  }

  std::cout << "\n";
  std::cout << std::setw(nameColWidth + timeColWidth) << "averages";
  for(unsigned j = 0; j < A; ++j) {
    const double errorAverage = Temple::average(errors.at(j));
    if(errorAverage < 1e-10) {
      std::cout << color::green << std::setw(nameColWidth) << 0 << color::reset;
    } else {
      std::cout << color::red << std::setw(nameColWidth) << errorAverage << color::reset;
    }

    std::cout << std::setw(timeColWidth) << static_cast<unsigned>(Temple::average(latencies.at(j)) / 1000);
    std::cout << std::setw(timeColWidth) << static_cast<unsigned>(Temple::average(speedups.at(j)));
  }
  std::cout << "\n";
}


// ====================================================================
// This file is part of FlexibleSUSY.
//
// FlexibleSUSY is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published
// by the Free Software Foundation, either version 3 of the License,
// or (at your option) any later version.
//
// FlexibleSUSY is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with FlexibleSUSY.  If not, see
// <http://www.gnu.org/licenses/>.
// ====================================================================

#ifndef SFERMIONS_H
#define SFERMIONS_H

#include <Eigen/Core>

namespace flexiblesusy {
namespace sfermions {

enum Sparticles {
   up = 0,
   down = 1,
   neutrino = 2,
   electron = 3,
   NUMBER_OF_MSSM_SPARTICLES
};

extern const double Isospin[NUMBER_OF_MSSM_SPARTICLES];
extern const double Hypercharge_left[NUMBER_OF_MSSM_SPARTICLES];
extern const double Hypercharge_right[NUMBER_OF_MSSM_SPARTICLES];

/**
 * data needed to fill 2 x 2 sfermion mass matrix
 */
struct Mass_data {
   double ml2;    ///< soft mass of left-handed sfermion
   double mr2;    ///< soft mass of right-handed sfermion
   double yf;     ///< Yukawa coupling
   double vd, vu; ///< Higgs VEVs
   double gY, g2; ///< gauge couplings (not GUT normalized)
   double Tyf;    ///< trilinear coupling
   double mu;     ///< Superpotential parameter
   double T3;     ///< weak isospin
   double Yl;     ///< Hypercharge of left-handed sfermion
   double Yr;     ///< Hypercharge of right-handed sfermion
};

double diagonalize_sfermions_2x2(const Mass_data&,
                                 Eigen::Array<double,2,1>&);

} // namespace sfermions
} // namespace flexiblesusy

#endif


/* Copyright 2016 Ramakrishnan Kannan */

#ifndef DISTNMF_AUNMF_HPP_
#define DISTNMF_AUNMF_HPP_

#include <mpi.h>
#include <armadillo>
#include <string>
#include <vector>
#include <cstdio>
#include <cmath>
#include "distnmf/distnmf.hpp"
#include "distnmf/mpicomm.hpp"

/**
 * There are totally prxpc process.
 * Each process will hold the following
 * An A of size \f$\frac{globalm}{p_r} \times \frac{globaln}{p_c}\f$
 * Here each process \f$m=\frac{globalm}{p_r} and n=\frac{globaln}{p_c}\f$
 * H of size \f$\frac{globaln}{p} \times k\f$
 * W of size \f${globalm}{p} \times k\f$
 * A is \f$m \times n\f$ matrix
 * H is \f$n \times k\f$ matrix
 */
namespace planc {

template <class INPUTMATTYPE>
class DistAUNMF : public DistNMF<INPUTMATTYPE> {
  // needed in derived algorithms to
  // call BPP routines
 protected:
  MAT HtH;    /// H is of size (globaln/p)*k;
  MAT WtW;    /// W is of size (globaln/p)*k;
  MAT AHtij;  /// AHtij is of size k*(globalm/p)
  MAT WtAij;  /// WtAij is of size k*(globaln/p)
  MAT Wt;     /// Wt is of size k*(globalm/p)
  MAT Ht;     /// Ht is of size k*(globaln/p)

  virtual void updateW() = 0;
  virtual void updateH() = 0;

 private:
  // Things needed while solving for W
  MAT localHtH;         /// H is of size (globaln/p)*k;
  MAT Hjt, Hj;          /// Hj is of size n*k;
  MAT AijHj, AijHjt;    /// AijHj is of size m*k;
  // INPUTMATTYPE A_ij_t;  /// n*m matrix. Transpose of A_ij
  // Things needed while solving for H
  MAT localWtW;        /// W is of size (globalm/p)*k;
  MAT Wit, Wi;         /// Wi is of size m*k;
  MAT WitAij, AijWit;  /// WijtAij is of size k*n;

  // needed for error computation
  MAT prevH;        // used for error computation
  MAT prevHtH;      // used for error computation
  MAT WtAijH;       /// global k*k matrix.
  MAT localWtAijH;  /// local k*k matrix
  MAT errMtx;
  MAT A_errMtx;

  // needed for symm regularization
  MAT crossFac;     // holds the appropriate row of W,H
  int paired_proc;  // processor to swap factors with

  // needed for block implementation to save memory
  MAT Ht_blk;
  MAT AHtij_blk;
  MAT Wt_blk;
  MAT WtAij_blk;

  // Gatherv and Reducescatter variables
  std::vector<int> gatherWtAcnts;
  std::vector<int> gatherWtAdisp;
  std::vector<int> scatterWtAcnts;

  std::vector<int> gatherAHcnts;
  std::vector<int> gatherAHdisp;
  std::vector<int> scatterAHcnts;

  int num_k_blocks;
  std::string relative_error_dir;
  int perk;

  /**
   * Allocates matrices
   */
  void allocateMatrices() {
    // collective call related initializations.
    // These initialization are for solving W.
    DISTPRINTINFO("k::" << this->k << "::perk::" << this->perk
                        << "::localm::" << this->m << "::localn::" << this->n
                        << "::globalm::" << this->globalm() << "::globaln::"
                        << this->globaln() << "::MPI_SIZE::" << MPI_SIZE);
    HtH.zeros(this->k, this->k);
    localHtH.zeros(this->k, this->k);
    Hj.zeros(this->n, this->perk);
    Hjt.zeros(this->perk, this->n);
    AijHj.zeros(this->m, this->perk);
    AijHjt.zeros(this->perk, this->m);
    AHtij.zeros(this->k, this->W.n_rows);
    this->scatterAHcnts.resize(NUMCOLPROCS);
    int fillsize = this->perk * (this->W.n_rows);
    fillVector<int>(fillsize, &scatterAHcnts);
    this->gatherAHcnts.resize(NUMROWPROCS);
    fillVector<int>(0, &gatherAHcnts);
    this->gatherAHdisp.resize(NUMROWPROCS);
    fillVector<int>(0, &gatherAHdisp);
#ifdef MPI_VERBOSE
    if (ISROOT) {
      INFO << "::recvAHsize::";
      printVector<int>(recvAHsize);
    }
#endif
    // allocated for block implementation.
    Ht_blk.zeros(this->perk, (this->W.n_rows));
    AHtij_blk.zeros(this->perk, this->W.n_rows);

    // These initialization are for solving H.
    Wt.zeros(this->k, this->W.n_rows);
    WtW.zeros(this->k, this->k);
    localWtW.zeros(this->k, this->k);
    Wi.zeros(this->m, this->perk);
    Wit.zeros(this->perk, this->m);
    WitAij.zeros(this->perk, this->n);
    AijWit.zeros(this->n, this->perk);
    WtAij.zeros(this->k, this->H.n_rows);
    this->scatterWtAcnts.resize(NUMROWPROCS);
    fillsize = this->perk * (this->H.n_rows);
    fillVector<int>(fillsize, &scatterWtAcnts);
    this->gatherWtAcnts.resize(NUMCOLPROCS);
    fillVector<int>(0, &gatherWtAcnts);
    this->gatherWtAdisp.resize(NUMCOLPROCS);
    fillVector<int>(0, &gatherWtAdisp);

    // allocated for block implementation
    Wt_blk.zeros(this->perk, this->W.n_rows);
    WtAij_blk.zeros(this->perk, this->H.n_rows);

    // allocated for symmetric regularisation
    if (this->symm_reg() > 0) {
      crossFac.zeros(this->k, this->H.n_rows);
    }
#ifdef MPI_VERBOSE
    if (ISROOT) {
      INFO << "::recvWtAsize::";
      printVector<int>(recvWtAsize);
    }
#endif
#ifndef BUILD_SPARSE
    if (this->is_compute_error()) {
      errMtx.zeros(this->m, this->n);
      A_errMtx.zeros(this->m, this->n);
    }
#endif
  }

  void freeMatrices() {
    HtH.clear();
    localHtH.clear();
    Hj.clear();
    Hjt.clear();
    AijHj.clear();
    AijHjt.clear();
    AHtij.clear();
    Wt.clear();
    WtW.clear();
    localWtW.clear();
    Wi.clear();
    Wit.clear();
    WitAij.clear();
    AijWit.clear();
    WtAij.clear();
    // A_ij_t.clear();
    if (this->is_compute_error()) {
      prevH.clear();
      prevHtH.clear();
      WtAijH.clear();
      localWtAijH.clear();
    }
    Ht_blk.clear();
    AHtij_blk.clear();
    Wt_blk.clear();
    WtAij_blk.clear();
    if (this->symm_reg() > 0) {
      crossFac.clear();
    }
    if (this->is_compute_error()) {
      errMtx.clear();
      A_errMtx.clear();
    }
  }

  /**
   * Sets up the communication pattern for the matrix multiplies
  */
  void setupCommcounts() {
    // WtA
    // Allgatherv counts
    gatherWtAcnts[0] = itersplit(this->A.n_rows, NUMCOLPROCS, 0) * this->perk;
    gatherWtAdisp[0] = 0;
    for (int i = 1; i < NUMCOLPROCS; i++) {
      gatherWtAcnts[i] = itersplit(this->A.n_rows,
                                   NUMCOLPROCS, i) * this->perk;
      gatherWtAdisp[i] = gatherWtAdisp[i-1] + gatherWtAcnts[i-1];
    }
    // Reducescatter counts
    for (int i = 0; i < NUMROWPROCS; i++) {
      scatterWtAcnts[i] = itersplit(this->A.n_cols,
                                    NUMROWPROCS, i) * this->perk;
    }
    // AH
    // Allgatherv counts
    gatherAHcnts[0] = itersplit(this->A.n_cols, NUMROWPROCS, 0) * this->perk;
    gatherAHdisp[0] = 0;
    for (int i = 1; i < NUMROWPROCS; i++) {
      gatherAHcnts[i] = itersplit(this->A.n_cols, NUMROWPROCS, i) * this->perk;
      gatherAHdisp[i] = gatherAHdisp[i-1] + gatherAHcnts[i-1];
    }
    // Reducescatter counts
    for (int i = 0; i < NUMCOLPROCS; i++) {
      scatterAHcnts[i] = itersplit(this->A.n_rows,
                                   NUMCOLPROCS, i) * this->perk;
    }
  }

 public:
  /**
   * Public constructor with local input matrix, local factors and communicator
   * @param[in] local input matrix of size \f$\frac{globalm}{p_r} \times \frac{globaln}{p_c}\f$.
   *            Each process owns \f$m=\frac{globalm}{p_r}\f$ and \f$n=\frac{globaln}{p_c}\f$
   * @param[in] local left low rank factor of size \f$\frac{globalm}{p} \times k \f$
   * @param[in] local right low rank factor of size \f$\frac{globaln}{p} \times k \f$
   * @param[in] MPICommunicator that has row and column communicators
   * @param[in] numkblks. the columns of the local factor can further be
   *            partitioned into numkblks
   */
  DistAUNMF(const INPUTMATTYPE &input, const MAT &leftlowrankfactor,
            const MAT &rightlowrankfactor, const MPICommunicator &communicator,
            const int numkblks, const std::string relerr_dir)
      : DistNMF<INPUTMATTYPE>(input, leftlowrankfactor, rightlowrankfactor,
                              communicator) {
    num_k_blocks = numkblks;
    relative_error_dir = relerr_dir;
    perk = this->k / num_k_blocks;
    allocateMatrices();
    setupCommcounts();
    this->Wt = leftlowrankfactor.t();
    this->Ht = rightlowrankfactor.t();
    // A_ij_t = input.t();
    if (this->symm_reg() >= 0) {
      // Get paired processor
      int coords[2];
      coords[0] = MPI_COL_RANK;
      coords[1] = MPI_ROW_RANK;
      MPI_Cart_rank(this->m_mpicomm.gridComm(), &coords[0], &paired_proc);
      DISTPRINTINFO("rank::" << MPI_RANK << "::paired_proc::" << paired_proc);
    }
    PRINTROOT("aunmf()::constructor succesful");
  }
  ~DistAUNMF() {
    // freeMatrices();
  }

  /**
   * This is a matrix multiplication routine based on
   * reduce_scatter.
   * A is mxn in column major ordering
   * W is mxk in column major ordering
   * AtW is nxk in column major ordering
   * There are totally p processes. Every process has
   * A_i as m_i * n
   * W_i as m_i * k
   * AtW_i as n_i * k
   * this->m_mpicomm.comm_subs()[0] is column communicator.
   * this->m_mpicomm.comm_subs()[1] is row communicator.
   */
  void distWtA() {
    for (int i = 0; i < num_k_blocks; i++) {
      int start_row = i * perk;
      int end_row = (i + 1) * perk - 1;
      Wt_blk = Wt.rows(start_row, end_row);
      distWtABlock();
      WtAij.rows(start_row, end_row) = WtAij_blk;
    }
  }
  void distWtABlock() {
#ifdef USE_PACOSS
    // Perform expand communication using Pacoss.
    memcpy(Wit.memptr(), Wt_blk.memptr(),
           Wt_blk.n_rows * Wt_blk.n_cols * sizeof(Wt_blk[0]));
    MPITIC;
    this->m_rowcomm->expCommBegin(Wit.memptr(), this->perk);
    this->m_rowcomm->expCommFinish(Wit.memptr(), this->perk);
#else
    int sendcnt = (this->W.n_rows) * this->perk;
    Wit.zeros();
    MPITIC;  // allgather WtA
    MPI_Allgatherv(Wt_blk.memptr(), sendcnt, MPI_DOUBLE, Wit.memptr(),
                  &(gatherWtAcnts[0]), &(gatherWtAdisp[0]), MPI_DOUBLE,
                  this->m_mpicomm.commSubs()[1]);
#endif
    double temp = MPITOC;  // allgather WtA
    PRINTROOT("n::" << this->n << "::k::" << this->k << PRINTMATINFO(Wt)
                    << PRINTMATINFO(Wit));
#ifdef MPI_VERBOSE
    DISTPRINTINFO(PRINTMAT(Wt_blk));
    DISTPRINTINFO(PRINTMAT(Wit));
#endif
    this->time_stats.communication_duration(temp);
    this->time_stats.allgather_duration(temp);
    MPITIC;  // mm WtA
    this->WitAij = this->Wit * this->A;
    // #if defined(MKL_FOUND) && defined(BUILD_SPARSE)
    //     // void ARMAMKLSCSCMM(const SRC &mklMat, const DESTN &Bt, const char
    //     transa,
    //     //               DESTN *Ct)
    //     ARMAMKLSCSCMM(this->A_ij_t, 'N', this->Wit, this->AijWit.memptr());
    // #ifdef MPI_VERBOSE
    //     DISTPRINTINFO(PRINTMAT(this->AijWit));
    // #endif
    //     this->WitAij = reshape(this->AijWit, this->k, this->n);
    // #else
    //     this->WitAij = this->Wit * this->A;
    // #endif
    temp = MPITOC;  // mm WtA
#ifdef MPI_VERBOSE
    DISTPRINTINFO(PRINTMAT(this->WitAij));
#endif
    PRINTROOT(PRINTMATINFO(this->A)
              << PRINTMATINFO(this->Wit) << PRINTMATINFO(this->WitAij));
    this->time_stats.compute_duration(temp);
    this->time_stats.mm_duration(temp);
    this->reportTime(temp, "WtA::");
#ifdef USE_PACOSS
    // Perform fold communication using Pacoss.
    MPITIC;
    this->m_colcomm->foldCommBegin(WitAij.memptr(), this->perk);
    this->m_colcomm->foldCommFinish(WitAij.memptr(), this->perk);
    temp = MPITOC;
    memcpy(WtAij_blk.memptr(), WitAij.memptr(),
           WtAij_blk.n_rows * WtAij_blk.n_cols * sizeof(WtAij_blk[0]));
#else
    WtAij_blk.zeros();
    MPITIC;  // reduce_scatter WtA
    MPI_Reduce_scatter(this->WitAij.memptr(), this->WtAij_blk.memptr(),
                       &(scatterWtAcnts[0]), MPI_DOUBLE, MPI_SUM,
                       this->m_mpicomm.commSubs()[0]);
    temp = MPITOC;  // reduce_scatter WtA
#endif
    this->time_stats.communication_duration(temp);
    this->time_stats.reducescatter_duration(temp);
  }
  /**
   * There are totally prxpc process.
   * Each process will hold the following
   * An A of size (m/pr) x (n/pc)
   * H of size (n/p)xk
   * find AHt kx(m/p) by reducing and scatter it using MPI_Reduce_scatter call.
   * That is, p process will hold a kx(m/p) matrix.
   * this->m_mpicomm.comm_subs()[0] is column communicator.
   * this->m_mpicomm.comm_subs()[1] is row communicator.
   * To preserve the memory for Hj, we collect only partial k
   */
  void distAH() {
    for (int i = 0; i < num_k_blocks; i++) {
      int start_row = i * perk;
      int end_row = (i + 1) * perk - 1;
      Ht_blk = Ht.rows(start_row, end_row);
      distAHBlock();
      AHtij.rows(start_row, end_row) = AHtij_blk;
    }
  }
  void distAHBlock() {
    /*
    DISTPRINTINFO("distAH::" << "::Acolst::" \
                  Acolst.n_rows<<"x"<<Acolst.n_cols \
                  << "::norm::" << arma::norm(Acolst, "fro"));
    DISTPRINTINFO("distAH::" << "::H::" \
                  << this->H.n_rows << "x" << this->H.n_cols);
    */
#ifdef USE_PACOSS
    // Perform expand communication using Pacoss.
    memcpy(Hjt.memptr(), Ht_blk.memptr(),
           Ht_blk.n_rows * Ht_blk.n_cols * sizeof(Ht_blk[0]));
    MPITIC;
    this->m_colcomm->expCommBegin(Hjt.memptr(), this->perk);
    this->m_colcomm->expCommFinish(Hjt.memptr(), this->perk);
#else
    int sendcnt = (this->H.n_rows) * this->perk;
    Hjt.zeros();
    MPITIC;  // allgather AH
    MPI_Allgatherv(this->Ht_blk.memptr(), sendcnt, MPI_DOUBLE,
                  this->Hjt.memptr(), &(gatherAHcnts[0]), &(gatherAHdisp[0]),
                  MPI_DOUBLE, this->m_mpicomm.commSubs()[0]);
#endif
    PRINTROOT("n::" << this->n << "::k::" << this->k << PRINTMATINFO(Ht)
                    << PRINTMATINFO(Hjt));
    double temp = MPITOC;  // allgather AH
#ifdef MPI_VERBOSE
    DISTPRINTINFO(PRINTMAT(Ht_blk));
    DISTPRINTINFO(PRINTMAT(Hjt));
    // DISTPRINTINFO(PRINTMAT(this->A_ij_t));
#endif
    this->time_stats.communication_duration(temp);
    this->time_stats.allgather_duration(temp);
    MPITIC;  // mm AH
/*
#ifdef BUILD_SPARSE
    this->Hj = this->Hjt.t();
    this->AijHj = this->A * this->Hj;
    this->AijHjt = this->AijHj.t();
#else
    this->AijHjt = this->Hjt * this->A.t();
#endif
More memory efficient to rewrite code as above since A.t() will construct a
copy of the A in the sparse case. However sparse x dense matmul is much slower
than in dense x sparse. Keeping current version for performance reasons.
*/
    this->AijHjt = this->Hjt * this->A.t();
    // #if defined(MKL_FOUND) && defined(BUILD_SPARSE)
    //     // void ARMAMKLSCSCMM(const SRC &mklMat, const DESTN &Bt, const char
    //     transa,
    //     //               DESTN *Ct)
    //     ARMAMKLSCSCMM(this->A, 'N', this->Hjt, this->AijHj.memptr());
    // #ifdef MPI_VERBOSE
    //     DISTPRINTINFO(PRINTMAT(this->AijHj));
    // #endif
    //     this->AijHjt = reshape(this->AijHj, this->k, this->m);
    // #else
    //     this->AijHjt = this->Hjt * this->A_ij_t;
    // #endif
    // #ifdef MPI_VERBOSE
    //     DISTPRINTINFO(PRINTMAT(this->AijHjt));
    // #endif
    temp = MPITOC;  // mm AH
    // PRINTROOT(PRINTMATINFO(this->A_ij_t)
    PRINTROOT(PRINTMATINFO(this->Hjt) << PRINTMATINFO(this->AijHjt));
    this->time_stats.compute_duration(temp);
    this->time_stats.mm_duration(temp);
    this->reportTime(temp, "AH::");
#ifdef USE_PACOSS
    // Perform fold communication using Pacoss.
    MPITIC;
    this->m_rowcomm->foldCommBegin(AijHjt.memptr(), this->perk);
    this->m_rowcomm->foldCommFinish(AijHjt.memptr(), this->perk);
    temp = MPITOC;
    memcpy(AHtij_blk.memptr(), AijHjt.memptr(),
           AHtij_blk.n_rows * AHtij_blk.n_cols * sizeof(AHtij_blk[0]));
#else
    AHtij_blk.zeros();
    MPITIC;  // reduce_scatter AH
    MPI_Reduce_scatter(this->AijHjt.memptr(), this->AHtij_blk.memptr(),
                       &(this->scatterAHcnts[0]), MPI_DOUBLE, MPI_SUM,
                       this->m_mpicomm.commSubs()[1]);
    temp = MPITOC;  // reduce_scatter AH
#endif
    this->time_stats.communication_duration(temp);
    this->time_stats.reducescatter_duration(temp);
  }
  /**
   * There are p processes.
   * Every process i has W in m_i * k
   * At the end of this call, all process will have
   * WtW of size k*k is symmetric. So not to worry
   * about column/row major formats.
   * @param[in] X is of size m_i x k
   * @param[out] XtX Every process owns the same kxk global gram matrix of X
   */
  void distInnerProduct(const MAT &X, MAT *XtX) {
    // each process computes its own kxk matrix
    MPITIC;  // gram
    localWtW = X.t() * X;
#ifdef MPI_VERBOSE
    DISTPRINTINFO("W::" << norm(X, "fro")
                        << "::localWtW::" << norm(this->localWtW, "fro"));
#endif
    double temp = MPITOC;  // gram
    this->time_stats.compute_duration(temp);
    this->time_stats.gram_duration(temp);
    (*XtX).zeros();
    if (X.n_rows == this->m) {
      this->reportTime(temp, "Gram::W::");
    } else {
      this->reportTime(temp, "Gram::H::");
    }
    MPITIC;  // allreduce gram
    MPI_Allreduce(localWtW.memptr(), (*XtX).memptr(), this->k * this->k,
                  MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
    temp = MPITOC;  // allreduce gram
    this->time_stats.communication_duration(temp);
    this->time_stats.allreduce_duration(temp);
  }
  /**
   * This is the main loop function
   * Refer Algorithm 1 in Page 3 of
   * the PPoPP HPC-NMF paper.
   */
  void computeNMF() {
    PRINTROOT("computeNMF started");
#ifdef MPI_VERBOSE
    DISTPRINTINFO(PRINTMAT(this->A));
#endif
    // error computation
    if (this->is_compute_error()) {
      prevH.zeros(size(this->H));
      prevHtH.zeros(this->k, this->k);
      WtAijH.zeros(this->k, this->k);
      localWtAijH.zeros(this->k, this->k);
    }
#ifdef __WITH__BARRIER__TIMING__
    MPI_Barrier(MPI_COMM_WORLD);
#endif
    for (unsigned int iter = 0; iter < this->num_iterations(); iter++) {
      // saving current instance for error computation.
      if (iter > 0 && this->is_compute_error()) {
        this->prevH = this->H;
        this->prevHtH = this->HtH;
      }
      MPITIC;  // total_d W&H
      // update H given WtW and WtA step 4 of the algorithm
      {
        // compute WtW
        this->distInnerProduct(this->W, &this->WtW);
        PRINTROOT(PRINTMATINFO(this->WtW));
        this->applyReg(this->regH(), &this->WtW);
#ifdef MPI_VERBOSE
        PRINTROOT(PRINTMAT(this->WtW));
#endif
        // compute WtA
        this->distWtA();
        if (this->symm_reg() > 0) {
          // Get the appropriate Wt from the transposed processor
          this->crossFac.zeros(arma::size(this->Ht));
          int recvsize = this->crossFac.n_elem;
          int sendsize = this->Wt.n_elem;
          MPITIC;
          MPI_Sendrecv(this->Wt.memptr(), sendsize, MPI_DOUBLE, paired_proc, 0,
              this->crossFac.memptr(), recvsize, MPI_DOUBLE, paired_proc, 0,
              this->m_mpicomm.gridComm(), MPI_STATUS_IGNORE);
          double temp = MPITOC;  // sendrecv
          this->time_stats.communication_duration(temp);
          this->time_stats.sendrecv_duration(temp);

          this->applySymmetricReg(this->symm_reg(), &this->WtW,
                  &this->crossFac, &this->WtAij);
        }
        // PRINTROOT(PRINTMATINFO(this->WtAij));
#ifdef MPI_VERBOSE
        DISTPRINTINFO(PRINTMAT(this->WtAij));
#endif
        MPITIC;  // nnls H
        // ensure both Ht and H are consistent after the update
        // some function find Ht and some H.
        updateH();
#ifdef MPI_VERBOSE
        DISTPRINTINFO("::it=" << iter << PRINTMAT(this->H));
#endif
        double temp = MPITOC;  // nnls H
        this->time_stats.compute_duration(temp);
        this->time_stats.nnls_duration(temp);
        this->reportTime(temp, "NNLS::H::");
      }
      // Update W given HtH and AH step 3 of the algorithm.
      {
        // compute HtH
        this->distInnerProduct(this->H, &this->HtH);
        PRINTROOT("HtH::" << PRINTMATINFO(this->HtH));
        this->applyReg(this->regW(), &this->HtH);
#ifdef MPI_VERBOSE
        PRINTROOT(PRINTMAT(this->HtH));
#endif
        // compute AH
        this->distAH();
        if (this->symm_reg() > 0) {
          // Get the appropriate Ht from the transposed processor
          this->crossFac.zeros(arma::size(this->Wt));
          int recvsize = this->crossFac.n_elem;
          int sendsize = this->Ht.n_elem;
          MPITIC;
          MPI_Sendrecv(this->Ht.memptr(), sendsize, MPI_DOUBLE, paired_proc, 0,
              this->crossFac.memptr(), recvsize, MPI_DOUBLE, paired_proc, 0,
              this->m_mpicomm.gridComm(), MPI_STATUS_IGNORE);
          double temp = MPITOC;  // sendrecv
          this->time_stats.communication_duration(temp);
          this->time_stats.sendrecv_duration(temp);

          this->applySymmetricReg(this->symm_reg(), &this->HtH,
                &this->crossFac, &this->AHtij);
        }
        // PRINTROOT(PRINTMATINFO(this->AHtij));
#ifdef MPI_VERBOSE
        DISTPRINTINFO(PRINTMAT(this->AHtij));
#endif
        MPITIC;  // nnls W
        // Update W given HtH and AH step 3 of the algorithm.
        // ensure W and Wt are consistent. As some algorithms
        // determine W and some Wt.
        updateW();
#ifdef MPI_VERBOSE
        DISTPRINTINFO("::it=" << iter << PRINTMAT(this->W));
#endif
        double temp = MPITOC;  // nnls W
        this->time_stats.compute_duration(temp);
        this->time_stats.nnls_duration(temp);
        this->reportTime(temp, "NNLS::W::");
      }
      this->time_stats.duration(MPITOC);  // total_d W&H
      if (iter > 0 && this->is_compute_error()) {
#ifdef BUILD_SPARSE
        this->computeError(iter);
#else
        this->computeError2(iter);
#endif
        std::string outfullName = relative_error_dir+"relerr_at_k"+std::to_string(this->k);
        PRINTROOT("it=" << iter << "::algo::" << this->m_algorithm << "::k::"
                        << this->k << "::err::" << sqrt(this->objective_err)
                        << "::relerr::"
                        << sqrt(this->objective_err / this->m_globalsqnormA));
        std::ofstream outfile;
        if (iter == this->num_iterations()-1 && this->m_mpicomm.rank() == 0) {
            outfile.open(outfullName.c_str(), std::ios_base::app);
            outfile<<  sqrt(this->objective_err / this->m_globalsqnormA) <<"\n";
            //PRINTROOT(sqrt(this->objective_err/this->m_globalsqnormA));
            outfile.close();
        }
        // Compute the difference between factor matrices
        if (this->symm_reg() > 0) {
          double localdiff = arma::norm(this->Wt-this->crossFac, "fro");
          double globaldiff = 0.0;

          // Compute global difference
          localdiff = localdiff * localdiff;
          MPI_Allreduce(&localdiff, &globaldiff, 1,
              MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);

          double localWnorm = arma::norm(this->Wt, "fro");
          double globalWnorm = 0.0;

          // Compute global W norm
          localWnorm = localWnorm * localWnorm;
          MPI_Allreduce(&localWnorm, &globalWnorm, 1,
              MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);

          PRINTROOT("it=" << iter << "::symmdiff::" << globaldiff
                    << "::reldiff::" << sqrt(globaldiff / globalWnorm));
        }
      }
      PRINTROOT("completed it=" << iter
                                << "::taken::" << this->time_stats.duration());
    }  // end for loop
    MPI_Barrier(MPI_COMM_WORLD);
    this->reportTime(this->time_stats.duration(), "total_d");
    this->reportTime(this->time_stats.communication_duration(), "total_comm");
    this->reportTime(this->time_stats.compute_duration(), "total_comp");
    this->reportTime(this->time_stats.allgather_duration(), "total_allgather");
    this->reportTime(this->time_stats.allreduce_duration(), "total_allreduce");
    this->reportTime(this->time_stats.reducescatter_duration(),
                     "total_reducescatter");
    this->reportTime(this->time_stats.gram_duration(), "total_gram");
    this->reportTime(this->time_stats.mm_duration(), "total_mm");
    this->reportTime(this->time_stats.nnls_duration(), "total_nnls");
    if (this->symm_reg() > 0) {
      this->reportTime(this->time_stats.sendrecv_duration(), "total_sendrecv");
    }
    if (this->is_compute_error()) {
      this->reportTime(this->time_stats.err_compute_duration(),
                       "total_err_compute");
      this->reportTime(this->time_stats.err_compute_duration(),
                       "total_err_communication");
    }
  }

  /**
   * We assume this error function will be called in
   * every iteration before updating the block to
   * compute the error from previous iteration
   * \f$\|A\|_F^2 - 2trace(H(A^TW))+trace((W^TW)*(HH^T))\f$
   * each process owns globalsqnormA will have \|A\|_F^2
   * each process owns WtAij is of size \f$k \times \frac{globaln}{p}\f$
   * each process owns H is of size \f${globaln}{p} \times k \f$
   * compute WtAij*H and do an MPI_ALL reduce to get the kxk matrix.
   * every process local computation
   */

  void computeError(const int it) {
    MPITIC;  // computeerror
    this->localWtAijH = this->WtAij * this->prevH;
#ifdef MPI_VERBOSE
    DISTPRINTINFO("::it=" << it << PRINTMAT(this->WtAij));
    DISTPRINTINFO("::it=" << it << PRINTMAT(this->localWtAijH));
    DISTPRINTINFO("::it=" << it << PRINTMAT(this->prevH));
    PRINTROOT("::it=" << it << PRINTMAT(this->WtW));
    PRINTROOT("::it=" << it << PRINTMAT(this->prevHtH));
#endif
    double temp = MPITOC;  // computererror
    this->time_stats.err_compute_duration(temp);
    MPITIC;  // coommunication error
    MPI_Allreduce(this->localWtAijH.memptr(), this->WtAijH.memptr(),
                  this->k * this->k, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
    temp = MPITOC;  // communication error
#ifdef MPI_VERBOSE
    DISTPRINTINFO(PRINTMAT(WtAijH));
#endif
    this->time_stats.err_communication_duration(temp);
    double tWtAijh = trace(this->WtAijH);
    double tWtWHtH = trace(this->WtW * this->prevHtH);
    PRINTROOT("::it=" << it << "normA::" << this->m_globalsqnormA
                      << "::tWtAijH::" << 2 * tWtAijh
                      << "::tWtWHtH::" << tWtWHtH);
    this->objective_err = this->m_globalsqnormA - 2 * tWtAijh + tWtWHtH;
  }
  /*
   * Compute error the old-fashioned way
   */
  void computeError2(const int it) {
    double local_sqerror = 0.0;
    PRINTROOT("::it=" << it << "::Calling compute error 2");
    MPITIC;
    // DISTPRINTINFO("::norm(Wi,fro)::" << norm(this->Wit, "fro") <<
    // "::norm(Hjt, fro)::" << norm(this->Hjt, "fro"));
    this->Wi = this->Wit.t();
    errMtx = this->Wi * this->Hjt;
    A_errMtx = this->A - errMtx;
    local_sqerror = norm(A_errMtx, "fro");
    local_sqerror *= local_sqerror;
    double temp = MPITOC;
    this->time_stats.err_compute_duration(temp);
    // DISTPRINTINFO("::it=" << it << "::local_sqerror::" << local_sqerror);
    MPITIC;
    MPI_Allreduce(&local_sqerror, &this->objective_err, 1, MPI_DOUBLE, MPI_SUM,
                  MPI_COMM_WORLD);
    temp = MPITOC;
    this->time_stats.err_communication_duration(temp);
  }

  // Set the LUC inner iterations for iterative LUC
  void set_luciters(int max_luciters) {}
};

}  // namespace planc

#endif  // DISTNMF_AUNMF_HPP_


#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <math.h>
#include <omp.h>
#include <iomanip>

#include "mvnormal.h"
#include "macau.h"
#include "chol.h"
#include "linop.h"
#include "noisemodels.h"
#include "truncnorm.h"
extern "C" {
  #include <sparse.h>
}

using namespace std;
using namespace Eigen;

void ILatentPrior::sample_latents(FixedGaussianNoise & noise, Eigen::MatrixXd &U, const Eigen::SparseMatrix<double> &mat,
                    double mean_value, const Eigen::MatrixXd &samples, const int num_latent) {
  this->sample_latents(U, mat, mean_value, samples, noise.alpha, num_latent);
}

void ILatentPrior::sample_latents(AdaptiveGaussianNoise & noise, Eigen::MatrixXd &U, const Eigen::SparseMatrix<double> &mat,
                    double mean_value, const Eigen::MatrixXd &samples, const int num_latent) {
  this->sample_latents(U, mat, mean_value, samples, noise.alpha, num_latent);
}

void ILatentPrior::sample_latents(FixedGaussianNoise & noiseModel, MatrixData & matrixData,
                                std::vector< std::unique_ptr<Eigen::MatrixXd> > & samples, int mode, const int num_latent) {
  if (mode == 0) {
    this->sample_latents(noiseModel, *samples[0], matrixData.Yt, matrixData.mean_value, *samples[1], num_latent);
  } else {
    this->sample_latents(noiseModel, *samples[1], matrixData.Y,  matrixData.mean_value, *samples[0], num_latent);
  }
}

void ILatentPrior::sample_latents(AdaptiveGaussianNoise & noiseModel, MatrixData & matrixData,
                                std::vector< std::unique_ptr<Eigen::MatrixXd> > & samples, int mode, const int num_latent) {
  if (mode == 0) {
    this->sample_latents(noiseModel, *samples[0], matrixData.Yt, matrixData.mean_value, *samples[1], num_latent);
  } else {
    this->sample_latents(noiseModel, *samples[1], matrixData.Y,  matrixData.mean_value, *samples[0], num_latent);
  }
}

void ILatentPrior::sample_latents(ProbitNoise & noiseModel, MatrixData & matrixData,
                                std::vector< std::unique_ptr<Eigen::MatrixXd> > & samples, int mode, const int num_latent) {
  if (mode == 0) {
    this->sample_latents(noiseModel, *samples[0], matrixData.Yt, matrixData.mean_value, *samples[1], num_latent);
  } else {
    this->sample_latents(noiseModel, *samples[1], matrixData.Y,  matrixData.mean_value, *samples[0], num_latent);
  }
}

void ILatentPrior::sample_latents(FixedGaussianNoise& noiseModel, TensorData & data,
                                std::vector< std::unique_ptr<Eigen::MatrixXd> > & samples, int mode, const int num_latent) {
  sample_latents(noiseModel.alpha, data, samples, mode, num_latent);
}

void ILatentPrior::sample_latents(AdaptiveGaussianNoise& noiseModel, TensorData & data,
                            std::vector< std::unique_ptr<Eigen::MatrixXd> > & samples, int mode, const int num_latent) {
  sample_latents(noiseModel.alpha, data, samples, mode, num_latent);
}

/** BPMFPrior */
void BPMFPrior::sample_latents(Eigen::MatrixXd &U, const Eigen::SparseMatrix<double> &mat, double mean_value,
                    const Eigen::MatrixXd &samples, double alpha, const int num_latent) {
  const int N = U.cols();

#pragma omp parallel for schedule(dynamic, 2)
  for(int n = 0; n < N; n++) {
    sample_latent_blas(U, n, mat, mean_value, samples, alpha, mu, Lambda, num_latent);
  }
}

void BPMFPrior::update_prior(const Eigen::MatrixXd &U) {
  tie(mu, Lambda) = CondNormalWishart(U, mu0, b0, WI, df);
}

void BPMFPrior::init(const int num_latent) {
  mu.resize(num_latent);
  mu.setZero();

  Lambda.resize(num_latent, num_latent);
  Lambda.setIdentity();
  Lambda *= 10;

  // parameters of Inv-Whishart distribution
  WI.resize(num_latent, num_latent);
  WI.setIdentity();
  mu0.resize(num_latent);
  mu0.setZero();
  b0 = 2;
  df = num_latent;
}

void BPMFPrior::sample_latents(ProbitNoise & noise, Eigen::MatrixXd &U, const Eigen::SparseMatrix<double> &mat,
                               double mean_value, const Eigen::MatrixXd &samples, const int num_latent) {
  const int N = U.cols();

#pragma omp parallel for schedule(dynamic, 2)
  for(int n = 0; n < N; n++) {
    sample_latent_blas_probit(U, n, mat, mean_value, samples, mu, Lambda, num_latent);
  }

}

void BPMFPrior::sample_latents(ProbitNoise& noiseModel, TensorData & data,
                               std::vector< std::unique_ptr<Eigen::MatrixXd> > & samples, int mode, const int num_latent) {
  // TODO
  throw std::runtime_error("Unimplemented: sample_latents");
}

void sample_latent_tensor(std::unique_ptr<Eigen::MatrixXd> &U,
                          int n,
                          std::unique_ptr<SparseMode> & sparseMode,
                          VectorView<Eigen::MatrixXd> & view,
                          double mean_value,
                          double alpha,
                          Eigen::VectorXd & mu,
                          Eigen::MatrixXd & Lambda) {
  const int nmodes1 = view.size();
  const int num_latent = U->rows();

  MatrixXd MM(num_latent, num_latent);
  MM = Lambda;
  VectorXd rr = VectorXd::Zero(mu.size());

  Eigen::VectorXi & row_ptr = sparseMode->row_ptr;
  Eigen::MatrixXi & indices = sparseMode->indices;
  Eigen::VectorXd & values  = sparseMode->values;

  Eigen::MatrixXd* S0 = view.get(0);

  for (int j = row_ptr(n); j < row_ptr(n + 1); j++) {
    VectorXd col = S0->col(indices(j, 0));
    for (int m = 1; m < nmodes1; m++) {
      col.noalias() = col.cwiseProduct(view.get(m)->col(indices(j, m)));
    }

    MM.triangularView<Eigen::Lower>() += alpha * col * col.transpose();
    rr.noalias() += col * ((values(j) - mean_value) * alpha);
  }

  Eigen::LLT<MatrixXd> chol = MM.llt();
  if(chol.info() != Eigen::Success) {
    throw std::runtime_error("Cholesky Decomposition failed!");
  }

  rr.noalias() += Lambda * mu;
  chol.matrixL().solveInPlace(rr);
  for (int i = 0; i < num_latent; i++) {
    rr[i] += randn0();
  }
  chol.matrixU().solveInPlace(rr);
  U->col(n).noalias() = rr;
}

void BPMFPrior::sample_latents(double noisePrecision,
                               TensorData & data,
                               std::vector< std::unique_ptr<Eigen::MatrixXd> > & samples,
                               const int mode,
                               const int num_latent) {
  auto& sparseMode = (*data.Y)[mode];
  auto& U = samples[mode];
  const int N = U->cols();
  VectorView<Eigen::MatrixXd> view(samples, mode);

#pragma omp parallel for schedule(dynamic, 2)
  for (int n = 0; n < N; n++) {
    sample_latent_tensor(U, n, sparseMode, view, data.mean_value, noisePrecision, mu, Lambda);
  }
}

/** MacauPrior */
template<class FType>
void MacauPrior<FType>::init(const int num_latent, std::unique_ptr<FType> &Fmat, bool comp_FtF) {
  mu.resize(num_latent);
  mu.setZero();

  Lambda.resize(num_latent, num_latent);
  Lambda.setIdentity();
  Lambda *= 10;

  // parameters of Inv-Whishart distribution
  WI.resize(num_latent, num_latent);
  WI.setIdentity();
  mu0.resize(num_latent);
  mu0.setZero();
  b0 = 2;
  df = num_latent;

  // side information
  F = std::move(Fmat);
  use_FtF = comp_FtF;
  if (use_FtF) {
    FtF.resize(F->cols(), F->cols());
    At_mul_A(FtF, *F);
  }

  Uhat.resize(num_latent, F->rows());
  Uhat.setZero();

  beta.resize(num_latent, F->cols());
  beta.setZero();

  // initial value (should be determined automatically)
  lambda_beta = 5.0;
  // Hyper-prior for lambda_beta (mean 1.0, var of 1e+3):
  lambda_beta_mu0 = 1.0;
  lambda_beta_nu0 = 1e-3;
}

template<class FType>
void MacauPrior<FType>::sample_latents(Eigen::MatrixXd &U, const Eigen::SparseMatrix<double> &mat, double mean_value,
                    const Eigen::MatrixXd &samples, double alpha, const int num_latent) {
  const int N = U.cols();
#pragma omp parallel for schedule(dynamic, 2)
  for(int n = 0; n < N; n++) {
    // TODO: try moving mu + Uhat.col(n) inside sample_latent for speed
    sample_latent_blas(U, n, mat, mean_value, samples, alpha, mu + Uhat.col(n), Lambda, num_latent);
  }
}

template<class FType>
void MacauPrior<FType>::sample_latents(ProbitNoise& noiseModel, TensorData & data,
                               std::vector< std::unique_ptr<Eigen::MatrixXd> > & samples, int mode, const int num_latent) {
  // TODO:
}

template<class FType>
void MacauPrior<FType>::sample_latents(double noisePrecision,
                                       TensorData & data,
                                       std::vector< std::unique_ptr<Eigen::MatrixXd> > & samples,
                                       int mode,
                                       const int num_latent) {
  auto& sparseMode = (*data.Y)[mode];
  auto& U = samples[mode];
  const int N = U->cols();
  VectorView<Eigen::MatrixXd> view(samples, mode);

#pragma omp parallel for schedule(dynamic, 2)
  for (int n = 0; n < N; n++) {
    Eigen::VectorXd mu2 = mu + Uhat.col(n);
    sample_latent_tensor(U, n, sparseMode, view, data.mean_value, noisePrecision, mu2, Lambda);
  }
}

template<class FType>
void MacauPrior<FType>::update_prior(const Eigen::MatrixXd &U) {
  // residual (Uhat is later overwritten):
  Uhat.noalias() = U - Uhat;
  MatrixXd BBt = A_mul_At_combo(beta);
  // sampling Gaussian
  tie(mu, Lambda) = CondNormalWishart(Uhat, mu0, b0, WI + lambda_beta * BBt, df + beta.cols());
  sample_beta(U);
  compute_uhat(Uhat, *F, beta);
  lambda_beta = sample_lambda_beta(beta, Lambda, lambda_beta_nu0, lambda_beta_mu0);
}

template<class FType>
double MacauPrior<FType>::getLinkNorm() {
  return beta.norm();
}

/** Update beta and Uhat */
template<class FType>
void MacauPrior<FType>::sample_beta(const Eigen::MatrixXd &U) {
  const int num_feat = beta.cols();
  // Ft_y = (U .- mu + Normal(0, Lambda^-1)) * F + sqrt(lambda_beta) * Normal(0, Lambda^-1)
  // Ft_y is [ D x F ] matrix
  MatrixXd tmp = (U + MvNormal_prec_omp(Lambda, U.cols())).colwise() - mu;
  MatrixXd Ft_y = A_mul_B(tmp, *F) + sqrt(lambda_beta) * MvNormal_prec_omp(Lambda, num_feat);

  if (use_FtF) {
    MatrixXd K(FtF.rows(), FtF.cols());
    K.triangularView<Eigen::Lower>() = FtF;
    for (int i = 0; i < K.cols(); i++) {
      K(i,i) += lambda_beta;
    }
    chol_decomp(K);
    chol_solve_t(K, Ft_y);
    beta = Ft_y;
  } else {
    // BlockCG
    solve_blockcg(beta, *F, lambda_beta, Ft_y, tol, 32, 8);
  }
}

template<class FType>
void MacauPrior<FType>::sample_latents(ProbitNoise & noise, Eigen::MatrixXd &U, const Eigen::SparseMatrix<double> &mat,
                                       double mean_value, const Eigen::MatrixXd &samples, const int num_latent) {
    const int N = U.cols();
#pragma omp parallel for schedule(dynamic, 2)
  for(int n = 0; n < N; n++) {
    // TODO: try moving mu + Uhat.col(n) inside sample_latent for speed
    sample_latent_blas_probit(U, n, mat, mean_value, samples, mu + Uhat.col(n), Lambda, num_latent);
  }

}

void BPMFPrior::saveModel(std::string prefix) {
  writeToCSVfile(prefix + "-latentmean.csv", mu);
}

template<class FType>
void MacauPrior<FType>::saveModel(std::string prefix) {
  writeToCSVfile(prefix + "-latentmean.csv", mu);
  writeToCSVfile(prefix + "-link.csv", beta);
}

std::pair<double,double> posterior_lambda_beta(Eigen::MatrixXd & beta, Eigen::MatrixXd & Lambda_u, double nu, double mu) {
  const int D = beta.rows();
  MatrixXd BB(D, D);
  A_mul_At_combo(BB, beta);
  double nux = nu + beta.rows() * beta.cols();
  double mux = mu * nux / (nu + mu * (BB.selfadjointView<Eigen::Lower>() * Lambda_u).trace() );
  double b   = nux / 2;
  double c   = 2 * mux / nux;
  return std::make_pair(b, c);
}

double sample_lambda_beta(Eigen::MatrixXd & beta, Eigen::MatrixXd & Lambda_u, double nu, double mu) {
  auto gamma_post = posterior_lambda_beta(beta, Lambda_u, nu, mu);
  return rgamma(gamma_post.first, gamma_post.second);
}

/** global function */
void sample_latent(MatrixXd &s, int mm, const SparseMatrix<double> &mat, double mean_rating,
    const MatrixXd &samples, double alpha, const VectorXd &mu_u, const MatrixXd &Lambda_u,
    const int num_latent)
{
  // TODO: add cholesky update version
  MatrixXd MM = MatrixXd::Zero(num_latent, num_latent);
  VectorXd rr = VectorXd::Zero(num_latent);
  for (SparseMatrix<double>::InnerIterator it(mat, mm); it; ++it) {
    auto col = samples.col(it.row());
    MM.noalias() += col * col.transpose();
    rr.noalias() += col * ((it.value() - mean_rating) * alpha);
  }

  Eigen::LLT<MatrixXd> chol = (Lambda_u + alpha * MM).llt();
  if(chol.info() != Eigen::Success) {
    throw std::runtime_error("Cholesky Decomposition failed!");
  }

  rr.noalias() += Lambda_u * mu_u;
  chol.matrixL().solveInPlace(rr);
  for (int i = 0; i < num_latent; i++) {
    rr[i] += randn0();
  }
  chol.matrixU().solveInPlace(rr);
  s.col(mm).noalias() = rr;
}

void sample_latent_blas(MatrixXd &s, int mm, const SparseMatrix<double> &mat, double mean_rating,
    const MatrixXd &samples, double alpha, const VectorXd &mu_u, const MatrixXd &Lambda_u,
    const int num_latent)
{
  MatrixXd MM = Lambda_u;
  VectorXd rr = VectorXd::Zero(num_latent);
  for (SparseMatrix<double>::InnerIterator it(mat, mm); it; ++it) {
    auto col = samples.col(it.row());
    MM.triangularView<Eigen::Lower>() += alpha * col * col.transpose();
    rr.noalias() += col * ((it.value() - mean_rating) * alpha);
  }

  Eigen::LLT<MatrixXd> chol = MM.llt();
  if(chol.info() != Eigen::Success) {
    throw std::runtime_error("Cholesky Decomposition failed!");
  }

  rr.noalias() += Lambda_u * mu_u;
  chol.matrixL().solveInPlace(rr);
  for (int i = 0; i < num_latent; i++) {
    rr[i] += randn0();
  }
  chol.matrixU().solveInPlace(rr);
  s.col(mm).noalias() = rr;
}

void sample_latent_blas_probit(MatrixXd &s, int mm, const SparseMatrix<double> &mat, double mean_rating,
    const MatrixXd &samples, const VectorXd &mu_u, const MatrixXd &Lambda_u,
    const int num_latent)
{
    MatrixXd MM = Lambda_u;
    VectorXd rr = VectorXd::Zero(num_latent);
    double z;
    auto u = s.col(mm);
    for (SparseMatrix<double>::InnerIterator it(mat, mm); it; ++it) {
      auto col = samples.col(it.row());
      MM.triangularView<Eigen::Lower>() += col * col.transpose();
			double y = 2 * it.value() - 1;
      z = y * rand_truncnorm(y * col.dot(u), 1.0, 0.0);
      rr.noalias() += col * z;
    }
  Eigen::LLT<MatrixXd> chol = MM.llt();
  if(chol.info() != Eigen::Success) {
    throw std::runtime_error("Cholesky Decomposition failed!");
  }

  rr.noalias() += Lambda_u * mu_u;
  chol.matrixL().solveInPlace(rr);
  for (int i = 0; i < num_latent; i++) {
    rr[i] += randn0();
  }
  chol.matrixU().solveInPlace(rr);
  s.col(mm).noalias() = rr;
}

/**
 * X = A * B
 */
Eigen::MatrixXd A_mul_B(Eigen::MatrixXd & A, Eigen::MatrixXd & B) {
  MatrixXd out(A.rows(), B.cols());
  A_mul_B_blas(out, A, B);
  return out;
}

Eigen::MatrixXd A_mul_B(Eigen::MatrixXd & A, SparseFeat & B) {
  MatrixXd out(A.rows(), B.cols());
  A_mul_Bt(out, B.Mt, A);
  return out;
}

Eigen::MatrixXd A_mul_B(Eigen::MatrixXd & A, SparseDoubleFeat & B) {
  MatrixXd out(A.rows(), B.cols());
  A_mul_Bt(out, B.Mt, A);
  return out;
}

MacauPrior<Eigen::MatrixXd>* make_dense_prior(int nlatent, double* ptr, int nrows, int ncols, bool colMajor, bool comp_FtF) {
	MatrixXd* Fmat = new MatrixXd(0, 0);
	if (colMajor) {
		*Fmat = Map<Matrix<double, Dynamic, Dynamic, ColMajor> >(ptr, nrows, ncols);
	} else {
		*Fmat = Map<Matrix<double, Dynamic, Dynamic, RowMajor> >(ptr, nrows, ncols);
	}
	unique_ptr<MatrixXd> Fmat_ptr = unique_ptr<MatrixXd>(Fmat);
	return new MacauPrior<MatrixXd>(nlatent, Fmat_ptr, comp_FtF);
}

template class MacauPrior<SparseFeat>;
template class MacauPrior<SparseDoubleFeat>;
template class MacauPrior<Eigen::MatrixXd>;


#include <memory>
#include <vector>

#include <boost/shared_ptr.hpp>

#include <CGAL/Straight_skeleton_2.h>
#include <CGAL/create_offset_polygons_2.h>
#include <CGAL/create_straight_skeleton_2.h>
#include <CGAL/create_straight_skeleton_from_polygon_with_holes_2.h>

#include <jlcxx/module.hpp>
#include <jlcxx/smart_pointers.hpp>

#include <julia.h>

#include "polygon_2.hpp"
#include "utils.hpp"

namespace jlcgal {

template<typename T>
std::shared_ptr<T>
to_std(boost::shared_ptr<T> ptr) {
  return std::shared_ptr<T>(ptr.get(), [ptr](T*) mutable { ptr.reset(); });
}

template<typename Polygon>
jlcxx::Array<Polygon>
to_poly_jlarr(std::vector<boost::shared_ptr<Polygon>> ps) {
  jlcxx::Array<Polygon> jlarr;
  for (const auto& poly : ps) jlarr.push_back(Polygon(*poly));
  return jlarr;
}

void wrap_straight_skeleton_2(jlcxx::Module& cgal) {
  typedef CGAL::Straight_skeleton_2<Kernel> Skeleton_2;

  typedef CGAL::HalfedgeDS_in_place_list_face<Skeleton_2::Face>         Face;
  typedef CGAL::HalfedgeDS_in_place_list_halfedge<Skeleton_2::Halfedge> Halfedge;
  typedef CGAL::HalfedgeDS_in_place_list_vertex<Skeleton_2::Vertex>     Vertex;

  const std::string skel_2_name = "StraightSkeleton2";

  auto ssf = cgal.add_type<Face>    (skel_2_name + "Face");
  auto ssh = cgal.add_type<Halfedge>(skel_2_name + "Halfedge");
  auto ssv = cgal.add_type<Vertex>  (skel_2_name + "Vertex");

  ssf
    // Access Functions
    .method("id", &Face::id)
    .method("halfedge", [](const Face& f) { return *f.halfedge(); })
    ;

  ssh
    // Access Functions
    .method("id",       &Halfedge::id)
    .method("opposite", [](const Halfedge& h) { return *h.opposite(); })
    .method("next",     [](const Halfedge& h) { return *h.next(); })
    .method("prev",     [](const Halfedge& h) { return *h.prev(); })
    .method("vertex",   [](const Halfedge& h) { return *h.vertex(); })
    .method("face",     [](const Halfedge& h) { return *h.face(); })
    .method("defining_contour_edge", [](const Halfedge& h) {
      return *h.defining_contour_edge();
    })
    // Predicates
    .method("has_null_segment",  &Halfedge::has_null_segment)
    .method("has_infinite_time", &Halfedge::has_infinite_time)
    .method("is_border",         &Halfedge::is_border)
    .method("is_bisector",       &Halfedge::is_bisector)
    .method("is_inner_bisector", &Halfedge::is_inner_bisector)
    .method("slope",             &Halfedge::slope)
    ;

  ssv
    // Access Functions
    .method("id",       &Vertex::id)
    .method("degree",   &Vertex::degree)
    .method("halfedge", [](const Vertex& v) { return *v.halfedge(); })
    .method("point",    &Vertex::point)
    ;
  cgal.set_override_module(jl_base_module);
  ssv
    .method("time", &Vertex::time)
    ;
  cgal.unset_override_module();
  ssv
    .method("primary_bisector", [](const Vertex& v) {
      return *v.primary_bisector();
    })
    // Queries
    .method("has_infinite_time", &Vertex::has_infinite_time)
    .method("has_null_point",    &Vertex::has_null_point)
    .method("is_contour",        &Vertex::is_contour)
    .method("is_skeleton",       &Vertex::is_skeleton)
    .method("is_split",          &Vertex::is_split)
    ;

  cgal.add_type<Skeleton_2>(skel_2_name)
    // Access Member Functions
    .method("size_of_faces",     &Skeleton_2::size_of_faces)
    .method("size_of_halfedges", &Skeleton_2::size_of_halfedges)
    .method("size_of_vertices",  &Skeleton_2::size_of_vertices)
    .method("faces", [](const Skeleton_2& s) {
      return collect(s.faces_begin(), s.faces_end());
    })
    .method("halfedges", [](const Skeleton_2& s) {
      return collect(s.halfedges_begin(), s.halfedges_end());
    })
    .method("vertices", [](const Skeleton_2& s) {
      return collect(s.vertices_begin(), s.vertices_end());
    })
    // Predicates
    .method("is_valid", &Skeleton_2::is_valid)
    ;

  cgal.method("create_exterior_straight_skeleton_2",
              [](const FT& max_offset, const Polygon_2& poly) {
    return to_std(CGAL::create_exterior_straight_skeleton_2(max_offset,
                                                            poly,
                                                            Kernel()));
  });
  cgal.method("create_exterior_straight_skeleton_2",
              [](const FT& max_offset, jlcxx::ArrayRef<Point_2> ps) {
    // because `bbox_2` uses operator-> from the input iterator... which
    // ArrayRef::iterator is missing, so we copy it into a vector.
    std::vector<Point_2> vps(ps.begin(), ps.end());
    return to_std(CGAL::create_exterior_straight_skeleton_2(max_offset,
                                                            vps.begin(),
                                                            vps.end(),
                                                            Kernel()));
  });
  cgal.method("create_interior_straight_skeleton_2", [](const Polygon_2& poly) {
    return to_std(CGAL::create_interior_straight_skeleton_2(poly, Kernel()));
  });
  cgal.method("create_interior_straight_skeleton_2",
              [](const Polygon_with_holes_2& poly) {
    return to_std(CGAL::create_interior_straight_skeleton_2(poly));
  });
  cgal.method("create_interior_straight_skeleton_2",
              [](jlcxx::ArrayRef<Point_2> contour) {
    return to_std(CGAL::create_interior_straight_skeleton_2(contour.begin(),
                                                            contour.end(),
                                                            Kernel()));
  });
  cgal.method("create_interior_straight_skeleton_2",
              [](jlcxx::ArrayRef<Point_2> contour,
                 jlcxx::ArrayRef<Polygon_2> holes) {
    return to_std(CGAL::create_interior_straight_skeleton_2(contour.begin(),
                                                            contour.end(),
                                                            holes.begin(),
                                                            holes.end(),
                                                            Kernel()));
  });

  cgal.method("create_offset_polygons_2", [](const FT& offset,
                                             const Skeleton_2& ss) {
    return to_poly_jlarr(CGAL::create_offset_polygons_2<Polygon_2>(offset,
                                                                   ss,
                                                                   Kernel()));
  });
}

} // jlcgal