#include <crab/config.h>

#include "../program_options.hpp"
#include "../common.hpp"
#ifdef HAVE_LDD
#include <crab/domains/ldd/ldd.hpp>
#endif
#include <boost/optional.hpp>
#include <crab/domains/boxes.hpp>
using namespace std;
using namespace ikos;
using namespace crab::cfg_impl;
using namespace crab::domain_impl;
#ifdef HAVE_LDD
using namespace crab::domains::ldd;
#endif

#define RATIONALS

#ifdef RATIONALS
typedef q_number number_t;
#else
typedef z_number number_t;
#endif

typedef linear_constraint<number_t, varname_t> linear_constraint_t;
typedef linear_expression<number_t, varname_t> linear_expression_t;
typedef interval <number_t> interval_t;

#ifdef HAVE_LDD

LddManager* create_ldd_man(size_t num_vars) {
  DdManager* cudd = Cudd_Init(0, 0, CUDD_UNIQUE_SLOTS, 127, 0);
  #ifndef RATIONALS
  theory_t * theory = tvpi_create_boxz_theory(num_vars);
  #else
  theory_t * theory = tvpi_create_box_theory(num_vars);
  #endif

  LddManager* ldd = Ldd_Init(cudd, theory);

  const bool dvo = true;
  if(dvo) Cudd_AutodynEnable(cudd, CUDD_REORDER_GROUP_SIFT);
  return ldd;
}

void destroy_ldd_man(LddManager* ldd) {
  DdManager *cudd = NULL;
  theory_t *theory = NULL;
  if (ldd) {
    cudd = Ldd_GetCudd(ldd);
    theory = Ldd_GetTheory(ldd);
    Ldd_Quit(ldd);
  }

  if (theory) tvpi_destroy_theory(theory);
  if (cudd) Cudd_Quit(cudd);
}

void dump(LddNodePtr val,  LddManager *man) {
  DdManager *cudd = Ldd_GetCudd(man);
  FILE *fp = Cudd_ReadStdout(cudd);
  Cudd_SetStdout(cudd,stderr);
  if (val.get() == Ldd_GetTrue(man)) crab::outs() << "true\n";
  else if (val.get() == Ldd_GetFalse(man)) crab::outs() << "false\n";
  else Ldd_PrintMinterm(man,val.get());
  Cudd_SetStdout(cudd,fp);
}

LddNodePtr top(LddManager *man)
{
  return lddPtr(man, Ldd_GetTrue(man));
}

LddNodePtr bot(LddManager *man)
{
  return lddPtr(man, Ldd_GetFalse(man));
}

bool isBot(LddManager * man, LddNodePtr v)
{
  return &*v == Ldd_GetFalse(man);
}

bool isTop(LddManager * man, LddNodePtr v)
{
  return &*v == Ldd_GetTrue(man);
}

LddNodePtr approx(LddManager *man, LddNodePtr v)
{
  return lddPtr(man, Ldd_TermMinmaxApprox(man, &*v));
}

LddNodePtr meet(LddManager *man, LddNodePtr n1, LddNodePtr n2)
{
  return lddPtr(man, Ldd_And(man, &*n1, &*n2));
}

LddNodePtr convex_join(LddManager *man, LddNodePtr n1, LddNodePtr n2)
{
  return approx(man, lddPtr(man, Ldd_Or(man, &*n1, &*n2)));
}

LddNodePtr join(LddManager *man, LddNodePtr n1, LddNodePtr n2)
{
  return lddPtr(man, Ldd_Or(man, &*n1, &*n2));
}

bool isLeq(LddManager *man, LddNodePtr n1, LddNodePtr n2)
{
  return Ldd_TermLeq(man, &(*n1), &(*n2));
}

LddNodePtr widen(LddManager *man, LddNodePtr n1, LddNodePtr n2)
{
  return lddPtr(man, Ldd_IntervalWiden(man, &*n1, &*n2));
}

int getVarId(varname_t v) {
  return v.index();
}

/** return term for variable v, neg for negation of variable */
linterm_t termForVal(LddManager* man, varname_t v, bool neg = false)
{
  int varId = getVarId(v);
  int sgn = neg ? -1 : 1;
  linterm_t term = Ldd_GetTheory(man)->create_linterm_sparse_si(&varId, &sgn, 1);
  return term;
}

/** v := k, where k is a constant */
LddNodePtr assign(LddManager *man, LddNodePtr n, varname_t v, number_t k) {
  if (isBot(man, n)) return n;

  q_number q(k);
  linterm_t t = termForVal(man, v);
  constant_t kkk = (constant_t) tvpi_create_cst(q.get_mpq_t());
  LddNodePtr newVal =
      lddPtr(man, Ldd_TermReplace(man, &(*n), t, NULL, NULL, kkk, kkk));
  Ldd_GetTheory(man)->destroy_cst(kkk);

  assert(!isBot(man, newVal));
  return newVal;
}

/** v := k, where k is an interval */
LddNodePtr assign(LddManager* man, LddNodePtr n, varname_t v, interval_t ival) {

  if (isBot(man, n)) return n;

  constant_t kmin = NULL, kmax = NULL;

  if (boost::optional <number_t> l = ival.lb().number()) {
    q_number q(*l);
    kmin = (constant_t)tvpi_create_cst(q.get_mpq_t());
  }

  if (boost::optional <number_t> u = ival.ub().number()) {
    q_number q(*u);
    kmax = (constant_t)tvpi_create_cst(q.get_mpq_t());
  }

  linterm_t t = termForVal(man, v);
  LddNodePtr newVal =
      lddPtr(man, Ldd_TermReplace(man, &(*n), t, NULL, NULL, kmin, kmax));

  if (kmin) Ldd_GetTheory(man)->destroy_cst(kmin);
  if (kmax) Ldd_GetTheory(man)->destroy_cst(kmax);

  assert(!isBot(man, newVal));
  return newVal;
}

/** v := a * u + k, where a, k are constants and u variable */
LddNodePtr apply(LddManager* man,
                  const LddNodePtr n,
                  varname_t v, varname_t u,
                  number_t a, number_t k)
{
  if (isTop(man, n) || isBot(man, n)) return n;

  linterm_t t = termForVal(man, v);
  linterm_t r = termForVal(man, u);

  q_number qa(a);
  q_number qk(k);

  constant_t aaa = (constant_t) tvpi_create_cst (qa.get_mpq_t());
  constant_t kkk = (constant_t) tvpi_create_cst(qk.get_mpq_t());
  LddNodePtr newVal =
      lddPtr(man, Ldd_TermReplace(man, &(*n), t, r, aaa, kkk, kkk));
  Ldd_GetTheory(man)->destroy_cst(aaa);
  Ldd_GetTheory(man)->destroy_cst(kkk);

  assert(!isBot(man, newVal));
  return newVal;
}

/** v := u */
LddNodePtr assign(LddManager* man, const LddNodePtr n,
                   varname_t v, varname_t u)
{
  return apply(man, n, v, u, 1, 0);
}
#endif

int main(int argc, char** argv )
{
#ifdef HAVE_LDD

  bool stats_enabled = false;
  if (!crab_tests::parse_user_options(argc,argv,stats_enabled)) {
      return 0;
  }

  LddManager* man = create_ldd_man(3000);
  Ldd_SanityCheck(man);

  variable_factory_t vfac;
  varname_t x = vfac["x"];
  varname_t y = vfac["y"];

  {
    LddNodePtr s1  = top(man);
    LddNodePtr s11 = assign(man, s1, x, 5.5);
    LddNodePtr s111 = assign(man, s11, y, 8.3);

    LddNodePtr s2  = top(man);
    LddNodePtr s22 = assign(man, s2, x, 5.5);
    LddNodePtr s222 = assign(man, s22, y, 12.2);

    dump(s111, man);
    dump(s222, man);
    crab::outs() << "Join \n";
    LddNodePtr s3 = join(man, s111, s222);
    dump(s3, man);
    crab::outs() << "Widening \n";
    LddNodePtr s4 = widen(man, s111, s222);
    dump(s4, man);
  }

  #if 1
  {
    // To reproduce an old bug with boxes.
    // Assertion failed: (level < (unsigned) cuddI(unique,Cudd_Regular(E)->index)), function cuddUniqueInter, file /Users/E30338/Repos/crab/build_crab2/ldd/src/ldd/cudd-2.4.2/cudd/cuddTable.c, line 1143.

    typedef boxes_domain< number_t, varname_t > boxes_domain_t;
    typedef variable <number_t, varname_t> var_t;

    boxes_domain_t s1;
    s1 += linear_constraint_t(var_t(x) >= number_t(0));
    s1 += linear_constraint_t(var_t(x) <= number_t(1,4));
    s1 += linear_constraint_t(var_t(y) >= number_t(1,4));
    s1 += linear_constraint_t(var_t(y) <= number_t(1,2));
    boxes_domain_t s2;
    s2 += linear_constraint_t(var_t(x) == number_t(0));
    s2 += linear_constraint_t(var_t(y) == number_t(1,2));

    boxes_domain_t s3 = s2 | s1;
    ///////////////////
    boxes_domain_t s4;
    s4 += linear_constraint_t(var_t(x) >= number_t(1,4));
    s4 += linear_constraint_t(var_t(x) <= number_t(1));
    s4 += linear_constraint_t(var_t(y) >= number_t(1,8));
    s4 += linear_constraint_t(var_t(y) <= number_t(1,4));

    boxes_domain_t s5 = s4 | s2;
    //////
    crab::outs() << "Widening of \n" << s3;
    crab::outs() << "and\n" << s5 << " = \n";
    boxes_domain_t s6 = s3 || s5;
    crab::outs() << s6 << "\n";

  }
  #endif

  // Ldd_NodeSanityCheck(man, &(*s4));

  // FIXME: seg fault here
  // destroy_ldd_man(man);
#endif
  return 0;
}


/* Siconos is a program dedicated to modeling, simulation and control
 * of non smooth dynamical systems.
 *
 * Copyright 2021 INRIA.
 *
 * 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.
*/

#include <boost/numeric/ublas/matrix_proxy.hpp>

#include "determinant.hpp"
#include "SiconosVector.hpp"
#include "SimpleMatrix.hpp"
#include "BlockMatrixIterators.hpp"
#include "BlockMatrix.hpp"

#include "SiconosAlgebra.hpp"
#include "SiconosException.hpp"

using namespace Siconos;

//=======================
//       get norm
//=======================

double SimpleMatrix::normInf() const
{
  if(_num == DENSE)
    return norm_inf(*mat.Dense);
  else if(_num == TRIANGULAR)
    return norm_inf(*mat.Triang);
  else if(_num == SYMMETRIC)
    return norm_inf(*mat.Sym);
  else if(_num == SPARSE)
    return norm_inf(*mat.Sparse);
  else if(_num == SPARSE_COORDINATE)
    return norm_inf(*mat.SparseCoordinate);
  else if(_num == BANDED)
    return norm_inf(*mat.Banded);
  else if(_num == ZERO)
    return 0;
  else if(_num == IDENTITY)
    return 1;

  THROW_EXCEPTION("Matrix type not supported");
  return std::numeric_limits<double>::infinity();
}

void SimpleMatrix::normInfByColumn(SP::SiconosVector vIn) const
{
  if(_num == DENSE)
  {
    if(vIn->size() != size(1))
      THROW_EXCEPTION("the given vector does not have the right length");
    DenseVect tmpV = DenseVect(size(0));
    for(unsigned int i = 0; i < size(1); i++)
    {
      ublas::noalias(tmpV) = ublas::column(*mat.Dense, i);
      (*vIn)(i) = norm_inf(tmpV);
    }
  }
  else
    THROW_EXCEPTION("not implemented for data other than DenseMat");
}
//=======================
//       determinant
//=======================

double SimpleMatrix::det() const
{
  if(_num == DENSE)
    return determinant(*mat.Dense);
  else if(_num == TRIANGULAR)
    return determinant(*mat.Triang);
  else if(_num == SYMMETRIC)
    return determinant(*mat.Sym);
  else if(_num == SPARSE)
    return determinant(*mat.Sparse);
  else if(_num == SPARSE_COORDINATE)
    return determinant(*mat.Sparse);
  else if(_num == BANDED)
    return determinant(*mat.Banded);
  else if(_num == ZERO)
    return 0;
  else  if(_num == IDENTITY)
    return 1;
  THROW_EXCEPTION("Matrix type not supported");
  return std::numeric_limits<double>::infinity();
}

void SimpleMatrix::trans()
{
  switch(_num)
  {
  case DENSE:
    *mat.Dense = ublas::trans(*mat.Dense);
    break;
  case TRIANGULAR:
    THROW_EXCEPTION("failed, the matrix is triangular matrix and can not be transposed in place.");
    break;
  case SYMMETRIC:
    break;
  case SPARSE:
    *mat.Sparse = ublas::trans(*mat.Sparse);
    break;
  case SPARSE_COORDINATE:
    *mat.Sparse = ublas::trans(*mat.Sparse);
    break;
  case BANDED:
    *mat.Banded = ublas::trans(*mat.Banded);
    break;
  case Siconos::ZERO:
    break;
  case Siconos::IDENTITY:
    break;
  default:
    THROW_EXCEPTION("Matrix type not supported");
  }
  resetFactorizationFlags();
}

void SimpleMatrix::trans(const SiconosMatrix &m)
{
  if(m.isBlock())
    THROW_EXCEPTION("not yet implemented for m being a BlockMatrix.");

  if(&m == this)
    trans();
  else
  {
    Siconos::UBLAS_TYPE numM = m.num();
    switch(numM)
    {
    case DENSE:
      if(_num != DENSE)
        THROW_EXCEPTION("try to transpose a dense matrix into another type.");
      noalias(*mat.Dense) = ublas::trans(*m.dense());
      break;
    case TRIANGULAR:
      if(_num != DENSE)
        THROW_EXCEPTION("try to transpose a triangular matrix into a non-dense one.");
      noalias(*mat.Dense) = ublas::trans(*m.triang());
      break;
    case SYMMETRIC:
      *this = m;
      break;
    case SPARSE:
      if(_num == DENSE)
        noalias(*mat.Dense) = ublas::trans(*m.sparse());
      else if(_num == SPARSE)
        noalias(*mat.Sparse) = ublas::trans(*m.sparse());
      else if(_num == SPARSE_COORDINATE)
        noalias(*mat.SparseCoordinate) = ublas::trans(*m.sparse());
      else
        THROW_EXCEPTION("try to transpose a sparse matrix into a forbidden type (not dense nor sparse).");
      break;
    case SPARSE_COORDINATE:
      if(_num == DENSE)
        noalias(*mat.Dense) = ublas::trans(*m.sparseCoordinate());
      else if(_num == SPARSE)
        noalias(*mat.Sparse) = ublas::trans(*m.sparseCoordinate());
      else if(_num == SPARSE_COORDINATE)
        noalias(*mat.SparseCoordinate) = ublas::trans(*m.sparseCoordinate());
      else
        THROW_EXCEPTION("try to transpose a sparse coordinate matrix into a forbidden type (not dense nor sparse coordinate).");
      break;
    case BANDED:
      if(_num == DENSE)
        noalias(*mat.Dense) = ublas::trans(*m.banded());
      else if(_num == BANDED)
        noalias(*mat.Banded) = ublas::trans(*m.banded());
      else
        THROW_EXCEPTION("try to transpose a banded matrix into a forbidden type (not dense nor banded).");
      break;
    case ZERO:
      *this = m;
      break;
    case IDENTITY:
      *this = m;
      break;
    default:
      THROW_EXCEPTION("");
    }
    resetFactorizationFlags();
  }
}

/*
The following code inverts the matrix input using LU-decomposition with backsubstitution of unit vectors. Reference: Numerical Recipies in C, 2nd ed., by Press, Teukolsky, Vetterling & Flannery.

you can solve Ax=b using three lines of ublas code:

permutation_matrix<> piv;
lu_factorize(A, piv);
lu_substitute(A, piv, x);

*/
#ifndef INVERT_MATRIX_HPP
#define INVERT_MATRIX_HPP

// REMEMBER to update "lu.hpp" header includes from boost-CVS
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/vector_proxy.hpp>
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/triangular.hpp>
#include <boost/numeric/ublas/lu.hpp>
#include <boost/numeric/ublas/io.hpp>

/* Matrix inversion routine.
Uses lu_factorize and lu_substitute in uBLAS to invert a matrix */
template<class T, class U, class V>
bool InvertMatrix(const ublas::matrix<T, U, V>& input, ublas::matrix<T, U, V>& inverse)
{
  using namespace boost::numeric::ublas;
  typedef permutation_matrix<std::size_t> pmatrix;
// create a working copy of the input
  matrix<T, U, V> A(input);
// create a permutation matrix for the LU-factorization
  pmatrix pm(A.size1());

// perform LU-factorization
  int res = lu_factorize(A,pm);
  if(res != 0) return false;

// create identity matrix of "inverse"
  inverse.assign(ublas::identity_matrix<T>(A.size1()));

// backsubstitute to get the inverse
  lu_substitute(A, pm, inverse);

  return true;
}

#endif //INVERT_MATRIX_HPP

// Note FP: never used. Comment before removal ?
// void invertMatrix(const SimpleMatrix& input, SimpleMatrix& output)
// {
//   InvertMatrix(*input.dense(), *output.dense());
// }


#include <fstream>
#include <sstream>
#include <iostream>
#include <iomanip>
#include <fmt/format.h>
#include <string>
#include <iomanip>
#include <boost/variant.hpp>
#include <unordered_map>
#include <algorithm>
#include <range/v3/all.hpp>
#include "matrix.h"
namespace views = ranges::views;
struct Displayer
{
  FixedMatrix<bool, 6, 50> screen_;
  void matrix(int y, int x) noexcept
  {
    for (int i = 0; i < x; ++i) {
      ranges::fill(screen_.row(i).begin(), screen_.row(i).begin() + y, true);
    }
  }
  void rotateRow(int rownum, int n) noexcept
  {
    auto row = screen_.row(rownum);
    std::rotate(row.begin(), screen_.row(rownum).begin() + screen_.colCount() - n, row.end());
  }
  void rotateCol(int colnum, int n) noexcept
  {
    auto col = screen_.column(colnum);
    std::rotate(col.begin(), screen_.column(colnum).begin() + screen_.rowCount() - n, col.end());
  }
  void print() const
  {
    for (int k = 0; k < 10; ++k) {
      int cols = k * 5;
      for (int i = 0; i < 6; ++i) {
        for (int j = cols; j < 5 + cols; ++j) {
          if (screen_(i, j)) {
            fmt::print("#");
          } else {
            fmt::print(".");
          }
        }
        fmt::print("\n");
      }

      fmt::print("\n");
    }
  }
};

int main(int argc, char **argv)
{
  if (argc > 1) {
    typedef void (Displayer::*MemFunc)(int, int) noexcept;
    Displayer d;
    std::map<std::string, MemFunc> m;
    m["rect"] = &Displayer::matrix;
    m["row"] = &Displayer::rotateRow;
    m["column"] = &Displayer::rotateCol;
    std::ifstream ifs(argv[1]);
    std::string s;
    std::string command;
    int x, y;

    for (; ifs >> command >> x >> y;) {
      auto f = m[command];
      (d.*f)(x, y);
    }
    std::cout << ranges::count(d.screen_.array(), true);
    d.print();
  }
}


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

#ifndef BOOST_MATH_INTERPOLATORS_BARYCENTRIC_RATIONAL_DETAIL_HPP
#define BOOST_MATH_INTERPOLATORS_BARYCENTRIC_RATIONAL_DETAIL_HPP

#include <vector>
#include <utility> // for std::move
#include <algorithm> // for std::is_sorted
#include <boost/lexical_cast.hpp>
#include <boost/math/special_functions/fpclassify.hpp>
#include <boost/core/demangle.hpp>
#include <boost/assert.hpp>

namespace boost{ namespace math{ namespace detail{

template<class Real>
class barycentric_rational_imp
{
public:
    template <class InputIterator1, class InputIterator2>
    barycentric_rational_imp(InputIterator1 start_x, InputIterator1 end_x, InputIterator2 start_y, size_t approximation_order = 3);

    barycentric_rational_imp(std::vector<Real>&& x, std::vector<Real>&& y, size_t approximation_order = 3);

    Real operator()(Real x) const;

    Real prime(Real x) const;

    // The barycentric weights are not really that interesting; except to the unit tests!
    Real weight(size_t i) const { return m_w[i]; }

    std::vector<Real>&& return_x()
    {
        return std::move(m_x);
    }

    std::vector<Real>&& return_y()
    {
        return std::move(m_y);
    }

private:

    void calculate_weights(size_t approximation_order);

    std::vector<Real> m_x;
    std::vector<Real> m_y;
    std::vector<Real> m_w;
};

template <class Real>
template <class InputIterator1, class InputIterator2>
barycentric_rational_imp<Real>::barycentric_rational_imp(InputIterator1 start_x, InputIterator1 end_x, InputIterator2 start_y, size_t approximation_order)
{
    std::ptrdiff_t n = std::distance(start_x, end_x);

    if (approximation_order >= (std::size_t)n)
    {
        throw std::domain_error("Approximation order must be < data length.");
    }

    // Big sad memcpy.
    m_x.resize(n);
    m_y.resize(n);
    for(unsigned i = 0; start_x != end_x; ++start_x, ++start_y, ++i)
    {
        // But if we're going to do a memcpy, we can do some error checking which is inexpensive relative to the copy:
        if(boost::math::isnan(*start_x))
        {
            std::string msg = std::string("x[") + boost::lexical_cast<std::string>(i) + "] is a NAN";
            throw std::domain_error(msg);
        }

        if(boost::math::isnan(*start_y))
        {
           std::string msg = std::string("y[") + boost::lexical_cast<std::string>(i) + "] is a NAN";
           throw std::domain_error(msg);
        }

        m_x[i] = *start_x;
        m_y[i] = *start_y;
    }
    calculate_weights(approximation_order);
}

template <class Real>
barycentric_rational_imp<Real>::barycentric_rational_imp(std::vector<Real>&& x, std::vector<Real>&& y,size_t approximation_order) : m_x(std::move(x)), m_y(std::move(y))
{
    BOOST_ASSERT_MSG(m_x.size() == m_y.size(), "There must be the same number of abscissas and ordinates.");
    BOOST_ASSERT_MSG(approximation_order < m_x.size(), "Approximation order must be < data length.");
    BOOST_ASSERT_MSG(std::is_sorted(m_x.begin(), m_x.end()), "The abscissas must be listed in increasing order x[0] < x[1] < ... < x[n-1].");
    calculate_weights(approximation_order);
}

template<class Real>
void barycentric_rational_imp<Real>::calculate_weights(size_t approximation_order)
{
    using std::abs;
    int64_t n = m_x.size();
    m_w.resize(n, 0);
    for(int64_t k = 0; k < n; ++k)
    {
        int64_t i_min = (std::max)(k - (int64_t) approximation_order, (int64_t) 0);
        int64_t i_max = k;
        if (k >= n - (std::ptrdiff_t)approximation_order)
        {
            i_max = n - approximation_order - 1;
        }

        for(int64_t i = i_min; i <= i_max; ++i)
        {
            Real inv_product = 1;
            int64_t j_max = (std::min)(static_cast<int64_t>(i + approximation_order), static_cast<int64_t>(n - 1));
            for(int64_t j = i; j <= j_max; ++j)
            {
                if (j == k)
                {
                    continue;
                }

                Real diff = m_x[k] - m_x[j];
                using std::numeric_limits;
                if (abs(diff) < (numeric_limits<Real>::min)())
                {
                   std::string msg = std::string("Spacing between  x[")
                      + boost::lexical_cast<std::string>(k) + std::string("] and x[")
                      + boost::lexical_cast<std::string>(i) + std::string("] is ")
                      + boost::lexical_cast<std::string>(diff) + std::string(", which is smaller than the epsilon of ")
                      + boost::core::demangle(typeid(Real).name());
                    throw std::logic_error(msg);
                }
                inv_product *= diff;
            }
            if (i % 2 == 0)
            {
                m_w[k] += 1/inv_product;
            }
            else
            {
                m_w[k] -= 1/inv_product;
            }
        }
    }
}

template<class Real>
Real barycentric_rational_imp<Real>::operator()(Real x) const
{
    Real numerator = 0;
    Real denominator = 0;
    for(size_t i = 0; i < m_x.size(); ++i)
    {
        // Presumably we should see if the accuracy is improved by using ULP distance of say, 5 here, instead of testing for floating point equality.
        // However, it has been shown that if x approx x_i, but x != x_i, then inaccuracy in the numerator cancels the inaccuracy in the denominator,
        // and the result is fairly accurate. See: http://epubs.siam.org/doi/pdf/10.1137/S0036144502417715
        if (x == m_x[i])
        {
            return m_y[i];
        }
        Real t = m_w[i]/(x - m_x[i]);
        numerator += t*m_y[i];
        denominator += t;
    }
    return numerator/denominator;
}

/*
 * A formula for computing the derivative of the barycentric representation is given in
 * "Some New Aspects of Rational Interpolation", by Claus Schneider and Wilhelm Werner,
 * Mathematics of Computation, v47, number 175, 1986.
 * http://www.ams.org/journals/mcom/1986-47-175/S0025-5718-1986-0842136-8/S0025-5718-1986-0842136-8.pdf
 * and reviewed in
 * Recent developments in barycentric rational interpolation
 * Jean-Paul Berrut, Richard Baltensperger and Hans D. Mittelmann
 *
 * Is it possible to complete this in one pass through the data?
 */

template<class Real>
Real barycentric_rational_imp<Real>::prime(Real x) const
{
    Real rx = this->operator()(x);
    Real numerator = 0;
    Real denominator = 0;
    for(size_t i = 0; i < m_x.size(); ++i)
    {
        if (x == m_x[i])
        {
            Real sum = 0;
            for (size_t j = 0; j < m_x.size(); ++j)
            {
                if (j == i)
                {
                    continue;
                }
                sum += m_w[j]*(m_y[i] - m_y[j])/(m_x[i] - m_x[j]);
            }
            return -sum/m_w[i];
        }
        Real t = m_w[i]/(x - m_x[i]);
        Real diff = (rx - m_y[i])/(x-m_x[i]);
        numerator += t*diff;
        denominator += t;
    }

    return numerator/denominator;
}
}}}
#endif


#pragma once

#include "../../util/assert.hh"
#include "../../util/meta.hh"

#include <Eigen/Core>
#include <stdexcept>
#include <sstream>

namespace bold
{
  template<typename T,int dim>
  struct LineSegment : Eigen::Matrix<T,dim,2>
  {
  public:
    typedef Eigen::Matrix<T,dim,2> Base;
    typedef Eigen::Matrix<T,dim,1> PointType;

    template<typename OtherDerived>
    LineSegment(Eigen::MatrixBase<OtherDerived> const& other)
      :  Base(other)
    {}

    LineSegment(PointType const& p1,
                PointType const& p2)
    {
      static_assert(std::is_arithmetic<T>::value, "Must be an arithmetic type");

      this->col(0) = p1;
      this->col(1) = p2;

      ASSERT((p2 - p1).cwiseAbs().maxCoeff() != 0);
    }

    template<typename OtherDerived>
    LineSegment& operator=(Eigen::MatrixBase<OtherDerived> const& other)
    {
      this->Base::operator=(other);
      return *this;
    }

    PointType p1() const { return this->col(0); }
    PointType p2() const { return this->col(1); }

    /** Returns the vector formed by p2() - p1() */
    PointType delta() const { return p2() - p1(); }

    /** Returns the vector formed by (p2() + p1()) / 2.0 */
    PointType mid() const { return (p2() + p1()) / 2.0; }

    T length() const { return delta().norm(); }

    double normalisedDot(LineSegment<T,dim> other) const
    {
      return delta().normalized().dot( other.delta().normalized() );
    }

    /** Calculate the smallest angle between the two line segments.
     *
     * The output will be between 0 and PI/2, inclusive.
     */
    double smallestAngleBetween(LineSegment<T,dim> other)
    {
      auto a = fabs(acos(normalisedDot(other)));
      if (a >= M_PI/2)
        a = M_PI - a;
      return a;
    }

    template<int newDim>
    LineSegment<T,newDim> to() const
    {
      auto newLineSegment = LineSegment<T,newDim>{LineSegment<T,newDim>::Zero()};
      newLineSegment.col(0).template head< meta::min<dim,newDim>::value >() =
        this->col(0).template head< meta::min<dim,newDim>::value >();
      newLineSegment.col(1).template head< meta::min<dim,newDim>::value >() =
        this->col(1).template head< meta::min<dim,newDim>::value >();
      return newLineSegment;
    }

    bool operator==(LineSegment<T,dim> const& other) const
    {
      const double epsilon = 0.0000004;
      return (this->array() - other.array()).abs().sum() < epsilon;
    }

    LineSegment<T,dim> operator+(PointType const& delta) const
    {
      return LineSegment<T,dim>{this->colwise() + delta};
    }

    LineSegment<T,dim> operator-(PointType const& delta) const
    {
      return LineSegment<T,dim>{this->colwise() - delta};
    }

    friend std::ostream& operator<<(std::ostream& stream, LineSegment<T,dim> const& lineSegment)
    {
      return stream << "LineSegment (P1=" << lineSegment.p1().transpose() << " P2=" << lineSegment.p2().transpose() << ")";
    }

  };

  typedef LineSegment<double,3> LineSegment3d;
  typedef LineSegment<int,3> LineSegment3i;
}


#include <boost/math/special_functions/round.hpp>
#include <cmath>
#include <jngl.hpp>

const int screenWidth = 800;
const int screenHeight = 600;

class Ball {
public:
	explicit Ball(const std::string& filename);
	void draw() const;
	void Move();
	void CheckMouse(int x, int y);

private:
	double x_, y_, xSpeed_, ySpeed_;
	const std::string filename_;
	const int width_, height_;
};

const double timePerFrame = 0.01; // 100 FPS

JNGL_MAIN_BEGIN {
	jngl::setScaleFactor(
	    std::floor(std::min(jngl::getDesktopWidth() / 800, jngl::getDesktopHeight() / 600)));
	jngl::showWindow("Throw Example", boost::math::iround(screenWidth * jngl::getScaleFactor()),
	                 boost::math::iround(screenHeight * jngl::getScaleFactor()));
	Ball ball("ball");
	double oldTime = jngl::getTime();
	while (jngl::running()) {
		if (jngl::getTime() - oldTime > timePerFrame) {
			oldTime += timePerFrame;
			ball.CheckMouse(static_cast<int>(jngl::getMousePos().x + screenWidth / 2.),
			                static_cast<int>(jngl::getMousePos().y + screenHeight / 2.));
			ball.Move();
		}
		else
		{
			jngl::updateInput();
			jngl::translate(-screenWidth / 2., -screenHeight / 2.);
			ball.draw();
			jngl::swapBuffers();
		}
	}
} JNGL_MAIN_END

Ball::Ball(const std::string& filename)
: x_(100), y_(100), xSpeed_(200), ySpeed_(200), filename_(filename),
  width_(jngl::getWidth(filename)), height_(jngl::getHeight(filename)) {
}

void Ball::draw() const {
	jngl::draw("ball", static_cast<int>(x_), static_cast<int>(y_));
}

void Ball::Move()
{
	x_ += xSpeed_ * timePerFrame; // 100 FPS also jedes mal ein hunderstel der Bewegung
	y_ += ySpeed_ * timePerFrame;
	xSpeed_ *= 0.99; // Die Geschwindigkeit wird 100 mal pro Sekunde um 1 % langsamer
	ySpeed_ *= 0.99;
	if(x_ + width_ > screenWidth)
	{
		xSpeed_ = -xSpeed_;
		x_ = screenWidth - width_;
	}
	if(y_ + height_ > screenHeight)
	{
		ySpeed_ = -ySpeed_;
		y_ = screenHeight - height_;
	}
	if(x_ < 0)
	{
		xSpeed_ = -xSpeed_;
		x_ = 0;
	}
	if(y_ < 0)
	{
		ySpeed_ = -ySpeed_;
		y_ = 0;
	}
}

void Ball::CheckMouse(const int x, const int y)
{
	if (jngl::mouseDown()) {
		xSpeed_ = (x - width_ / 2. - x_) * 10; // Geschwindigkeit vom Ball auf Entfernung von
		ySpeed_ = (y - height_ / 2. - y_) * 10; // Mausposition und Ballmittelpunkt mal 10 setzen
	}
}


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

#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <boost/multi_array.hpp>
#include <fftw3.h>
#include "dvs_reconstruction/poisson_solver/laplace.h"

#ifdef TIME_REPORT
#include <time.h>
#include <sys/time.h>
namespace pde
{

double stoptime(void)
{
   struct timeval t;
   gettimeofday(&t,NULL);
   return (double) t.tv_sec + t.tv_usec/1000000.0;
}
}
#endif // TIME_REPORT

namespace
{

// convenience function: square
inline double sqr(double x)
{
   return x*x;
}

} // namespace

// 2d-array utility functions
// --------------------------
namespace arr
{

// helper function: [0,1] uniform random generator
double runif()
{
   return (double) rand()/(RAND_MAX+1.0);
}
// fill vector with random values
void runif(std::vector<double>& X)
{
   size_t n=X.size();
   for(size_t i=0; i<n; i++) {
      X[i]=runif();
   }
}
// fill 2d-array with random values
void runif(boost::multi_array<double,2>& X)
{
   size_t n1=X.shape()[0];
   size_t n2=X.shape()[1];
   for(size_t i=0; i<n1; i++) {
      for(size_t j=0; j<n2; j++) {
         X[i][j]=runif();
      }
   }
}

// print the matrix
void print(const boost::multi_array<double,2>& X)
{
   size_t n1=X.shape()[0];
   size_t n2=X.shape()[1];
   for(size_t i=0; i<n1; i++) {
      for(size_t j=0; j<n2; j++) {
         printf("%9.5f ", X[i][j]);
      }
      printf("\n");
   }
}

// L^2 difference between two vectors
double diff(const std::vector<double>& X,
            const std::vector<double>& Y,
            bool allow_shift)
{

   assert(X.size()==Y.size());
   size_t n=X.size();

   double c=0.0;
   if(allow_shift) {
      c=X[1]-Y[1];
   }
   double sum=0.0;
   for(size_t i=0; i<n; i++) {
      sum+=sqr(X[i]-Y[i]-c);
   }
   return sqrt(sum);
}

// L^2 difference between vector and a constant
double diff(const std::vector<double>& X, double c)
{
   size_t n=X.size();
   double sum=0.0;
   for(size_t i=0; i<n; i++) {
      sum+=sqr(X[i]-c);
   }
   return sqrt(sum);
}

// L^2 difference between two 2d-arrays
// allow_shift: removes a constant shift between 2 matrices,
//    i.e.  comparing X[][] and Y[][] + (X[1][1]-Y[1][1])
// ignore_corner: ignores all 4 corners, X[0][0], ...
double diff(const boost::multi_array<double,2>& X,
            const boost::multi_array<double,2>& Y,
            bool allow_shift, bool ignore_corner)
{

   size_t n1=X.shape()[0];
   size_t n2=X.shape()[1];

   assert(Y.shape()[0]==n1 && n1>1);
   assert(Y.shape()[1]==n2 && n2>1);

   double c=0.0;
   if(allow_shift) {
      c=X[1][1]-Y[1][1];
   }
   double sum=0.0;
   for(size_t i=0; i<n1; i++) {
      for(size_t j=0; j<n2; j++) {
         if( !(ignore_corner==true && (i==0 || i==n1-1) && (j==0 || j==n2-1)) )
            sum+=sqr(X[i][j]-Y[i][j]-c);
      }
   }
   return sqrt(sum);
}

} // namespace arr

namespace pde
{

// sets the number of threads for the fftw solver
int fftw_threads(int n)
{
   int ret=fftw_init_threads();
   fftw_plan_with_nthreads(n);
   return ret;
}
// cleans up memory used by fftw for threads
void fftw_clean()
{
   fftw_cleanup();
   fftw_cleanup_threads();
}

// sets boundary vectors to a uniform value
void set_boundary_to_uniform(
   std::vector<double>& bd1a, std::vector<double>& bd1b,
   std::vector<double>& bd2a, std::vector<double>& bd2b,
   double value, size_t n1, size_t n2)
{

   bd1a.resize(n2);
   bd1b.resize(n2);
   bd2a.resize(n1);
   bd2b.resize(n1);
   for(size_t i=0; i<n2; i++) {
      bd1a[i]=value;
      bd1b[i]=value;
   }
   for(size_t i=0; i<n1; i++) {
      bd2a[i]=value;
      bd2b[i]=value;
   }
}

// saves the 4 boundaries/edges of a 2-d array as individual vectors
// note, the 4 corners/vertices are not saved and the resulting
// vectors are of size n-2
void get_boundary(std::vector<double>& bd1a, std::vector<double>& bd1b,
                  std::vector<double>& bd2a, std::vector<double>& bd2b,
                  const boost::multi_array<double,2>& X,
                  double h1, double h2, types::boundary boundary)
{
   size_t n1=X.shape()[0];
   size_t n2=X.shape()[1];
   assert(n1>2 && n2>2);

   // allocating memory for boundary vectors
   bd1a.resize(n2-2);
   bd1b.resize(n2-2);
   bd2a.resize(n1-2);
   bd2b.resize(n1-2);

   // reading boundary values
   if(boundary==types::Dirichlet) {
      for(size_t i=1; i<n2-1; i++) {
         bd1a[i-1] = X[0][i];
         bd1b[i-1] = X[n1-1][i];
      }
      for(size_t i=1; i<n1-1; i++) {
         bd2a[i-1] = X[i][0];
         bd2b[i-1] = X[i][n2-1];
      }
   } else if(boundary==types::Neumann) {
      for(size_t i=1; i<n2-1; i++) {
         bd1a[i-1] = (X[0][i]-X[2][i]) / (2.0*h1);
         bd1b[i-1] = (X[n1-1][i]-X[n1-3][i]) / (2.0*h1);
      }
      for(size_t i=1; i<n1-1; i++) {
         bd2a[i-1] = (X[i][0]-X[i][2]) / (2.0*h2);
         bd2b[i-1] = (X[i][n2-1]-X[i][n2-3]) / (2.0*h2);
      }
   } else {
      assert(false);
   }
}

// set boundary of matrix X according to conditions
// can either add boundaries to the array (resize), or overwrite the
// current boundary values
// boundary vectors shall not contain the 4 corners (vertices), ie need to
// be of size n-2 of the resulting array
void set_boundary(boost::multi_array<double,2>& X, double h1, double h2,
                  const std::vector<double>& bd1a, const std::vector<double>& bd1b,
                  const std::vector<double>& bd2a, const std::vector<double>& bd2b,
                  types::boundary boundary, bool add)
{

   size_t n1=X.shape()[0];
   size_t n2=X.shape()[1];
   assert(n1>0 && n2>0);

   if(add==true) {
      // add boundary --> grow array
      X.resize(boost::extents[n1+2][n2+2]);  // element preserving resize
      n1=X.shape()[0];
      n2=X.shape()[1];
      // shifting numbers inside the array to make space for boundary
      for(int i=n1-3; i>=0; i--) {
         for(int j=n2-3; j>=0; j--) {
            X[i+1][j+1]=X[i][j];
         }
      }
   }

   // check the boundary vectors are of the correct size
   assert( bd1a.size()==bd1b.size() && bd1a.size()==n2-2 );
   assert( bd2a.size()==bd2b.size() && bd2a.size()==n1-2 );

   // setting boundary values
   if(boundary==types::Dirichlet) {
      for(int i=0; i<(int)n2; i++) {
         // include corner cases, which are otherwise undetermined
         int idx = std::min(std::max(i-1,0),(int)n2-3);
         X[0][i]    = bd1a[idx];
         X[n1-1][i] = bd1b[idx];
      }
      for(size_t i=1; i<n1-1; i++) {
         X[i][0]    = bd2a[i-1];
         X[i][n2-1] = bd2b[i-1];
      }
   } else if(boundary==types::Neumann) {
      for(int i=0; i<(int)n2; i++) {
         // include corner cases, which are otherwise undetermined
         int idx = std::min(std::max(i-1,0),(int)n2-3);
         X[0][i]    = X[2][i]    + 2.0*h1*bd1a[idx];
         X[n1-1][i] = X[n1-3][i] + 2.0*h1*bd1b[idx];
      }
      for(size_t i=1; i<n1-1; i++) {
         X[i][0]    = X[i][2]    + 2.0*h2*bd2a[i-1];
         X[i][n2-1] = X[i][n2-3] + 2.0*h2*bd2b[i-1];
      }
   } else {
      assert(false);
   }
}
// specialised version where all boundaries have the same value
void set_boundary(boost::multi_array<double,2>& X, double h1, double h2,
                  double bdvalue, types::boundary bdcond, bool add)
{
   size_t n1=X.shape()[0];
   size_t n2=X.shape()[1];
   // set the correct size of the boundary vectors
   size_t m1,m2;
   if(add==true) {
      m1=n1;
      m2=n2;
   } else {
      assert(n1>=2 && n2>=2);
      m1=n1-2;
      m2=n2-2;
   }
   // fill boundary vectors
   std::vector<double> bd1a,bd1b,bd2a,bd2b;
   set_boundary_to_uniform(bd1a,bd1b,bd2a,bd2b,bdvalue,m1,m2);

   // call the general routine
   set_boundary(X,h1,h2,bd1a,bd1b,bd2a,bd2b,bdcond,add);
}

// PDE functions to solve the discrete Poisson's Equation
// ------------------------------------------------------

// discrete Gradient operator on a uniform grid
// given u(x,y), it returns DX = a_1 du/dx, DY = a_2 du/dy,
// using one-sided finite differences
// U[i][j]  = u(x_i, y_j), i=0...n-1
// DX[i][j] = a1*(U[i+1][j]-U[i][j])/h, i=0...n-2
void grad(boost::multi_array<double,2>& DX, boost::multi_array<double,2>& DY,
          const boost::multi_array<double,2>& U,
          double a1, double a2, double h1, double h2)
{
   assert(U.shape()[0]>1 && U.shape()[1]>1);
   size_t n1=U.shape()[0]-1;
   size_t n2=U.shape()[1]-1;
   DX.resize(boost::extents[n1][n2]);
   DY.resize(boost::extents[n1][n2]);
   for(size_t i=0; i<n1; i++) {
      for(size_t j=0; j<n2; j++) {
         DX[i][j] = (a1/h1) * (U[i+1][j]-U[i][j]);
         DY[i][j] = (a2/h2) * (U[i][j+1]-U[i][j]);
      }
   }
}
// discrete Gradient operator on a uniform grid
// same as above, but the input U is assumed to contain only inner grid points
// and the boundary grid points are inferred from the boundary condition
// U[i][j] = u(x_i, y_j), i=0...n-1
// DX[i][j] = a1*(U[i][j]-U[i-1][j])/h, i=0...n+1
void grad(boost::multi_array<double,2>& DX, boost::multi_array<double,2>& DY,
          const boost::multi_array<double,2>& U,
          double a1, double a2, double h1, double h2,
          const std::vector<double>& bd1a, const std::vector<double>& bd1b,
          const std::vector<double>& bd2a, const std::vector<double>& bd2b,
          types::boundary boundary)
{
   // this is lazy and very memory inefficient
   // TODO: replace with a direct calculation
   boost::multi_array<double,2> V=U;
   pde::set_boundary(V,h1,h2,bd1a,bd1b,bd2a,bd2b,boundary,true);
   grad(DX,DY,V,a1,a2,h1,h2);
}
// simplified version where only one boundary value can be specified
void grad(boost::multi_array<double,2>& DX, boost::multi_array<double,2>& DY,
          const boost::multi_array<double,2>& U,
          double a1, double a2, double h1, double h2,
          double bdvalue, types::boundary bdcond)
{
   size_t n1=U.shape()[0];
   size_t n2=U.shape()[1];

   // fill boundary vectors
   std::vector<double> bd1a,bd1b,bd2a,bd2b;
   set_boundary_to_uniform(bd1a,bd1b,bd2a,bd2b,bdvalue,n1,n2);

   // call the general routine
   grad(DX,DY,U,a1,a2,h1,h2,bd1a,bd1b,bd2a,bd2b,bdcond);
}

// discrete Divergence operator on a uniform grid
// given u(x,y), v(x,y), it returns f = a_1 du/dx + a_2 dv/dx,
// using one-sided finite differences
void div(boost::multi_array<double,2>& F,
         const boost::multi_array<double,2>& U,
         const boost::multi_array<double,2>& V,
         double a1, double a2, double h1, double h2)
{
   assert(U.shape()[0]==V.shape()[0]);
   assert(U.shape()[1]==V.shape()[1]);
   assert(U.shape()[0]>1 && U.shape()[1]>1);
   size_t n1=U.shape()[0]-1;
   size_t n2=V.shape()[1]-1;
   F.resize(boost::extents[n1][n2]);
   for(size_t i=0; i<n1; i++) {
      for(size_t j=0; j<n2; j++) {
         F[i][j] = (a1/h1) * (U[i+1][j+1]-U[i][j+1])
                   + (a2/h2) * (V[i+1][j+1]-V[i+1][j]);
      }
   }
}

// discrete Laplace operator on a uniform grid
// given u, it returns f = a_1 u_xx + a_2 u_yy
// note, the Laplace operator can only be applied to inner grid points
// and so the result F has two grid points less in each dimension, i.e.
// U[i][j] = u(x_i, y_j)
// F[i][j] = f(x_{i+1}, y_{i+1})
void laplace(boost::multi_array<double,2>& F,
             const boost::multi_array<double,2>& U,
             double a1, double a2, double h1, double h2)
{
   assert(U.shape()[0]>2 && U.shape()[1]>2);
   size_t n1=U.shape()[0]-2;
   size_t n2=U.shape()[1]-2;
   F.resize(boost::extents[n1][n2]);
   for(size_t i=0; i<n1; i++) {
      for(size_t j=0; j<n2; j++) {
         // index [i][j] in F corresponds to [i+1][j+1] in U
         F[i][j] = a1 * (U[i][j+1] - 2.0*U[i+1][j+1] + U[i+2][j+1]) / (h1*h1)
                   + a2 * (U[i+1][j] - 2.0*U[i+1][j+1] + U[i+1][j+2]) / (h2*h2);
      }
   }
}
// discrete Laplace operator on a uniform grid
// the input U is assumed to contain only inner grid points and the
// boundary grid points are inferred from the boundary condition
// U[i][j] = u(x_i, y_j)
// F[i][j] = f(x_i, y_i)
void laplace(boost::multi_array<double,2>& F,
             const boost::multi_array<double,2>& U,
             double a1, double a2, double h1, double h2,
             const std::vector<double>& bd1a, const std::vector<double>& bd1b,
             const std::vector<double>& bd2a, const std::vector<double>& bd2b,
             types::boundary boundary)
{
   size_t n1=U.shape()[0];
   size_t n2=U.shape()[1];
   assert(n1>0 && n2>0);
   assert(bd1a.size()==bd1b.size() && bd1a.size()==n2);
   assert(bd2a.size()==bd2b.size() && bd2a.size()==n1);

   F.resize(boost::extents[n1][n2]);
   for(int i=0; i<(int)n1; i++) {
      for(int j=0; j<(int)n2; j++) {
         double Um1, U0, Up1;       // U[i-1], U[i], U[i+1]

         // first dimension
         Um1 = U[std::max(i-1,0)][j];
         U0  = U[i][j];
         Up1 = U[std::min(i+1,(int)n1-1)][j];
         if(i==0) {
            if(boundary==types::Dirichlet) {
               Um1 = bd1a[j];
            } else if(boundary==types::Neumann) {
               Um1 = U[1][j] + 2.0*h1*bd1a[j];
            }
         } else if(i==(int)n1-1) {
            if(boundary==types::Dirichlet) {
               Up1 = bd1b[j];
            } else if(boundary==types::Neumann) {
               Up1 = U[n1-2][j] + 2.0*h1*bd1b[j];
            }
         }
         F[i][j] = a1 * (Um1 - 2.0*U0 + Up1) / (h1*h1);

         // second dimension
         Um1 = U[i][std::max(j-1,0)];
         Up1 = U[i][std::min(j+1,(int)n2-1)];
         if(j==0) {
            if(boundary==types::Dirichlet) {
               Um1 = bd2a[i];
            } else if(boundary==types::Neumann) {
               Um1 = U[i][1] + 2.0*h2*bd2a[i];
            }
         } else if(j==(int)n2-1) {
            if(boundary==types::Dirichlet) {
               Up1 = bd2b[i];
            } else if(boundary==types::Neumann) {
               Up1 = U[i][n2-2] + 2.0*h2*bd2b[i];
            }
         }
         F[i][j] += a2 * (Um1 - 2.0*U0 + Up1) / (h2*h2);
      }
   }
}
// simplified version where only one boundary value can be specified
void laplace(boost::multi_array<double,2>& F,
             const boost::multi_array<double,2>& U,
             double a1, double a2, double h1, double h2,
             double bdvalue, types::boundary bdcond)
{

   size_t n1=U.shape()[0];
   size_t n2=U.shape()[1];
   assert(n1>0 && n2>0);

   // fill boundary vectors
   std::vector<double> bd1a,bd1b,bd2a,bd2b;
   set_boundary_to_uniform(bd1a,bd1b,bd2a,bd2b,bdvalue,n1,n2);

   // call the general routine
   laplace(F,U,a1,a2,h1,h2,bd1a,bd1b,bd2a,bd2b,bdcond);
}

// Neumann boundary condition with a general right hand side F will not
// have a solution, however by looking at F we can say what the minimum
// L2-norm error will be for 0-Neumann boundary conditions:
// return = min ||Laplace U - F||
// min is taken over all U consistent with 0-Neumann condition
double neumann_error(const boost::multi_array<double,2>& F)
{
   size_t n1=F.shape()[0];
   size_t n2=F.shape()[1];
   assert(n1>1 && n2>1);

   // assuming 0 boundary conditions, otherwise would have to adjust
   // right hand side F as in poisolve()

   // calculate min L2 error ||Laplace U - F|| = |\hat F[0]| * ||EV[0]||
   // where \hat F, are the coordinates of F in EV-space,
   // and can simply be calculated as follows
   double sum=0.0;
   double fac=1.0;
   for(size_t i=0; i<n1; i++) {
      for(size_t j=0; j<n2; j++) {
         fac=1.0;
         if(j==0 || j==n2-1)
            fac*=0.5;
         if(i==0 || i==n1-1)
            fac*=0.5;
         sum+=fac*F[i][j];
      }
   }
   double F00 = sum/((n1-1)*(n2-1));      // \hat F [0][0] (EV space)
   double norm_ev=sqrt((double)(n1*n2));  // EV[0]=(1,...,1) --> norm=sqrt(n)
   double l2_error=F00*norm_ev;
   return l2_error;
}

// given a right hand side F, we can find a constant Neumann-boundary
// value which will have a solution U, so that Laplace U = F,
// returns the boundary value
double neumann_compat(const boost::multi_array<double,2>& F,
                      double a1, double a2, double h1, double h2)
{
   size_t n1=F.shape()[0];
   size_t n2=F.shape()[1];

   double l2_error=neumann_error(F);
   double norm_ev=sqrt((double)(n1*n2));  // EV[0]=(1,...,1) --> norm=sqrt(n)
   double F00 = l2_error/norm_ev;         // \hat F [0][0] (EV space)

   // with non-zero Neumann boundary condition, rhs F is modified,
   // as in poisolve(), so we can calculate the exact boundary
   // value to make the l2_error zero
   double bd = F00 / (2.0*a1/(h1*(n1-1)) + 2.0*a2/(h2*(n2-1)));

   return bd;
}

// solves the 2D Poisson equation: a1 u_xx + a2 u_yy = f
//
// discretised with uniform grid, x_{i+1} = x_i + h1, y_{i+1} = y_i + h2
// input: rhs       F[i][j] = f(x_i, y_j)
// output: solution  U[i][j] = u(x_i, y_j), ie only inner points by default
//
//
//
// boundary condition:
// - lower and upper boundary in dimension 1: bd1a, bd1b
// - lower and upper boundary in dimension 2: bd2a, bd2b
// - e.g. bd1a[i] refers to U[-1][i]
double poisolve( boost::multi_array<double,2>& U,
                 const boost::multi_array<double,2>& F,
                 double a1, double a2, double h1, double h2,
                 const std::vector<double>& bd1a, const std::vector<double>& bd1b,
                 const std::vector<double>& bd2a, const std::vector<double>& bd2b,
                 types::boundary boundary, bool add_boundary_to_solution)
{

#ifdef TIME_REPORT
   double t0=stoptime();
   double t1=t0;
   double t2;
#endif

   size_t n1=F.shape()[0];
   size_t n2=F.shape()[1];
   assert(n1>0 && n2>0);

   // adjust right hand side F with boundary condition (nothing to do for =0)
   boost::multi_array<double,2> rhs = F;
   assert( bd1a.size()==bd1b.size() && bd1a.size()==n2 );
   assert( bd2a.size()==bd2b.size() && bd2a.size()==n1 );
   {
      // factors for boundary adjustment to rhs
      double c1 = 0.0, c2 = 0.0;
      if(boundary==types::Dirichlet) {
         c1=a1/sqr(h1);
         c2=a2/sqr(h2);
      } else if(boundary==types::Neumann) {
         c1=2.0*a1/h1;
         c2=2.0*a2/h2;
      } else {
         assert(false);
      }
      // adjust rhs with boundary conditions
      for(size_t i=0; i<n2; i++) {
         rhs[0][i]    -= c1 * bd1a[i];
         rhs[n1-1][i] -= c1 * bd1b[i];
      }
      for(size_t i=0; i<n1; i++) {
         rhs[i][0]    -= c2 * bd2a[i];
         rhs[i][n2-1] -= c2 * bd2b[i];
      }
   }

#ifdef TIME_REPORT
   t2=stoptime();
   printf("poisolve(): %5.0f ms: rhs boundary conditions\n",(t2-t1)*1000.0);
   t1=t2;
#endif

   // transform rhs into EV space (inverse fft, in-place)
   {
      fftw_plan p;
      double fft_norm=0.0;  // i.e. FFT(FFT(x)) = fft_norm * x

      if(boundary==types::Dirichlet) {
         // DST-I = FFTW_RODFT00 (note, additional factor of 2 in fftw)
         p=fftw_plan_r2r_2d(n1, n2, &(rhs[0][0]), &(rhs[0][0]),
                            FFTW_RODFT00, FFTW_RODFT00, FFTW_ESTIMATE);
         fftw_execute(p);
         fftw_destroy_plan(p);
         fft_norm = 4.0*((n1+1)*(n2+1));

      } else if(boundary==types::Neumann) {
         // EV space is similar to DCT-I but without the factor of 1/2
         // for the 1st and last element

         // DCT-I = REDFT00 (note, additional factor of 2 in fftw)
         // so that DCT-I*DCT-I = 2(n-1) I (and not the usual (n-1)/2)
         p=fftw_plan_r2r_2d(n1, n2, &(rhs[0][0]), &(rhs[0][0]),
                            FFTW_REDFT00, FFTW_REDFT00, FFTW_ESTIMATE);
         fftw_execute(p);
         fftw_destroy_plan(p);

         for(size_t i=0; i<n1; i++) {
            rhs[i][0]    *= 0.5;
            rhs[i][n2-1] *= 0.5;
         }
         for(size_t j=0; j<n2; j++) {
            rhs[0][j]    *= 0.5;
            rhs[n1-1][j] *= 0.5;
         }
         fft_norm = 4.0*((n1-1)*(n2-1));

      } else {
         assert(false);
      }
      // scale so we get the inverse fft
      for(size_t i=0; i<n1; i++) {
         for(size_t j=0; j<n2; j++) {
            rhs[i][j] *= (1.0/fft_norm);      // div is more expensive than mul
         }
      }
   }

#ifdef TIME_REPORT
   t2=stoptime();
   printf("poisolve(): %5.0f ms: rhs to EV space (inv fft)\n",(t2-t1)*1000.0);
   t1=t2;
#endif

   // calculate eigenvalues of the linear operators L
   std::vector<double> lambda1(n1);
   std::vector<double> lambda2(n2);
   if(boundary==types::Dirichlet) {
      for(size_t i=0; i<n1; i++)
         lambda1[i] = -4.0*sqr( sin((M_PI*(i+1))/(2.0*(n1+1))) );
      for(size_t i=0; i<n2; i++)
         lambda2[i] = -4.0*sqr( sin((M_PI*(i+1))/(2.0*(n2+1))) );
   } else if(boundary==types::Neumann) {
      for(size_t i=0; i<n1; i++)
         lambda1[i] = -4.0*sqr( sin((M_PI*i)/(2.0*(n1-1))) );
      for(size_t i=0; i<n2; i++)
         lambda2[i] = -4.0*sqr( sin((M_PI*i)/(2.0*(n2-1))) );
   } else {
      assert(false);
   }

   // solve the equation for U in EV space
   double error=0.0;
   U.resize(boost::extents[n1][n2]);
   for(size_t i=0; i<n1; i++) {
      for(size_t j=0; j<n2; j++) {
         double div = (a1*lambda1[i]/(h1*h1) + a2*lambda2[j]/(h2*h2));
         if(div==0.0) {
            // here we need rhs[i][j] == 0 or else there is no solution
            // calculate the L2-norm of the error, need to know norm of EV
            // we know div==0.0, only for Neumann and EV 0, ||EV[0]||^2=n
            // however, since fftw has an extra factor 2, ||EV[0]||^2=4n
            double norm_ev = 16.0*n1*n2;     // TODO: make this more general
            error+=sqr(rhs[i][j])*norm_ev;
            // U[i][j] is arbitrary here
            U[i][j] = 0.0;
         } else {
            U[i][j] = rhs[i][j] / div;
         }
      }
   }

   // free rhs memory as it's no longer needed
   rhs.resize(boost::extents[0][0]);

#ifdef TIME_REPORT
   t2=stoptime();
   printf("poisolve(): %5.0f ms: solve equation in EV space\n",(t2-t1)*1000.0);
   t1=t2;
#endif

   // transform U from EV space into canonical space (fft, in-place)
   {
      fftw_plan p = nullptr;

      if(boundary==types::Dirichlet) {
         // DST-I = FFTW_RODFT00 (note, additional factor of 2 in fftw)
         // so that DST-I*DST-I = 2(n+1) I (and not the usual (n+1)/2)
         p=fftw_plan_r2r_2d(n1, n2, &(U[0][0]), &(U[0][0]),
                            FFTW_RODFT00, FFTW_RODFT00, FFTW_ESTIMATE);
         fftw_execute(p);

      } else if(boundary==types::Neumann) {
         // EV space is similar to DCT-I but without the factor of 1/2
         // for the 1st and last element

         // apply factor 2
         for(size_t i=0; i<n1; i++) {
            U[i][0]    *= 2.0;
            U[i][n2-1] *= 2.0;
         }
         for(size_t j=0; j<n2; j++) {
            U[0][j]    *= 2.0;
            U[n1-1][j] *= 2.0;
         }
         // DCT-I = REDFT00 (note, additional factor of 2 in fftw)
         p=fftw_plan_r2r_2d(n1, n2, &(U[0][0]), &(U[0][0]),
                            FFTW_REDFT00, FFTW_REDFT00, FFTW_ESTIMATE);
         fftw_execute(p);

      } else {
         assert(false);
      }
      fftw_destroy_plan(p);
   }

#ifdef TIME_REPORT
   t2=stoptime();
   printf("poisolve(): %5.0f ms: solution to normal space (fft)\n",(t2-t1)*1000.0);
   t1=t2;
#endif

   // by default, U only contains the inner grid points of the solution
   // however, we can also add the boundary if needed (inefficient)
   if(add_boundary_to_solution) {
      set_boundary(U,h1,h2,bd1a,bd1b,bd2a,bd2b,boundary,true);
   }

#ifdef TIME_REPORT
   t2=stoptime();
   printf("poisolve(): %5.0f ms: total time\n", (t2-t0)*1000.0);
#endif

   return sqrt(error);
}

// Poisson solver, assuming uniform boundary
double poisolve(boost::multi_array<double,2>& U,
                const boost::multi_array<double,2>& F,
                double a1, double a2, double h1, double h2,
                double bound_value,
                types::boundary boundarycondition,
                bool add_boundary_to_solution)
{

   size_t n1=F.shape()[0];
   size_t n2=F.shape()[1];

   // fill boundary vectors
   std::vector<double> bd1a,bd1b,bd2a,bd2b;
   set_boundary_to_uniform(bd1a,bd1b,bd2a,bd2b,bound_value,n1,n2);

   // call the general solver
   return poisolve(U,F,a1,a2,h1,h2,bd1a,bd1b,bd2a,bd2b,
                   boundarycondition,add_boundary_to_solution);
}

} // namespace pde


#ifndef EIGEN_SOLVER_H
#define EIGEN_SOLVER_H

#include <cassert>
#include <armadillo>

namespace willow { namespace qcmol {

class EigenSolver
{
public:

  EigenSolver (): m_is_init(false), m_eig_vals(), m_eig_vecs() {}

  // Sinvh.n_rows = nbf
  // Sinvh.n_cols = nmo
  // F.n_rows = nbf
  // F.n_cols = nbf
  // solve WF(nmo, nmo) = Sinv^T(nmo,nbf) F(nbf,nbf) Sinv(nbf,nmo)
  // dsyevd(WF) or eig_sym ()
  // save m_eig_vecs (nbf,nmo) = Sinvh(nbf,nmo)*Cvec(nmo,nmo)
  //
  EigenSolver  (const arma::mat Sinvh, const arma::mat F):
    m_is_init(true),
    m_eig_vals(Sinvh.n_cols),
    m_eig_vecs(Sinvh.n_rows, Sinvh.n_cols) {

    compute (Sinvh, F);

  }

  void compute (const arma::mat Sinvh, const arma::mat F) {

    const int nbf = Sinvh.n_rows;
    const int nmo = Sinvh.n_cols;

    if (!m_is_init) {
      m_is_init = true;
      m_eig_vals.set_size (nmo);
      m_eig_vecs.set_size (nbf, nmo);
    }

    // F'
    arma::mat wf = Sinvh.t()*F*Sinvh; // (nmo, nmo)

    // Update Orbitals and Energies
    arma::mat eg_vecs;
    bool ok = arma::eig_sym (m_eig_vals, eg_vecs, wf);

    // Transform back to non-orthogonal basis
    m_eig_vecs = Sinvh*eg_vecs; // (nbf, nmol)

  }

  arma::vec eigenvalues() const {
    assert (m_is_init);
    return m_eig_vals;
  }

  arma::mat eigenvectors() const {
    assert (m_is_init);
    return m_eig_vecs;
  }

protected:
  bool      m_is_init;
  arma::vec m_eig_vals;
  arma::mat m_eig_vecs;
};

} } // namespace willow::qcmol

#endif


#include "lapnig_hierarchy.h"

#include <google/protobuf/stubs/casts.h>

#include <Eigen/Dense>
#include <algorithm>
#include <cmath>
#include <stan/math/prim/prob.hpp>
#include <vector>

#include "algorithm_state.pb.h"
#include "hierarchy_prior.pb.h"
#include "ls_state.pb.h"
#include "src/utils/rng.h"

unsigned int LapNIGHierarchy::accepted_ = 0;
unsigned int LapNIGHierarchy::iter_ = 0;

void LapNIGHierarchy::set_state_from_proto(
    const google::protobuf::Message &state_) {
  auto &statecast = downcast_state(state_);
  state.mean = statecast.uni_ls_state().mean();
  state.scale = statecast.uni_ls_state().var();
  set_card(statecast.cardinality());
}

void LapNIGHierarchy::set_hypers_from_proto(
    const google::protobuf::Message &hypers_) {
  auto &hyperscast = downcast_hypers(hypers_).lapnig_state();
  hypers->mean = hyperscast.mean();
  hypers->var = hyperscast.var();
  hypers->scale = hyperscast.scale();
  hypers->shape = hyperscast.shape();
  hypers->mh_mean_var = hyperscast.mh_mean_var();
  hypers->mh_log_scale_var = hyperscast.mh_log_scale_var();
}

double LapNIGHierarchy::like_lpdf(const Eigen::RowVectorXd &datum) const {
  return stan::math::double_exponential_lpdf(datum(0), state.mean,
                                             state.scale);
}

std::shared_ptr<bayesmix::AlgorithmState::ClusterState>
LapNIGHierarchy::get_state_proto() const {
  bayesmix::UniLSState state_;
  state_.set_mean(state.mean);
  state_.set_var(state.scale);

  auto out = std::make_shared<bayesmix::AlgorithmState::ClusterState>();
  out->mutable_uni_ls_state()->CopyFrom(state_);
  return out;
}

std::shared_ptr<bayesmix::AlgorithmState::HierarchyHypers>
LapNIGHierarchy::get_hypers_proto() const {
  bayesmix::LapNIGState hypers_;
  hypers_.set_mean(hypers->mean);
  hypers_.set_var(hypers->var);
  hypers_.set_shape(hypers->shape);
  hypers_.set_scale(hypers->scale);
  hypers_.set_mh_mean_var(hypers->mh_mean_var);
  hypers_.set_mh_log_scale_var(hypers->mh_log_scale_var);

  auto out = std::make_shared<bayesmix::AlgorithmState::HierarchyHypers>();
  out->mutable_lapnig_state()->CopyFrom(hypers_);
  return out;
}

void LapNIGHierarchy::clear_summary_statistics() {
  cluster_data_values.clear();
  sum_abs_diff_curr = 0;
  sum_abs_diff_prop = 0;
}

void LapNIGHierarchy::initialize_hypers() {
  if (prior->has_fixed_values()) {
    // Set values
    hypers->mean = prior->fixed_values().mean();
    hypers->var = prior->fixed_values().var();
    hypers->shape = prior->fixed_values().shape();
    hypers->scale = prior->fixed_values().scale();
    hypers->mh_mean_var = prior->fixed_values().mh_mean_var();
    hypers->mh_log_scale_var = prior->fixed_values().mh_log_scale_var();
    // Check validity
    if (hypers->var <= 0) {
      throw std::invalid_argument("Variance-scaling parameter must be > 0");
    }
    if (hypers->shape <= 0) {
      throw std::invalid_argument("Shape parameter must be > 0");
    }
    if (hypers->scale <= 0) {
      throw std::invalid_argument("Scale parameter must be > 0");
    }
    if (hypers->mh_mean_var <= 0) {
      throw std::invalid_argument("mh_mean_var parameter must be > 0");
    }
    if (hypers->mh_log_scale_var <= 0) {
      throw std::invalid_argument("mh_log_scale_var parameter must be > 0");
    }
  } else {
    throw std::invalid_argument("Unrecognized hierarchy prior");
  }
}

void LapNIGHierarchy::initialize_state() {
  state.mean = hypers->mean;
  state.scale = hypers->scale / (hypers->shape + 1);  // mode of Inv-Gamma
}

void LapNIGHierarchy::update_hypers(
    const std::vector<bayesmix::AlgorithmState::ClusterState> &states) {
  auto &rng = bayesmix::Rng::Instance().get();
  if (prior->has_fixed_values()) {
    return;
  } else {
    throw std::invalid_argument("Unrecognized hierarchy prior");
  }
}

LapNIG::State LapNIGHierarchy::draw(const LapNIG::Hyperparams &params) {
  auto &rng = bayesmix::Rng::Instance().get();
  LapNIG::State out{};
  out.scale = stan::math::inv_gamma_rng(params.shape, 1. / params.scale, rng);
  out.mean = stan::math::normal_rng(params.mean, sqrt(params.var), rng);
  return out;
}

void LapNIGHierarchy::update_summary_statistics(
    const Eigen::RowVectorXd &datum, bool add) {
  if (add) {
    sum_abs_diff_curr += std::abs(state.mean - datum(0, 0));
    cluster_data_values.push_back(datum);
  } else {
    sum_abs_diff_curr -= std::abs(state.mean - datum(0, 0));
    auto it = std::find(cluster_data_values.begin(), cluster_data_values.end(),
                        datum);
    cluster_data_values.erase(it);
  }
}

void LapNIGHierarchy::sample_full_cond(bool update_params) {
  if (this->card == 0) {
    // No posterior update possible
    this->sample_prior();
  } else {
    // Number of iterations to compute the acceptance rate
    ++iter_;

    // Random generator
    auto &rng = bayesmix::Rng::Instance().get();

    // Candidate mean and candidate log_scale
    Eigen::VectorXd curr_unc_params(2);
    curr_unc_params << state.mean, std::log(state.scale);

    Eigen::VectorXd prop_unc_params = propose_rwmh(curr_unc_params);

    double log_target_prop =
        eval_prior_lpdf_unconstrained(prop_unc_params) +
        eval_like_lpdf_unconstrained(prop_unc_params, false);

    double log_target_curr =
        eval_prior_lpdf_unconstrained(curr_unc_params) +
        eval_like_lpdf_unconstrained(curr_unc_params, true);

    double log_a_rate = log_target_prop - log_target_curr;

    if (std::log(stan::math::uniform_rng(0, 1, rng)) < log_a_rate) {
      ++accepted_;
      state.mean = prop_unc_params(0);
      state.scale = std::exp(prop_unc_params(1));
      sum_abs_diff_curr = sum_abs_diff_prop;
    }
  }
}

Eigen::VectorXd LapNIGHierarchy::propose_rwmh(
    const Eigen::VectorXd &curr_vals) {
  auto &rng = bayesmix::Rng::Instance().get();
  double candidate_mean =
      curr_vals(0) + stan::math::normal_rng(0, sqrt(hypers->mh_mean_var), rng);
  double candidate_log_scale =
      curr_vals(1) +
      stan::math::normal_rng(0, sqrt(hypers->mh_log_scale_var), rng);
  Eigen::VectorXd proposal(2);
  proposal << candidate_mean, candidate_log_scale;
  return proposal;
}

double LapNIGHierarchy::eval_prior_lpdf_unconstrained(
    Eigen::VectorXd unconstrained_parameters) {
  double mu = unconstrained_parameters(0);
  double log_scale = unconstrained_parameters(1);
  double scale = std::exp(log_scale);
  return stan::math::normal_lpdf(mu, hypers->mean, std::sqrt(hypers->var)) +
         stan::math::inv_gamma_lpdf(scale, hypers->shape, hypers->scale) +
         log_scale;
}

double LapNIGHierarchy::eval_like_lpdf_unconstrained(
    Eigen::VectorXd unconstrained_parameters, bool is_current) {
  double mean = unconstrained_parameters(0);
  double log_scale = unconstrained_parameters(1);
  double scale = std::exp(log_scale);
  double diff_sum = 0;  // Sum of absolute values of data - candidate_mean
  if (is_current) {
    diff_sum = sum_abs_diff_curr;
  } else {
    for (auto &elem : cluster_data_values) {
      diff_sum += std::abs(elem(0, 0) - mean);
    }
    sum_abs_diff_prop = diff_sum;
  }
  return std::log(0.5 / scale) + (-0.5 / scale * diff_sum);
}