/*******************************************************************************
 * Copyright (c) 2018-, UT-Battelle, LLC.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the MIT License
 * which accompanies this distribution.
 *
 * Contributors:
 *   Alexander J. McCaskey - initial API and implementation
 *   Thien Nguyen - implementation
 *******************************************************************************/
#include "qcor_observable.hpp"

#include "ObservableTransform.hpp"
#include "xacc.hpp"
// #include "xacc_quantum_gate_api.hpp"
#include <spdlog/fmt/fmt.h>
#include <Eigen/Dense>

#include <algorithm>
#include <cassert>

#include "CompositeInstruction.hpp"
#include "FermionOperator.hpp"
#include "PauliOperator.hpp"
#include "qalloc.hpp"
#include "qcor_ir.hpp"
#include "qcor_pimpl_impl.hpp"
#include "xacc_internal_compiler.hpp"
#include "xacc_quantum_gate_api.hpp"
#include "xacc_service.hpp"

namespace qcor {

// ---------------- Operator ---------------------- //

// Internal hidden implementation
class Operator::OperatorImpl
    : public commutative_ring<OperatorImpl>,
      public equality_comparable<OperatorImpl>,
      public commutative_multipliable<OperatorImpl, double>,
      public commutative_multipliable<OperatorImpl, std::complex<double>> {
  friend class Operator;

 private:
  enum operation { PlusEquals, MinusEqual, StarEqual };

  // This function performs sub-class specific algebraic operation
  // with another operator instance, or scalar value (double, complex)
  template <typename T>
  struct cast_and_apply {
    void operator()(std::shared_ptr<xacc::Observable> obs, operation _op,
                    const OperatorImpl &other) {
      auto input_casted = std::dynamic_pointer_cast<T>(obs);
      auto other_casted = std::dynamic_pointer_cast<T>(other.op);
      assert(input_casted && other_casted &&
             "Invalid Operator sub-types for arithmetic operation.");

      switch (_op) {
        case operation::PlusEquals:
          input_casted->operator+=(*other_casted.get());
          break;
        case operation::MinusEqual:
          input_casted->operator-=(*other_casted.get());
          break;
        case operation::StarEqual:
          input_casted->operator*=(*other_casted.get());
          break;
        default:
          std::cout << "Invalid Selection\n";
          exit(0);
          break;
      }
    }

    void operator()(std::shared_ptr<xacc::Observable> obs,
                    const double &value) {
      auto input_casted = std::dynamic_pointer_cast<T>(obs);
      assert(input_casted &&
             "Invalid Operator sub-types for arithmetic operation.");
      input_casted->operator*=(value);
    }

    void operator()(std::shared_ptr<xacc::Observable> obs,
                    const std::complex<double> &value) {
      auto input_casted = std::dynamic_pointer_cast<T>(obs);
      assert(input_casted &&
             "Invalid Operator sub-types for arithmetic operation.");
      input_casted->operator*=(value);
    }
  };

  // Apply the desired operation between this operator an the given other operator
  OperatorImpl &apply(const operation &_op, const OperatorImpl &other) {
    if (type == "pauli") {
      cast_and_apply<xacc::quantum::PauliOperator>()(op, _op, other);
    } else if (type == "fermion") {
      cast_and_apply<xacc::quantum::FermionOperator>()(op, _op, other);
    } else {
      // FIXME / TODO As we add more Operator types, update
      // this if/else section
      error("Invalid Operator type for the given algebraic operation.");
    }
    return *this;
  }

  // Apply the desired operation between this operator an the given scalar double
  OperatorImpl &apply(const double &d) {
    if (type == "pauli") {
      cast_and_apply<xacc::quantum::PauliOperator>()(op, d);
    } else if (type == "fermion") {
      cast_and_apply<xacc::quantum::FermionOperator>()(op, d);
    } else {
      // FIXME / TODO As we add more Operator types, update
      // this if/else section
      error("Invalid Operator type for the given algebraic operation.");
    }
    return *this;
  }

  // Apply the desired operation between this operator an the given scalar complex
  OperatorImpl &apply(const std::complex<double> &d) {
    if (type == "pauli") {
      cast_and_apply<xacc::quantum::PauliOperator>()(op, d);
    } else if (type == "fermion") {
      cast_and_apply<xacc::quantum::FermionOperator>()(op, d);
    } else {
      // FIXME / TODO As we add more Operator types, update
      // this if/else section
      error("Invalid Operator type for the given algebraic operation.");
    }
    return *this;
  }

  std::string type;
  std::shared_ptr<xacc::Observable> op;

 public:
  // Internal impl constructed from HetMap of options, or string-like expression
  OperatorImpl() = default;
  OperatorImpl(const std::string &_type, const std::string &expr)
      : type(_type) {
    op = xacc::quantum::getObservable(type, expr);
  }

  OperatorImpl(const std::string &name, xacc::HeterogeneousMap &options) {
    auto tmp_op = xacc::quantum::getObservable(name, options);
    auto obs_str = tmp_op->toString();

    if (obs_str.find("^") != std::string::npos) {

      op = xacc::quantum::getObservable("fermion", obs_str);
      type = "fermion";

    } else if (obs_str.find("X") != std::string::npos ||
               obs_str.find("Y") != std::string::npos ||
               obs_str.find("Z") != std::string::npos) {
      op = xacc::quantum::getObservable("pauli", obs_str);
      type = "pauli";
    }
  }

  OperatorImpl(const std::string &_type, std::shared_ptr<xacc::Observable> _op)
      : type(_type), op(_op) {}
  OperatorImpl(const OperatorImpl &other) : type(other.type), op(other.op) {}

  // Implement internal Algebraic API
  OperatorImpl &operator+=(const OperatorImpl &v) noexcept {
    return apply(operation::PlusEquals, v);
  }

  OperatorImpl &operator-=(const OperatorImpl &v) noexcept {
    return apply(operation::MinusEqual, v);
  }
  OperatorImpl &operator*=(const OperatorImpl &v) noexcept {
    return apply(operation::StarEqual, v);
  }
  OperatorImpl &operator*=(const double v) noexcept { return apply(v); }
  OperatorImpl &operator*=(const std::complex<double> v) noexcept {
    return apply(v);
  }

  bool operator==(const OperatorImpl &v) noexcept {
    if (type == "pauli") {
      auto casted = std::dynamic_pointer_cast<xacc::quantum::PauliOperator>(op);
      auto other_casted =
          std::dynamic_pointer_cast<xacc::quantum::PauliOperator>(v.op);
      assert(casted && other_casted && "Invalid types for Operator == check.");
      return casted->operator==(*other_casted.get());
    } else if (type == "fermion") {
      auto casted =
          std::dynamic_pointer_cast<xacc::quantum::FermionOperator>(op);
      auto other_casted =
          std::dynamic_pointer_cast<xacc::quantum::FermionOperator>(v.op);
      assert(casted && other_casted && "Invalid types for Operator == check.");
      return casted->operator==(*other_casted.get());
    } else {
      // FIXME / TODO As we add more Operator types, update
      // this if/else section
    }
    return false;
  }

  std::vector<Operator> getSubTerms() {
    std::vector<Operator> ret;
    for (auto sub_term : op->getSubTerms()) {
      ret.emplace_back(OperatorImpl(type, sub_term));
    }
    return ret;
  }

  std::vector<Operator> getNonIdentitySubTerms() {
    std::vector<Operator> ret;
    for (auto sub_term : op->getNonIdentitySubTerms()) {
      ret.emplace_back(OperatorImpl(type, sub_term));
    }
    return ret;
  }

  Operator getIdentitySubTerm() {
    auto id_term = op->getIdentitySubTerm();
    if (!id_term) {
      // THROW AN ERROR.
     error("There is no identity sub term. exiting.");
    }
    return Operator(OperatorImpl(type, op->getIdentitySubTerm()));
  }

  bool hasIdentitySubTerm() { return op->getIdentitySubTerm() != nullptr; }
  std::complex<double> coefficient() { return op->coefficient(); }

  std::vector<SparseElement> to_sparse_matrix() {
    auto sp_el = op->to_sparse_matrix();
    std::vector<SparseElement> ret;
    for (auto el : sp_el) {
      ret.emplace_back(el.row(), el.col(), el.coeff());
    }
    return ret;
  }

  Operator commutator(Operator &other) {
    return Operator(OperatorImpl(type, op->commutator(other.m_internal->op)));
  }

  std::pair<std::vector<int>, std::vector<int>> toBinaryVectors(
      const int nQubits) {
    assert(type == "pauli" && "toBinaryVectors only works for pauli operators");
    return std::dynamic_pointer_cast<xacc::quantum::PauliOperator>(op)
        ->toBinaryVectors(nQubits);
  }

  void mapQubitSites(std::map<int, int> &siteMap) {
    assert(type == "pauli" && "mapQubitSites only works for pauli operators");
    std::dynamic_pointer_cast<xacc::quantum::PauliOperator>(op)->mapQubitSites(
        siteMap);
  }
};

Operator &Operator::operator=(const Operator &other) {
  m_internal->op = other.m_internal->op;
  m_internal->type = other.m_internal->type;
  return *this;
}

Operator::Operator(const std::string &name, xacc::HeterogeneousMap &options)
    : m_internal(name, options) {}
Operator::Operator(const std::string &type, const std::string &expr)
    : m_internal(type, expr) {}
Operator::Operator(const Operator &op)
    : m_internal(op.m_internal.operator->()->type,
                 op.m_internal.operator->()->op->toString()) {}
Operator::Operator(const OperatorImpl &&impl)
    : m_internal(impl.type, impl.op) {}

Operator::Operator() = default;
Operator::~Operator() = default;

std::shared_ptr<xacc::Identifiable> Operator::get_as_opaque() {
  return std::dynamic_pointer_cast<xacc::Identifiable>(m_internal->op);
}

Operator &Operator::operator+=(const Operator &v) noexcept {
  OperatorImpl *other = v.m_internal.operator->();
  m_internal->operator+=(*other);
  return *this;
}

Operator &Operator::operator-=(const Operator &v) noexcept {
  OperatorImpl *other = v.m_internal.operator->();
  m_internal->operator-=(*other);
  return *this;
}
Operator &Operator::operator*=(const Operator &v) noexcept {
  OperatorImpl *other = v.m_internal.operator->();
  m_internal->operator*=(*other);
  return *this;
}
bool Operator::operator==(const Operator &v) noexcept {
  OperatorImpl *other = v.m_internal.operator->();
  return m_internal->operator==(*other);
}
Operator &Operator::operator*=(const double v) noexcept {
  m_internal->operator*=(v);
  return *this;
}
Operator &Operator::operator*=(const std::complex<double> v) noexcept {
  m_internal->operator*=(v);
  return *this;
}

int Operator::nQubits() { return m_internal->op->nBits(); }

std::pair<std::vector<int>, std::vector<int>> Operator::toBinaryVectors(
    const int nQubits) {
  return m_internal->toBinaryVectors(nQubits);
}
void Operator::mapQubitSites(std::map<int, int> &siteMap) {
  return m_internal->mapQubitSites(siteMap);
}

std::vector<std::shared_ptr<CompositeInstruction>> Operator::observe(
    std::shared_ptr<CompositeInstruction> program) {
  auto as_xacc = program->as_xacc();
  auto cis = m_internal->op->observe(as_xacc);
  std::vector<std::shared_ptr<CompositeInstruction>> ret;
  for (auto ci : cis) {
    ret.emplace_back(std::make_shared<CompositeInstruction>(
        std::dynamic_pointer_cast<xacc::Identifiable>(ci)));
  }

  return ret;
}

std::vector<Operator> Operator::getSubTerms() {
  return m_internal->getSubTerms();
}

std::vector<Operator> Operator::getNonIdentitySubTerms() {
  return m_internal->getNonIdentitySubTerms();
}
bool Operator::hasIdentitySubTerm() { return m_internal->hasIdentitySubTerm(); }

std::string Operator::toString() const { return m_internal->op->toString(); }
Operator Operator::getIdentitySubTerm() {
  return m_internal->getIdentitySubTerm();
}
std::complex<double> Operator::coefficient() {
  return m_internal->coefficient();
}

std::vector<Operator::SparseElement> Operator::to_sparse_matrix() {
  return m_internal->to_sparse_matrix();
}

Operator Operator::commutator(Operator &op) {
  return m_internal->commutator(op);
}

Operator Operator::transform(const std::string &type,
                             xacc::HeterogeneousMap m) {
  auto transformed =
      xacc::getService<xacc::ObservableTransform>(type)->transform(
          m_internal->op);
  auto new_type =
      std::dynamic_pointer_cast<xacc::quantum::FermionOperator>(transformed)
          ? "fermion"
          : "pauli";
  return Operator(OperatorImpl(new_type, transformed));
}

void __internal_exec_observer(
    xacc::AcceleratorBuffer *b,
    std::vector<std::shared_ptr<CompositeInstruction>> v) {
  // auto vv = v.get()->operator->()->program;
  std::vector<std::shared_ptr<xacc::CompositeInstruction>> tmp;
  std::transform(v.begin(), v.end(), std::back_inserter(tmp),
                 [](std::shared_ptr<CompositeInstruction> c)
                     -> std::shared_ptr<xacc::CompositeInstruction> {
                   return c->as_xacc();
                 });
  xacc::internal_compiler::execute(b, tmp);
}

Operator operator+(double coeff, Operator op) {
  return Operator("pauli", fmt::format("{}", coeff)) + op;
}
Operator operator+(Operator op, double coeff) {
  return op + Operator("pauli", fmt::format("{}", coeff));
}

Operator operator-(double coeff, Operator op) {
  return Operator("pauli", fmt::format("{}", coeff)) - op;
}

Operator operator-(Operator op, double coeff) {
  return op - Operator("pauli", fmt::format("{}", coeff));
}

Operator adag(int idx) {
  return Operator("fermion", fmt::format("1.0 {}^", idx));
}
Operator a(int idx) { return Operator("fermion", fmt::format("1.0 {}", idx)); }

Operator X(int idx) { return Operator("pauli", fmt::format("X{}", idx)); }
Operator Y(int idx) { return Operator("pauli", fmt::format("Y{}", idx)); }
Operator Z(int idx) { return Operator("pauli", fmt::format("Z{}", idx)); }

Operator allZs(const int nQubits) {
  auto ret = Z(0);
  for (int i = 1; i < nQubits; i++) {
    ret *= Z(i);
  }
  return ret;
}

Operator SP(int idx) {
  std::complex<double> imag(0.0, 1.0);
  return X(idx) + imag * Y(idx);
}

Operator SM(int idx) {
  std::complex<double> imag(0.0, 1.0);
  return X(idx) - imag * Y(idx);
}

Eigen::MatrixXcd get_dense_matrix(Operator &op) {
  auto mat_el = op.to_sparse_matrix();
  auto size = std::pow(2, op.nBits());
  Eigen::MatrixXcd mat = Eigen::MatrixXcd::Zero(size, size);
  for (auto el : mat_el) {
    mat(el.row(), el.col()) = el.coeff();
  }
  return mat;
}

Operator createOperator(const std::string &repr) {
  if (!xacc::isInitialized())
    xacc::internal_compiler::compiler_InitializeXACC();
  return qcor::Operator("pauli", repr);
}

Operator createOperator(const std::string &name, const std::string &repr) {
  if (!xacc::isInitialized())
    xacc::internal_compiler::compiler_InitializeXACC();
  return qcor::Operator(name, repr);
}

Operator createOperator(const std::string &name, HeterogeneousMap &&options) {
  return createOperator(name, options);
}
Operator createOperator(const std::string &name, HeterogeneousMap &options) {
  return qcor::Operator(name, options);
}

Operator createObservable(const std::string &repr) {
  return createOperator(repr);
}

Operator createObservable(const std::string &name, const std::string &repr) {
  return createOperator(name, repr);
}

Operator operatorTransform(const std::string &type, Operator &op) {
  return op.transform(type);
}

Operator _internal_python_createObservable(const std::string &name,
                                           const std::string &repr) {
  return createOperator(name, repr);
}

namespace __internal__ {
std::map<std::size_t, Operator> cached_observables = {};

std::vector<std::shared_ptr<CompositeInstruction>> observe(
    Operator &obs, std::shared_ptr<CompositeInstruction> program) {
  return obs.observe(program);
}
}  // namespace __internal__

double observe(std::shared_ptr<CompositeInstruction> program, Operator &obs,
               xacc::internal_compiler::qreg &q) {
  // Observe the program
  auto v = obs.observe(program);

  std::vector<std::shared_ptr<xacc::CompositeInstruction>> tmp;
  std::transform(v.begin(), v.end(), std::back_inserter(tmp),
                 [](std::shared_ptr<CompositeInstruction> c)
                     -> std::shared_ptr<xacc::CompositeInstruction> {
                   return c->as_xacc();
                 });
  xacc::internal_compiler::execute(q.results(), tmp);

  // We want to contract q children buffer
  // exp-val-zs with obs term coeffs
  return q.weighted_sum(
      std::dynamic_pointer_cast<xacc::Observable>(obs.get_as_opaque()).get());
}

}  // namespace qcor

std::ostream &operator<<(std::ostream &os, qcor::Operator const &m) {
  return os << m.toString();
}


#include <kr_serial_interface/comm_types.h>
#include <kr_serial_interface/decode_msgs.h>

#include <Eigen/Geometry>

namespace kr_mav_msgs
{
bool decodeOutputData(const std::vector<uint8_t> &data, kr_mav_msgs::OutputData &output)
{
  struct OUTPUT_DATA output_data;
  if(data.size() != sizeof(output_data))
    return false;

  memcpy(&output_data, &data[0], sizeof(output_data));
  output.loop_rate = output_data.loop_rate;
  output.voltage = output_data.voltage / 1e3f;

  const double roll = output_data.roll / 1e2f * M_PI / 180;
  const double pitch = output_data.pitch / 1e2f * M_PI / 180;
  const double yaw = output_data.yaw / 1e2f * M_PI / 180;
  // Asctec (2012 firmware) uses  Z-Y-X convention
  Eigen::Quaternionf q = Eigen::AngleAxisf(yaw, Eigen::Vector3f::UnitZ()) *
                         Eigen::AngleAxisf(pitch, Eigen::Vector3f::UnitY()) *
                         Eigen::AngleAxisf(roll, Eigen::Vector3f::UnitX());
  output.orientation.w = q.w();
  output.orientation.x = q.x();
  output.orientation.y = q.y();
  output.orientation.z = q.z();

  output.angular_velocity.x = output_data.ang_vel[0] * 0.0154f * M_PI / 180;
  output.angular_velocity.y = output_data.ang_vel[1] * 0.0154f * M_PI / 180;
  output.angular_velocity.z = output_data.ang_vel[2] * 0.0154f * M_PI / 180;

  output.linear_acceleration.x = output_data.acc[0] / 1e3f * 9.81f;
  output.linear_acceleration.y = output_data.acc[1] / 1e3f * 9.81f;
  output.linear_acceleration.z = output_data.acc[2] / 1e3f * 9.81f;

  output.pressure_dheight = output_data.dheight / 1e3f;
  output.pressure_height = output_data.height / 1e3f;

  output.magnetic_field.x = output_data.mag[0] / 2500.0f;
  output.magnetic_field.y = output_data.mag[1] / 2500.0f;
  output.magnetic_field.z = output_data.mag[2] / 2500.0f;

  for(int i = 0; i < 8; i++)
  {
    output.radio_channel[i] = output_data.radio[i];
  }

  // Asctec firmware uses the following rotor numbering convention:
  //   *1*    Front
  // 3     4
  //    2
  //
  // But we want:
  //   *1*    Front
  // 2     4
  //    3
  const int motors_map[] = {0, 2, 1, 3};
  for(int i = 0; i < 4; i++)
  {
    // The following conversion is from
    // http://wiki.asctec.de/display/AR/List+of+all+predefined+variables%2C+commands+and+parameters
    //
    // motorRPM = 1075+m*37.625
    //
    // Note: If m == 0, the motors are not spinning.
    int m = output_data.rpm[motors_map[i]];
    output.motor_rpm[i] = (m == 0 ? 0 : 1075) + m * 37.625;
  }

  output.seq = output_data.seq;

  return true;
}

bool decodeStatusData(const std::vector<uint8_t> &data, kr_mav_msgs::StatusData &status)
{
  struct STATUS_DATA status_data;
  if(data.size() != sizeof(status_data))
    return false;
  memcpy(&status_data, &data[0], sizeof(status_data));

  status.loop_rate = status_data.loop_rate;
  status.voltage = status_data.voltage / 1e3f;
  status.seq = status_data.seq;

  return true;
}

}  // namespace kr_mav_msgs


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_ARCH_COMMON_SIMD_FUNCTION_POW_HPP_INCLUDED
#define BOOST_SIMD_ARCH_COMMON_SIMD_FUNCTION_POW_HPP_INCLUDED
#include <boost/simd/detail/overload.hpp>

#include <boost/simd/constant/nan.hpp>
#include <boost/simd/constant/one.hpp>
#include <boost/simd/function/any.hpp>
#include <boost/simd/function/if_else.hpp>
#include <boost/simd/function/is_flint.hpp>
#include <boost/simd/function/is_ltz.hpp>
#include <boost/simd/function/is_odd.hpp>
#include <boost/simd/function/logical_and.hpp>
#include <boost/simd/function/logical_andnot.hpp>
#include <boost/simd/function/multiplies.hpp>
#include <boost/simd/function/if_neg.hpp>
#include <boost/simd/function/pow_abs.hpp>
#include <boost/simd/function/shift_right.hpp>
#include <boost/simd/function/sqr.hpp>

namespace boost { namespace simd { namespace ext
{
  namespace bd = boost::dispatch;
  namespace bs = boost::simd;
  BOOST_DISPATCH_OVERLOAD_IF( pow_
                           , (typename A0,typename X)
                           , (detail::is_native<X>)
                           , bd::cpu_
                           , bs::pack_<bd::floating_<A0>,X>
                           , bs::pack_<bd::floating_<A0>,X>
                           )
  {
    BOOST_FORCEINLINE A0 operator()( const A0& a0, const A0& a1) BOOST_NOEXCEPT
    {
      auto nega0 = is_negative(a0);
      A0 z = pow_abs(a0, a1);
      z =  if_neg(logical_and(is_odd(a1), nega0), z);
      auto invalid =  logical_andnot(nega0, logical_or(is_flint(a1), is_inf(a1)));
      z = if_else(invalid, Nan<A0>(), z);
      return z;
    }
  };

  BOOST_DISPATCH_OVERLOAD_IF(pow_
                          , (typename A0,typename A1,typename X)
                          , (detail::is_native<X>)
                          , bd::cpu_
                          , bs::pack_<bd::arithmetic_<A0>, X>
                          , bs::pack_<bd::uint_<A1>, X>
                          )
   {
      BOOST_FORCEINLINE A0 operator()(const A0& a0, const A1& a1 ) const BOOST_NOEXCEPT
      {
        A0 base = a0;
        A1 exp = a1;
        A0 result = One<A0>();
        while(bs::any(exp))
        {
          result *= if_else(is_odd(exp), base, One<A0>());
          exp =  exp >> 1;
          base = sqr(base);
        }
        return result;
       }
   };

  BOOST_DISPATCH_OVERLOAD_IF ( pow_
                          , (typename A0, typename A1, typename X)
                          , (detail::is_native<X>)
                          , bd::cpu_
                          , bs::pack_< bd::arithmetic_<A0>, X>
                          , bd::constant_< bd::uint_<A1>>
                          )
  {
    using result_type = A0;

    BOOST_FORCEINLINE result_type operator() ( A0 const& a0, A1) const BOOST_NOEXCEPT
    {
      return pow_expander<A1::value>::call(a0);
    }
  };

  BOOST_DISPATCH_OVERLOAD_IF ( pow_
                          , (typename A0, typename A1, typename X)
                          , (detail::is_native<X>)
                          , bd::cpu_
                          , bs::pack_< bd::floating_<A0>, X>
                          , bd::constant_< bd::int_<A1>>
                          )
  {
    using result_type = A0;

    BOOST_FORCEINLINE result_type operator() ( A0 const& a0, A1) const BOOST_NOEXCEPT
    {
      return eval(a0, boost::mpl::bool_<(A1::value >= 0)>());
    }

    BOOST_FORCEINLINE result_type eval( A0 const& a0, boost::mpl::true_) const BOOST_NOEXCEPT
    {
      return pow_expander<A1::value>::call(a0);
    }

    BOOST_FORCEINLINE result_type eval( A0 const& a0, boost::mpl::false_) const BOOST_NOEXCEPT
    {
      return pow_expander<-A1::value>::call(rec(a0));
    }
  };

  BOOST_DISPATCH_OVERLOAD_IF ( pow_
                          , (typename A0, typename A1, typename X)
                          , (detail::is_native<X>)
                          , bd::cpu_
                          , bs::pack_< bd::arithmetic_<A0>, X>
                          , bd::scalar_< bd::uint_<A1>>
                          )
  {
    using result_type = A0;

    A0 operator() ( A0 const& a0, A1 const& a1) const BOOST_NOEXCEPT
    {
      A0 base = a0;
      A1 exp = a1;

      result_type result = One<result_type>();
      while(exp)
      {
        if(is_odd(exp))
          result *= base;
        exp >>= 1;
        base = sqr(base);
      }

      return result;
    }
  };

  BOOST_DISPATCH_OVERLOAD_IF ( pow_
                          , (typename A0, typename A1, typename X)
                          , (detail::is_native<X>)
                          , bd::cpu_
                          , bs::pack_< bd::integer_<A0>, X>
                          , bs::pack_< bd::int_<A1>, X>
                          )
  {
    using result_type = A0;

    BOOST_FORCEINLINE A0 operator() ( A0 const& a0, A1 const& a1) const BOOST_NOEXCEPT
    {
      BOOST_ASSERT_MSG( boost::simd::assert_all(a1 >= 0), "integral pow with signed exponent" );

      using u_t =  bd::as_integer_t<A1, unsigned>;
      return pow(a0, bitwise_cast<u_t>(a1));
    }
  };

  BOOST_DISPATCH_OVERLOAD_IF ( pow_
                          , (typename A0, typename A1, typename X)
                          , (detail::is_native<X>)
                          , bd::cpu_
                          , bs::pack_< bd::floating_<A0>, X>
                          , bs::pack_< bd::int_<A1>, X>
                          )
  {
    using result_type = A0;

    A0 operator() ( A0 const& a0, A1 const& a1) const BOOST_NOEXCEPT
    {
      using u_t =  bd::as_integer_t<A1, unsigned>;
      auto ltza1 = is_ltz(a1);
      A0 p = pow(a0, bitwise_cast<u_t>(if_neg(ltza1, a1)));
      return if_else(ltza1, rec(p), p);
    }
  };

  BOOST_DISPATCH_OVERLOAD_IF( pow_
                            , (typename A0,typename X)
                            , (detail::is_native<X>)
                            , bs::raw_tag
                            , bd::cpu_
                            , bs::pack_<bd::floating_<A0>,X>
                            , bs::pack_<bd::floating_<A0>,X>
                           )
  {
    BOOST_FORCEINLINE A0 operator()(const raw_tag &,
                                    const A0& a0, const A0& a1) BOOST_NOEXCEPT
    {
      auto nega0 = is_negative(a0);
      A0 z = raw_(pow_abs(a0, a1));
      z =  if_neg(logical_and(is_odd(a1), nega0), z);
      auto invalid =  logical_andnot(nega0, logical_or(is_flint(a1), is_inf(a1)));
      z = if_else(invalid, Nan<A0>(), z);
      return z;
    }
  };
} } }

#endif


//
// Created by David Wise on 24/03/2017.
//

#include <iostream>
#include "donorClass.h"
#include <stdio.h>
#include <Eigen/Dense>
//#include "gnuplot-iostream/gnuplot-iostream.h"
#include <fstream>

using namespace std;

int main() {

// Initialise a donor with required nuclear spin and hyperfine coupling
	Donor phos;
    phos.initialise(0.5, 7.29e-26);

// Create a savefile for the eigen values
    std::ofstream saveFile("donorsgotest.txt");

// Create matrices to store eigenvalues (need (nucSpin + 0.5) * 4)
    MatrixXcd eig1;
    MatrixXcd eig2;
    MatrixXcd eig3;
    MatrixXcd eig4;
    MatrixXcd fields;

// Loop through a number of steps in magnetic field finding eigenvalues at each step
    int numSteps = 1000;
    for (double incr = 0; incr <numSteps; ++incr) {
        double maxField = 1;
        eig1.resize(numSteps, 1);
        eig2.resize(numSteps, 1);
        eig3.resize(numSteps, 1);
        eig4.resize(numSteps, 1);
        fields.resize(numSteps, 1);
        MatrixXcd Eigs = phos.getEigs(incr);
        eig1(incr) = Eigs(0);
        eig2(incr) = Eigs(1);
        eig4(incr) = Eigs(2);
        eig3(incr) = Eigs(3);
        fields(incr) = (maxField/numSteps)*incr;
    };

// Save eigenvalues to file
    saveFile << "Fields \n" << fields << "\n Eigs 1 \n" << eig1 << "\n Eigs 2 \n" << eig2 << "\n Eigs 3 \n" << eig3 << "\n Eigs 4 \n" << eig4;
    saveFile.close();

	return 0;
}


/*

PICCANTE
The hottest HDR imaging library!
http://vcg.isti.cnr.it/piccante

Copyright (C) 2014
Visual Computing Laboratory - ISTI CNR
http://vcg.isti.cnr.it
First author: Francesco Banterle

This Source Code Form is subject to the terms of the Mozilla Public
License, v. 2.0. If a copy of the MPL was not distributed with this
file, You can obtain one at http://mozilla.org/MPL/2.0/.

*/

#ifndef PIC_ALGORITHMS_POISSON_IMAGE_EDITING_HPP
#define PIC_ALGORITHMS_POISSON_IMAGE_EDITING_HPP

#include <vector>

#include "../base.hpp"
#include "../image.hpp"
#include "../util/std_util.hpp"
#include "../filtering/filter_laplacian.hpp"

#ifndef PIC_DISABLE_EIGEN

#ifndef PIC_EIGEN_NOT_BUNDLED
    #include "../externals/Eigen/Sparse"
    #include "../externals/Eigen/src/SparseCore/SparseMatrix.h"
#else
    #include <Eigen/Sparse>
    #include <Eigen/src/SparseCore/SparseMatrix.h>
#endif

#endif

namespace pic {

#ifndef PIC_DISABLE_EIGEN
/**
 * @brief computePoissonImageEditing
 * @param source
 * @param target
 * @param mask
 * @param ret
 * @return
 */
PIC_INLINE Image *computePoissonImageEditing(Image *source, Image *target, bool *mask, Image *ret = NULL)
{
    if((source == NULL) || (target == NULL) || (mask == NULL)) {
        return NULL;
    }

    //allocate the output
    if(ret == NULL) {
        ret = target->clone();
    }

    int width  = target->width;
    int height = target->height;

    #ifdef PIC_DEBUG
        printf("Init matrix...");
    #endif

    Image *lap_source = FilterLaplacian::execute(source, NULL);

    std::vector< Eigen::Triplet< double > > tL;

    //indices pass
    int *index = new int[width * height];
    int count = 0;
    for(int i = 0; i < height; i++) {
        int tmpI = i * width;

        for(int j = 0; j < width; j++) {
            int indI = tmpI + j;

            if(mask[indI]) {
                index[indI] = count;
                count++;
            } else {
                index[indI] = 0;
            }
        }
    }

    //matrix A pass
    count = 0;
    for(int i = 0; i < height; i++) {
        int tmpI = i * width;

        for(int j = 0; j < width; j++) {
            int indI = tmpI + j;

            if(mask[indI]) {
                if((j + 1) < (width - 1)) {
                    if(mask[indI + 1]) {
                        tL.push_back(Eigen::Triplet< double > (count, index[indI + 1], -1.0));
                    }
                }

                if((j - 1) > -1) {
                    if(mask[indI - 1]) {
                        tL.push_back(Eigen::Triplet< double > (count, index[indI - 1], -1.0));
                    }
                }

                if((i + 1) < (height - 1)) {
                    if(mask[indI + width]) {
                        tL.push_back(Eigen::Triplet< double > (count, index[indI + width], -1.0));
                    }
                }

                if((i - 1) > -1) {
                    if(mask[indI - width]) {
                        tL.push_back(Eigen::Triplet< double > (count, index[indI - width], -1.0));
                    }
                }

                tL.push_back(Eigen::Triplet< double > (count, count , 4.0));

                count++;
            }
        }
    }

    int tot = count;
    Eigen::SparseMatrix<double> A = Eigen::SparseMatrix<double>(tot, tot);
    A.setFromTriplets(tL.begin(), tL.end());

    #ifdef PIC_DEBUG
        printf("Ok\n");
    #endif

    //solve the linear system for each color channel
    Eigen::SimplicialCholesky<Eigen::SparseMatrix<double> > solver(A);

    for(int k=0; k< target->channels; k++) {

        Eigen::VectorXd b, x;
        b = Eigen::VectorXd::Zero(tot);

        //assign values to b
        int count = 0;
        for(int i = 0; i < height; i++) {
            int tmpI = i * width;

            for(int j = 0; j < width; j++) {
                int indI = (tmpI + j);

                if(mask[indI]) {

                    b[count] = -(*lap_source)(j, i)[k];

                    if((j + 1) < (width - 1)) {
                        if(!mask[indI + 1]) {
                            b[count] += (*target)(j + 1, i)[k];
                        }
                    }

                    if((j - 1) > -1) {
                        if(!mask[indI - 1]) {
                            b[count] += (*target)(j - 1, i)[k];
                        }
                    }

                    if((i + 1) < (height - 1)) {
                        if(!mask[indI + width]) {
                            b[count] += (*target)(j, i + 1)[k];
                        }
                    }

                    if((i - 1) > -1) {
                        if(!mask[indI - width]) {
                           b[count] += (*target)(j, i - 1)[k];
                        }
                    }

                    count++;
                }
            }
        }

        x = solver.solve(b);

        if(solver.info() != Eigen::Success) {
            #ifdef PIC_DEBUG
                printf("SOLVER FAILED!\n");
            #endif

            return NULL;
        }

        #ifdef PIC_DEBUG
            printf("SOLVER SUCCESS!\n");
        #endif

        count = 0;
        for(int i = 0; i < height; i++) {
            int tmpI = i * width;

            for(int j = 0; j < width; j++) {
                int indI = (tmpI + j);

                if(mask[indI]) {
                    float val = float(x(count));
                    (*ret)(j, i)[k] = val > 0.0f ? val : 0.0f;
                    count++;
                }
            }
        }
    }

    delete_s(lap_source);
    delete_vec_s(index);

    return ret;
}
#endif

} // end namespace pic

#endif /* PIC_ALGORITHMS_POISSON_IMAGE_EDITING_HPP */


// Copyright 2004 The Trustees of Indiana University.

// Use, modification and distribution is subject to the Boost Software
// License, Version 1.0. (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)

//  Authors: Douglas Gregor
//           Andrew Lumsdaine
#ifndef BOOST_GRAPH_BETWEENNESS_CENTRALITY_CLUSTERING_HPP
#define BOOST_GRAPH_BETWEENNESS_CENTRALITY_CLUSTERING_HPP

#include <boost/graph/betweenness_centrality.hpp>
#include <boost/graph/graph_traits.hpp>
#include <boost/pending/indirect_cmp.hpp>
#include <algorithm>
#include <vector>
#include <boost/property_map.hpp>

namespace boost {

/** Threshold termination function for the betweenness centrality
 * clustering algorithm.
 */
template<typename T>
struct bc_clustering_threshold
{
  typedef T centrality_type;

  /// Terminate clustering when maximum absolute edge centrality is
  /// below the given threshold.
  explicit bc_clustering_threshold(T threshold)
    : threshold(threshold), dividend(1.0) {}

  /**
   * Terminate clustering when the maximum edge centrality is below
   * the given threshold.
   *
   * @param threshold the threshold value
   *
   * @param g the graph on which the threshold will be calculated
   *
   * @param normalize when true, the threshold is compared against the
   * normalized edge centrality based on the input graph; otherwise,
   * the threshold is compared against the absolute edge centrality.
   */
  template<typename Graph>
  bc_clustering_threshold(T threshold, const Graph& g, bool normalize = true)
    : threshold(threshold), dividend(1.0)
  {
    if (normalize) {
      typename graph_traits<Graph>::vertices_size_type n = num_vertices(g);
      dividend = T((n - 1) * (n - 2)) / T(2);
    }
  }

  /** Returns true when the given maximum edge centrality (potentially
   * normalized) falls below the threshold.
   */
  template<typename Graph, typename Edge>
  bool operator()(T max_centrality, Edge, const Graph&)
  {
    return (max_centrality / dividend) < threshold;
  }

 protected:
  T threshold;
  T dividend;
};

/** Graph clustering based on edge betweenness centrality.
 *
 * This algorithm implements graph clustering based on edge
 * betweenness centrality. It is an iterative algorithm, where in each
 * step it compute the edge betweenness centrality (via @ref
 * brandes_betweenness_centrality) and removes the edge with the
 * maximum betweenness centrality. The @p done function object
 * determines when the algorithm terminates (the edge found when the
 * algorithm terminates will not be removed).
 *
 * @param g The graph on which clustering will be performed. The type
 * of this parameter (@c MutableGraph) must be a model of the
 * VertexListGraph, IncidenceGraph, EdgeListGraph, and Mutable Graph
 * concepts.
 *
 * @param done The function object that indicates termination of the
 * algorithm. It must be a ternary function object thats accepts the
 * maximum centrality, the descriptor of the edge that will be
 * removed, and the graph @p g.
 *
 * @param edge_centrality (UTIL/OUT) The property map that will store
 * the betweenness centrality for each edge. When the algorithm
 * terminates, it will contain the edge centralities for the
 * graph. The type of this property map must model the
 * ReadWritePropertyMap concept. Defaults to an @c
 * iterator_property_map whose value type is
 * @c Done::centrality_type and using @c get(edge_index, g) for the
 * index map.
 *
 * @param vertex_index (IN) The property map that maps vertices to
 * indices in the range @c [0, num_vertices(g)). This type of this
 * property map must model the ReadablePropertyMap concept and its
 * value type must be an integral type. Defaults to
 * @c get(vertex_index, g).
 */
template<typename MutableGraph, typename Done, typename EdgeCentralityMap,
         typename VertexIndexMap>
void
betweenness_centrality_clustering(MutableGraph& g, Done done,
                                  EdgeCentralityMap edge_centrality,
                                  VertexIndexMap vertex_index)
{
  typedef typename property_traits<EdgeCentralityMap>::value_type
    centrality_type;
  typedef typename graph_traits<MutableGraph>::edge_iterator edge_iterator;
  typedef typename graph_traits<MutableGraph>::edge_descriptor edge_descriptor;
  typedef typename graph_traits<MutableGraph>::vertices_size_type
    vertices_size_type;

  if (edges(g).first == edges(g).second) return;

  // Function object that compares the centrality of edges
  indirect_cmp<EdgeCentralityMap, std::less<centrality_type> >
    cmp(edge_centrality);

  bool is_done;
  do {
    brandes_betweenness_centrality(g,
                                   edge_centrality_map(edge_centrality)
                                   .vertex_index_map(vertex_index));
    edge_descriptor e = *max_element(edges(g).first, edges(g).second, cmp);
    centrality_type max_centrality = get(edge_centrality, e);
    is_done = done(get(edge_centrality, e), e, g);
    if (!is_done) remove_edge(e, g);
  } while (!is_done && edges(g).first != edges(g).second);
}

/**
 * \overload
 */
template<typename MutableGraph, typename Done, typename EdgeCentralityMap>
void
betweenness_centrality_clustering(MutableGraph& g, Done done,
                                  EdgeCentralityMap edge_centrality)
{
  betweenness_centrality_clustering(g, done, edge_centrality,
                                    get(vertex_index, g));
}

/**
 * \overload
 */
template<typename MutableGraph, typename Done>
void
betweenness_centrality_clustering(MutableGraph& g, Done done)
{
  typedef typename Done::centrality_type centrality_type;
  std::vector<centrality_type> edge_centrality(num_edges(g));
  betweenness_centrality_clustering(g, done,
    make_iterator_property_map(edge_centrality.begin(), get(edge_index, g)),
    get(vertex_index, g));
}

} // end namespace boost

#endif // BOOST_GRAPH_BETWEENNESS_CENTRALITY_CLUSTERING_HPP


/* dos_estimate_quartic.cc

Copyright 2018 Grant M. Rotskoff

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

 */

#include <armadillo>
#include "particles.h"
#include "quartic.h"

int main(int argc, char **argv)
{

  arma_rng::set_seed(45632426);
  quartic g;
  int dim = atoi(argv[1]);
  double gamma = atof(argv[2]);
  g.initialize(dim, gamma);
  double hmin = 0.;
  double hmax = 2. * dim;
  double hres = 0.5;

  g.qmin = -pow(4*hmax,0.25);
  g.qmax = pow(4*hmax,0.25);

  // set the integration parameters
  double dt = 1e-3;
  double tol = 1e-6;
  int max_iter = 0;
  bool use_mc = false;

  int n_traj = atoi(argv[3]);
  g.initialize_integration_variables(dt, gamma);
  g.initialize_estimator(hmax, hres, hmin, n_traj);
  // run the estimation trajectories
  for (int i=0; i<n_traj; i++) {
    g.run_estimation_trajectory(tol, hmax, max_iter, i, langevin, use_mc);
  }
  char dos_filename[80];
  sprintf(dos_filename, "quartic_dim=%03d_gamma=%05.3f_ntraj=%04d.dat", dim, gamma, n_traj);
  FILE *dosf = fopen(dos_filename, "w");
  g.dump_dos_estimate(n_traj, dosf);
}


#include "tasktorrent/tasktorrent.hpp"
#ifdef USE_MKL
#include <mkl_cblas.h>
#include <mkl_lapacke.h>
#else
#include <cblas.h>
#include <lapacke.h>
#endif
#include <Eigen/Core>
#include <Eigen/Cholesky>
#include <fstream>
#include <array>
#include <random>
#include <mutex>
#include <iostream>
#include <map>
#include <memory>
#include <set>

#include <mpi.h>
#include <cxxopts.hpp>

using namespace std;
using namespace Eigen;
using namespace ttor;

typedef array<int, 2> int2;
typedef array<int, 3> int3;

/*
Parametrized priorities for cholesky:
0. No priority, only enforces potrf>trsm>gemm
1. Row-based priority, prioritize tasks with smaller row number in addition to priority 0.
2. Critical path priority, prioritize tasks with longest distance to the exit task. For references, check out the paper
    Beaumont, Olivier, et al. "A Makespan Lower Bound for the Scheduling of the Tiled Cholesky Factorization based on ALAP Schedule." (2020).
3. Critical path and row priority, prioritize tasks with smaller row number in addition to priority 2. We also enforces potrf>trsm>gemm
*/

enum PrioKind { no = 0, row = 1, cp = 2, cp_row = 3};

void cholesky(const int n_threads, const int verb, const int block_size, const int num_blocks, const int nprows, const int npcols,
              const PrioKind prio_kind, const bool log, const bool deps_log, const bool test, const int accumulate_parallel, const int upper_block_size)
{
    const int rank = comm_rank();
    const int n_ranks = comm_size();
    const int matrix_size = block_size * num_blocks;
    assert(nprows * npcols == n_ranks);
    std::atomic<long long int> potrf_us_t(0);
    std::atomic<long long int> trsm_us_t(0);
    std::atomic<long long int> gemm_us_t(0);
    std::atomic<long long int> accu_us_t(0);

    // Warmup MKL
    {
        Eigen::MatrixXd A = Eigen::MatrixXd::Identity(256,256);
        Eigen::MatrixXd B = Eigen::MatrixXd::Identity(256,256);
        Eigen::MatrixXd C = Eigen::MatrixXd::Identity(256,256);
        for(int i = 0; i < 10; i++) {
            cblas_dgemm(CblasColMajor, CblasNoTrans, CblasNoTrans, 256, 256, 256, 1.0, A.data(), 256, B.data(), 256, 1.0, C.data(), 256);
        }
    }

    // Compute random sizes
    std::mt19937 gen(2020);
    assert(upper_block_size <= 2*block_size);
    const int lower_block_size = 2*block_size - upper_block_size;
    std::uniform_int_distribution<> distrib(lower_block_size,upper_block_size); // average is block_size
    if(rank == 0) printf("lower_block_size %d, upper_block_size %d\n", lower_block_size, upper_block_size);
    std::vector<int> block_sizes(num_blocks, block_size);
    {
        int n = 0;
        for(int i = 0; i < num_blocks-1; i++) {
            int bs = std::min(matrix_size - n, distrib(gen));
            n += bs;
            block_sizes[i] = bs;
        }
        assert(matrix_size - n >= 0);
        block_sizes[num_blocks-1] = matrix_size - n;
    }
    int total = std::accumulate(block_sizes.begin(), block_sizes.end(), 0);
    assert(total == matrix_size);
    std::vector<int> block_displ(num_blocks+1, 0);
    for(int i = 1; i < num_blocks+1; i++) {
        block_displ[i] = block_displ[i-1] + block_sizes[i-1];
    }
    std::vector<int> block_sizes_lda(num_blocks, 0);
    for(int i = 0; i < num_blocks; i++) {
        block_sizes_lda[i] = std::max(1, block_sizes[i]);
    }
    assert(block_displ[num_blocks] == matrix_size);
    if(rank == 0) {
        printf("block sizes: ");
        for(int i = 0; i < num_blocks; i++) {
            assert(block_sizes[i] >= 0);
            printf("%d ", block_sizes[i]);
        };
        printf("\n");
        printf("block displ: ");
        for(int i = 0; i < num_blocks+1; i++) {
            assert(block_displ[i] >= 0);
            printf("%d ", block_displ[i]);
        };
        printf("\n");
    }

    // Map tasks to ranks
    auto block_2_rank = [&](int i, int j) {
        assert(i >= 0 && i < num_blocks);
        assert(j >= 0 && j < num_blocks);
        int r = (j % npcols) * nprows + (i % nprows);
        assert(r >= 0 && r < n_ranks);
        return r;
    };

    const int rank_row = (rank % nprows);
    const int rank_col = (rank / nprows);
    auto block_2_rank_row = [&](int i, int j) {
        return i % nprows;
    };
    auto block_2_rank_col = [&](int i, int j) {
        return j % npcols;
    };

    // Map threads to ranks
    auto block_2_thread = [&](int i, int j) {
        int ii = i / nprows;
        int jj = j / npcols;
        int num_blocksit = num_blocks / nprows;
        return (ii + jj * num_blocksit) % n_threads;
    };

    // Initializes the matrix
    auto val = [&](int i, int j) { return 1/(double)((i-j)*(i-j)+1); };
    vector<unique_ptr<MatrixXd>> blocks(num_blocks*num_blocks);
    for (int ii=0; ii<num_blocks; ii++) {
        for (int jj=0; jj<num_blocks; jj++) {
            auto val_loc = [&](int i, int j) { return val(block_displ[ii]+i,block_displ[jj]+j); };
            if(ii >= jj) {
                if(block_2_rank(ii,jj) == rank) {
                    blocks[ii+jj*num_blocks]=make_unique<MatrixXd>(block_sizes[ii],block_sizes[jj]);
                    *blocks[ii+jj*num_blocks]=MatrixXd::NullaryExpr(block_sizes[ii],block_sizes[jj], val_loc);
                } else {
                    blocks[ii+jj*num_blocks]=make_unique<MatrixXd>(0,0);
                }
            }
        }
    }

    // Holds the temporary matrices result of gemm to be accumulated by accu
    // Each block holds data to be accumulated into a given block[ii+jj*num_blocks]
    struct acc_data {
        std::map<int, std::unique_ptr<MatrixXd>> to_accumulate; // to_accumulate[k] holds matrix result of gemm(k,i,j)
        std::mutex mtx; // Protects that map
    };
    std::vector<acc_data> gemm_results(num_blocks*num_blocks); // gemm_results[ii+jj*num_blocks] holds the data to be accumulated into blocks[ii+jj*num_blocks]

    // Set priorities
    auto potf_block_2_prio = [&](int j) {
        if (prio_kind == PrioKind::cp_row) {
            return (double)(9*(num_blocks - j)-1) + 18 * num_blocks * num_blocks;
        }
        else if(prio_kind == PrioKind::cp) {
            return (double)(9*(num_blocks - j)-1);
        }
        else if(prio_kind == PrioKind::row) {
            return 3.0*(double)(num_blocks-j);
        }
        else {
            return 3.0;
        }
    };
    auto trsm_block_2_prio = [&](int2 ij) {
        if (prio_kind == PrioKind::cp_row) {
            return (double)((num_blocks - ij[0]) + num_blocks * (9.0 * num_blocks - 9.0 * ij[1] - 2.0) + 9 * num_blocks * num_blocks);
        }
        else if(prio_kind == PrioKind::cp) {
            return (double)(9*(num_blocks - ij[1])-2);
        }
        else if(prio_kind == PrioKind::row) {
            return 2.0*(double)(num_blocks - ij[0]);
        }
        else {
            return 2.0;
        }
    };
    auto gemm_block_2_prio = [&](int3 kij) {
        if (prio_kind == PrioKind::cp_row) {
            if (accumulate_parallel) {
                return (double)(num_blocks - kij[1]) + num_blocks * (9.0 * num_blocks - 9.0 * kij[2] - 2.0);
            }
            else {
                return (double)(num_blocks - kij[1]) + num_blocks * (9.0 * num_blocks - 3.0 * kij[2] - 6.0 * kij[0] - 2.0);
            }
        }
        else if(prio_kind == PrioKind::cp) {
            return (double)(9*num_blocks-9*kij[2]-2);
        }
        else if(prio_kind == PrioKind::row) {
            return (double)(num_blocks - kij[1]);
        }
        else {
            return 1.0;
        }
    };
    // Names
    auto potrf_name = [](int j) {
        return "POTRF_" + to_string(j);
    };
    auto trsm_name = [](int2 ij) {
        return "TRSM_" + to_string(ij[0]) + "_" + to_string(ij[1]);
    };
    auto gemm_name = [](int3 kij) {
        return "GEMM_" + to_string(kij[0]) + "_" + to_string(kij[1]) + "_" + to_string(kij[2]);
    };
    auto accu_name = [](int3 kij) {
        return "ACCU_" + to_string(kij[0]) + "_" + to_string(kij[1]) + "_" + to_string(kij[2]);
    };

    const int num_blocksmax = 15;
    MPI_Barrier(MPI_COMM_WORLD);
    if(comm_rank() == 0) {
        printf("Block -> Rank\n");
        for(int i = 0; i < min(num_blocksmax, num_blocks); i++) {
            for(int j = 0; j < min(num_blocksmax, num_blocks); j++) {
                if(i >= j) {
                    printf("%2d ", block_2_rank(i, j));
                }
            }
            printf("\n");
        }
        printf("Potf/trsm -> Priority\n");
        for(int i = 0; i < min(num_blocksmax, num_blocks); i++) {
            for(int j = 0; j < min(num_blocksmax, num_blocks); j++) {
                if(i == j) {
                    printf("%5f ", potf_block_2_prio(i));
                } else if (i > j) {
                    printf("%5f ", trsm_block_2_prio({i,j}));
                };
            }
            printf("\n");
        }
        printf("Gemm -> Priority\n");
        for(int k = 0; k < min(num_blocksmax, num_blocks); k++) {
            printf("k = %d\n", k);
            for(int i = 0; i < min(num_blocksmax, num_blocks); i++) {
                for(int j = 0; j < min(num_blocksmax, num_blocks); j++) {
                    if(i >= j) {
                        if(k < j) {
                            printf("%5f ", gemm_block_2_prio({k,i,j}));
                        } else {
                            printf(".     ");
                        }
                    }
                }
                printf("\n");
            }
        }
    }
    for(int r = 0; r < ttor::comm_size(); r++) {
        if(r == comm_rank()) {
            printf("[%d] Block -> thread\n", r);
            for(int i = 0; i < min(num_blocksmax, num_blocks); i++) {
                for(int j = 0; j < min(num_blocksmax, num_blocks); j++) {
                    if(i >= j && block_2_rank(i,j) == r) {
                        printf("%2d ", block_2_thread(i, j));
                    } else {
                        printf(" . ");
                    }
                }
                printf("\n");
            }
        }
        MPI_Barrier(MPI_COMM_WORLD);
    }

    // Initialize the communicator structure
    Communicator comm(MPI_COMM_WORLD, verb);

    // Initialize the runtime structures
    Threadpool tp(n_threads, &comm, verb, "Wk_Chol_" + to_string(rank) + "_");
    Taskflow<int>  potrf(&tp, verb);
    Taskflow<int2> trsm(&tp, verb);
    Taskflow<int3> gemm(&tp, verb);
    Taskflow<int3> accu(&tp, verb);

    Logger logger(1000000);
    if(log) {
        tp.set_logger(&logger);
        comm.set_logger(&logger);
    }

    DepsLogger dlog(1000000);

    // Send a potrf'ed pivot A(k,k) and trigger trsms below requiring A(k,k)
    auto am_trsm = comm.make_large_active_msg(
            [&](int& j) {
                int off = (nprows + rank_row - block_2_rank_row(j,j)) % nprows;
                assert(off > 0); // Can't be me
                assert(off < nprows);
                for (int i = j + off; i < num_blocks; i += nprows) {
                    assert(block_2_rank(i,j) == rank);
                    trsm.fulfill_promise({i,j});
                }
            },
            [&](int& j) {
                blocks[j+j*num_blocks]->resize(block_sizes[j],block_sizes[j]);
                return blocks[j+j*num_blocks]->data();
            },
            [&](int&){
                return;
            });

    /**
     * j is the pivot's position at A(j,j)
     */
    potrf.set_task([&](int j) { // A[j,j] -> A[j,j]
            assert(block_2_rank(j,j) == rank);
            timer t_ = wctime();
            LAPACKE_dpotrf(LAPACK_COL_MAJOR, 'L', block_sizes[j], blocks[j+j*num_blocks]->data(), block_sizes_lda[j]);
            timer t__ = wctime();
            potrf_us_t += 1e6 * elapsed(t_, t__);
        })
        .set_fulfill([&](int j) { // Triggers all trsms on rows i > j, A[i,j]
            assert(block_2_rank(j,j) == rank);
            if(deps_log) {
                for(int i = j+1; i < num_blocks; i++) {
                    dlog.add_event(make_unique<DepsEvent>(potrf.name(j), trsm.name({i,j})));
                }
            }
            // Trigger myself
            for (int i = j + nprows; i < num_blocks; i += nprows) {
                assert(block_2_rank(i,j) == rank);
                trsm.fulfill_promise({i,j});
            }
            // Send to other procs in column
            auto Ljjv = view<double>(blocks[j+j*num_blocks]->data(), block_sizes[j]*block_sizes[j]);
            for(int p = 0; p < nprows; p++) {
                if(j+p >= num_blocks) break;
                int dest = block_2_rank(j+p,j);
                if(dest != rank) {
                    am_trsm->send_large(dest, Ljjv, j);
                }
            }

        })
        .set_indegree([&](int j) {
            assert(block_2_rank(j,j) == rank);
            if(accumulate_parallel) {
                return j == 0 ? 1 : j; // Need j accumulations into (j,j) to trigger the potf
            } else {
                return 1;
            }
        })
        .set_priority(potf_block_2_prio)
        .set_mapping([&](int j) {
            assert(block_2_rank(j,j) == rank);
            return block_2_thread(j, j);
        })
        .set_name([&](int j) { // This is just for debugging and profiling
            return potrf_name(j);
        });

    // Sends a panel (trsm'ed block A(i,j)) and trigger gemms requiring A(i,j)
    auto am_gemm = comm.make_large_active_msg(
        [&](int& i, int& j) {
            if(block_2_rank_row(i,j) == rank_row) {
                const int off_right = (npcols + rank_col - block_2_rank_col(i,j)) % npcols;
                assert(off_right > 0); // Can't be me
                assert(off_right < npcols);
                assert(i >= 0 && i < num_blocks);
                assert(j >= 0 && j < num_blocks);
                for (int k = j + off_right; k < i; k += npcols) {
                    assert(block_2_rank(i,k) == rank);
                    gemm.fulfill_promise({j,i,k});
                }
            }
            if(block_2_rank_col(i,i) == rank_col) {
                const int off_below = (nprows + rank_row - block_2_rank_row(i,i)) % nprows;
                assert(off_below >= 0); // Could be me
                assert(off_below < nprows);
                for (int k = i + off_below; k < num_blocks; k += nprows) {
                    assert(block_2_rank(k,i) == rank);
                    gemm.fulfill_promise({j,k,i});
                }
            }
        },
        [&](int& i, int& j) {
            blocks[i+j*num_blocks]->resize(block_sizes[i],block_sizes[j]);
            return blocks[i+j*num_blocks]->data();
        },
        [&](int& i, int& j) {
            return;
        });

    /**
     * ij is (Row, Col) of the block in the matrix at A(i,j)
     **/
    trsm.set_task([&](int2 ij) { // A[j,j] & A[i,j] -> A[i,j]
            int i=ij[0];
            int j=ij[1];
            assert(block_2_rank(i,j) == rank);
            assert(i > j);
            timer t_ = wctime();
            cblas_dtrsm(CblasColMajor, CblasRight, CblasLower, CblasTrans, CblasNonUnit,
                block_sizes[i], block_sizes[j], 1.0, blocks[j+j*num_blocks]->data(), block_sizes_lda[j], blocks[i+j*num_blocks]->data(), block_sizes_lda[i]);
            timer t__ = wctime();
            trsm_us_t += 1e6 * elapsed(t_, t__);
        })
        .set_fulfill([&](int2 ij) {
            int i=ij[0];
            int j=ij[1];
            assert(block_2_rank(i,j) == rank);
            assert(i > j);
            if(deps_log) {
                for(int k = j+1; k < num_blocks; k++) {
                    dlog.add_event(make_unique<DepsEvent>(trsm.name(ij), gemm.name({j,std::max(i,k),std::min(i,k)})));
                }
            }
            // Local
            // Careful to not count the pivot (syrk) twice
            for (int k = j + npcols; k < i; k += npcols) {
                assert(block_2_rank(i,k) == rank);
                gemm.fulfill_promise({j,i,k});
            }
            if(block_2_rank_col(i,i) == rank_col) {
                int off_below = (nprows + rank_row - block_2_rank_row(i,i)) % nprows;
                for (int k = i + off_below; k < num_blocks; k += nprows) {
                    assert(block_2_rank(k,i) == rank);
                    gemm.fulfill_promise({j,k,i});
                }
            }
            // Remote
            auto Lijv = view<double>(blocks[i+j*num_blocks]->data(), block_sizes[i]*block_sizes[j]);
            std::set<int> dests;
            for (int c = 0; c < npcols; c++) {
                if(j+c >= num_blocks) break;
                int dest = block_2_rank(i,j+c);
                if(dest != rank) dests.insert(dest);
            }
            for (int r = 0; r < nprows; r++) {
                if(i+r >= num_blocks) break;
                int dest = block_2_rank(i+r,i);
                if(dest != rank) dests.insert(dest);
            }
            for(auto& dest: dests) {
                am_gemm->send_large(dest, Lijv, i, j);
            }
        })
        .set_indegree([&](int2 ij) {
            assert(block_2_rank(ij[0],ij[1]) == rank);
            if(accumulate_parallel) {
                return 1 + ij[1]; // Potrf above and all gemms before
            } else {
                return 1 + (ij[1] == 0 ? 0 : 1); // Potrf and last gemm before
            }
        })
        .set_priority(trsm_block_2_prio)
        .set_mapping([&](int2 ij) {
            assert(block_2_rank(ij[0],ij[1]) == rank);
            return block_2_thread(ij[0], ij[1]);
        })
        .set_name([&](int2 ij) { // This is just for debugging and profiling
            return trsm_name(ij);
        });

    /**
     * k is the step (the pivot's position), ij are Row and Column, at A(i,j)
     **/
    gemm.set_task([&](int3 kij) {
            assert(block_2_rank(kij[1],kij[2]) == rank);
            const int k=kij[0];
            const int i=kij[1];
            const int j=kij[2];
            assert(j <= i);
            assert(k < j);
            std::unique_ptr<MatrixXd> Atmp;
            MatrixXd* Aij;
            double beta = 1.0;
            if(accumulate_parallel) {
                beta = 0.0;
                Atmp = make_unique<MatrixXd>(block_sizes[i], block_sizes[j]); // The matrix is allocated with garbage. The 0 in the BLAS call make sure its overwritten by 0's before doing any math
                Aij = Atmp.get();
            } else {
                beta = 1.0;
                Aij = blocks[i+j*num_blocks].get();
            }
            assert(Aij->rows() == block_sizes[i] && Aij->cols() == block_sizes[j]);
            timer t_ = wctime();
            if (i == j) {
                cblas_dsyrk(CblasColMajor, CblasLower, CblasNoTrans,
                    block_sizes[i], block_sizes[k], -1.0, blocks[i+k*num_blocks]->data(), block_sizes_lda[i], beta, Aij->data(), block_sizes_lda[i]);
            } else {
                cblas_dgemm(CblasColMajor, CblasNoTrans, CblasTrans,
                    block_sizes[i], block_sizes[j], block_sizes[k], -1.0, blocks[i+k*num_blocks]->data(), block_sizes_lda[i], blocks[j+k*num_blocks]->data(), block_sizes_lda[j], beta, Aij->data(), block_sizes_lda[i]);
            }
            timer t__ = wctime();
            gemm_us_t += 1e6 * elapsed(t_, t__);
            if(accumulate_parallel) {
                lock_guard<mutex> lock(gemm_results[i+j*num_blocks].mtx);
                gemm_results[i+j*num_blocks].to_accumulate[k] = move(Atmp);
            }
        })
        .set_fulfill([&](int3 kij) {
            const int k=kij[0];
            const int i=kij[1];
            const int j=kij[2];
            assert(block_2_rank(kij[1],kij[2]) == rank);
            if(accumulate_parallel) {
                if(deps_log) {
                    dlog.add_event(make_unique<DepsEvent>(gemm.name(kij), accu.name(kij)));
                }
                accu.fulfill_promise(kij);
            } else {
                if (k < j-1) {
                    if(deps_log) {
                        dlog.add_event(make_unique<DepsEvent>(gemm.name(kij), gemm.name({k+1, i, j})));
                    }
                    gemm.fulfill_promise({k+1, i, j});
                } else {
                    if (i == j) {
                        if(deps_log) {
                            dlog.add_event(make_unique<DepsEvent>(gemm.name(kij), potrf.name(i)));
                        }
                        potrf.fulfill_promise(i);
                    } else {
                        if(deps_log) {
                            dlog.add_event(make_unique<DepsEvent>(gemm.name(kij), trsm.name({i,j})));
                        }
                        trsm.fulfill_promise({i,j});
                    }
                }
            }
        })
        .set_indegree([&](int3 kij) {
            assert(block_2_rank(kij[1],kij[2]) == rank);
            if(accumulate_parallel) {
                return kij[1] == kij[2] ? 1 : 2; // Either one potf or two trsms
            } else {
                return (kij[1] == kij[2] ? 1 : 2) + (kij[0] == 0 ? 0 : 1); // one potrf or two trsms + the gemm before
            }
        })
        .set_priority(gemm_block_2_prio)
        .set_mapping([&](int3 kij) {
            assert(block_2_rank(kij[1],kij[2]) == rank);
            return block_2_thread(kij[1], kij[2]); // IMPORTANT if accumulate_parallel is true
        })
        .set_binding([&](int3 kij) {
            return false; // If we accumulate in parallel, there is no order for the gemm so it doesnt matter ; If we don't then we do the gemm in sequence anyway
        }).set_name([&](int3 kij) { // This is just for debugging and profiling
            return gemm_name(kij);
        });

    /**
     * k is the step (the pivot's position), ij are Row and Column, at A(i,j)
     **/
    accu.set_task([&](int3 kij) {
            assert(block_2_rank(kij[1],kij[2]) == rank);
            int k=kij[0]; // Step (gemm's pivot)
            int i=kij[1]; // Row
            int j=kij[2]; // Col
            assert(j <= i);
            assert(k < j);
            std::unique_ptr<Eigen::MatrixXd> Atmp;
            {
                lock_guard<mutex> lock(gemm_results[i+j*num_blocks].mtx);
                Atmp = move(gemm_results[i+j*num_blocks].to_accumulate[k]);
                gemm_results[i+j*num_blocks].to_accumulate.erase(k);
            }
            timer t_ = wctime();
            *blocks[i+j*num_blocks] += (*Atmp);
            timer t__ = wctime();
            accu_us_t += 1e6 * elapsed(t_, t__);
        })
        .set_fulfill([&](int3 kij) {
            assert(block_2_rank(kij[1],kij[2]) == rank);
            int k=kij[0];
            int i=kij[1];
            int j=kij[2];
            assert(j <= i);
            assert(k < j);
            if(i == j) {
                if(deps_log) {
                    dlog.add_event(make_unique<DepsEvent>(accu.name(kij), potrf.name(i)));
                }
                potrf.fulfill_promise(i);
            } else {
                if(deps_log) {
                    dlog.add_event(make_unique<DepsEvent>(accu.name(kij), trsm.name({i,j})));
                }
                trsm.fulfill_promise({i,j});
            }
        })
        .set_indegree([&](int3 kij) {
            assert(block_2_rank(kij[1],kij[2]) == rank);
            return 1;
        })
        .set_mapping([&](int3 kij) {
            assert(block_2_rank(kij[1],kij[2]) == rank);
            return block_2_thread(kij[1], kij[2]); // IMPORTANT. Every (i,j) should map to a given fixed thread
        })
        .set_priority(gemm_block_2_prio)
        .set_binding([&](int3 kij) {
            assert(block_2_rank(kij[1],kij[2]) == rank);
            return true; // IMPORTANT
        })
        .set_name([&](int3 kij) { // This is just for debugging and profiling
            return accu_name(kij);
        });

    printf("Starting Cholesky factorization...\n");
    MPI_Barrier(MPI_COMM_WORLD);
    timer t0 = wctime();
    if (rank == 0){
        potrf.fulfill_promise(0);
    }
    tp.join();
    MPI_Barrier(MPI_COMM_WORLD);
    timer t1 = wctime();
    double total_time = elapsed(t0, t1);
    printf("Done with Cholesky factorization...\n");
    printf("Elapsed time: %e\n", total_time);
    printf("Potrf time: %e\n", potrf_us_t.load() * 1e-6);
    printf("Trsm time: %e\n", trsm_us_t.load() * 1e-6);
    printf("Gemm time: %e\n", gemm_us_t.load() * 1e-6);
    printf("Accu time: %e\n", accu_us_t.load() * 1e-6);

    printf("++++rank nranks n_threads matrix_size block_size num_blocks priority_kind accumulate upper_block_size total_time\n");
    printf("[%d]>>>>%d %d %d %d %d %d %d %d %d %e\n",rank,rank,n_ranks,n_threads,matrix_size,block_size,num_blocks,(int)prio_kind,(int)accumulate_parallel,upper_block_size,total_time);

    if(log) {
        std::ofstream logfile;
        string filename = "ttor_dist_"+to_string(block_size)+"_"+to_string(num_blocks)+"_"+ to_string(n_ranks)+"_"+to_string(n_threads)+"_"+to_string(prio_kind)+".log."+to_string(rank);
        logfile.open(filename);
        logfile << logger;
        logfile.close();
    }

    if(deps_log) {
        std::ofstream depsfile;
        string depsfilename = "deps_ttor_dist_"+to_string(block_size)+"_"+to_string(num_blocks)+"_"+ to_string(n_ranks)+"_"+to_string(n_threads)+"_"+to_string(prio_kind)+".dot."+to_string(rank);
        depsfile.open(depsfilename);
        depsfile << dlog;
        depsfile.close();
    }

    if(test) {
        printf("Starting sending matrix to rank 0...\n");
        MatrixXd A = MatrixXd::NullaryExpr(matrix_size,matrix_size,val);
        MatrixXd L = MatrixXd::Zero(matrix_size,matrix_size);
        // Send the matrix to rank 0
        for (int ii=0; ii<num_blocks; ii++) {
            for (int jj=0; jj<num_blocks; jj++) {
                if (jj<=ii)  {
                    int owner = block_2_rank(ii,jj);
                    MPI_Status status;
                    if (rank == 0 && rank != owner) { // Careful with deadlocks here
                        blocks[ii+jj*num_blocks] = make_unique<Eigen::MatrixXd>(block_sizes[ii], block_sizes[jj]);
                        MPI_Recv(blocks[ii+jj*num_blocks]->data(), block_sizes[ii]*block_sizes[jj], MPI_DOUBLE, owner, 0, MPI_COMM_WORLD, &status);
                    } else if (rank != 0 && rank == owner) {
                        MPI_Send(blocks[ii+jj*num_blocks]->data(), block_sizes[ii]*block_sizes[jj], MPI_DOUBLE, 0, 0, MPI_COMM_WORLD);
                    }
                }
            }
        }

        // Rank 0 test
        if(rank == 0) {
            printf("Starting test on rank 0...\n");
            for (int ii=0; ii<num_blocks; ii++) {
                for (int jj=0; jj<num_blocks; jj++) {
                    if (jj<=ii) {
                        L.block(block_displ[ii],block_displ[jj],block_sizes[ii],block_sizes[jj])=*blocks[ii+jj*num_blocks];
                    }
                }
            }
            auto L1=L.triangularView<Lower>();
            VectorXd x = VectorXd::Random(matrix_size);
            VectorXd b = A*x;
            VectorXd bref = b;
            L1.solveInPlace(b);
            L1.transpose().solveInPlace(b);
            double error = (b - x).norm() / x.norm();
            printf("\n=> Error solve %e\n\n", error);
            if(error > 1e-6) {
                printf("\n\nERROR: error is too large!\n\n");
                exit(1);
            }
        }
    }
}

int main(int argc, const char **argv)
{
    int req = MPI_THREAD_FUNNELED;
    int prov = -1;

    MPI_Init_thread(NULL, NULL, req, &prov);

    assert(prov == req);

    std::stringstream sstr;
    sstr << comm_size();
    const std::string comm_size_str = sstr.str();

    cxxopts::Options options("2d_cholesky", "2D dense cholesky using TaskTorrent");
    options.add_options()
        ("help", "Print help")
        ("n_threads", "Number of threads", cxxopts::value<int>()->default_value("2"))
        ("verb", "Verbosity level", cxxopts::value<int>()->default_value("0"))
        ("block_size", "Block size", cxxopts::value<int>()->default_value("5"))
        ("num_blocks", "Number of blocks", cxxopts::value<int>()->default_value("10"))
        ("nprows", "Number of processors accross rows", cxxopts::value<int>()->default_value("1"))
        ("npcols", "Number of processors accross columns", cxxopts::value<int>()->default_value(comm_size_str.c_str()))
        ("kind", "Priority kind", cxxopts::value<int>()->default_value("0"))
        ("log", "Enable logging", cxxopts::value<bool>()->default_value("false"))
        ("depslog", "Enable dependency logging", cxxopts::value<bool>()->default_value("false"))
        ("test", "Test or not", cxxopts::value<bool>()->default_value("true"))
        ("accumulate", "Accumulate block GEMMs in parallel", cxxopts::value<bool>()->default_value("false"))
        ("upper_block_size", "Upper block size", cxxopts::value<int>()->default_value("-1"));
    auto result = options.parse(argc, argv);

    const int n_threads = result["n_threads"].as<int>();
    const int verb = result["verb"].as<int>();
    const int block_size = result["block_size"].as<int>();
    const int num_blocks = result["num_blocks"].as<int>();
    const int nprows = result["nprows"].as<int>();
    const int npcols = result["npcols"].as<int>();
    const PrioKind kind = (PrioKind) result["kind"].as<int>();
    const bool log = result["log"].as<bool>();
    const bool depslog = result["depslog"].as<bool>();
    const bool test = result["test"].as<bool>();
    const bool accumulate = result["accumulate"].as<bool>();
    const int upper_block_size = (result["upper_block_size"].as<int>() == -1 ? block_size : result["upper_block_size"].as<int>());

    assert(block_size > 0);
    assert(num_blocks > 0);
    assert(n_threads > 0);
    assert(verb >= 0);
    assert(nprows >= 0);
    assert(npcols >= 0);
    assert(upper_block_size >= block_size && upper_block_size <= 2 * block_size);

    if (result.count("help")) {
        std::cout << options.help({"", "Group"}) << endl;
        exit(0);
    }
    if(comm_rank() == 0) printf("Arguments: block_size (size of blocks) %d\nnum_blocks (# of blocks) %d\nn_threads %d\nverb %d\nnprows %d\nnpcols %d\nkind %d\nlog %d\ndeplog %d\ntest %d\naccumulate %d\nupper_block_size %d\n", block_size, num_blocks, n_threads, verb, nprows, npcols, (int)kind, log, depslog, test, accumulate, upper_block_size);

    cholesky(n_threads, verb, block_size, num_blocks, nprows, npcols, kind, log, depslog, test, accumulate, upper_block_size);

    MPI_Finalize();
}


#include <boost/math/special_functions/next.hpp>
#include <boost/random.hpp>

#include <limits>
#include <vector>

#include "caffe/common.hpp"
#include "caffe/util/math_functions.hpp"
#include "caffe/util/rng.hpp"

namespace caffe {

template<>
void caffe_cpu_gemm<float,float>(const CBLAS_TRANSPOSE TransA,
    const CBLAS_TRANSPOSE TransB, const int M, const int N, const int K,
    const float alpha, const float* A, const float* B, const float beta,
    float* C) {
  int lda = (TransA == CblasNoTrans) ? K : M;
  int ldb = (TransB == CblasNoTrans) ? N : K;
  cblas_sgemm(CblasRowMajor, TransA, TransB, M, N, K, alpha, A, lda, B,
      ldb, beta, C, N);
}

template<>
void caffe_cpu_gemm<double,double>(const CBLAS_TRANSPOSE TransA,
    const CBLAS_TRANSPOSE TransB, const int M, const int N, const int K,
    const double alpha, const double* A, const double* B, const double beta,
    double* C) {
  int lda = (TransA == CblasNoTrans) ? K : M;
  int ldb = (TransB == CblasNoTrans) ? N : K;
  cblas_dgemm(CblasRowMajor, TransA, TransB, M, N, K, alpha, A, lda, B,
      ldb, beta, C, N);
}

#ifndef CPU_ONLY
template<>
void caffe_cpu_gemm<float16,float16>(const CBLAS_TRANSPOSE TransA,
    const CBLAS_TRANSPOSE TransB, const int M, const int N, const int K,
    const float16 alpha, const float16* A, const float16* B, const float16 beta,
    float16* C) {
}
template<>
void caffe_cpu_gemm<float16,float>(const CBLAS_TRANSPOSE TransA,
    const CBLAS_TRANSPOSE TransB, const int M, const int N, const int K,
    const float alpha, const float16* A, const float16* B, const float beta,
    float16* C) {
  if (M <= 0 || N <= 0 || K <= 0) {
    return;
  }
  std::vector<float> a(M*K), b(K*N), c(M*N);
  caffe_cpu_convert(a.size(), A, &a.front());
  caffe_cpu_convert(b.size(), B, &b.front());
  caffe_cpu_convert(c.size(), C, &c.front());
  const int lda = (TransA == CblasNoTrans) ? K : M;
  const int ldb = (TransB == CblasNoTrans) ? N : K;
  cblas_sgemm(CblasRowMajor, TransA, TransB, M, N, K, alpha, &a.front(), lda, &b.front(),
      ldb, beta, &c.front(), N);
  caffe_cpu_convert(c.size(), &c.front(), C);
}
#endif

template <>
void caffe_cpu_gemv<float,float>(const CBLAS_TRANSPOSE TransA, const int M,
    const int N, const float alpha, const float* A, const float* x,
    const float beta, float* y) {
  cblas_sgemv(CblasRowMajor, TransA, M, N, alpha, A, N, x, 1, beta, y, 1);
}

template <>
void caffe_cpu_gemv<double,double>(const CBLAS_TRANSPOSE TransA, const int M,
    const int N, const double alpha, const double* A, const double* x,
    const double beta, double* y) {
  cblas_dgemv(CblasRowMajor, TransA, M, N, alpha, A, N, x, 1, beta, y, 1);
}

#ifndef CPU_ONLY
template <>
void caffe_cpu_gemv<float16,float>(const CBLAS_TRANSPOSE TransA, const int M,
    const int N, const float alpha, const float16* A, const float16* x,
    const float beta, float16* y) {
  if (M <= 0 || N <= 0) {
    return;
  }
  const int lx = (TransA == CblasNoTrans) ? N : M;
  const int ly = (TransA == CblasNoTrans) ? M : N;
  std::vector<float> a(M*N), xv(lx), yv(ly);
  caffe_cpu_convert(a.size(), A, &a.front());
  caffe_cpu_convert(xv.size(), x, &xv.front());
  caffe_cpu_convert(yv.size(), y, &yv.front());
  cblas_sgemv(CblasRowMajor, TransA, M, N, alpha, &a.front(), N, &xv.front(), 1, beta, &yv.front(), 1);
  caffe_cpu_convert(yv.size(), &yv.front(), y);
}

template <>
void caffe_cpu_gemv<float16,float16>(const CBLAS_TRANSPOSE TransA, const int M,
    const int N, const float16 alpha, const float16* A, const float16* xv,
    const float16 beta, float16* y) {
  //  cblas_hgemv(CblasRowMajor, TransA, M, N, alpha, A, N, x, 1, beta, y, 1);
}

#endif

template <>
void caffe_axpy<float,float>(const int N, const float alpha, const float* X,
    float* Y) { cblas_saxpy(N, alpha, X, 1, Y, 1); }

template <>
void caffe_axpy<double,double>(const int N, const double alpha, const double* X,
    double* Y) { cblas_daxpy(N, alpha, X, 1, Y, 1); }

#ifndef CPU_ONLY
// TODO Consider CUDA
template<>
void caffe_axpy<float16,float>(const int N, const float alpha, const float16* X, float16* Y) {
  for (int i = 0; i < N; ++i) {
    Y[i] = Get<float16>(alpha * Get<float>(X[i]) + Get<float>(Y[i]));
  }
}
template<>
void caffe_axpy<float16,float16>(const int N, const float16 alpha, const float16* X, float16* Y) {
  for (int i = 0; i < N; ++i) {
    Y[i] = Get<float16>(Get<float>(alpha) * Get<float>(X[i]) + Get<float>(Y[i]));
  }
}
#endif

template <typename Dtype>
void caffe_set(const int N, const Dtype alpha, Dtype* Y) {
  if (alpha == 0) {
    memset(Y, 0, sizeof(Dtype) * N);  // NOLINT(caffe/alt_fn)
    return;
  }
  for (int i = 0; i < N; ++i) {
    Y[i] = alpha;
  }
}

template void caffe_set<int>(const int N, const int alpha, int* Y);
template void caffe_set<float>(const int N, const float alpha, float* Y);
template void caffe_set<double>(const int N, const double alpha, double* Y);

#ifndef CPU_ONLY
template void caffe_set<float16>(const int N, const float16 alpha, float16* Y);
#endif

template <>
void caffe_add_scalar(const int N, const float alpha, float* Y) {
  for (int i = 0; i < N; ++i) {
    Y[i] += alpha;
  }
}

template <>
void caffe_add_scalar(const int N, const double alpha, double* Y) {
  for (int i = 0; i < N; ++i) {
    Y[i] += alpha;
  }
}

#ifndef CPU_ONLY
template <>
void caffe_add_scalar(const int N, const float alpha, float16* Y) {
  for (int i = 0; i < N; ++i) {
    Y[i] = Get<float16>( Get<float>(Y[i]) + alpha );
  }
}
template <>
void caffe_add_scalar(const int N, const float16 alpha, float16* Y) {
  for (int i = 0; i < N; ++i) {
    Y[i] = Get<float16>( Get<float16>(Y[i]) + alpha );
  }
}
#endif

template <typename Dtype, typename Mtype>
void caffe_copy(const int N, const Dtype* X, Dtype* Y) {
  if (X != Y) {
    if (Caffe::mode() == Caffe::GPU) {
#ifndef CPU_ONLY
      // NOLINT_NEXT_LINE(caffe/alt_fn)
      CUDA_CHECK(cudaMemcpy(Y, X, sizeof(Dtype) * N, cudaMemcpyDefault));
#else
      NO_GPU;
#endif
    } else {
      memcpy(Y, X, sizeof(Dtype) * N);  // NOLINT(caffe/alt_fn)
    }
  }
}

template void caffe_copy<int,int>(const int N, const int* X, int* Y);
template void caffe_copy<unsigned int, unsigned int>(const int N, const unsigned int* X,
    unsigned int* Y);
template void caffe_copy<float,float>(const int N, const float* X, float* Y);
template void caffe_copy<double,double>(const int N, const double* X, double* Y);

#ifndef CPU_ONLY
template void caffe_copy<float16,float>(const int N, const float16* X, float16* Y);
template void caffe_copy<float16,float16>(const int N, const float16* X, float16* Y);
#endif

template <>
void caffe_scal<float,float>(const int N, const float alpha, float *X) {
  cblas_sscal(N, alpha, X, 1);
}

template <>
void caffe_scal<double,double>(const int N, const double alpha, double *X) {
  cblas_dscal(N, alpha, X, 1);
}

#ifndef CPU_ONLY
template <>
void caffe_scal<float16,float>(const int N, const float alpha, float16 *X) {
  for (int i = 0; i < N; ++i) {
    X[i] = Get<float16>( alpha * Get<float>(X[i]) );
  }
}

template <>
void caffe_scal<float16,float16>(const int N, const float16 alpha, float16 *X) {
  // cblas_hscal(N, alpha, X, 1);
}
#endif

template <>
void caffe_cpu_axpby<float,float>(const int N, const float alpha, const float* X,
                            const float beta, float* Y) {
  cblas_saxpby(N, alpha, X, 1, beta, Y, 1);
}

template <>
void caffe_cpu_axpby<double,double>(const int N, const double alpha, const double* X,
                             const double beta, double* Y) {
  cblas_daxpby(N, alpha, X, 1, beta, Y, 1);
}

#ifndef CPU_ONLY
template <>
void caffe_cpu_axpby<float16,float16>(const int N, const float16 alpha, const float16* X,
				    const float16 beta, float16* Y) {}
template <>
void caffe_cpu_axpby<float16,float>(const int N, const float alpha, const float16* X,
                             const float beta, float16* Y) {
  for (int i=0; i<N; i++) {
    Y[i] = Get<float16>( alpha * Get<float>(X[i]) + beta * Get<float>(Y[i]) );
  }
}
#endif

template <>
void caffe_add<float,float>(const int n, const float* a, const float* b,
    float* y) {
  vsAdd(n, a, b, y);
}

template <>
void caffe_add<double,double>(const int n, const double* a, const double* b,
    double* y) {
  vdAdd(n, a, b, y);
}

#ifndef CPU_ONLY
template <>
void caffe_add<float16,float>(const int n, const float16* a, const float16* b,
    float16* y) {
  for (int i=0; i<n; i++) {
    y[i] = Get<float16>( Get<float>(a[i]) + Get<float>(b[i]) );
  }
}
template <>
void caffe_add<float16,float16>(const int n, const float16* a, const float16* b,
    float16* y) {
  for (int i=0; i<n; i++) {
    y[i] = Get<float16>( Get<float16>(a[i]) + Get<float16>(b[i]) );
  }
}
#endif

template <>
void caffe_sub<float,float>(const int n, const float* a, const float* b,
    float* y) {
  vsSub(n, a, b, y);
}

template <>
void caffe_sub<double,double>(const int n, const double* a, const double* b,
    double* y) {
  vdSub(n, a, b, y);
}

#ifndef CPU_ONLY
template <>
void caffe_sub<float16,float>(const int n, const float16* a, const float16* b,
    float16* y) {
  for (int i=0; i<n; i++) {
    y[i] = Get<float16>( Get<float>(a[i]) - Get<float>(b[i]) );
  }
}
template <>
void caffe_sub<float16,float16>(const int n, const float16* a, const float16* b,
    float16* y) {
  //  vhSub(n, a, b, y);
}
#endif

template <>
void caffe_mul<float,float>(const int n, const float* a, const float* b,
    float* y) {
  vsMul(n, a, b, y);
}

template <>
void caffe_mul<double,double>(const int n, const double* a, const double* b,
    double* y) {
  vdMul(n, a, b, y);
}

#ifndef CPU_ONLY
template <>
void caffe_mul<float16,float>(const int n, const float16* a, const float16* b,
    float16* y) {
  for (int i=0; i<n; i++) {
    y[i] = Get<float16>( Get<float>(a[i]) * Get<float>(b[i]) );
  }
}
template <>
void caffe_mul<float16,float16>(const int n, const float16* a, const float16* b,
    float16* y) {
  //  vhMul(n, a, b, y);
}
#endif

template <>
void caffe_div<float,float>(const int n, const float* a, const float* b,
    float* y) {
  vsDiv(n, a, b, y);
}

template <>
void caffe_div<double,double>(const int n, const double* a, const double* b,
    double* y) {
  vdDiv(n, a, b, y);
}

#ifndef CPU_ONLY
template <>
void caffe_div<float16,float>(const int n, const float16* a, const float16* b,
    float16* y)
{
  for (int i=0; i<n; i++) {
    y[i] = Get<float16>( Get<float>(a[i]) / Get<float>(b[i]) );
  }
}

template <>
void caffe_div<float16,float16>(const int n, const float16* a, const float16* b,
    float16* y) {
  //  vhDiv(n, a, b, y);
}
#endif

template <>
void caffe_powx<float,float>(const int n, const float* a, const float b,
    float* y) {
  vsPowx(n, a, b, y);
}

template <>
void caffe_powx<double,double>(const int n, const double* a, const double b,
    double* y) {
  vdPowx(n, a, b, y);
}

#ifndef CPU_ONLY
template <>
void caffe_powx<float16,float16>(const int n, const float16* a, const float16 b,
    float16* y) {
  //  vhPowx(n, a, b, y);
}
template <>
void caffe_powx<float16,float>(const int n, const float16* a, const float b, float16* y) {
  for (int i=0; i<n; i++) {
    y[i] = Get<float16>( pow(Get<float>(a[i]), b) );
  }
}
#endif

template <>
void caffe_sqr<float,float>(const int n, const float* a, float* y) {
  vsSqr(n, a, y);
}

template <>
void caffe_sqr<double,double>(const int n, const double* a, double* y) {
  vdSqr(n, a, y);
}

#ifndef CPU_ONLY
template <>
void caffe_sqr<float16,float16>(const int n, const float16* a, float16* y) {
  vhSqr(n, a, y);
}
template <>
void caffe_sqr<float16,float>(const int n, const float16* a, float16* y) {
  float f;
  for (int i = 0; i < n; ++i) {
    f = Get<float>(a[i]);
    y[i] = Get<float16>(f * f);
  }
}
#endif

template <>
void caffe_exp<float,float>(const int n, const float* a, float* y) {
  vsExp(n, a, y);
}

template <>
void caffe_exp<double,double>(const int n, const double* a, double* y) {
  vdExp(n, a, y);
}

template <>
void caffe_log<float>(const int n, const float* a, float* y) {
  vsLn(n, a, y);
}

template <>
void caffe_log<double>(const int n, const double* a, double* y) {
  vdLn(n, a, y);
}

#ifndef CPU_ONLY
template <>
void caffe_exp<float16, float16>(const int n, const float16* a, float16* y) {
  vhExp(n, a, y);
}
template <>
void caffe_log<float16>(const int n, const float16* a, float16* y) {
  vhLn(n, a, y);
}
#endif // ! CPU_ONLY

template <>
void caffe_abs<float>(const int n, const float* a, float* y) {
    vsAbs(n, a, y);
}

template <>
void caffe_abs<double>(const int n, const double* a, double* y) {
    vdAbs(n, a, y);
}

#ifndef CPU_ONLY
template <>
void caffe_exp<float16,float>(const int n, const float16* a, float16* y) {
  for (int i=0; i<n; i++) {
    y[i] = Get<float16>( exp(Get<float>(a[i])) );
  }
}
#endif

#ifndef CPU_ONLY
template <>
void caffe_abs<float16>(const int n, const float16* a, float16* y) {
  for (int i=0; i<n; i++) {
    y[i] = Get<float16>( fabs(Get<float>(a[i])) );
  }
}
#endif

unsigned int caffe_rng_rand() {
  return (*caffe_rng())();
}

template <typename Dtype>
Dtype caffe_nextafter(const Dtype b) {
  return boost::math::nextafter<Dtype>(
      b, std::numeric_limits<Dtype>::max());
}

template
float caffe_nextafter(const float b);

template
double caffe_nextafter(const double b);

template <typename Dtype, typename Mtype>
void caffe_rng_uniform(const int n, const Mtype a, const Mtype b, Dtype* r) {
  CHECK_GE(n, 0);
  CHECK(r);
  CHECK_LE(a, b);
  boost::uniform_real<Mtype> random_distribution(a, caffe_nextafter<Mtype>(b));
  boost::variate_generator<caffe::rng_t*, boost::uniform_real<Mtype> >
      variate_generator(caffe_rng(), random_distribution);
  for (int i = 0; i < n; ++i) {
    r[i] = Get<Dtype>(variate_generator());
  }
}

template
void caffe_rng_uniform<float,float>(const int n, const float a, const float b,
                              float* r);

template
void caffe_rng_uniform<double,double>(const int n, const double a, const double b,
                               double* r);

#ifndef CPU_ONLY
template
void caffe_rng_uniform<float16,float>(const int n, const float a, const float b,
                               float16* r);
  template<>
void caffe_rng_uniform<float16,float16>(const int n, const float16 a, const float16 b,
					float16* r) {}
#endif

template <typename Dtype, typename Mtype>
void caffe_rng_gaussian(const int n, const Mtype a,
                        const Mtype sigma, Dtype* r) {
  CHECK_GE(n, 0);
  CHECK(r);
  CHECK_GT(sigma, 0);
  boost::normal_distribution<Mtype> random_distribution(a, sigma);
  boost::variate_generator<caffe::rng_t*, boost::normal_distribution<Mtype> >
      variate_generator(caffe_rng(), random_distribution);
  for (int i = 0; i < n; ++i) {
    r[i] = Get<Dtype>(variate_generator());
  }
}

template
void caffe_rng_gaussian<float,float>(const int n, const float mu,
                               const float sigma, float* r);

template
void caffe_rng_gaussian<double,double>(const int n, const double mu,
                                const double sigma, double* r);

#ifndef CPU_ONLY
template
void caffe_rng_gaussian<float16,float>(const int n, const float mu,
                                const float sigma, float16* r);
  template <>
void caffe_rng_gaussian<float16,float16>(const int n, const float16 mu,
					 const float16 sigma, float16* r) {}
#endif

template <typename Dtype, typename Mtype>
void caffe_rng_bernoulli(const int n, const Mtype p, int* r) {
  CHECK_GE(n, 0);
  CHECK(r);
  CHECK_GE(p, 0);
  CHECK_LE(p, 1);
  boost::bernoulli_distribution<Mtype> random_distribution(p);
  boost::variate_generator<caffe::rng_t*, boost::bernoulli_distribution<Mtype> >
      variate_generator(caffe_rng(), random_distribution);
  for (int i = 0; i < n; ++i) {
    r[i] = Get<int>(variate_generator());
  }
}

template
void caffe_rng_bernoulli<double,double>(const int n, const double p, int* r);

template
void caffe_rng_bernoulli<float,float>(const int n, const float p, int* r);

#ifndef CPU_ONLY
template
void caffe_rng_bernoulli<float16,float>(const int n, const float p, int* r);
template
void caffe_rng_bernoulli<float16,float16>(const int n, const float16 p, int* r);
#endif

template <typename Dtype, typename Mtype>
void caffe_rng_bernoulli(const int n, const Mtype p, unsigned int* r) {
  CHECK_GE(n, 0);
  CHECK(r);
  CHECK_GE(p, 0);
  CHECK_LE(p, 1);
  boost::bernoulli_distribution<Mtype> random_distribution(p);
  boost::variate_generator<caffe::rng_t*, boost::bernoulli_distribution<Mtype> >
      variate_generator(caffe_rng(), random_distribution);
  for (int i = 0; i < n; ++i) {
    r[i] = static_cast<unsigned int>(variate_generator());
  }
}

template
void caffe_rng_bernoulli<double,double>(const int n, const double p, unsigned int* r);

template
void caffe_rng_bernoulli<float,float>(const int n, const float p, unsigned int* r);

#ifndef CPU_ONLY
template
void caffe_rng_bernoulli<float16,float>(const int n, const float p, unsigned int* r);
template
void caffe_rng_bernoulli<float16,float16>(const int n, const float16 p, unsigned int* r);
#endif

template <>
float caffe_cpu_strided_dot<float,float>(const int n, const float* x, const int incx,
    const float* y, const int incy) {
  return cblas_sdot(n, x, incx, y, incy);
}

template <>
double caffe_cpu_strided_dot<double,double>(const int n, const double* x,
    const int incx, const double* y, const int incy) {
  return cblas_ddot(n, x, incx, y, incy);
}

#ifndef CPU_ONLY
template <>
float caffe_cpu_strided_dot<float16,float>(const int n, const float16* x,
    const int incx, const float16 *y, const int incy) {
  float sum = 0.0f;
  int idx_x, idx_y;
  for (int i = 0; i < n; ++i) {
    idx_x = i*incx;
    idx_y = i*incy;
    sum += Get<float>(x[idx_x]) * Get<float>(y[idx_y]);
  }
  return sum;
}
// TODO Consider CUDA
template <>
float16 caffe_cpu_strided_dot<float16,float16>(const int n, const float16* x,
    const int incx, const float16 *y, const int incy) {
  float sum = 0.0f;
  int idx_x, idx_y;
  for (int i = 0; i < n; ++i) {
    idx_x = i*incx;
    idx_y = i*incy;
    sum += Get<float>(x[idx_x]) * Get<float>(y[idx_y]);
  }
  return Get<float16>(sum);
}
#endif

template <typename Dtype, typename Mtype>
Mtype caffe_cpu_dot(const int n, const Dtype* x, const Dtype* y) {
  return caffe_cpu_strided_dot<Dtype,Mtype>(n, x, 1, y, 1);
}

template
float caffe_cpu_dot<float,float>(const int n, const float* x, const float* y);
template
double caffe_cpu_dot<double,double>(const int n, const double* x, const double* y);

#ifndef CPU_ONLY
template
float caffe_cpu_dot<float16,float>(const int n, const float16* x, const float16* y);
template
float16 caffe_cpu_dot<float16,float16>(const int n, const float16* x, const float16* y);
#endif

template <>
int caffe_cpu_hamming_distance<float>(const int n, const float* x,
                                  const float* y) {
  int dist = 0;
  for (int i = 0; i < n; ++i) {
    dist += __builtin_popcount(static_cast<uint32_t>(x[i]) ^
                               static_cast<uint32_t>(y[i]));
  }
  return dist;
}

template <>
int caffe_cpu_hamming_distance<double>(const int n, const double* x,
                                   const double* y) {
  int dist = 0;
  for (int i = 0; i < n; ++i) {
    dist += __builtin_popcountl(static_cast<uint64_t>(x[i]) ^
                                static_cast<uint64_t>(y[i]));
  }
  return dist;
}

#ifndef CPU_ONLY
template <>
int caffe_cpu_hamming_distance<float16>(const int n, const float16* x,
                                        const float16* y) {
  int dist = 0;
  for (int i = 0; i < n; ++i) {
    dist += __builtin_popcount(static_cast<uint16_t>(Get<float>(x[i])) ^
                               static_cast<uint16_t>(Get<float>(y[i])));
  }
  return dist;
}
#endif

template <>
float caffe_cpu_asum<float,float>(const int n, const float* x) {
  return cblas_sasum(n, x, 1);
}

template <>
double caffe_cpu_asum<double,double>(const int n, const double* x) {
  return cblas_dasum(n, x, 1);
}

#ifndef CPU_ONLY
template <>
float caffe_cpu_asum<float16,float>(const int n, const float16 *x) {
  float sum = 0.0f;
  for (int i = 0; i < n; ++i) {
    sum += fabs(Get<float>(x[i]));
  }
  return sum;
}

template <>
float16 caffe_cpu_asum<float16,float16>(const int n, const float16 *x) {
  float sum = 0.0f;
  for (int i = 0; i < n; ++i) {
    sum += fabs(Get<float>(x[i]));
  }
  return Get<float16>(sum);
}
#endif

template <>
void caffe_cpu_scale<float,float>(const int n, const float alpha, const float *x,
                            float* y) {
  cblas_scopy(n, x, 1, y, 1);
  cblas_sscal(n, alpha, y, 1);
}

template <>
void caffe_cpu_scale<double,double>(const int n, const double alpha, const double *x,
                             double* y) {
  cblas_dcopy(n, x, 1, y, 1);
  cblas_dscal(n, alpha, y, 1);
}

#ifndef CPU_ONLY
template <>
void caffe_cpu_scale<float16,float16>(const int n, const float16 alpha, const float16 *x,
    float16 *y) {
}

template <>
void caffe_cpu_scale<float16,float>(const int n, const float alpha, const float16 *x,
    float16 *y) {
  for (int i=0; i<n; i++) {
    y[i] = Get<float16>( alpha * Get<float>(x[i]) );
  }
}
#endif

}  // namespace caffe


// Copyright (c) 2020 Marcus Valtonen Örnhag
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#include "get_valtonenornhag_arxiv_2020a.hpp"
#include <float.h>  // DBL_MAX
#include <Eigen/Dense>
#include <Eigen/Geometry>
#include <cmath>  // abs
#include "solver_valtonenornhag_arxiv_2020a_fHf.hpp"
#include "normalize2dpts.hpp"

namespace HomLib {
namespace ValtonenOrnhagArxiv2020A {
    inline double get_algebraic_error_floor_fHf(const Eigen::VectorXd &data);

    HomLib::PoseData get_fHf(
        const Eigen::MatrixXd &p1,
        const Eigen::MatrixXd &p2,
        const Eigen::Matrix3d &R1,
        const Eigen::Matrix3d &R2
    ) {
        int nbr_coeffs = 30;
        int nbr_unknowns = 6;

        // Save copies of the inverse rotation
        Eigen::Matrix3d R1T = R1.transpose();
        Eigen::Matrix3d R2T = R2.transpose();

        // Compute normalization matrix
        double scale = normalize2dpts(p1);
        Eigen::Vector3d s;
        s << scale, scale, 1.0;
        Eigen::DiagonalMatrix<double, 3> S = s.asDiagonal();

        // Normalize data
        Eigen::Matrix3d x1;
        Eigen::Matrix3d x2;
        x1 = p1.colwise().homogeneous();
        x2 = p2.colwise().homogeneous();

        x1 = S * x1;
        x2 = S * x2;

        Eigen::MatrixXd x1t(2, 3);
        x1t << x1.colwise().hnormalized();
        Eigen::MatrixXd x2t(2, 3);
        x2t << x2.colwise().hnormalized();

        // Wrap input data to expected format
        Eigen::VectorXd input(nbr_coeffs);
        input << x1t.col(0),
                 x2t.col(0),
                 x1t.col(1),
                 x2t.col(1),
                 x1t.col(2),
                 x2t.col(2),
                 Eigen::Map<Eigen::VectorXd>(R1T.data(), 9),
                 Eigen::Map<Eigen::VectorXd>(R2T.data(), 9);

        // Extract solution
        Eigen::MatrixXcd sols = HomLib::ValtonenOrnhagArxiv2020A::solver_fHf(input);

        // Pre-processing: Remove complex-valued solutions
        double thresh = 1e-5;
        Eigen::ArrayXd real_sols(7);
        real_sols = sols.imag().cwiseAbs().colwise().sum();
        int nbr_real_sols = (real_sols <= thresh).count();

        // Allocate space for putative (real) homographies
        Eigen::MatrixXd best_homography(3, 3);
        double best_focal_length;
        double best_algebraic_error = DBL_MAX;
        double algebraic_error;

        // Since this is a 2.5 pt solver, use the last
        // (previously unused) constraint, to discard
        // false solutions.
        Eigen::ArrayXd xx(6);
        Eigen::VectorXd input_algebraic(nbr_coeffs + nbr_unknowns);

        for (int i = 0; i < real_sols.size(); i++) {
            if (real_sols(i) <= thresh) {
                // Compute algebraic error, and compare to other solutions.
                xx = sols.col(i).real();
                input_algebraic << xx, input;
                algebraic_error = HomLib::ValtonenOrnhagArxiv2020A::get_algebraic_error_floor_fHf(input_algebraic);

                if (algebraic_error < best_algebraic_error) {
                    best_algebraic_error = algebraic_error;
                    best_homography << xx[0], xx[2], xx[1],
                                           0, xx[3],     0,
                                      -xx[1], xx[4], xx[0];
                    best_focal_length = xx[5];
                }
            }
        }
        // Construct homography
        Eigen::Matrix3d K, Ki, H;
        K = Eigen::Vector3d(best_focal_length, best_focal_length, 1).asDiagonal();
        Ki = Eigen::Vector3d(1, 1, best_focal_length).asDiagonal();
        H = S.inverse() * K * R2 * best_homography * R1.transpose() * Ki * S;

        // Package output
        HomLib::PoseData posedata;
        posedata.homography = H;
        posedata.focal_length = best_focal_length / scale;

        return posedata;
    }

    // Function that utilizes the last equation of the DLT system to discard false solutions
    inline double get_algebraic_error_floor_fHf(const Eigen::VectorXd &data) {
        const double* d = data.data();

        // Compute algebraic error
        double error;
        error = -d[0]*std::pow(d[5], 2)*d[24]*d[34] - d[0]*d[5]*d[14]*d[18]*d[34] - d[0]*d[5]*d[15]*d[21]*d[34]
            - d[0]*d[5]*d[16]*d[24]*d[28] - d[0]*d[5]*d[17]*d[24]*d[31] - d[0]*d[14]*d[16]*d[18]*d[28]
            - d[0]*d[14]*d[17]*d[18]*d[31] - d[0]*d[15]*d[16]*d[21]*d[28] - d[0]*d[15]*d[17]*d[21]*d[31]
            - d[1]*std::pow(d[5], 2)*d[26]*d[34] - d[1]*d[5]*d[14]*d[20]*d[34] - d[1]*d[5]*d[15]*d[23]*d[34]
            - d[1]*d[5]*d[16]*d[26]*d[28] - d[1]*d[5]*d[17]*d[26]*d[31] - d[1]*d[14]*d[16]*d[20]*d[28]
            - d[1]*d[14]*d[17]*d[20]*d[31] - d[1]*d[15]*d[16]*d[23]*d[28] - d[1]*d[15]*d[17]*d[23]*d[31]
            - d[2]*std::pow(d[5], 2)*d[25]*d[34] - d[2]*d[5]*d[14]*d[19]*d[34] - d[2]*d[5]*d[15]*d[22]*d[34]
            - d[2]*d[5]*d[16]*d[25]*d[28] - d[2]*d[5]*d[17]*d[25]*d[31] - d[2]*d[14]*d[16]*d[19]*d[28]
            - d[2]*d[14]*d[17]*d[19]*d[31] - d[2]*d[15]*d[16]*d[22]*d[28] - d[2]*d[15]*d[17]*d[22]*d[31]
            + d[3]*std::pow(d[5], 2)*d[25]*d[33] + d[3]*d[5]*d[14]*d[19]*d[33] + d[3]*d[5]*d[15]*d[22]*d[33]
            + d[3]*d[5]*d[16]*d[25]*d[27] + d[3]*d[5]*d[17]*d[25]*d[30] + d[3]*d[14]*d[16]*d[19]*d[27]
            + d[3]*d[14]*d[17]*d[19]*d[30] + d[3]*d[15]*d[16]*d[22]*d[27] + d[3]*d[15]*d[17]*d[22]*d[30];
        return abs(error);
    }
}  // namespace ValtonenOrnhagArxiv2020A
}  // namespace HomLib