#include <iostream>
#include <Eigen/Core>
#include <Eigen/Dense>
#include <Eigen/IterativeLinearSolvers>

class MatrixReplacement;
template<typename Rhs> class MatrixReplacement_ProductReturnType;

namespace Eigen {
namespace internal {
  template<>
  struct traits<MatrixReplacement> :  Eigen::internal::traits<Eigen::SparseMatrix<double> >
  {};

  template <typename Rhs>
  struct traits<MatrixReplacement_ProductReturnType<Rhs> > {
    // The equivalent plain objet type of the product. This type is used if the product needs to be evaluated into a temporary.
    typedef Eigen::Matrix<typename Rhs::Scalar, Eigen::Dynamic, Rhs::ColsAtCompileTime> ReturnType;
  };
}
}

// Inheriting EigenBase should not be needed in the future.
class MatrixReplacement : public Eigen::EigenBase<MatrixReplacement> {
public:
  // Expose some compile-time information to Eigen:
  typedef double Scalar;
  typedef double RealScalar;
  enum {
    ColsAtCompileTime = Eigen::Dynamic,
    RowsAtCompileTime = Eigen::Dynamic,
    MaxColsAtCompileTime = Eigen::Dynamic,
    MaxRowsAtCompileTime = Eigen::Dynamic
  };

  Index rows() const { return 4; }
  Index cols() const { return 4; }

  void resize(Index a_rows, Index a_cols)
  {
    // This method should not be needed in the future.
    assert(a_rows==0 && a_cols==0 || a_rows==rows() && a_cols==cols());
  }

  // In the future, the return type should be Eigen::Product<MatrixReplacement,Rhs>
  template<typename Rhs>
  MatrixReplacement_ProductReturnType<Rhs> operator*(const Eigen::MatrixBase<Rhs>& x) const {
    return MatrixReplacement_ProductReturnType<Rhs>(*this, x.derived());
  }

};

// The proxy class representing the product of a MatrixReplacement with a MatrixBase<>
template<typename Rhs>
class MatrixReplacement_ProductReturnType : public Eigen::ReturnByValue<MatrixReplacement_ProductReturnType<Rhs> > {
public:
  typedef MatrixReplacement::Index Index;

  // The ctor store references to the matrix and right-hand-side object (usually a vector).
  MatrixReplacement_ProductReturnType(const MatrixReplacement& matrix, const Rhs& rhs)
    : m_matrix(matrix), m_rhs(rhs)
  {}

  Index rows() const { return m_matrix.rows(); }
  Index cols() const { return m_rhs.cols(); }

  // This function is automatically called by Eigen. It must evaluate the product of matrix * rhs into y.
  template<typename Dest>
  void evalTo(Dest& y) const
  {
    y.setZero(4);

    y(0) += 2 * m_rhs(0); y(1) += 1 * m_rhs(0);
    y(0) += 1 * m_rhs(1); y(1) += 2 * m_rhs(1); y(2) += 1 * m_rhs(1);
    y(1) += 1 * m_rhs(2); y(2) += 2 * m_rhs(2); y(3) += 1 * m_rhs(2);
    y(2) += 1 * m_rhs(3); y(3) += 2 * m_rhs(3);
  }

protected:
  const MatrixReplacement& m_matrix;
  typename Rhs::Nested m_rhs;
};

/*****/

// This class simply warp a diagonal matrix as a Jacobi preconditioner.
// In the future such simple and generic wrapper should be shipped within Eigen itsel.
template <typename _Scalar>
class MyJacobiPreconditioner
{
    typedef _Scalar Scalar;
    typedef Eigen::Matrix<Scalar,Eigen::Dynamic,1> Vector;
    typedef typename Vector::Index Index;

  public:
    // this typedef is only to export the scalar type and compile-time dimensions to solve_retval
    typedef Eigen::Matrix<Scalar,Eigen::Dynamic,Eigen::Dynamic> MatrixType;

    MyJacobiPreconditioner() : m_isInitialized(false) {}

    void setInvDiag(const Eigen::VectorXd &invdiag) {
      m_invdiag=invdiag;
      m_isInitialized=true;
    }

    Index rows() const { return m_invdiag.size(); }
    Index cols() const { return m_invdiag.size(); }

    template<typename MatType>
    MyJacobiPreconditioner& analyzePattern(const MatType& ) { return *this; }

    template<typename MatType>
    MyJacobiPreconditioner& factorize(const MatType& mat) { return *this; }

    template<typename MatType>
    MyJacobiPreconditioner& compute(const MatType& mat) { return *this; }

    template<typename Rhs, typename Dest>
    void _solve(const Rhs& b, Dest& x) const
    {
      x = m_invdiag.array() * b.array() ;
    }

    template<typename Rhs> inline const Eigen::internal::solve_retval<MyJacobiPreconditioner, Rhs>
    solve(const Eigen::MatrixBase<Rhs>& b) const
    {
      eigen_assert(m_isInitialized && "MyJacobiPreconditioner is not initialized.");
      eigen_assert(m_invdiag.size()==b.rows()
                && "MyJacobiPreconditioner::solve(): invalid number of rows of the right hand side matrix b");
      return Eigen::internal::solve_retval<MyJacobiPreconditioner, Rhs>(*this, b.derived());
    }

  protected:
    Vector m_invdiag;
    bool m_isInitialized;
};

namespace Eigen {
namespace internal {

template<typename _MatrixType, typename Rhs>
struct solve_retval<MyJacobiPreconditioner<_MatrixType>, Rhs>
  : solve_retval_base<MyJacobiPreconditioner<_MatrixType>, Rhs>
{
  typedef MyJacobiPreconditioner<_MatrixType> Dec;
  EIGEN_MAKE_SOLVE_HELPERS(Dec,Rhs)

  template<typename Dest> void evalTo(Dest& dst) const
  {
    dec()._solve(rhs(),dst);
  }
};

}
}

/*****/

int main()
{
  MatrixReplacement A;
  Eigen::VectorXd b(4), x;
  b << 1, 1, 1, 1;

  // solve Ax = b using CG with matrix-free version:
  Eigen::ConjugateGradient < MatrixReplacement, Eigen::Lower|Eigen::Upper, MyJacobiPreconditioner<double> > cg;

  Eigen::VectorXd invdiag(4);
  invdiag << 1./3., 1./4., 1./4., 1./3.;

  cg.preconditioner().setInvDiag(invdiag);
  cg.compute(A);
  x = cg.solve(b);

  std::cout << "#iterations: " << cg.iterations() << std::endl;
  std::cout << "estimated error: " << cg.error() << std::endl;
}


//////////////////////////////////////////////////////////////////////
// This file is distributed under the University of Illinois/NCSA Open Source
// License.  See LICENSE file in top directory for details.
//
// Copyright (c) 2016 Jeongnim Kim and QMCPACK developers.
//
// File developed by:
// Miguel A. Morales, moralessilva2@llnl.gov
//    Lawrence Livermore National Laboratory
//
// File created by:
// Miguel A. Morales, moralessilva2@llnl.gov
//    Lawrence Livermore National Laboratory
////////////////////////////////////////////////////////////////////////////////

#ifndef QMCPLUSPLUS_AFQMC_EXCITATIONS_HPP
#define QMCPLUSPLUS_AFQMC_EXCITATIONS_HPP

#include <boost/iterator/iterator_facade.hpp>
#include <map>
#include "AFQMC/config.h"
#include "AFQMC/Matrix/array_of_sequences.hpp"

namespace qmcplusplus
{

namespace afqmc
{

// CHEAT!!!
template<class TP, class integer>
size_t get_index(TP const& tp_, integer loc)
{
  if(loc == 0) return size_t(std::get<0>(tp_));
  else if(loc == 1) return size_t(std::get<1>(tp_));
  else throw std::runtime_error(" Error in qmcplusplus::afqmc::get_index<TP,integer>(). \n");
}

template<class Vector>
void push_excitation(Vector const& abij, Vector& v)
{
  if(abij.size()==0) return;
  assert(v.size()%abij.size()==0);
  size_t n = abij.size();
  for(typename Vector::iterator it=v.begin(); it<v.end(); it+=n)
    if( std::equal(abij.begin(),abij.end(),it) ) return;
  v.insert(v.end(),abij.begin(),abij.end());
}

template<class Vector>
size_t find_excitation(Vector const& abij, Vector& v)
{
  if(abij.size()==0) return 0;  // this assumes that the reference is 0
  assert(v.size()%abij.size()==0);
  size_t n = abij.size();
  size_t loc=0;
  for(typename Vector::iterator it=v.begin(); it<v.end(); it+=n, loc++)
    if( std::equal(abij.begin(),abij.end(),it) ) return loc;
  APP_ABORT("Error: Sequence not found in find_excitation.\n");
  return 0;
}

template<class excitations>
std::map<int,int> find_active_space(bool single_list, excitations const& abij, int NMO, int NAEA, int NAEB)
{
  std::map<int,int> mo2active;
  for(int i=0; i<2*NMO; i++) mo2active[i]=-1;
  std::vector<size_t> count(2*NMO);
  // reference first
  auto refc = abij.reference_configuration();
  for(int i=0; i<NAEA+NAEB; i++, ++refc) ++count[*refc];
  auto nex = abij.maximum_excitation_number();
  for(int n=1; n<nex[0]; ++n) {
    auto it = abij.alpha_begin(n);
    auto itend = abij.alpha_end(n);
    for(; it<itend; ++it) {
      auto exct = *it+n; // skip locations
      for(int ak=0; ak<n; ++ak, ++exct)
        ++count[*exct];
    }
  }
  for(int n=1; n<nex[1]; ++n) {
    auto it = abij.beta_begin(n);
    auto itend = abij.beta_end(n);
    for(; it<itend; ++it) {
      auto exct = *it+n; // skip locations
      for(int ak=0; ak<n; ++ak, ++exct)
        ++count[*exct];
    }
  }
  if(not single_list) {
    int ik=0;
    for(int i=0; i<NMO; ++i)
      if(count[i]>0 || count[i+NMO]>0) {
        if(count[i]>0) mo2active[i] = ik;
        if(count[i+NMO]>0) mo2active[i+NMO] = ik;
        ++ik;
      }
  } else {
    int ik=0;
    for(int i=0; i<NMO; ++i)
      if(count[i]>0) mo2active[i] = ik++;
    ik=0;
    for(int i=NMO; i<2*NMO; ++i)
      if(count[i]>0) mo2active[i] = ik++;
  }
  return mo2active;
}

/*
 * - exct stores, for each electron excitaion, the location of the orbital in the reference
 *    being excited and the index of the excited orbital.
 */
template<class Vector, class T>
int get_excitation_number(bool getIndx, Vector& refc, Vector& confg, Vector& exct, T& ci, Vector& Iwork) {
  int NE = refc.size();
  exct.clear();
  int cnt=0;
  if(getIndx) std::copy(refc.begin(),refc.end(),Iwork.begin());
  auto it = refc.data();
  assert(Iwork.size() >= refc.size());
  for(int i=0; i<NE; i++,it++)
    if(!std::binary_search(confg.begin(),confg.end(),*it)) {
      if(getIndx) {
        // store the location, NOT the index!!!
        //exct.emplace_back(*it);
        exct.emplace_back(i);
      }
      cnt++;
    }
  if(!getIndx)
    return cnt;
  it = confg.data();
  int cnt2=0;
  for(int i=0; i<NE; i++,it++)
    if(!std::binary_search(refc.begin(),refc.end(),*it)) {
      exct.emplace_back(*it);
      Iwork[exct[cnt2]]=*it;
      cnt2++;
    }
  assert(cnt==cnt2);
  // sort Iwork and count number of exchanges to determine permutation sign
  // sooo slow but sooo simple too
  for(int i=0; i<NE; i++)
    for(int j=i+1; j<NE; j++)
    {
      if(Iwork[j] < Iwork[i])
      {
        ci *= T(-1.0);
        std::swap(Iwork[i],Iwork[j]);
      }
    }
  return cnt;
}

template<typename intT, class csr>
inline std::vector<size_t> get_nnz(csr const& PsiT_MO, intT* refc, size_t N, size_t shift)
{
  std::vector<size_t> res(N);
  for(size_t i=0; i<N; i++)
    res[i] = PsiT_MO.num_non_zero_elements(*(refc+i)-shift);
  return res;
}

// try putting this in shared memory later on
template<class I = int,
	 class VType = std::complex<double>,
         class Alloc = shared_allocator<I>,
         class is_root = ma::sparse::is_root>
struct ph_excitations
{
  public:
  using integer_type = I;
  using configuration_type = std::tuple<int,int,VType>;

  private:
  using IAllocator = Alloc;
  using CAllocator = typename Alloc::template rebind<configuration_type>::other;
  using confg_aos = ma::sparse::array_of_sequences<configuration_type,int,CAllocator,is_root>;
  using index_aos = ma::sparse::array_of_sequences<integer_type,int,IAllocator,is_root>;

  IAllocator i_allocator_;
  CAllocator c_allocator_;

  template< typename Integer>
  class Iterator: public boost::iterator_facade< Iterator<Integer>,
                                               Integer*,
                                               std::random_access_iterator_tag,
                                               Integer*,
                                               std::ptrdiff_t
                                             >
  {
    public:
      using difference_type = std::ptrdiff_t;
      using reference = Integer *;
      using const_reference = Integer const*;
      using value_tupe = Integer*;

      Iterator(Integer* index, size_t d_) : p_index(index),D(d_) {}

        // What we implement is determined by the boost::forward_traversal_tag
    private:
      friend class boost::iterator_core_access;

      void increment() { p_index+=2*D; }

      bool equal(Iterator const& other) const
      {
        return this->p_index == other.p_index;
      }

      reference dereference() const
      {
        return reference(p_index);
      }

      void decrement() { p_index-=2*D; }

      void advance(int n) { p_index+=2*D*n; }

      difference_type distance_to(Iterator const& z) const { return ((z.p_index-p_index)/2/D); }

    private:
      Integer* p_index;
      size_t D;
  };

  template< typename Integer>
  class Iterator_const: public boost::iterator_facade< Iterator_const<Integer>,
                                               Integer const*,
                                               std::random_access_iterator_tag,
                                               Integer const*,
                                               std::ptrdiff_t
                                             >
  {
    public:
      using difference_type = std::ptrdiff_t;
      using reference = Integer const*;
      using const_reference = Integer const*;
      using value_tupe = Integer*;

      Iterator_const(Integer * index, size_t d_) : p_index(index),D(d_) {}
      Iterator_const(Integer const* index, size_t d_) : p_index(index),D(d_) {}

        // What we implement is determined by the boost::forward_traversal_tag
    private:
      friend class boost::iterator_core_access;

      void increment() { p_index+=2*D; }

      bool equal(Iterator_const const& other) const
      {
        return this->p_index == other.p_index;
      }

      reference dereference() const
      {
        return reference(p_index);
      }

      void decrement() { p_index-=2*D; }

      void advance(int n) { p_index+=2*D*n; }

      difference_type distance_to(Iterator_const const& z) const { return ((z.p_index-p_index)/2/D); }

    private:
      Integer* p_index;
      size_t D;
  };

  public:

  using Excitation_Iterator = Iterator<integer_type>;
  using Excitation_const_Iterator = Iterator_const<integer_type>;

  ph_excitations() = delete;

  // Note: terms_per_excitation[0] has special meaning, the number of electrons in the calculation.
  // coefficients[0] will store the reference configuration itself.
  ph_excitations(size_t number_of_configurations,
                 int na_, int nb_,
                 std::vector<size_t>& unique_alpha_counts,
                 std::vector<size_t>& unique_beta_counts,
                 Alloc alloc_ = Alloc{}):
                 i_allocator_(alloc_),
                 c_allocator_(alloc_),
                 NAEA(na_),NAEB(nb_),
                 configurations(1,number_of_configurations,c_allocator_),
                 reference(1,NAEA+NAEB,i_allocator_),
                 unique_alpha(unique_alpha_counts.size(),unique_alpha_counts,i_allocator_),
                 unique_beta(unique_beta_counts.size(),unique_beta_counts,i_allocator_)
  {
    size_t emax = std::max(unique_alpha_counts.size(),unique_beta_counts.size());
    sum_of_exct.resize(emax+1);
    sum_of_exct[0] = {0,0};
    sum_of_exct[1] = {1,1};
    for(size_t n=1; n<unique_alpha.size(); ++n)
      sum_of_exct[n+1][0] = sum_of_exct[n][0]+unique_alpha_counts[n]/size_t(2)/n;
    for(size_t n=unique_alpha.size()+1; n<=emax; n++)
      sum_of_exct[n][0]=sum_of_exct[n-1][0];
    for(size_t n=1; n<unique_beta.size(); ++n)
      sum_of_exct[n+1][1] = sum_of_exct[n][1]+unique_beta_counts[n]/size_t(2)/n;
    for(size_t n=unique_beta.size()+1; n<=emax; n++)
      sum_of_exct[n][1]=sum_of_exct[n-1][1];
  }

  ph_excitations(ph_excitations const& other) = delete;
  //ph_excitations(ph_excitations && other) = default;
  ph_excitations(ph_excitations && other):
                 i_allocator_(other.i_allocator_),
                 c_allocator_(other.c_allocator_),
                 NAEA(other.NAEA),NAEB(other.NAEB),
                 configurations(std::move(other.configurations)),
                 reference(std::move(other.reference)),
                 unique_alpha(std::move(other.unique_alpha)),
                 unique_beta(std::move(other.unique_beta)),
                 sum_of_exct(std::move(other.sum_of_exct))
  {}

  ph_excitations& operator=(ph_excitations const& other) = delete;
  ph_excitations& operator=(ph_excitations && other) = default;

  std::array<size_t,2> maximum_excitation_number() const
  { return {unique_alpha.size(),unique_beta.size()}; }
  size_t number_of_unique_alpha_excitations(int n) const
  {
    if(n==0) return 1;
    return unique_alpha.num_elements(n)/2/n;
  }
  size_t number_of_unique_beta_excitations(int n) const
  {
    if(n==0) return 1;
    return unique_beta.num_elements(n)/2/n;
  }
  std::array<size_t,2> number_of_unique_excitations(int n) const
  {
    if(n==0) return {1,1};
    std::array<size_t,2> res{0,0};
    if(n < unique_alpha.size()) res[0] = unique_alpha.num_elements(n)/2/n;
    if(n < unique_beta.size()) res[1] = unique_beta.num_elements(n)/2/n;
    return res;
  }
  std::array<size_t,2> number_of_unique_excitations() const
  {
    return sum_of_exct.back();
  }

  size_t number_of_configurations() const { return configurations.num_elements(0); }

  // returns the number of unique excitations with particle number less than n
  std::array<size_t,2> number_of_unique_smaller_than(int n) const { return sum_of_exct[n]; }

  template<class intIt>
  void add_alpha(size_t n, intIt indx) {
    assert(n<unique_alpha.size());
    assert(n>0);
    for(int i=0; i<2*n; i++, ++indx)
      unique_alpha.emplace_back(n,static_cast<integer_type>(*indx));
  }

  template<class intIt>
  void add_beta(size_t n, intIt indx) {
    assert(n<unique_beta.size());
    assert(n>0);
    for(int i=0; i<2*n; i++, ++indx)
      unique_beta.emplace_back(n,static_cast<integer_type>(*indx));
  }

  template<class Vector>
  void add_reference(Vector& refa, Vector& refb) {
    for(auto k: refa)
      reference.emplace_back(0,static_cast<integer_type>(k));
    for(auto k: refb)
      reference.emplace_back(0,static_cast<integer_type>(k));
  }

  // index=0 is reserved for the reference!!!
  template<typename integer, typename value>
  void add_configuration(integer alpha_indx, integer beta_index, value ci) {
    configurations.emplace_back(0,configuration_type{alpha_indx,beta_index,ci});
  }

  typename Excitation_Iterator::const_reference reference_configuration(int spin=0) const{
    return to_address(reference.values(0)) + (spin==0?0:NAEA);
  }

  typename Excitation_Iterator::reference reference_configuration(int spin=0) {
    return to_address(reference.values(0)) + (spin==0?0:NAEA);
  }

  configuration_type const* configurations_begin() const {
    return to_address(configurations.values(0));
  }

  configuration_type const* configurations_end() const {
    return to_address(configurations.values()) + (*configurations.pointers_end(0));
  }

  configuration_type const* configuration(int i) const {
    return to_address(configurations.values()) + i;
  }

  Excitation_Iterator alpha_begin(int n) {
    assert(n>0);
    if(n<unique_alpha.size()) {
      return Excitation_Iterator(to_address(unique_alpha.values(n)),n);
    } else
      return alpha_end(n);
  }

  Excitation_Iterator alpha_end(int n) {
    assert(n>0);
    if(n<unique_alpha.size())
      return Excitation_Iterator(to_address(unique_alpha.values())
                                    +(*unique_alpha.pointers_end(n)),n);
    else
      return Excitation_Iterator(to_address(unique_alpha.values())
                                    +(*unique_alpha.pointers_end(unique_alpha.size()-1)),1);
  }

  Excitation_const_Iterator alpha_begin(int n) const {
    assert(n>0);
    if(n<unique_alpha.size()) {
      return Excitation_const_Iterator(to_address(unique_alpha.values(n)),n);
    } else
      return alpha_end(n);
  }

  Excitation_const_Iterator alpha_end(int n) const {
    assert(n>0);
    if(n<unique_alpha.size())
      return Excitation_const_Iterator(to_address(unique_alpha.values())
                                    +(*unique_alpha.pointers_end(n)),n);
    else
      return Excitation_const_Iterator(to_address(unique_alpha.values())
                                    +(*unique_alpha.pointers_end(unique_alpha.size()-1)),1);
  }

  Excitation_Iterator beta_begin(int n) {
    assert(n>0);
    if(n<unique_beta.size())
      return Excitation_Iterator(to_address(unique_beta.values(n)),n);
    else
      return beta_end(n);
  }

  Excitation_Iterator beta_end(int n) {
    assert(n>0);
    if(n<unique_beta.size())
      return Excitation_Iterator(to_address(unique_beta.values())
                                    +(*unique_beta.pointers_end(n)),n);
    else
      return Excitation_Iterator(to_address(unique_beta.values())
                                    +(*unique_beta.pointers_end(unique_beta.size()-1)),1);
  }

  Excitation_const_Iterator beta_begin(int n) const {
    assert(n>0);
    if(n<unique_beta.size())
      return Excitation_const_Iterator(to_address(unique_beta.values(n)),n);
    else
      return beta_end(n);
  }

  Excitation_const_Iterator beta_end(int n) const {
    assert(n>0);
    if(n<unique_beta.size())
      return Excitation_const_Iterator(to_address(unique_beta.values())
                                    +(*unique_beta.pointers_end(n)),n);
    else
      return Excitation_const_Iterator(to_address(unique_beta.values())
                                    +(*unique_beta.pointers_end(unique_beta.size()-1)),1);
  }

  // for generic access
  std::array<Excitation_Iterator,2> unique_begin(int n)
  {  return std::array<Excitation_Iterator,2>{alpha_begin(n),beta_begin(n)}; }
  std::array<Excitation_Iterator,2> unique_end(int n)
  {  return std::array<Excitation_Iterator,2>{alpha_end(n),beta_end(n)}; }
  std::array<Excitation_const_Iterator,2> unique_begin(int n) const
  {  return std::array<Excitation_const_Iterator,2>{alpha_begin(n),beta_begin(n)}; }
  std::array<Excitation_const_Iterator,2> unique_end(int n) const
  {  return std::array<Excitation_const_Iterator,2>{alpha_end(n),beta_end(n)}; }

  template<class Vector>
  void get_alpha_configuration(size_t index, Vector& confg) const{
    assert(confg.size() >= NAEA);
    std::copy_n(to_address(reference.values(0)),NAEA,confg.data());
    if(index==0) return;
    // could use lower bound
    for(int i=1; i<unique_alpha.size(); i++) {
      if(index >= sum_of_exct[i][0] && index < sum_of_exct[i+1][0]) {
        size_t dn = index-sum_of_exct[i][0];
        auto exct = unique_alpha.values(i) + 2*i*dn;
        for(int n=0; n<i; n++)
          confg[ exct[n] ] = exct[n+i];
        return;
      }
    }
    APP_ABORT(" Error in ph_excitations::get_alpha_configuration() \n");
  }

  template<class Vector>
  void get_beta_configuration(size_t index, Vector& confg) const{
    assert(confg.size() >= NAEB);
    std::copy_n(to_address(reference.values(0))+NAEA,NAEB,confg.data());
    if(index==0) return;
    // could use lower bound
    for(int i=1; i<unique_beta.size(); i++) {
      if(index >= sum_of_exct[i][1] && index < sum_of_exct[i+1][1]) {
        size_t dn = index-sum_of_exct[i][1];
        auto exct = unique_beta.values(i) + 2*i*dn;
        for(int n=0; n<i; n++)
          confg[ exct[n] ] = exct[n+i];
        return;
      }
    }
    APP_ABORT(" Error in ph_excitations::get_beta_configuration() \n");
  }

  template<class Vector>
  void get_configuration(int spin, size_t index, Vector& confg) const{
    if(spin==0)
      get_alpha_configuration(index,confg);
    else
      get_beta_configuration(index,confg);
  }

  private:

  // using array_of_seq until I switch to Boost.Multi to be able to use shared_allocator
  int NAEA, NAEB;
  confg_aos configurations;
  index_aos reference;
  index_aos unique_alpha;
  index_aos unique_beta;
  std::vector<std::array<size_t,2>> sum_of_exct;

};

}

}

#endif


// Copyright (c) 2016 The UUV Simulator Authors.
// All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// This source code is from rotors_simulator
//   (https://github.com/ethz-asl/rotors_simulator)
// * Copyright 2015 Fadri Furrer, ASL, ETH Zurich, Switzerland
// * Copyright 2015 Michael Burri, ASL, ETH Zurich, Switzerland
// * Copyright 2015 Mina Kamel, ASL, ETH Zurich, Switzerland
// * Copyright 2015 Janosch Nikolic, ASL, ETH Zurich, Switzerland
// * Copyright 2015 Markus Achtelik, ASL, ETH Zurich, Switzerland
// This source code is licensed under the Apache-2.0 license found in the
// open_source_licenses.txt file in the root directory of this source tree.
//
// The original code was modified by the UUV Simulator Authors to adhere to
// Gazebo's coding standards.

#ifndef UUV_SENSOR_PLUGINS_COMMON_H_
#define UUV_SENSOR_PLUGINS_COMMON_H_

#include <string>
#include <Eigen/Dense>
#include <gazebo/gazebo.hh>

namespace gazebo {

/**
 * \brief Obtains a parameter from sdf.
 * \param[in] sdf Pointer to the sdf object.
 * \param[in] name Name of the parameter.
 * \param[out] param Param Variable to write the parameter to.
 * \param[in] default_value Default value, if the parameter not available.
 * \param[in] verbose If true, gzerror if the parameter is not available.
 */
template<class T>
bool getSdfParam(sdf::ElementPtr sdf, const std::string& name, T& param,
                 const T& default_value, const bool& verbose = false)
{
  if (sdf->HasElement(name))
  {
    param = sdf->GetElement(name)->Get<T>();
    return true;
  }
  else
  {
    param = default_value;
    if (verbose)
      gzerr << "[uuv_sensor_plugins] Please specify a value for parameter \""
            << name << "\".\n";
  }
  return false;
}
}  // namespace gazebo

template <typename T>
class FirstOrderFilter {
/*
This class can be used to apply a first order filter on a signal.
It allows different acceleration and deceleration time constants.

Short reveiw of discrete time implementation of firest order system:
Laplace:
    X(s)/U(s) = 1/(tau*s + 1)
continous time system:
    dx(t) = (-1/tau)*x(t) + (1/tau)*u(t)
discretized system (ZoH):
    x(k+1) = exp(samplingTime*(-1/tau))*x(k) + (1 - exp(samplingTime*(-1/tau))) * u(k)
*/

  public:
    FirstOrderFilter(double timeConstantUp, double timeConstantDown,
                     T initialState):
      timeConstantUp_(timeConstantUp),
      timeConstantDown_(timeConstantDown),
      previousState_(initialState) {}

    T updateFilter(T inputState, double samplingTime)
    {
      /*
      This method will apply a first order filter on the inputState.
      */
      T outputState;
      if (inputState > previousState_)
      {
        // Calcuate the outputState if accelerating.
        double alphaUp = exp(- samplingTime / timeConstantUp_);
        // x(k+1) = Ad*x(k) + Bd*u(k)
        outputState = alphaUp * previousState_ + (1 - alphaUp) * inputState;
      }
      else
      {
        // Calculate the outputState if decelerating.
        double alphaDown = exp(- samplingTime / timeConstantDown_);
        outputState = alphaDown * previousState_ + (1 - alphaDown) * inputState;
      }
      previousState_ = outputState;
      return outputState;
    }
    ~FirstOrderFilter() {}

  protected:
    double timeConstantUp_;
    double timeConstantDown_;
    T previousState_;
};

/// Computes a quaternion from the 3-element small angle approximation theta.
template<class Derived>
Eigen::Quaternion<typename Derived::Scalar> QuaternionFromSmallAngle(
        const Eigen::MatrixBase<Derived> & theta)
{
  typedef typename Derived::Scalar Scalar;
  EIGEN_STATIC_ASSERT_FIXED_SIZE(Derived);
  EIGEN_STATIC_ASSERT_VECTOR_SPECIFIC_SIZE(Derived, 3);
  const Scalar q_squared = theta.squaredNorm() / 4.0;

  if (q_squared < 1)
  {
    return Eigen::Quaternion<Scalar>(sqrt(1 - q_squared), theta[0] * 0.5,
            theta[1] * 0.5, theta[2] * 0.5);
  }
  else
  {
    const Scalar w = 1.0 / sqrt(1 + q_squared);
    const Scalar f = w * 0.5;
    return Eigen::Quaternion<Scalar>(w, theta[0] * f, theta[1] * f,
            theta[2] * f);
  }
}

template<class In, class Out>
void copyPosition(const In& in, Out* out)
{
  out->x = in.x;
  out->y = in.y;
  out->z = in.z;
}

#endif  // UUV_SENSOR_PLUGINS_COMMON_H_


// The quad_float module is derived from the doubledouble
// library originally developed by Keith Briggs:
//    http://keithbriggs.info/doubledouble.html
// I attach the original copyright notice.

/*

Copyright (C) 1997 Keith Martin Briggs

This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License as published by the Free Software Foundation; either
version 2.1 of the License, or (at your option) any later version.

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public
License along with this library; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

*/

#if (defined(__GNUC__) && __FAST_MATH__)
#error "do not compile quad_float.cpp with -ffast-math!!"
#endif

#ifdef __INTEL_COMPILER
#pragma float_control(precise,on)
#endif

// NOTE: the above will force the Intel compiler to adhere to
// language standards, which it does not do by default

#include <NTL/quad_float.h>
#include <NTL/RR.h>

#include <cfloat>

#include <NTL/new.h>

NTL_START_IMPL

#if (NTL_EXT_DOUBLE && defined(__GNUC__) && (defined(__i386__) || defined(__x86_64__)))

#if (!defined(NTL_X86_FIX) && !defined(NTL_NO_X86_FIX))

#define NTL_X86_FIX

#endif

#endif

#if (NTL_EXT_DOUBLE && !defined(NTL_X86_FIX))

#define DOUBLE volatile double

#else

#define DOUBLE double

#endif

#ifdef NTL_X86_FIX

#define START_FIX \
  volatile unsigned short __old_cw, __new_cw; \
  asm volatile ("fnstcw %0":"=m" (__old_cw)); \
  __new_cw = (__old_cw & ~0x300) | 0x200; \
  asm volatile ("fldcw %0": :"m" (__new_cw));

#define END_FIX  asm volatile ("fldcw %0": :"m" (__old_cw));

#else

#define START_FIX
#define END_FIX

#endif

static
void normalize(quad_float& z, const double& xhi, const double& xlo)
{
START_FIX
   DOUBLE u, v;

   u = xhi + xlo;
   v = xhi - u;
   v = v + xlo;

   z.hi = u;
   z.lo = v;
END_FIX
}

#if (NTL_BITS_PER_LONG >= NTL_DOUBLE_PRECISION)

quad_float to_quad_float(long n)
{
   DOUBLE xhi, xlo;

   xhi = TrueDouble(n);

   // Because we are assuming 2's compliment integer
   // arithmetic, the following prevents long(xhi) from overflowing.

   if (n > 0)
      xlo = TrueDouble(n+long(-xhi));
   else
      xlo = TrueDouble(n-long(xhi));

   // renormalize...just to be safe

   quad_float z;
   normalize(z, xhi, xlo);
   return z;
}

quad_float to_quad_float(unsigned long n)
{
   DOUBLE xhi, xlo, t;

   const double bnd = double(1L << (NTL_BITS_PER_LONG-2))*4.0;

   xhi = TrueDouble(n);

   if (xhi >= bnd)
      t = xhi - bnd;
   else
      t = xhi;

   // we use the "to_long" function here to be as portable as possible.
   long llo = to_long(n - (unsigned long)(t));
   xlo = TrueDouble(llo);

   quad_float z;
   normalize(z, xhi, xlo);
   return z;
}
#endif

NTL_CHEAP_THREAD_LOCAL
long quad_float::oprec = 10;

void quad_float::SetOutputPrecision(long p)
{
   if (p < 1) p = 1;

   if (NTL_OVERFLOW(p, 1, 0))
      ResourceError("quad_float: output precision too big");

   oprec = p;
}

quad_float operator +(const quad_float& x, const quad_float& y ) {
START_FIX
        DOUBLE    H, h, T, t, S, s, e, f;
        DOUBLE    t1;

        S = x.hi + y.hi;
        T = x.lo + y.lo;
        e = S - x.hi;
        f = T - x.lo;

        t1 = S-e;
        t1 = x.hi-t1;
        s = y.hi-e;
        s = s + t1;

        t1 = T-f;
        t1 = x.lo-t1;
        t = y.lo-f;
        t = t + t1;

        s = s + T;
        H = S + s;
        h = S - H;
        h = h + s;

        h = h + t;
        e = H + h;
        f = H - e;
        f = f + h;
END_FIX
        return quad_float(e, f);
}

quad_float& operator +=(quad_float& x, const quad_float& y ) {
START_FIX
        DOUBLE    H, h, T, t, S, s, e, f;
        DOUBLE    t1;

        S = x.hi + y.hi;
        T = x.lo + y.lo;
        e = S - x.hi;
        f = T - x.lo;

        t1 = S-e;
        t1 = x.hi-t1;
        s = y.hi-e;
        s = s + t1;

        t1 = T-f;
        t1 = x.lo-t1;
        t = y.lo-f;
        t = t + t1;

        s = s + T;
        H = S + s;
        h = S - H;
        h = h + s;

        h = h + t;
        e = H + h;
        f = H - e;
        f = f + h;

        x.hi = e;
        x.lo = f;
END_FIX
        return x;
}

quad_float operator -(const quad_float& x, const quad_float& y ) {
START_FIX
        DOUBLE    H, h, T, t, S, s, e, f;
        DOUBLE    t1, yhi, ylo;

        yhi = -y.hi;
        ylo = -y.lo;

        S = x.hi + yhi;
        T = x.lo + ylo;
        e = S - x.hi;
        f = T - x.lo;

        t1 = S-e;
        t1 = x.hi-t1;
        s = yhi-e;
        s = s + t1;

        t1 = T-f;
        t1 = x.lo-t1;
        t = ylo-f;
        t = t + t1;

        s = s + T;
        H = S + s;
        h = S - H;
        h = h + s;

        h = h + t;
        e = H + h;
        f = H - e;
        f = f + h;

END_FIX
        return quad_float(e, f);
}

quad_float& operator -=(quad_float& x, const quad_float& y ) {
START_FIX
        DOUBLE    H, h, T, t, S, s, e, f;
        DOUBLE    t1, yhi, ylo;

        yhi = -y.hi;
        ylo = -y.lo;

        S = x.hi + yhi;
        T = x.lo + ylo;
        e = S - x.hi;
        f = T - x.lo;

        t1 = S-e;
        t1 = x.hi-t1;
        s = yhi-e;
        s = s + t1;

        t1 = T-f;
        t1 = x.lo-t1;
        t = ylo-f;
        t = t + t1;

        s = s + T;
        H = S + s;
        h = S - H;
        h = h + s;

        h = h + t;
        e = H + h;
        f = H - e;
        f = f + h;

        x.hi = e;
        x.lo = f;
END_FIX
        return x;
}

quad_float operator -(const quad_float& x)
{
START_FIX
   DOUBLE xhi, xlo, u, v;

   xhi = -x.hi;
   xlo = -x.lo;

   // it is a good idea to renormalize here, just in case
   // the rounding rule depends on sign, and thus we will
   // maintain the "normal form" for quad_float's.

   u = xhi + xlo;
   v = xhi - u;
   v = v + xlo;

END_FIX
   return quad_float(u, v);
}

#if (NTL_FMA_DETECTED)

double quad_float_zero = 0;

static inline
double Protect(double x) { return x + quad_float_zero; }

#else

static inline
double Protect(double x) { return x; }

#endif

// NOTE: this is really sick: some compilers will issue FMA
// (fused mul add) instructions which will break correctness.
// C99 standard is supposed to prevent this across separate
// statements, but C++ standard doesn't guarantee much at all.
// In any case, gcc does not even implement the C99 standard
// correctly.  One could disable this by compiling with
// an appropriate flag: -mno-fma works for gcc, while -no-fma works
// for icc.  icc and MSVC++ also support pragmas to do this:
// #pragma fp_contract(off).  There is also a compiler flag for
// gcc: -ffp-contract=off, but -mno-fma seems more widely supported.
// These flags work for clang, as well.
//
// But in any case, I'd rather not mess with getting these flags right.
// By calling Protect(a*b), this has the effect of forcing the
// compiler to compute a*b + 0.  Assuming the compiler otherwise
// does not perform any re-association, this should do the trick.
// There is a small performance penalty, but it should be reasonable.

quad_float operator *(const quad_float& x,const quad_float& y ) {
START_FIX
  DOUBLE hx, tx, hy, ty, C, c;
  DOUBLE t1, t2;

  C = Protect(NTL_QUAD_FLOAT_SPLIT*x.hi);
  hx = C-x.hi;
  c = Protect(NTL_QUAD_FLOAT_SPLIT*y.hi);
  hx = C-hx;
  tx = x.hi-hx;
  hy = c-y.hi;
  C = Protect(x.hi*y.hi);
  hy = c-hy;
  ty = y.hi-hy;

  // c = ((((hx*hy-C)+hx*ty)+tx*hy)+tx*ty)+(x.hi*y.lo+x.lo*y.hi);

  t1 = Protect(hx*hy);
  t1 = t1-C;
  t2 = Protect(hx*ty);
  t1 = t1+t2;
  t2 = Protect(tx*hy);
  t1 = t1+t2;
  t2 = Protect(tx*ty);
  c = t1+t2;
  t1 = Protect(x.hi*y.lo);
  t2 = Protect(x.lo*y.hi);
  t1 = t1+t2;
  c = c + t1;

  hx = C+c;
  tx = C-hx;
  tx = tx+c;

END_FIX
  return quad_float(hx, tx);
}

quad_float& operator *=(quad_float& x,const quad_float& y ) {
START_FIX
  DOUBLE hx, tx, hy, ty, C, c;
  DOUBLE t1, t2;

  C = Protect(NTL_QUAD_FLOAT_SPLIT*x.hi);
  hx = C-x.hi;
  c = Protect(NTL_QUAD_FLOAT_SPLIT*y.hi);
  hx = C-hx;
  tx = x.hi-hx;
  hy = c-y.hi;
  C = Protect(x.hi*y.hi);
  hy = c-hy;
  ty = y.hi-hy;

  // c = ((((hx*hy-C)+hx*ty)+tx*hy)+tx*ty)+(x.hi*y.lo+x.lo*y.hi);

  t1 = Protect(hx*hy);
  t1 = t1-C;
  t2 = Protect(hx*ty);
  t1 = t1+t2;
  t2 = Protect(tx*hy);
  t1 = t1+t2;
  t2 = Protect(tx*ty);
  c = t1+t2;
  t1 = Protect(x.hi*y.lo);
  t2 = Protect(x.lo*y.hi);
  t1 = t1+t2;
  c = c + t1;

  hx = C+c;
  tx = C-hx;
  tx = tx+c;

  x.hi = hx;
  x.lo = tx;
END_FIX
  return x;
}

quad_float operator /(const quad_float& x, const quad_float& y ) {
START_FIX
  DOUBLE hc, tc, hy, ty, C, c, U, u;
  DOUBLE t1;

  C = x.hi/y.hi;
  c = Protect(NTL_QUAD_FLOAT_SPLIT*C);
  hc = c-C;
  u = Protect(NTL_QUAD_FLOAT_SPLIT*y.hi);
  hc = c-hc;
  tc = C-hc;
  hy = u-y.hi;
  U = Protect(C * y.hi);
  hy = u-hy;
  ty = y.hi-hy;

  // u = (((hc*hy-U)+hc*ty)+tc*hy)+tc*ty;

  u = Protect(hc*hy);
  u = u-U;
  t1 = Protect(hc*ty);
  u = u+t1;
  t1 = Protect(tc*hy);
  u = u+t1;
  t1 = Protect(tc*ty);
  u = u+t1;

  // c = ((((x.hi-U)-u)+x.lo)-C*y.lo)/y.hi;

  c = x.hi-U;
  c = c-u;
  c = c+x.lo;
  t1 = Protect(C*y.lo);
  c = c - t1;
  c = c/y.hi;

  hy = C+c;
  ty = C-hy;
  ty = ty+c;

END_FIX
  return quad_float(hy, ty);
}

quad_float& operator /=(quad_float& x, const quad_float& y ) {
START_FIX
  DOUBLE hc, tc, hy, ty, C, c, U, u;
  DOUBLE t1;

  C = x.hi/y.hi;
  c = Protect(NTL_QUAD_FLOAT_SPLIT*C);
  hc = c-C;
  u = Protect(NTL_QUAD_FLOAT_SPLIT*y.hi);
  hc = c-hc;
  tc = C-hc;
  hy = u-y.hi;
  U = Protect(C * y.hi);
  hy = u-hy;
  ty = y.hi-hy;

  // u = (((hc*hy-U)+hc*ty)+tc*hy)+tc*ty;

  u = Protect(hc*hy);
  u = u-U;
  t1 = Protect(hc*ty);
  u = u+t1;
  t1 = Protect(tc*hy);
  u = u+t1;
  t1 = Protect(tc*ty);
  u = u+t1;

  // c = ((((x.hi-U)-u)+x.lo)-C*y.lo)/y.hi;

  c = x.hi-U;
  c = c-u;
  c = c+x.lo;
  t1 = Protect(C*y.lo);
  c = c - t1;
  c = c/y.hi;

  hy = C+c;
  ty = C-hy;
  ty = ty+c;

  x.hi = hy;
  x.lo = ty;
END_FIX
  return x;
}

quad_float sqrt(const quad_float& y) {
  if (y.hi < 0.0)
    ArithmeticError("quad_float: square root of negative number");
  if (y.hi == 0.0) return quad_float(0.0,0.0);

  double c;
  c = sqrt(y.hi);
  ForceToMem(&c);  // This is fairly paranoid, but it doesn't cost too much.

START_FIX

  DOUBLE p,q,hx,tx,u,uu,cc;
  DOUBLE t1;

  p = Protect(NTL_QUAD_FLOAT_SPLIT*c);
  hx = (c-p);
  hx = hx+p;
  tx = c-hx;
  p = Protect(hx*hx);
  q = Protect(hx*tx);
  q = q+q;

  u = p+q;
  uu = p-u;
  uu = uu+q;
  t1 = Protect(tx*tx);
  uu = uu+t1;

  cc = y.hi-u;
  cc = cc-uu;
  cc = cc+y.lo;
  t1 = c+c;
  cc = cc/t1;

  hx = c+cc;
  tx = c-hx;
  tx = tx+cc;
END_FIX
  return quad_float(hx, tx);
}

void power(quad_float& z, const quad_float& a, long e)
{
   quad_float res, u;
   unsigned long k;

   if (e < 0)
      k = -((unsigned long) e);
   else
      k = e;

   res = 1.0;
   u = a;

   while (k) {
      if (k & 1)
         res = res * u;

      k = k >> 1;
      if (k)
         u = u * u;
   }

   if (e < 0)
      z = 1.0/res;
   else
      z = res;
}

void power2(quad_float& z, long e)
{
   z.hi = _ntl_ldexp(1.0, e);
   z.lo = 0;
}

long to_long(const quad_float& x)
{
   double fhi, flo;

   fhi = floor(x.hi);

   if (fhi == x.hi)
      flo = floor(x.lo);
   else
      flo = 0;

   // the following code helps to prevent unnecessary integer overflow,
   // and guarantees that to_long(to_quad_float(a)) == a, for all long a,
   // provided long's are not too wide.

   if (fhi > 0)
      return long(flo) - long(-fhi);
   else
      return long(fhi) + long(flo);
}

// This version of ZZ to quad_float coversion relies on the
// precise rounding rules implemented by the ZZ to double conversion.

void conv(quad_float& z, const ZZ& a)
{
   double xhi, xlo;

   conv(xhi, a);

   if (!IsFinite(&xhi)) {
      z.hi = xhi;
      z.lo = 0;
      return;
   }

   NTL_ZZRegister(t);

   conv(t, xhi);
   sub(t, a, t);

   conv(xlo, t);

   normalize(z, xhi, xlo);

   // The following is just paranoia.
   if (fabs(z.hi) < NTL_FDOUBLE_PRECISION && z.lo != 0)
      LogicError("internal error: ZZ to quad_float conversion");
}

void conv(ZZ& z, const quad_float& x)
{
   NTL_ZZRegister(t1);
   NTL_ZZRegister(t2);
   NTL_ZZRegister(t3);

   double fhi, flo;

   fhi = floor(x.hi);

   if (fhi == x.hi) {
      flo = floor(x.lo);

      conv(t1, fhi);
      conv(t2, flo);

      add(z, t1, t2);
   }
   else
      conv(z, fhi);
}

ostream& operator<<(ostream& s, const quad_float& a)
{
   quad_float aa = a;

   if (!IsFinite(&aa)) {
      s << "NaN";
      return s;
   }

   RRPush push;
   RROutputPush opush;

   RR::SetPrecision(long(3.33*quad_float::oprec) + 10);
   RR::SetOutputPrecision(quad_float::oprec);

   NTL_TLS_LOCAL(RR, t);

   conv(t, a);
   s << t;

   return s;
}

istream& operator>>(istream& s, quad_float& x)
{
   RRPush push;
   RR::SetPrecision(4*NTL_DOUBLE_PRECISION);

   NTL_TLS_LOCAL(RR, t);
   NTL_INPUT_CHECK_RET(s, s >> t);
   conv(x, t);

   return s;
}

void random(quad_float& x)
{
   RRPush push;
   RR::SetPrecision(4*NTL_DOUBLE_PRECISION);

   NTL_TLS_LOCAL(RR, t);
   random(t);
   conv(x, t);
}

quad_float random_quad_float()
{
   quad_float x;
   random(x);
   return x;
}

long IsFinite(quad_float *x)
{
   return IsFinite(&x->hi) && IsFinite(&x->lo);
}

long PrecisionOK()
{
START_FIX
   long k;
   DOUBLE l1 = (double)1;
   DOUBLE lh = 1/(double)2;
   DOUBLE epsilon;
   DOUBLE fudge, oldfudge;

   epsilon = l1;
   fudge = l1+l1;

   k = 0;

   do {
      k++;
      epsilon = epsilon * lh;
      oldfudge = fudge;
      fudge = l1 + epsilon;
   } while (fudge > l1 && fudge < oldfudge);

END_FIX
   return k == NTL_DOUBLE_PRECISION;
}

quad_float floor(const quad_float& x)
{
   double fhi = floor(x.hi);

   if (fhi != x.hi)
      return quad_float(fhi, 0.0);
   else {
      double flo = floor(x.lo);
      quad_float z;
      normalize(z, fhi, flo);
      return z;
   }
}

quad_float ceil(const quad_float& x) {
  return -floor(-x);
}

quad_float trunc(const quad_float& x) {
  if (x>=0.0) return floor(x); else return -floor(-x);
}

long compare(const quad_float& x, const quad_float& y)
{
   if (x.hi > y.hi)
      return 1;
   else if (x.hi < y.hi)
      return -1;
   else if (x.lo > y.lo)
      return 1;
   else if (x.lo < y.lo)
      return -1;
   else
      return 0;
}

quad_float fabs(const quad_float& x)
{ if (x.hi>=0.0) return x; else return -x; }

quad_float ldexp(const quad_float& x, long exp) { // x*2^exp
   double xhi, xlo;
   quad_float z;

   xhi = _ntl_ldexp(x.hi, exp);
   xlo = _ntl_ldexp(x.lo, exp);

   normalize(z, xhi, xlo);
   return z;
}

quad_float exp(const quad_float& x) { // New version 97 Aug 05
/*
!  Calculate a quadruple-precision exponential
!  Method:
!   x    x.log2(e)    nint[x.log2(e)] + frac[x.log2(e)]
!  e  = 2          = 2
!
!                     iy    fy
!                  = 2   . 2
!  Then
!   fy    y.loge(2)
!  2   = e
!
!  Now y.loge(2) will be less than 0.3466 in absolute value.
!  This is halved and a Pade aproximation is used to approximate e^x over
!  the region (-0.1733, +0.1733).   This approximation is then squared.
*/
  if (x.hi<DBL_MIN_10_EXP*2.302585092994045684017991)
    return to_quad_float(0.0);
  if (x.hi>DBL_MAX_10_EXP*2.302585092994045684017991) {
    ResourceError("exp(quad_float): overflow");
  }

  static const quad_float Log2 = to_quad_float("0.6931471805599453094172321214581765680755");
  // GLOBAL (assumes C++11 thread-safe init)

  quad_float y,temp,ysq,sum1,sum2;
  long iy;
  y=x/Log2;
  temp = floor(y+0.5);
  iy = to_long(temp);
  y=(y-temp)*Log2;
  y=ldexp(y,-1L);
  ysq=y*y;
  sum1=y*((((ysq+3960.0)*ysq+2162160.0)*ysq+302702400.0)*ysq+8821612800.0);
  sum2=(((90.0*ysq+110880.0)*ysq+30270240.0)*ysq+2075673600.0)*ysq+17643225600.0;
/*
!                     sum2 + sum1         2.sum1
! Now approximation = ----------- = 1 + ----------- = 1 + 2.temp
!                     sum2 - sum1       sum2 - sum1
!
! Then (1 + 2.temp)^2 = 4.temp.(1 + temp) + 1
*/
  temp=sum1/(sum2-sum1);
  y=temp*(temp+1);
  y=ldexp(y,2L);
  return ldexp(y+1,iy);
}

quad_float log(const quad_float& t) { // Newton method. See Bailey, MPFUN
  if (t.hi <= 0.0) {
    ArithmeticError("log(quad_float): argument must be positive");
  }
  double s1 = log(t.hi);
  ForceToMem(&s1);  // Again, this is fairly paranoid.
  quad_float s;
  s = s1;
  quad_float e;
  e=exp(s);
  return s+(t-e)/e;  // Newton step
}

long operator> (const quad_float& x, const quad_float& y) {
   return (x.hi> y.hi) || (x.hi==y.hi && x.lo> y.lo); }
long operator>=(const quad_float& x, const quad_float& y) {
   return (x.hi>y.hi) || (x.hi==y.hi && x.lo>=y.lo); }
long operator< (const quad_float& x, const quad_float& y) {
   return (x.hi< y.hi) || (x.hi==y.hi && x.lo< y.lo); }
long operator<=(const quad_float& x, const quad_float& y) {
   return (x.hi<y.hi) || (x.hi==y.hi && x.lo<=y.lo); }
long operator==(const quad_float& x, const quad_float& y)
   { return x.hi==y.hi && x.lo==y.lo; }
long operator!=(const quad_float& x, const quad_float& y)
   { return x.hi!=y.hi || x.lo!=y.lo; }

NTL_END_IMPL


// Copyright (c) by respective owners including Yahoo!, Microsoft, and
// individual contributors. All rights reserved. Released under a BSD (revised)
// license as described in the file LICENSE.
#ifdef _WIN32
#  pragma warning(disable : 4996)  // generated by inner_product use
#endif
#include <fstream>
#include <vector>
#include <queue>
#include <algorithm>
#include <numeric>
#include <cmath>
#include "correctedMath.h"
#include "vw_versions.h"
#include "vw.h"
#include "mwt.h"

#include <cstring>
#include <cstdio>
#include <cassert>
#include "no_label.h"
#include "gd.h"
#include "rand48.h"
#include "reductions.h"
#include "array_parameters.h"
#include "vw_exception.h"

#include "io/logger.h"
#include "shared_data.h"

#include <boost/version.hpp>
#include <boost/math/special_functions/digamma.hpp>
#include <boost/math/special_functions/gamma.hpp>

#if BOOST_VERSION >= 105600
#  include <boost/align/is_aligned.hpp>
#endif

using namespace VW::config;

namespace logger = VW::io::logger;

enum lda_math_mode
{
  USE_SIMD,
  USE_PRECISE,
  USE_FAST_APPROX
};

class index_feature
{
public:
  uint32_t document;
  feature f;
  bool operator<(const index_feature b) const { return f.weight_index < b.f.weight_index; }
};

struct lda
{
  size_t topics;
  float lda_alpha;
  float lda_rho;
  float lda_D;
  float lda_epsilon;
  size_t minibatch;
  lda_math_mode mmode;

  size_t finish_example_count;

  v_array<float> Elogtheta;
  v_array<float> decay_levels;
  v_array<float> total_new;
  v_array<example *> examples;
  v_array<float> total_lambda;
  v_array<int> doc_lengths;
  v_array<float> digammas;
  v_array<float> v;
  std::vector<index_feature> sorted_features;

  bool compute_coherence_metrics;

  // size by 1 << bits
  std::vector<uint32_t> feature_counts;
  std::vector<std::vector<size_t>> feature_to_example_map;

  bool total_lambda_init;

  double example_t;
  vw *all;  // regressor, lda

  static constexpr float underflow_threshold = 1.0e-10f;
  inline float digamma(float x);
  inline float lgamma(float x);
  inline float powf(float x, float p);
  inline void expdigammify(vw &all, float *gamma);
  inline void expdigammify_2(vw &all, float *gamma, float *norm);
};

// #define VW_NO_INLINE_SIMD

namespace
{
inline bool is_aligned16(void *ptr)
{
#if BOOST_VERSION >= 105600
  return boost::alignment::is_aligned(16, ptr);
#else
  return ((reinterpret_cast<uintptr_t>(ptr) & 0x0f) == 0);
#endif
}
}  // namespace

namespace ldamath
{
inline float fastlog2(float x)
{
  uint32_t mx;
  memcpy(&mx, &x, sizeof(uint32_t));
  mx = (mx & 0x007FFFFF) | (0x7e << 23);

  float mx_f;
  memcpy(&mx_f, &mx, sizeof(float));

  uint32_t vx;
  memcpy(&vx, &x, sizeof(uint32_t));

  float y = static_cast<float>(vx);
  y *= 1.0f / (float)(1 << 23);

  return y - 124.22544637f - 1.498030302f * mx_f - 1.72587999f / (0.3520887068f + mx_f);
}

inline float fastlog(float x) { return 0.69314718f * fastlog2(x); }

inline float fastpow2(float p)
{
  float offset = (p < 0) * 1.0f;
  float clipp = (p < -126.0) ? -126.0f : p;
  int w = (int)clipp;
  float z = clipp - w + offset;
  uint32_t approx = (uint32_t)((1 << 23) * (clipp + 121.2740838f + 27.7280233f / (4.84252568f - z) - 1.49012907f * z));

  float v;
  memcpy(&v, &approx, sizeof(uint32_t));
  return v;
}

inline float fastexp(float p) { return fastpow2(1.442695040f * p); }

inline float fastpow(float x, float p) { return fastpow2(p * fastlog2(x)); }

inline float fastlgamma(float x)
{
  float logterm = fastlog(x * (1.0f + x) * (2.0f + x));
  float xp3 = 3.0f + x;

  return -2.081061466f - x + 0.0833333f / xp3 - logterm + (2.5f + x) * fastlog(xp3);
}

inline float fastdigamma(float x)
{
  float twopx = 2.0f + x;
  float logterm = fastlog(twopx);

  return -(1.0f + 2.0f * x) / (x * (1.0f + x)) - (13.0f + 6.0f * x) / (12.0f * twopx * twopx) + logterm;
}

#if !defined(VW_NO_INLINE_SIMD)

#  if defined(__SSE2__) || defined(__SSE3__) || defined(__SSE4_1__)

// Include headers for the various SSE versions:
#    if defined(__SSE2__)
#      include <emmintrin.h>
#    endif
#    if defined(__SSE3__)
#      include <tmmintrin.h>
#    endif
#    if defined(__SSE4_1__)
#      include <smmintrin.h>
#    endif

#    define HAVE_SIMD_MATHMODE

typedef __m128 v4sf;
typedef __m128i v4si;

inline v4sf v4si_to_v4sf(v4si x) { return _mm_cvtepi32_ps(x); }

inline v4si v4sf_to_v4si(v4sf x) { return _mm_cvttps_epi32(x); }

// Extract v[idx]
template <const int idx>
float v4sf_index(const v4sf x)
{
#    if defined(__SSE4_1__)
  float ret;
  uint32_t val;

  val = _mm_extract_ps(x, idx);
  // Portably convert uint32_t bit pattern to float. Optimizers will generally
  // make this disappear.
  memcpy(&ret, &val, sizeof(uint32_t));
  return ret;
#    else
  return _mm_cvtss_f32(_mm_shuffle_ps(x, x, _MM_SHUFFLE(idx, idx, idx, idx)));
#    endif
}

// Specialization for the 0'th element
template <>
float v4sf_index<0>(const v4sf x)
{
  return _mm_cvtss_f32(x);
}

inline v4sf v4sfl(const float x) { return _mm_set1_ps(x); }

inline v4si v4sil(const uint32_t x) { return _mm_set1_epi32(x); }

#    ifdef _WIN32

inline __m128 operator+(const __m128 a, const __m128 b) { return _mm_add_ps(a, b); }

inline __m128 operator-(const __m128 a, const __m128 b) { return _mm_sub_ps(a, b); }

inline __m128 operator*(const __m128 a, const __m128 b) { return _mm_mul_ps(a, b); }

inline __m128 operator/(const __m128 a, const __m128 b) { return _mm_div_ps(a, b); }

#    endif

inline v4sf vfastpow2(const v4sf p)
{
  v4sf ltzero = _mm_cmplt_ps(p, v4sfl(0.0f));
  v4sf offset = _mm_and_ps(ltzero, v4sfl(1.0f));
  v4sf lt126 = _mm_cmplt_ps(p, v4sfl(-126.0f));
  v4sf clipp = _mm_andnot_ps(lt126, p) + _mm_and_ps(lt126, v4sfl(-126.0f));
  v4si w = v4sf_to_v4si(clipp);
  v4sf z = clipp - v4si_to_v4sf(w) + offset;

  const v4sf c_121_2740838 = v4sfl(121.2740838f);
  const v4sf c_27_7280233 = v4sfl(27.7280233f);
  const v4sf c_4_84252568 = v4sfl(4.84252568f);
  const v4sf c_1_49012907 = v4sfl(1.49012907f);

  v4sf v = v4sfl(1 << 23) * (clipp + c_121_2740838 + c_27_7280233 / (c_4_84252568 - z) - c_1_49012907 * z);

  return _mm_castsi128_ps(v4sf_to_v4si(v));
}

inline v4sf vfastexp(const v4sf p)
{
  const v4sf c_invlog_2 = v4sfl(1.442695040f);

  return vfastpow2(c_invlog_2 * p);
}

inline v4sf vfastlog2(v4sf x)
{
  v4si vx_i = _mm_castps_si128(x);
  v4sf mx_f = _mm_castsi128_ps(_mm_or_si128(_mm_and_si128(vx_i, v4sil(0x007FFFFF)), v4sil(0x3f000000)));
  v4sf y = v4si_to_v4sf(vx_i) * v4sfl(1.1920928955078125e-7f);

  const v4sf c_124_22551499 = v4sfl(124.22551499f);
  const v4sf c_1_498030302 = v4sfl(1.498030302f);
  const v4sf c_1_725877999 = v4sfl(1.72587999f);
  const v4sf c_0_3520087068 = v4sfl(0.3520887068f);

  return y - c_124_22551499 - c_1_498030302 * mx_f - c_1_725877999 / (c_0_3520087068 + mx_f);
}

inline v4sf vfastlog(v4sf x)
{
  const v4sf c_0_69314718 = v4sfl(0.69314718f);

  return c_0_69314718 * vfastlog2(x);
}

inline v4sf vfastdigamma(v4sf x)
{
  v4sf twopx = v4sfl(2.0f) + x;
  v4sf logterm = vfastlog(twopx);

  return (v4sfl(-48.0f) + x * (v4sfl(-157.0f) + x * (v4sfl(-127.0f) - v4sfl(30.0f) * x))) /
      (v4sfl(12.0f) * x * (v4sfl(1.0f) + x) * twopx * twopx) +
      logterm;
}

void vexpdigammify(vw &all, float *gamma, const float underflow_threshold)
{
  float extra_sum = 0.0f;
  v4sf sum = v4sfl(0.0f);
  float *fp;
  const float *fpend = gamma + all.lda;

  // Iterate through the initial part of the array that isn't 128-bit SIMD
  // aligned.
  for (fp = gamma; fp < fpend && !is_aligned16(fp); ++fp)
  {
    extra_sum += *fp;
    *fp = fastdigamma(*fp);
  }

  // Rip through the aligned portion...
  for (; is_aligned16(fp) && fp + 4 < fpend; fp += 4)
  {
    v4sf arg = _mm_load_ps(fp);
    sum = sum + arg;
    arg = vfastdigamma(arg);
    _mm_store_ps(fp, arg);
  }

  for (; fp < fpend; ++fp)
  {
    extra_sum += *fp;
    *fp = fastdigamma(*fp);
  }

#    if defined(__SSE3__) || defined(__SSE4_1__)
  // Do two horizontal adds on sum, extract the total from the 0 element:
  sum = _mm_hadd_ps(sum, sum);
  sum = _mm_hadd_ps(sum, sum);
  extra_sum += v4sf_index<0>(sum);
#    else
  extra_sum += v4sf_index<0>(sum) + v4sf_index<1>(sum) + v4sf_index<2>(sum) + v4sf_index<3>(sum);
#    endif

  extra_sum = fastdigamma(extra_sum);
  sum = v4sfl(extra_sum);

  for (fp = gamma; fp < fpend && !is_aligned16(fp); ++fp) { *fp = fmax(underflow_threshold, fastexp(*fp - extra_sum)); }

  for (; is_aligned16(fp) && fp + 4 < fpend; fp += 4)
  {
    v4sf arg = _mm_load_ps(fp);
    arg = arg - sum;
    arg = vfastexp(arg);
    arg = _mm_max_ps(v4sfl(underflow_threshold), arg);
    _mm_store_ps(fp, arg);
  }

  for (; fp < fpend; ++fp) { *fp = fmax(underflow_threshold, fastexp(*fp - extra_sum)); }
}

void vexpdigammify_2(vw &all, float *gamma, const float *norm, const float underflow_threshold)
{
  float *fp = gamma;
  const float *np;
  const float *fpend = gamma + all.lda;

  for (np = norm; fp < fpend && !is_aligned16(fp); ++fp, ++np)
    *fp = fmax(underflow_threshold, fastexp(fastdigamma(*fp) - *np));

  for (; is_aligned16(fp) && fp + 4 < fpend; fp += 4, np += 4)
  {
    v4sf arg = _mm_load_ps(fp);
    arg = vfastdigamma(arg);
    v4sf vnorm = _mm_loadu_ps(np);
    arg = arg - vnorm;
    arg = vfastexp(arg);
    arg = _mm_max_ps(v4sfl(underflow_threshold), arg);
    _mm_store_ps(fp, arg);
  }

  for (; fp < fpend; ++fp, ++np) *fp = fmax(underflow_threshold, fastexp(fastdigamma(*fp) - *np));
}

#  else
// PLACEHOLDER for future ARM NEON code
// Also remember to define HAVE_SIMD_MATHMODE
#  endif

#endif  // !VW_NO_INLINE_SIMD

// Templates for common code shared between the three math modes (SIMD, fast approximations
// and accurate).
//
// The generic template takes a type and a specialization flag, mtype.
//
// mtype == USE_PRECISE: Use the accurate computation for lgamma, digamma.
// mtype == USE_FAST_APPROX: Use the fast approximations for lgamma, digamma.
// mtype == USE_SIMD: Use CPU SIMD instruction
//
// The generic template is specialized for the particular accuracy setting.

// Log gamma:
template <typename T, const lda_math_mode mtype>
inline T lgamma(T /* x */)
{
  BOOST_STATIC_ASSERT_MSG(true, "ldamath::lgamma is not defined for this type and math mode.");
}

// Digamma:
template <typename T, const lda_math_mode mtype>
inline T digamma(T /* x */)
{
  BOOST_STATIC_ASSERT_MSG(true, "ldamath::digamma is not defined for this type and math mode.");
}

// Exponential
template <typename T, lda_math_mode mtype>
inline T exponential(T /* x */)
{
  BOOST_STATIC_ASSERT_MSG(true, "ldamath::exponential is not defined for this type and math mode.");
}

// Powf
template <typename T, lda_math_mode mtype>
inline T powf(T /* x */, T /* p */)
{
  BOOST_STATIC_ASSERT_MSG(true, "ldamath::powf is not defined for this type and math mode.");
}

// High accuracy float specializations:

template <>
inline float lgamma<float, USE_PRECISE>(float x)
{
  return boost::math::lgamma(x);
}
template <>
inline float digamma<float, USE_PRECISE>(float x)
{
  return boost::math::digamma(x);
}
template <>
inline float exponential<float, USE_PRECISE>(float x)
{
  return correctedExp(x);
}
template <>
inline float powf<float, USE_PRECISE>(float x, float p)
{
  return std::pow(x, p);
}

// Fast approximation float specializations:

template <>
inline float lgamma<float, USE_FAST_APPROX>(float x)
{
  return fastlgamma(x);
}
template <>
inline float digamma<float, USE_FAST_APPROX>(float x)
{
  return fastdigamma(x);
}
template <>
inline float exponential<float, USE_FAST_APPROX>(float x)
{
  return fastexp(x);
}
template <>
inline float powf<float, USE_FAST_APPROX>(float x, float p)
{
  return fastpow(x, p);
}

// SIMD specializations:

template <>
inline float lgamma<float, USE_SIMD>(float x)
{
  return lgamma<float, USE_FAST_APPROX>(x);
}
template <>
inline float digamma<float, USE_SIMD>(float x)
{
  return digamma<float, USE_FAST_APPROX>(x);
}
template <>
inline float exponential<float, USE_SIMD>(float x)
{
  return exponential<float, USE_FAST_APPROX>(x);
}
template <>
inline float powf<float, USE_SIMD>(float x, float p)
{
  return powf<float, USE_FAST_APPROX>(x, p);
}

template <typename T, const lda_math_mode mtype>
inline void expdigammify(vw &all, T *gamma, T threshold, T initial)
{
  T sum = digamma<T, mtype>(std::accumulate(gamma, gamma + all.lda, initial));

  std::transform(gamma, gamma + all.lda, gamma,
      [sum, threshold](T g) { return fmax(threshold, exponential<T, mtype>(digamma<T, mtype>(g) - sum)); });
}
template <>
inline void expdigammify<float, USE_SIMD>(vw &all, float *gamma, float threshold, float)
{
#if defined(HAVE_SIMD_MATHMODE)
  vexpdigammify(all, gamma, threshold);
#else
  // Do something sensible if SIMD math isn't available:
  expdigammify<float, USE_FAST_APPROX>(all, gamma, threshold, 0.0);
#endif
}

template <typename T, const lda_math_mode mtype>
inline void expdigammify_2(vw &all, float *gamma, T *norm, const T threshold)
{
  std::transform(gamma, gamma + all.lda, norm, gamma,
      [threshold](float g, float n) { return fmax(threshold, exponential<T, mtype>(digamma<T, mtype>(g) - n)); });
}
template <>
inline void expdigammify_2<float, USE_SIMD>(vw &all, float *gamma, float *norm, const float threshold)
{
#if defined(HAVE_SIMD_MATHMODE)
  vexpdigammify_2(all, gamma, norm, threshold);
#else
  // Do something sensible if SIMD math isn't available:
  expdigammify_2<float, USE_FAST_APPROX>(all, gamma, norm, threshold);
#endif
}

}  // namespace ldamath

float lda::digamma(float x)
{
  switch (mmode)
  {
    case USE_FAST_APPROX:
      return ldamath::digamma<float, USE_FAST_APPROX>(x);
    case USE_PRECISE:
      return ldamath::digamma<float, USE_PRECISE>(x);
    case USE_SIMD:
      return ldamath::digamma<float, USE_SIMD>(x);
    default:
      // Should not happen.
      logger::errlog_critical("lda::digamma: Trampled or invalid math mode, aborting");
      abort();
      return 0.0f;
  }
}

float lda::lgamma(float x)
{
  switch (mmode)
  {
    case USE_FAST_APPROX:
      return ldamath::lgamma<float, USE_FAST_APPROX>(x);
    case USE_PRECISE:
      return ldamath::lgamma<float, USE_PRECISE>(x);
    case USE_SIMD:
      return ldamath::lgamma<float, USE_SIMD>(x);
    default:
      logger::errlog_critical("lda::lgamma: Trampled or invalid math mode, aborting");
      abort();
      return 0.0f;
  }
}

float lda::powf(float x, float p)
{
  switch (mmode)
  {
    case USE_FAST_APPROX:
      return ldamath::powf<float, USE_FAST_APPROX>(x, p);
    case USE_PRECISE:
      return ldamath::powf<float, USE_PRECISE>(x, p);
    case USE_SIMD:
      return ldamath::powf<float, USE_SIMD>(x, p);
    default:
      logger::errlog_critical("lda::powf: Trampled or invalid math mode, aborting");
      abort();
      return 0.0f;
  }
}

void lda::expdigammify(vw &all_, float *gamma)
{
  switch (mmode)
  {
    case USE_FAST_APPROX:
      ldamath::expdigammify<float, USE_FAST_APPROX>(all_, gamma, underflow_threshold, 0.0f);
      break;
    case USE_PRECISE:
      ldamath::expdigammify<float, USE_PRECISE>(all_, gamma, underflow_threshold, 0.0f);
      break;
    case USE_SIMD:
      ldamath::expdigammify<float, USE_SIMD>(all_, gamma, underflow_threshold, 0.0f);
      break;
    default:
      logger::errlog_critical("lda::expdigammify: Trampled or invalid math mode, aborting");
      abort();
  }
}

void lda::expdigammify_2(vw &all_, float *gamma, float *norm)
{
  switch (mmode)
  {
    case USE_FAST_APPROX:
      ldamath::expdigammify_2<float, USE_FAST_APPROX>(all_, gamma, norm, underflow_threshold);
      break;
    case USE_PRECISE:
      ldamath::expdigammify_2<float, USE_PRECISE>(all_, gamma, norm, underflow_threshold);
      break;
    case USE_SIMD:
      ldamath::expdigammify_2<float, USE_SIMD>(all_, gamma, norm, underflow_threshold);
      break;
    default:
      logger::errlog_critical("lda::expdigammify_2: Trampled or invalid math mode, aborting");
      abort();
  }
}

static inline float average_diff(vw &all, float *oldgamma, float *newgamma)
{
  float sum;
  float normalizer;

  // This warps the normal sense of "inner product", but it accomplishes the same
  // thing as the "plain old" for loop. clang does a good job of reducing the
  // common subexpressions.
  sum = std::inner_product(
      oldgamma, oldgamma + all.lda, newgamma, 0.0f, [](float accum, float absdiff) { return accum + absdiff; },
      [](float old_g, float new_g) { return std::abs(old_g - new_g); });

  normalizer = std::accumulate(newgamma, newgamma + all.lda, 0.0f);
  return sum / normalizer;
}

// Returns E_q[log p(\theta)] - E_q[log q(\theta)].
float theta_kl(lda &l, v_array<float> &Elogtheta, float *gamma)
{
  float gammasum = 0;
  Elogtheta.clear();
  for (size_t k = 0; k < l.topics; k++)
  {
    Elogtheta.push_back(l.digamma(gamma[k]));
    gammasum += gamma[k];
  }
  float digammasum = l.digamma(gammasum);
  gammasum = l.lgamma(gammasum);
  float kl = -(l.topics * l.lgamma(l.lda_alpha));
  kl += l.lgamma(l.lda_alpha * l.topics) - gammasum;
  for (size_t k = 0; k < l.topics; k++)
  {
    Elogtheta[k] -= digammasum;
    kl += (l.lda_alpha - gamma[k]) * Elogtheta[k];
    kl += l.lgamma(gamma[k]);
  }

  return kl;
}

static inline float find_cw(lda &l, float *u_for_w, float *v)
{
  return 1.0f / std::inner_product(u_for_w, u_for_w + l.topics, v, 0.0f);
}

namespace
{
// Effectively, these are static and not visible outside the compilation unit.
v_array<float> new_gamma = v_init<float>();
v_array<float> old_gamma = v_init<float>();
}  // namespace

// Returns an estimate of the part of the variational bound that
// doesn't have to do with beta for the entire corpus for the current
// setting of lambda based on the document passed in. The value is
// divided by the total number of words in the document This can be
// used as a (possibly very noisy) estimate of held-out likelihood.
float lda_loop(lda &l, v_array<float> &Elogtheta, float *v, example *ec, float)
{
  parameters &weights = l.all->weights;
  new_gamma.clear();
  old_gamma.clear();

  for (size_t i = 0; i < l.topics; i++)
  {
    new_gamma.push_back(1.f);
    old_gamma.push_back(0.f);
  }
  size_t num_words = 0;
  for (features &fs : *ec) num_words += fs.size();

  float xc_w = 0;
  float score = 0;
  float doc_length = 0;
  do
  {
    memcpy(v, new_gamma.begin(), sizeof(float) * l.topics);
    l.expdigammify(*l.all, v);

    memcpy(old_gamma.begin(), new_gamma.begin(), sizeof(float) * l.topics);
    memset(new_gamma.begin(), 0, sizeof(float) * l.topics);

    score = 0;
    size_t word_count = 0;
    doc_length = 0;
    for (features &fs : *ec)
    {
      for (features::iterator &f : fs)
      {
        float *u_for_w = &(weights[f.index()]) + l.topics + 1;
        float c_w = find_cw(l, u_for_w, v);
        xc_w = c_w * f.value();
        score += -f.value() * log(c_w);
        size_t max_k = l.topics;
        for (size_t k = 0; k < max_k; k++, ++u_for_w) new_gamma[k] += xc_w * *u_for_w;
        word_count++;
        doc_length += f.value();
      }
    }
    for (size_t k = 0; k < l.topics; k++) new_gamma[k] = new_gamma[k] * v[k] + l.lda_alpha;
  } while (average_diff(*l.all, old_gamma.begin(), new_gamma.begin()) > l.lda_epsilon);

  ec->pred.scalars.clear();
  ec->pred.scalars.resize_but_with_stl_behavior(l.topics);
  memcpy(ec->pred.scalars.begin(), new_gamma.begin(), l.topics * sizeof(float));

  score += theta_kl(l, Elogtheta, new_gamma.begin());

  return score / doc_length;
}

size_t next_pow2(size_t x)
{
  int i = 0;
  x = x > 0 ? x - 1 : 0;
  while (x > 0)
  {
    x >>= 1;
    i++;
  }
  return ((size_t)1) << i;
}

struct initial_weights
{
  weight _initial;
  weight _initial_random;
  bool _random;
  uint32_t _lda;
  uint32_t _stride;
};

void save_load(lda &l, io_buf &model_file, bool read, bool text)
{
  vw &all = *(l.all);
  uint64_t length = (uint64_t)1 << all.num_bits;
  if (read)
  {
    initialize_regressor(all);
    initial_weights init{all.initial_t, static_cast<float>(l.lda_D / all.lda / all.length() * 200.f),
        all.random_weights, all.lda, all.weights.stride()};

    auto initial_lda_weight_initializer = [init](weight *weights, uint64_t index) {
      uint32_t lda = init._lda;
      weight initial_random = init._initial_random;
      if (init._random)
      {
        for (size_t i = 0; i != lda; ++i, ++index)
        { weights[i] = static_cast<float>(-std::log(merand48(index) + 1e-6) + 1.0f) * initial_random; }
      }
      weights[lda] = init._initial;
    };

    all.weights.set_default(initial_lda_weight_initializer);
  }
  if (model_file.num_files() != 0)
  {
    uint64_t i = 0;
    std::stringstream msg;
    size_t brw = 1;

    do
    {
      brw = 0;
      size_t K = all.lda;
      if (!read && text) msg << i << " ";

      if (!read || all.model_file_ver >= VERSION_FILE_WITH_HEADER_ID)
        brw += bin_text_read_write_fixed(model_file, (char *)&i, sizeof(i), "", read, msg, text);
      else
      {
        // support 32bit build models
        uint32_t j;
        brw += bin_text_read_write_fixed(model_file, (char *)&j, sizeof(j), "", read, msg, text);
        i = j;
      }

      if (brw != 0)
      {
        weight *w = &(all.weights.strided_index(i));
        for (uint64_t k = 0; k < K; k++)
        {
          weight *v = w + k;
          if (!read && text) msg << *v + l.lda_rho << " ";
          brw += bin_text_read_write_fixed(model_file, (char *)v, sizeof(*v), "", read, msg, text);
        }
      }
      if (text)
      {
        if (!read) msg << "\n";
        brw += bin_text_read_write_fixed(model_file, nullptr, 0, "", read, msg, text);
      }
      if (!read) ++i;
    } while ((!read && i < length) || (read && brw > 0));
  }
}

void return_example(vw &all, example &ec)
{
  all.sd->update(ec.test_only, true, ec.loss, ec.weight, ec.num_features);
  for (auto &sink : all.final_prediction_sink) { MWT::print_scalars(sink.get(), ec.pred.scalars, ec.tag); }

  if (all.sd->weighted_examples() >= all.sd->dump_interval && !all.logger.quiet)
    all.sd->print_update(*all.trace_message, all.holdout_set_off, all.current_pass, "none", 0, ec.num_features,
        all.progress_add, all.progress_arg);
  VW::finish_example(all, ec);
}

void learn_batch(lda &l)
{
  parameters &weights = l.all->weights;

  assert(l.finish_example_count == (l.examples.size() - 1));

  if (l.sorted_features.empty())  // FAST-PASS for real "true"
  {
    // This can happen when the socket connection is dropped by the client.
    // If l.sorted_features is empty, then l.sorted_features[0] does not
    // exist, so we should not try to take its address in the beginning of
    // the for loops down there. Since it seems that there's not much to
    // do in this case, we just return.
    for (size_t d = 0; d < l.examples.size(); d++)
    {
      l.examples[d]->pred.scalars.clear();
      l.examples[d]->pred.scalars.resize_but_with_stl_behavior(l.topics);
      memset(l.examples[d]->pred.scalars.begin(), 0, l.topics * sizeof(float));

      l.examples[d]->pred.scalars.clear();

      if (l.finish_example_count > 0)
      {
        return_example(*l.all, *l.examples[d]);
        l.finish_example_count--;
      }
    }
    l.examples.clear();
    return;
  }

  float eta = -1;
  float minuseta = -1;

  if (l.total_lambda.empty())
  {
    for (size_t k = 0; k < l.all->lda; k++) l.total_lambda.push_back(0.f);
    // This part does not work with sparse parameters
    size_t stride = weights.stride();
    for (size_t i = 0; i <= weights.mask(); i += stride)
    {
      weight *w = &(weights[i]);
      for (size_t k = 0; k < l.all->lda; k++) l.total_lambda[k] += w[k];
    }
  }

  l.example_t++;
  l.total_new.clear();
  for (size_t k = 0; k < l.all->lda; k++) l.total_new.push_back(0.f);

  size_t batch_size = l.examples.size();

  sort(l.sorted_features.begin(), l.sorted_features.end());

  eta = l.all->eta * l.powf((float)l.example_t, -l.all->power_t);
  minuseta = 1.0f - eta;
  eta *= l.lda_D / batch_size;
  l.decay_levels.push_back(l.decay_levels.back() + log(minuseta));

  l.digammas.clear();
  float additional = (float)(l.all->length()) * l.lda_rho;
  for (size_t i = 0; i < l.all->lda; i++) l.digammas.push_back(l.digamma(l.total_lambda[i] + additional));

  auto last_weight_index = std::numeric_limits<uint64_t>::max();
  for (index_feature *s = &l.sorted_features[0]; s <= &l.sorted_features.back(); s++)
  {
    if (last_weight_index == s->f.weight_index) continue;
    last_weight_index = s->f.weight_index;
    // float *weights_for_w = &(weights[s->f.weight_index]);
    float *weights_for_w = &(weights[s->f.weight_index & weights.mask()]);
    float decay_component =
        l.decay_levels.end()[-2] - l.decay_levels.end()[(int)(-1 - l.example_t + *(weights_for_w + l.all->lda))];
    float decay = fmin(1.0f, correctedExp(decay_component));
    float *u_for_w = weights_for_w + l.all->lda + 1;

    *(weights_for_w + l.all->lda) = (float)l.example_t;
    for (size_t k = 0; k < l.all->lda; k++)
    {
      weights_for_w[k] *= decay;
      u_for_w[k] = weights_for_w[k] + l.lda_rho;
    }

    l.expdigammify_2(*l.all, u_for_w, l.digammas.begin());
  }

  for (size_t d = 0; d < batch_size; d++)
  {
    float score = lda_loop(l, l.Elogtheta, &(l.v[d * l.all->lda]), l.examples[d], l.all->power_t);
    if (l.all->audit) GD::print_audit_features(*l.all, *l.examples[d]);
    // If the doc is empty, give it loss of 0.
    if (l.doc_lengths[d] > 0)
    {
      l.all->sd->sum_loss -= score;
      l.all->sd->sum_loss_since_last_dump -= score;
    }

    if (l.finish_example_count > 0)
    {
      return_example(*l.all, *l.examples[d]);
      l.finish_example_count--;
    }
  }

  // -t there's no need to update weights (especially since it's a noop)
  if (eta != 0)
  {
    for (index_feature *s = &l.sorted_features[0]; s <= &l.sorted_features.back();)
    {
      index_feature *next = s + 1;
      while (next <= &l.sorted_features.back() && next->f.weight_index == s->f.weight_index) next++;

      float *word_weights = &(weights[s->f.weight_index]);
      for (size_t k = 0; k < l.all->lda; k++, ++word_weights)
      {
        float new_value = minuseta * *word_weights;
        *word_weights = new_value;
      }

      for (; s != next; s++)
      {
        float *v_s = &(l.v[s->document * l.all->lda]);
        float *u_for_w = &(weights[s->f.weight_index]) + l.all->lda + 1;
        float c_w = eta * find_cw(l, u_for_w, v_s) * s->f.x;
        word_weights = &(weights[s->f.weight_index]);
        for (size_t k = 0; k < l.all->lda; k++, ++u_for_w, ++word_weights)
        {
          float new_value = *u_for_w * v_s[k] * c_w;
          l.total_new[k] += new_value;
          *word_weights += new_value;
        }
      }
    }

    for (size_t k = 0; k < l.all->lda; k++)
    {
      l.total_lambda[k] *= minuseta;
      l.total_lambda[k] += l.total_new[k];
    }
  }
  l.sorted_features.resize(0);

  l.examples.clear();
  l.doc_lengths.clear();
}

void learn(lda &l, VW::LEARNER::single_learner &, example &ec)
{
  uint32_t num_ex = (uint32_t)l.examples.size();
  l.examples.push_back(&ec);
  l.doc_lengths.push_back(0);
  for (features &fs : ec)
  {
    for (features::iterator &f : fs)
    {
      index_feature temp = {num_ex, feature(f.value(), f.index())};
      l.sorted_features.push_back(temp);
      l.doc_lengths[num_ex] += (int)f.value();
    }
  }
  if (++num_ex == l.minibatch) learn_batch(l);
}

void learn_with_metrics(lda &l, VW::LEARNER::single_learner &base, example &ec)
{
  if (l.all->passes_complete == 0)
  {
    // build feature to example map
    uint64_t stride_shift = l.all->weights.stride_shift();
    uint64_t weight_mask = l.all->weights.mask();

    for (features &fs : ec)
    {
      for (features::iterator &f : fs)
      {
        uint64_t idx = (f.index() & weight_mask) >> stride_shift;
        l.feature_counts[idx] += (uint32_t)f.value();
        l.feature_to_example_map[idx].push_back(ec.example_counter);
      }
    }
  }

  learn(l, base, ec);
}

// placeholder
void predict(lda &l, VW::LEARNER::single_learner &base, example &ec) { learn(l, base, ec); }
void predict_with_metrics(lda &l, VW::LEARNER::single_learner &base, example &ec) { learn_with_metrics(l, base, ec); }

struct word_doc_frequency
{
  // feature/word index
  uint64_t idx;
  // document count
  uint32_t count;
};

// cooccurence of 2 features/words
struct feature_pair
{
  // feature/word 1
  uint64_t f1;
  // feature/word 2
  uint64_t f2;

  feature_pair(uint64_t _f1, uint64_t _f2) : f1(_f1), f2(_f2) {}
};

template <class T>
void get_top_weights(vw *all, int top_words_count, int topic, std::vector<feature> &output, T &weights)
{
  uint64_t length = (uint64_t)1 << all->num_bits;

  // get top features for this topic
  auto cmp = [](feature left, feature right) { return left.x > right.x; };
  std::priority_queue<feature, std::vector<feature>, decltype(cmp)> top_features(cmp);
  typename T::iterator iter = weights.begin();

  for (uint64_t i = 0; i < std::min(static_cast<uint64_t>(top_words_count), length); i++, ++iter)
    top_features.push({(&(*iter))[topic], iter.index()});

  for (uint64_t i = top_words_count; i < length; i++, ++iter)
  {
    weight v = (&(*iter))[topic];
    if (v > top_features.top().x)
    {
      top_features.pop();
      top_features.push({v, i});
    }
  }

  // extract idx and sort descending
  output.resize(top_features.size());
  for (int i = (int)top_features.size() - 1; i >= 0; i--)
  {
    output[i] = top_features.top();
    top_features.pop();
  }
}

void get_top_weights(vw *all, int top_words_count, int topic, std::vector<feature> &output)
{
  if (all->weights.sparse)
    get_top_weights(all, top_words_count, topic, output, all->weights.sparse_weights);
  else
    get_top_weights(all, top_words_count, topic, output, all->weights.dense_weights);
}

template <class T>
void compute_coherence_metrics(lda &l, T &weights)
{
  uint64_t length = (uint64_t)1 << l.all->num_bits;

  std::vector<std::vector<feature_pair>> topics_word_pairs;
  topics_word_pairs.resize(l.topics);

  int top_words_count = 10;  // parameterize and check

  for (size_t topic = 0; topic < l.topics; topic++)
  {
    // get top features for this topic
    auto cmp = [](feature &left, feature &right) { return left.x > right.x; };
    std::priority_queue<feature, std::vector<feature>, decltype(cmp)> top_features(cmp);
    typename T::iterator iter = weights.begin();
    for (uint64_t i = 0; i < std::min(static_cast<uint64_t>(top_words_count), length); i++, ++iter)
      top_features.push(feature((&(*iter))[topic], iter.index()));

    for (typename T::iterator v = weights.begin(); v != weights.end(); ++v)
      if ((&(*v))[topic] > top_features.top().x)
      {
        top_features.pop();
        top_features.push(feature((&(*v))[topic], v.index()));
      }

    // extract idx and sort descending
    std::vector<uint64_t> top_features_idx;
    top_features_idx.resize(top_features.size());
    for (int i = (int)top_features.size() - 1; i >= 0; i--)
    {
      top_features_idx[i] = top_features.top().weight_index;
      top_features.pop();
    }

    auto &word_pairs = topics_word_pairs[topic];
    for (size_t i = 0; i < top_features_idx.size(); i++)
      for (size_t j = i + 1; j < top_features_idx.size(); j++)
        word_pairs.emplace_back(top_features_idx[i], top_features_idx[j]);
  }

  // compress word pairs and create record for storing frequency
  std::map<uint64_t, std::vector<word_doc_frequency>> coWordsDFSet;
  for (auto &vec : topics_word_pairs)
  {
    for (auto &wp : vec)
    {
      auto f1 = wp.f1;
      auto f2 = wp.f2;
      auto wdf = coWordsDFSet.find(f1);

      if (wdf != coWordsDFSet.end())
      {
        // http://stackoverflow.com/questions/5377434/does-stdmapiterator-return-a-copy-of-value-or-a-value-itself
        // if (wdf->second.find(f2) == wdf->second.end())

        if (std::find_if(wdf->second.begin(), wdf->second.end(),
                [&f2](const word_doc_frequency &v) { return v.idx == f2; }) != wdf->second.end())
        {
          wdf->second.push_back({f2, 0});
          // printf(" add %d %d\n", f1, f2);
        }
      }
      else
      {
        std::vector<word_doc_frequency> tmp_vec = {{f2, 0}};
        coWordsDFSet.insert(std::make_pair(f1, tmp_vec));
      }
    }
  }

  // this.GetWordPairsDocumentFrequency(coWordsDFSet);
  for (auto &pair : coWordsDFSet)
  {
    auto &examples_for_f1 = l.feature_to_example_map[pair.first];
    for (auto &wdf : pair.second)
    {
      auto &examples_for_f2 = l.feature_to_example_map[wdf.idx];

      // assumes examples_for_f1 and examples_for_f2 are orderd
      size_t i = 0;
      size_t j = 0;
      while (i < examples_for_f1.size() && j < examples_for_f2.size())
      {
        if (examples_for_f1[i] == examples_for_f2[j])
        {
          wdf.count++;
          i++;
          j++;
        }
        else if (examples_for_f2[j] < examples_for_f1[i])
          j++;
        else
          i++;
      }
    }
  }

  float epsilon = 1e-6f;  // TODO
  float avg_coherence = 0;
  for (size_t topic = 0; topic < l.topics; topic++)
  {
    float coherence = 0;

    for (auto &pairs : topics_word_pairs[topic])
    {
      auto f1 = pairs.f1;
      if (l.feature_counts[f1] == 0) continue;

      auto f2 = pairs.f2;
      auto &co_feature = coWordsDFSet[f1];
      auto co_feature_df = std::find_if(
          co_feature.begin(), co_feature.end(), [&f2](const word_doc_frequency &v) { return v.idx == f2; });

      if (co_feature_df != co_feature.end())
      {
        // printf("(%d:%d + eps)/(%d:%d)\n", f2, co_feature_df->count, f1, l.feature_counts[f1]);
        coherence += logf((co_feature_df->count + epsilon) / l.feature_counts[f1]);
      }
    }

    printf("Topic %3d coherence: %f\n", (int)topic, coherence);

    // TODO: expose per topic coherence

    // TODO: good vs. bad topics
    avg_coherence += coherence;
  }

  avg_coherence /= l.topics;

  printf("Avg topic coherence: %f\n", avg_coherence);
}

void compute_coherence_metrics(lda &l)
{
  if (l.all->weights.sparse)
    compute_coherence_metrics(l, l.all->weights.sparse_weights);
  else
    compute_coherence_metrics(l, l.all->weights.dense_weights);
}

void end_pass(lda &l)
{
  if (!l.examples.empty()) learn_batch(l);

  if (l.compute_coherence_metrics && l.all->passes_complete == l.all->numpasses)
  {
    compute_coherence_metrics(l);
    // FASTPASS return;
  }
}

template <class T>
void end_examples(lda &l, T &weights)
{
  for (typename T::iterator iter = weights.begin(); iter != weights.end(); ++iter)
  {
    float decay_component =
        l.decay_levels.back() - l.decay_levels.end()[(int)(-1 - l.example_t + (&(*iter))[l.all->lda])];
    float decay = fmin(1.f, correctedExp(decay_component));

    weight *wp = &(*iter);
    for (size_t i = 0; i < l.all->lda; ++i) wp[i] *= decay;
  }
}

void end_examples(lda &l)
{
  if (l.all->weights.sparse)
    end_examples(l, l.all->weights.sparse_weights);
  else
    end_examples(l, l.all->weights.dense_weights);
}

void finish_example(vw &all, lda &l, example &e)
{
  if (l.minibatch <= 1) { return return_example(all, e); }

  if (l.examples.size() > 0)
  {
    // if there's still examples to be queued, inc only to finish later
    l.finish_example_count++;
  }
  else
  {
    // return now since it has been processed (example size = 0)
    return_example(all, e);
  }

  assert(l.finish_example_count <= l.minibatch);
}

std::istream &operator>>(std::istream &in, lda_math_mode &mmode)
{
  std::string token;
  in >> token;
  if (token == "simd")
    mmode = USE_SIMD;
  else if (token == "accuracy" || token == "precise")
    mmode = USE_PRECISE;
  else if (token == "fast-approx" || token == "approx")
    mmode = USE_FAST_APPROX;
  else
    THROW_EX(VW::vw_unrecognised_option_exception, token);
  return in;
}

VW::LEARNER::base_learner *lda_setup(options_i &options, vw &all)
{
  auto ld = scoped_calloc_or_throw<lda>();
  option_group_definition new_options("Latent Dirichlet Allocation");
  int math_mode;
  new_options.add(make_option("lda", ld->topics).keep().necessary().help("Run lda with <int> topics"))
      .add(make_option("lda_alpha", ld->lda_alpha)
               .keep()
               .default_value(0.1f)
               .help("Prior on sparsity of per-document topic weights"))
      .add(make_option("lda_rho", ld->lda_rho)
               .keep()
               .default_value(0.1f)
               .help("Prior on sparsity of topic distributions"))
      .add(make_option("lda_D", ld->lda_D).default_value(10000.0f).help("Number of documents"))
      .add(make_option("lda_epsilon", ld->lda_epsilon).default_value(0.001f).help("Loop convergence threshold"))
      .add(make_option("minibatch", ld->minibatch).default_value(1).help("Minibatch size, for LDA"))
      .add(make_option("math-mode", math_mode).default_value(USE_SIMD).help("Math mode: simd, accuracy, fast-approx"))
      .add(make_option("metrics", ld->compute_coherence_metrics).help("Compute metrics"));

  if (!options.add_parse_and_check_necessary(new_options)) return nullptr;

  // Convert from int to corresponding enum value.
  ld->mmode = static_cast<lda_math_mode>(math_mode);

  ld->finish_example_count = 0;

  all.lda = (uint32_t)ld->topics;
  ld->sorted_features = std::vector<index_feature>();
  ld->total_lambda_init = false;
  ld->all = &all;
  ld->example_t = all.initial_t;
  if (ld->compute_coherence_metrics)
  {
    ld->feature_counts.resize((uint32_t)(UINT64_ONE << all.num_bits));
    ld->feature_to_example_map.resize((uint32_t)(UINT64_ONE << all.num_bits));
  }

  float temp = ceilf(logf((float)(all.lda * 2 + 1)) / logf(2.f));

  all.weights.stride_shift((size_t)temp);
  all.random_weights = true;
  all.add_constant = false;

  if (all.eta > 1.)
  {
    logger::errlog_warn("your learning rate is too high, setting it to 1");
    all.eta = std::min(all.eta, 1.f);
  }

  size_t minibatch2 = next_pow2(ld->minibatch);

  //should num_parse_threads be specifiable or just 1? tell user this? write in the wiki page you make?
  int num_parse_threads = 1;
  if (minibatch2 > all.example_parser->ring_size)
  {
    bool previous_strict_parse = all.example_parser->strict_parse;
    delete all.example_parser;
    all.example_parser = new parser{minibatch2, previous_strict_parse, num_parse_threads};
    all.example_parser->_shared_data = all.sd;
  }

  ld->v.resize_but_with_stl_behavior(all.lda * ld->minibatch);

  ld->decay_levels.push_back(0.f);

  all.example_parser->lbl_parser = no_label::no_label_parser;

  VW::LEARNER::learner<lda, example> &l = init_learner(ld, ld->compute_coherence_metrics ? learn_with_metrics : learn,
      ld->compute_coherence_metrics ? predict_with_metrics : predict, UINT64_ONE << all.weights.stride_shift(),
      prediction_type_t::scalars, all.get_setupfn_name(lda_setup), true);

  l.set_save_load(save_load);
  l.set_finish_example(finish_example);
  l.set_end_examples(end_examples);
  l.set_end_pass(end_pass);

  return make_base(l);
}


/*
 * This file belongs to the Galois project, a C++ library for exploiting
 * parallelism. The code is being released under the terms of the 3-Clause BSD
 * License (a copy is located in LICENSE.txt at the top-level directory).
 *
 * Copyright (C) 2018, The University of Texas at Austin. All rights reserved.
 * UNIVERSITY EXPRESSLY DISCLAIMS ANY AND ALL WARRANTIES CONCERNING THIS
 * SOFTWARE AND DOCUMENTATION, INCLUDING ANY WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR ANY PARTICULAR PURPOSE, NON-INFRINGEMENT AND WARRANTIES OF
 * PERFORMANCE, AND ANY WARRANTY THAT MIGHT OTHERWISE ARISE FROM COURSE OF
 * DEALING OR USAGE OF TRADE.  NO WARRANTY IS EITHER EXPRESS OR IMPLIED WITH
 * RESPECT TO THE USE OF THE SOFTWARE OR DOCUMENTATION. Under no circumstances
 * shall University be liable for incidental, special, indirect, direct or
 * consequential damages or loss of profits, interruption of business, or
 * related expenses which may arise from use of Software or Documentation,
 * including but not limited to those resulting from defects in Software and/or
 * Documentation, or loss or inaccuracy of data of any kind.
 */

/*  This file is part of libDAI - http://www.libdai.org/
 *
 *  Copyright (c) 2006-2011, The libDAI authors. All rights reserved.
 *
 *  Use of this source code is governed by a BSD-style license that can be found
 * in the LICENSE file.
 */

#include <algorithm>
#include <cmath>
#include <boost/dynamic_bitset.hpp>
#include <dai/regiongraph.h>
#include <dai/factorgraph.h>
#include <dai/clustergraph.h>

namespace dai {

using namespace std;

void RegionGraph::construct(
    const FactorGraph& fg, const std::vector<VarSet>& ors,
    const std::vector<Region>& irs,
    const std::vector<std::pair<size_t, size_t>>& edges) {
  // Copy factor graph structure
  FactorGraph::operator=(fg);

  // Copy inner regions
  _IRs = irs;

  // Construct outer regions (giving them counting number 1.0)
  _ORs.clear();
  _ORs.reserve(ors.size());
  diaforeach(const VarSet& alpha, ors)
      _ORs.push_back(FRegion(Factor(alpha, 1.0), 1.0));

  // For each factor, find an outer region that subsumes that factor.
  // Then, multiply the outer region with that factor.
  _fac2OR.clear();
  _fac2OR.reserve(nrFactors());
  for (size_t I = 0; I < nrFactors(); I++) {
    size_t alpha;
    for (alpha = 0; alpha < nrORs(); alpha++)
      if (OR(alpha).vars() >> factor(I).vars()) {
        _fac2OR.push_back(alpha);
        break;
      }
    DAI_ASSERT(alpha != nrORs());
  }
  recomputeORs();

  // Create bipartite graph
  _G.construct(nrORs(), nrIRs(), edges.begin(), edges.end());
}

void RegionGraph::constructCVM(const FactorGraph& fg,
                               const std::vector<VarSet>& cl, size_t verbose) {
  if (verbose)
    cerr << "constructCVM called (" << fg.nrVars() << " vars, "
         << fg.nrFactors() << " facs, " << cl.size() << " clusters)" << endl;

  // Retain only maximal clusters
  if (verbose)
    cerr << "  Constructing ClusterGraph" << endl;
  ClusterGraph cg(cl);
  if (verbose)
    cerr << "  Erasing non-maximal clusters" << endl;
  cg.eraseNonMaximal();

  // Create inner regions - first pass
  if (verbose)
    cerr << "  Creating inner regions (first pass)" << endl;
  set<VarSet> betas;
  for (size_t alpha = 0; alpha < cg.nrClusters(); alpha++)
    for (size_t alpha2 = alpha; (++alpha2) != cg.nrClusters();) {
      VarSet intersection = cg.cluster(alpha) & cg.cluster(alpha2);
      if (intersection.size() > 0)
        betas.insert(intersection);
    }

  // Create inner regions - subsequent passes
  if (verbose)
    cerr << "  Creating inner regions (next passes)" << endl;
  set<VarSet> new_betas;
  do {
    new_betas.clear();
    for (set<VarSet>::const_iterator gamma = betas.begin();
         gamma != betas.end(); gamma++)
      for (set<VarSet>::const_iterator gamma2 = gamma;
           (++gamma2) != betas.end();) {
        VarSet intersection = (*gamma) & (*gamma2);
        if ((intersection.size() > 0) && (betas.count(intersection) == 0))
          new_betas.insert(intersection);
      }
    betas.insert(new_betas.begin(), new_betas.end());
  } while (new_betas.size());

  // Create inner regions - final phase
  if (verbose)
    cerr << "  Creating inner regions (final phase)" << endl;
  vector<Region> irs;
  irs.reserve(betas.size());
  for (set<VarSet>::const_iterator beta = betas.begin(); beta != betas.end();
       beta++)
    irs.push_back(Region(*beta, 0.0));

  // Create edges
  if (verbose)
    cerr << "  Creating edges" << endl;
  vector<pair<size_t, size_t>> edges;
  for (size_t beta = 0; beta < irs.size(); beta++)
    for (size_t alpha = 0; alpha < cg.nrClusters(); alpha++)
      if (cg.cluster(alpha) >> irs[beta])
        edges.push_back(pair<size_t, size_t>(alpha, beta));

  // Construct region graph
  if (verbose)
    cerr << "  Constructing region graph" << endl;
  construct(fg, cg.clusters(), irs, edges);

  // Calculate counting numbers
  if (verbose)
    cerr << "  Calculating counting numbers" << endl;
  calcCVMCountingNumbers();

  if (verbose)
    cerr << "Done." << endl;
}

void RegionGraph::calcCVMCountingNumbers() {
  // Calculates counting numbers of inner regions based upon counting numbers of
  // outer regions

  vector<vector<size_t>> ancestors(nrIRs());
  boost::dynamic_bitset<> assigned(nrIRs());
  for (size_t beta = 0; beta < nrIRs(); beta++) {
    IR(beta).c() = 0.0;
    for (size_t beta2 = 0; beta2 < nrIRs(); beta2++)
      if ((beta2 != beta) && IR(beta2) >> IR(beta))
        ancestors[beta].push_back(beta2);
  }

  bool new_counting;
  do {
    new_counting = false;
    for (size_t beta = 0; beta < nrIRs(); beta++) {
      if (!assigned[beta]) {
        bool has_unassigned_ancestor = false;
        for (vector<size_t>::const_iterator beta2 = ancestors[beta].begin();
             (beta2 != ancestors[beta].end()) && !has_unassigned_ancestor;
             beta2++)
          if (!assigned[*beta2])
            has_unassigned_ancestor = true;
        if (!has_unassigned_ancestor) {
          Real c = 1.0;
          diaforeach(const Neighbor& alpha, nbIR(beta)) c -= OR(alpha).c();
          for (vector<size_t>::const_iterator beta2 = ancestors[beta].begin();
               beta2 != ancestors[beta].end(); beta2++)
            c -= IR(*beta2).c();
          IR(beta).c() = c;
          assigned.set(beta, true);
          new_counting = true;
        }
      }
    }
  } while (new_counting);
}

bool RegionGraph::checkCountingNumbers() const {
  // Checks whether the counting numbers satisfy the fundamental relation

  bool all_valid = true;
  for (vector<Var>::const_iterator n = vars().begin(); n != vars().end(); n++) {
    Real c_n = 0.0;
    for (size_t alpha = 0; alpha < nrORs(); alpha++)
      if (OR(alpha).vars().contains(*n))
        c_n += OR(alpha).c();
    for (size_t beta = 0; beta < nrIRs(); beta++)
      if (IR(beta).contains(*n))
        c_n += IR(beta).c();
    if (fabs(c_n - 1.0) > 1e-15) {
      all_valid = false;
      cerr << "WARNING: counting numbers do not satisfy relation for " << *n
           << "(c_n = " << c_n << ")." << endl;
    }
  }

  return all_valid;
}

void RegionGraph::recomputeORs() {
  for (size_t alpha = 0; alpha < nrORs(); alpha++)
    OR(alpha).fill(1.0);
  for (size_t I = 0; I < nrFactors(); I++)
    if (fac2OR(I) != -1U)
      OR(fac2OR(I)) *= factor(I);
}

void RegionGraph::recomputeORs(const VarSet& ns) {
  for (size_t alpha = 0; alpha < nrORs(); alpha++)
    if (OR(alpha).vars().intersects(ns))
      OR(alpha).fill(1.0);
  for (size_t I = 0; I < nrFactors(); I++)
    if (fac2OR(I) != -1U)
      if (OR(fac2OR(I)).vars().intersects(ns))
        OR(fac2OR(I)) *= factor(I);
}

void RegionGraph::recomputeOR(size_t I) {
  DAI_ASSERT(I < nrFactors());
  if (fac2OR(I) != -1U) {
    size_t alpha = fac2OR(I);
    OR(alpha).fill(1.0);
    for (size_t J = 0; J < nrFactors(); J++)
      if (fac2OR(J) == alpha)
        OR(alpha) *= factor(J);
  }
}

/// Send RegionGraph to output stream
ostream& operator<<(ostream& os, const RegionGraph& rg) {
  os << "digraph RegionGraph {" << endl;
  os << "node[shape=box];" << endl;
  for (size_t alpha = 0; alpha < rg.nrORs(); alpha++)
    os << "\ta" << alpha << " [label=\"a" << alpha << ": "
       << rg.OR(alpha).vars() << ", c=" << rg.OR(alpha).c() << "\"];" << endl;
  os << "node[shape=ellipse];" << endl;
  for (size_t beta = 0; beta < rg.nrIRs(); beta++)
    os << "\tb" << beta << " [label=\"b" << beta << ": " << (VarSet)rg.IR(beta)
       << ", c=" << rg.IR(beta).c() << "\"];" << endl;
  for (size_t alpha = 0; alpha < rg.nrORs(); alpha++)
    diaforeach(const Neighbor& beta, rg.nbOR(alpha)) os
        << "\ta" << alpha << " -> b" << beta << ";" << endl;
  os << "}" << endl;
  return os;
}

} // end of namespace dai