#include<iostream>
#include <Eigen/Eigenvalues>
#include"numerics.hpp"
#include"reddm.hpp"
#include"tpoperators.hpp"
#include "files.hpp"
#include"FTLanczos.hpp"
#include<omp.h>
#include <boost/program_options.hpp>
using namespace Many_Body;
using Mat= Operators::Mat;

using namespace boost::program_options;
int main(int argc, char *argv[])
{

  int M{};
  int L{};
  double t0{};
  double omega{};
  double gamma{};
  bool PB{0};
   int runs={};
   int Ldim={};
    double T{};
    double err{};
  std::string sLdim={};
  std::string sruns={};
std::string sT{};
std::string sM{};
std::string sL{};
std::string st0{};
std::string somega{};
std::string sgamma{};
std::string sPB{};
 std::string serr{};
 std::string filename="ompFTLM";

  try
  {
    options_description desc{"Options"};
    desc.add_options()
      ("help,h", "Help screen")
      ("L", value(&L)->default_value(4), "L")
      ("M", value(&M)->default_value(2), "M")
      ("r", value(&runs)->default_value(20), "r")
      ("Ld", value(&Ldim)->default_value(20), "Ld")
      ("t0", value(&t0)->default_value(1.), "t0")
      ("gam", value(&gamma)->default_value(1.), "gamma")
      ("omg", value(&omega)->default_value(1.), "omega")
      ("T", value(&T)->default_value(1.), "T")
      ("pb", value(&PB)->default_value(false), "PB")
      ("err", boost::program_options::value(&err)->default_value(1E-9), "err");
    boost::program_options::variables_map vm;
    boost::program_options::store(parse_command_line(argc, argv, desc), vm);
    boost::program_options::notify(vm);
 if (vm.count("help"))
      {std::cout << desc << '\n'; return 0;}
    else{
      if (vm.count("L"))
      {
  	std::cout << "L: " << vm["L"].as<int>() << '\n';
	sL="L"+std::to_string(vm["L"].as<int>());
	filename+=sL;

      }
     if (vm.count("M"))
      {
  	std::cout << "M: " << vm["M"].as<int>() << '\n';
	sM="M"+std::to_string(vm["M"].as<int>());
	filename+=sM;

      }
     if (vm.count("r"))
      {
  	std::cout << "runs: " << vm["r"].as<int>() << '\n';
	sruns="r"+std::to_string(vm["r"].as<int>());
	filename+=sruns;

      }
if (vm.count("Ld"))
      {
  	std::cout << "lanczos dim: " << vm["Ld"].as<int>() << '\n';

      }
      if (vm.count("t"))
      {
      	std::cout << "t0: " << vm["t"].as<double>() << '\n';
      	      	st0="t0"+std::to_string(vm["t0"].as<double>()).substr(0, 3);
      	filename+=st0;
      }
       if (vm.count("omg"))
      {
      	std::cout << "omega: " << vm["omg"].as<double>() << '\n';
      	      	somega="omg"+std::to_string(vm["omg"].as<double>()).substr(0, 3);
      	filename+=somega;
      }
       if (vm.count("gam"))
      {
      	std::cout << "gamma: " << vm["gam"].as<double>() << '\n';
      	sgamma="gam"+std::to_string(vm["gam"].as<double>()).substr(0, 3);
      	filename+=sgamma;
      }
       if (vm.count("T"))
      {
      	std::cout << "T: " << vm["T"].as<double>() << '\n';
      }if (vm.count("pb"))
      {
      	std::cout << "PB: " << vm["pb"].as<bool>() << '\n';
      	sPB="PB"+std::to_string(vm["pb"].as<bool>());
      	filename+=sPB;
      }
      		     if (vm.count("err"))
      {      std::cout << "error: " << vm["err"].as<double>() << '\n';
      	 std::stringstream ss;
      	 ss<<vm["err"].as<double>();
      	 serr="err"+ss.str();
      	filename+=serr;
      }
    }
  }
  catch (const error &ex)
  {
    std::cerr << ex.what() << '\n';
    return 0;
  }

     using HolsteinBasis= TensorProduct<ElectronBasis, PhononBasis>;

  filename+=".bin";
     std::vector<double> Tem;
     std::vector<double> beta;
     for(int i=1; i<11; i++)
       {
	 	  Tem.push_back((0.1*i));
	  beta.push_back(1./(0.1*i));
       }

    Mat H;
    Mat N;
    Mat EK;
    Mat X;
  std::vector<Mat> obs;
  {

    ElectronBasis e( L, 1);
  PhononBasis ph(L, M);
  HolsteinBasis TP(e, ph);
  std::cout<< TP.dim<<std::endl;
        Mat E1=Operators::EKinOperatorL(TP, e, t0,PB);
       Mat Ebdag=Operators::NBosonCOperator(TP, ph, gamma, PB);
       Mat Eb=Operators::NBosonDOperator(TP, ph, gamma, PB);
      Mat Eph=Operators::NumberOperator(TP, ph, omega,  PB);

      Eigen::VectorXd eigenVals(TP.dim);
       H=E1+Eb+Eph+Ebdag;
        N=Eph/omega;
	EK=E1;
	X=(Ebdag+Eb)/gamma;
       obs.push_back(H);
       obs.push_back(N);
       obs.push_back(E1);
       obs.push_back(X);
  }

     Eigen::MatrixXd Astot=Eigen::MatrixXd::Zero(beta.size(), obs.size());
Eigen::VectorXd Zstot=Eigen::VectorXd::Zero(beta.size());
#pragma omp parallel
  {

#pragma omp for
 for(int i=0; i<runs; i++)
   {
     //     std::cout<< " om num thread " <<omp_get_thread_num()<<std::endl;
     auto [Observables, SUMs]=calculate_lanczFT(obs[0], obs, beta, Ldim, err);

	#pragma omp critical
	{
	  Astot+=Observables;
	  Zstot+=SUMs;
	}
	//	std::cout << " got meanZ "<< SUMs.mean() << std::endl;

   }
  }

  std::cout<< "Astot  "<<std::endl<< Astot<< std::endl;
  for(int i=0; i<beta.size(); i++)
    {
      Astot.row(i)/=Zstot(i);
  	   std::cout<<" T "<< 1./beta[i] << "  "<<Astot(i, 0)<<" SUM "<< Astot(i, 1)+Astot(i, 2)+Astot(i, 3)*gamma<<std::endl;
    }
  bin_write("E"+filename, Eigen::VectorXd(Astot.col(0)));
  bin_write("Nph"+filename,  Eigen::VectorXd(Astot.col(1)));
  bin_write("EK"+filename, Eigen::VectorXd(Astot.col(2)));
  bin_write("nX"+filename, Eigen::VectorXd(Astot.col(3)));
  bin_write("temp"+filename, Tem);

  return 0;
}


#pragma once
#include <mtao/types.h>
#include <mtao/eigen_utils.h>
#include <Eigen/Sparse>
#include <boost/hana.hpp>

//Use DECMesh by inheriting it and manually implementing boundary / cell volumes / dual cell volumes / the cell counts per type
//
//CachedDECMesh caches the boundary computatoins / cell volume information for when it's expensive to compute these things
//
//DECMeshCore implements DEC h(hodge dual) d(exterior derivative) cod(codifferential operator as well as
//  the various laplacians (d cod / cod d / Laplace deRham (cod d + d cod) / weak laplacian (down) d cod h / weak laplacian (up) h cod d )
//  Note: standard cotan laplacian is weak_laplacian_up<0>()

namespace mtao { namespace geometry { namespace mesh {
template <typename Derived, typename T, int EmbeddedDim, int Dim_>
class DECMeshCore {
    public:
        using types = mtao::embedded_types<T,EmbeddedDim>;
        constexpr static int Dim = Dim_;
        using Scalar = typename types::Scalar;
        using VecX = typename types::VectorX;
        using SparseMatrix = Eigen::SparseMatrix<Scalar>;

        Derived& derived() { return *static_cast<Derived*>(this); }
        const Derived& derived() const { return *static_cast<const Derived*>(this); }

        template <int D>
        SparseMatrix boundary() const { return derived().template boundary<D>(); }
        template <int D>
        VecX volume() const { return derived().template volume<D>(); }
        template <int D>
        VecX dual_volume() const { return derived().template dual_volume<D>(); }

        //base operators
        template <int D>
        VecX h() const;
        template <int D>
        VecX hi() const;
        template <int D>
        SparseMatrix d() const;
        template <int D>
        SparseMatrix cod() const;

        template <int D>
        VecX h(const VecX& o) const { return h<D>().asDiagonal() * o; }
        //inverse hodge
        template <int D>
        VecX hi(const VecX& o) const { return hi<D>().asDiagonal() * o; }
        template <int D>
        VecX d(const VecX& o) const { return d<D>() * o; }
        //hi<D>dt<D-1>h<D>
        template <int D>
        VecX cod(const VecX& o) const { return cod<D>() * o; }

        //laplacians
        //d<D-1>hi<D>dt<D-1>h<D>
        template <int D>
        SparseMatrix laplacian_down() const;
        //hi<D>dt<D>h<D+1>d<D>
        template <int D>
        SparseMatrix laplacian_up() const;
        //d<D-1>hi<D>dt<D-1>h<D> + hi<D>dt<D>h<D+1>d<D>
        template <int D>
        SparseMatrix laplacian() const;
        //d<D-1>hi<D>dt<D-1>   (h<D>)
        template <int D>
        SparseMatrix weak_laplacian_down() const;
        //(hi<D>)    dt<D>h<D+1>d<D>
        template <int D>
        SparseMatrix weak_laplacian_up() const;

        template <int D>
        VecX weak_poisson_down(const VecX& o) const;
        template <int D>
        VecX weak_poisson_up(const VecX& o) const;
        template <int D>
        VecX weak_poisson_down_rhs(const VecX& o) const;
        template <int D>
        VecX weak_poisson_up_rhs(const VecX& o) const;

        //Find the ___ part according to HH decomposition
        //i.e o = da + hdhb + g

        //a
        template <int D>
        VecX coexact_down(const VecX& o) const ;
        //b
        template <int D>
        VecX exact_up(const VecX& o) const ;
        //da
        template <int D>
        VecX exact(const VecX& o) const ;
        //hdhb
        template <int D>
        VecX coexact(const VecX& o) const ;
        //da + g i.e o - hdhb
        template <int D>
        VecX closed(const VecX& o) const ;
        //hdhb + g i.e o - da
        template <int D>
        VecX coclosed(const VecX& o) const ;

};

template <typename Derived, typename T, int EmbeddedDim, int Dim_>
class DECMesh: public DECMeshCore<DECMesh<Derived, T, EmbeddedDim, Dim_>, T, EmbeddedDim, Dim_> {
    public:
        using types = mtao::embedded_types<T,EmbeddedDim>;
        constexpr static int Dim = Dim_;
        using Scalar = typename types::Scalar;
        using VecX = typename types::VectorX;
        using SparseMatrix = Eigen::SparseMatrix<Scalar>;
        using Mesh = DECMeshCore<DECMesh<Derived,T,EmbeddedDim,Dim>,T,EmbeddedDim,Dim>;

        Derived& derived() { return *static_cast<Derived*>(this); }
        const Derived& derived() const { return *static_cast<const Derived*>(this); }

        template <int D>
        SparseMatrix boundary() const { return derived().template boundary<D>(); }
        /*
        template <int C>
        auto volume(int i) const { return derived().template volume<C>(i); }
        template <int C>
        auto dual_volume(int i) const { return derived().template dual_volume<C>(i); }
        */
        template <int C>
        int cell_count() const { return derived().template cell_count<C>(); }
        template <int C>
        int dual_cell_count() const { return derived().template dual_cell_count<C>(); }
        template <int D>
        VecX volume() const { return derived().template volume<D>(); }
        template <int D>
        VecX dual_volume() const { return derived().template dual_volume<D>(); }
};
template <typename T, int EmbeddedDim, int Dim_>
class CachedDECMesh: public DECMesh<CachedDECMesh<T, EmbeddedDim,Dim_>, T, EmbeddedDim, Dim_> {
    public:
        using types = mtao::embedded_types<T,EmbeddedDim>;
        constexpr static int Dim = Dim_;
        using Scalar = typename types::Scalar;
        using VecX = typename types::VectorX;
        using SparseMatrix = Eigen::SparseMatrix<Scalar>;
        template <typename Derived>
            CachedDECMesh(const DECMesh<Derived,T,EmbeddedDim,Dim>& d) {
                using namespace boost;
                boost::hana::for_each(hana::make_range(hana::int_c<0>,hana::int_c<Dim+1>), [&](auto t) {
                        constexpr static int D = decltype(t)::value;
                        m_volume[D] = d.template volume<D>();
                        m_dual_volume[D] = d.template dual_volume<D>();
                        });
                boost::hana::for_each(hana::make_range(hana::int_c<1>,hana::int_c<Dim+1>), [&](auto t) {
                        constexpr static int D = decltype(t)::value;
                        m_boundary[D-1] = d.template boundary<D>();

                        });
            }

        template <int C>
        const SparseMatrix& boundary() const { return m_boundary[C-1]; }
        template <int C>
        auto volume(int i) const { return volume<C>()(i); }
        template <int C>
        auto dual_volume(int i) const { return dual_volume<C>()(i); }
        template <int C>
        int cell_count() const { return m_volume[C].rows();}
        template <int C>
        int dual_cell_count() const { return m_dual_volume[C].rows();}
        template <int C>
        const VecX& volume() const { return m_volume[C]; }
        template <int C>
        const VecX& dual_volume() const { return m_dual_volume[C]; }
    private:
        std::array<SparseMatrix, Dim> m_boundary;
        std::array<VecX, Dim+1> m_volume;
        std::array<VecX, Dim+1> m_dual_volume;

};

/*
template <typename Derived, typename T, int EmbeddedDim, int Dim>
template <int D>
auto DECMesh<Derived,T,EmbeddedDim,Dim>::volume() const -> VecX {
    int s = cell_count<D>();
    VecX v = VecX::Zero(s);
    for(int i = 0; i < s; ++i) {
        v(i) = volume<D>(i);
    }
    return v;
}
template <typename Derived, typename T, int EmbeddedDim, int Dim>
template <int D>
auto DECMesh<Derived,T,EmbeddedDim,Dim>::dual_volume() const -> VecX {
    int s = dual_cell_count<D>();
    VecX v = VecX::Zero(s);
    for(int i = 0; i < s; ++i) {
        v(i) = dual_volume<D>(i);
    }
    return v;
}
*/

//Base operators
template <typename Derived, typename T, int EmbeddedDim, int Dim> template <int D>
auto DECMeshCore<Derived,T,EmbeddedDim,Dim>::h() const -> VecX {
    auto p = volume<D>();
    auto d = dual_volume<D>();

    VecX ret = d.cwiseQuotient(p);
    return mtao::eigen::finite(ret);
}
template <typename Derived, typename T, int EmbeddedDim, int Dim> template <int D>
auto DECMeshCore<Derived,T,EmbeddedDim,Dim>::hi() const -> VecX {
    auto p = volume<D>();
    auto d = dual_volume<D>();

    VecX ret = p.cwiseQuotient(d);
    return mtao::eigen::finite(ret);
}
template <typename Derived, typename T, int EmbeddedDim, int Dim> template <int D>
auto DECMeshCore<Derived,T,EmbeddedDim,Dim>::d() const -> SparseMatrix {
    return boundary<D+1>().transpose();
}

// * d *
template <typename Derived, typename T, int EmbeddedDim, int Dim> template <int D>
auto DECMeshCore<Derived,T,EmbeddedDim,Dim>::cod() const -> SparseMatrix {
    return hi<D-1>().asDiagonal() * d<D-1>().transpose() * h<D>().asDiagonal();
}

//Laplacians
//d * d *
template <typename Derived, typename T, int EmbeddedDim, int Dim> template <int D>
auto DECMeshCore<Derived,T,EmbeddedDim,Dim>:: laplacian_down() const -> SparseMatrix {
    return d<D-1>() * cod<D>();
}
//* d * d
template <typename Derived, typename T, int EmbeddedDim, int Dim> template <int D>
auto DECMeshCore<Derived,T,EmbeddedDim,Dim>:: laplacian_up() const -> SparseMatrix {
    //h<1>
    return cod<D+1>() * d<D>();
}
template <typename Derived, typename T, int EmbeddedDim, int Dim> template <int D>
auto DECMeshCore<Derived,T,EmbeddedDim,Dim>:: laplacian() const -> SparseMatrix {
    if constexpr(D == 0) {
        return laplacian_up<D>();
    } else if(D == Dim) {
        return laplacian_down<D>();
    } else {
        return laplacian_down<D>() + laplacian_down<D>();
    }
}
//d * d that goes up down a dimension
template <typename Derived, typename T, int EmbeddedDim, int Dim> template <int D>
auto DECMeshCore<Derived,T,EmbeddedDim,Dim>:: weak_laplacian_down() const -> SparseMatrix {
    auto a = d<D-1>();
    return a * hi<D-1>().asDiagonal() * a.transpose();
}
//d * d that goes up a dimension
template <typename Derived, typename T, int EmbeddedDim, int Dim> template <int D>
auto DECMeshCore<Derived,T,EmbeddedDim,Dim>:: weak_laplacian_up() const -> SparseMatrix {
    auto a = d<D>();
    return a.transpose() * h<D+1>().asDiagonal() * a;
}
//d x
template <typename Derived, typename T, int EmbeddedDim, int Dim> template <int D>
auto DECMeshCore<Derived,T,EmbeddedDim,Dim>::weak_poisson_down_rhs(const VecX& o) const -> VecX {
    return d<D>(o);
}
//d * x
template <typename Derived, typename T, int EmbeddedDim, int Dim> template <int D>
auto DECMeshCore<Derived,T,EmbeddedDim,Dim>::weak_poisson_up_rhs(const VecX& o) const -> VecX {
    return d<D-1>().transpose() * h<D>(o);
}
//(d * d )^{-1} ( d x )
template <typename Derived, typename T, int EmbeddedDim, int Dim> template <int D>
auto DECMeshCore<Derived,T,EmbeddedDim,Dim>::weak_poisson_down(const VecX& rhs) const -> VecX {
    auto L = weak_laplacian_down<D>();
    Eigen::ConjugateGradient<SparseMatrix, Eigen::Upper|Eigen::Lower> solver(L);
    VecX p = solver.solveWithGuess(rhs,VecX::Zero(rhs.rows()));
    return p;
}
//(d * d )^{-1} ( d * x )
template <typename Derived, typename T, int EmbeddedDim, int Dim> template <int D>
auto DECMeshCore<Derived,T,EmbeddedDim,Dim>::weak_poisson_up(const VecX& rhs) const -> VecX {
    auto L = weak_laplacian_up<D>();
    Eigen::ConjugateGradient<SparseMatrix, Eigen::Upper|Eigen::Lower> solver(L);
    VecX p = solver.solveWithGuess(rhs,VecX::Zero(rhs.rows()));
    return p;
}

//Hodge decomposition components
template <typename Derived, typename T, int EmbeddedDim, int Dim> template <int D>
auto DECMeshCore<Derived,T,EmbeddedDim,Dim>::exact_up(const VecX& o) const -> VecX {
    //d (u - HdHp) = 0
    //dHdHp = du
    //d^THdHp = du
    //Hp =(d^THd)^{-1} du

    //d<D> | hi<D>dt<D>   (h<D+1>)

    return hi<D+1>(weak_poisson_down<D+1>(weak_poisson_down_rhs<D>(o)));
}
template <typename Derived, typename T, int EmbeddedDim, int Dim> template <int D>
auto DECMeshCore<Derived,T,EmbeddedDim,Dim>::coexact_down(const VecX& o) const -> VecX {
    //HdH (u - dp) = 0
    //HdHdp = HdHu
    //dHd^Tp = d^THu
    //p =(dHd^T)^{-1} dHu

    //(hi<D>)    dt<D-1>h<D> | d<D-1>
    return weak_poisson_up<D-1>(weak_poisson_up_rhs<D>(o));
}
template <typename Derived, typename T, int EmbeddedDim, int Dim> template <int D>
auto DECMeshCore<Derived,T,EmbeddedDim,Dim>::exact(const VecX& o) const -> VecX {
    return d<D-1>(coexact_down<D>(o));
}
template <typename Derived, typename T, int EmbeddedDim, int Dim> template <int D>
auto DECMeshCore<Derived,T,EmbeddedDim,Dim>::coexact(const VecX& o) const -> VecX {
    return cod<D+1>(exact_up<D>(o));
}
template <typename Derived, typename T, int EmbeddedDim, int Dim> template <int D>
auto DECMeshCore<Derived,T,EmbeddedDim,Dim>::closed(const VecX& o) const -> VecX {
    return o - coexact<D>(o);
}
template <typename Derived, typename T, int EmbeddedDim, int Dim> template <int D>
auto DECMeshCore<Derived,T,EmbeddedDim,Dim>::coclosed(const VecX& o) const -> VecX {
    return o - exact<D>(o);
}
}}}


// This file is part of the pyMOR project (http://www.pymor.org).
// Copyright 2013-2018 pyMOR developers and contributors. All rights reserved.
// License: BSD 2-Clause License (http://opensource.org/licenses/BSD-2-Clause)

#include "elasticity.hh"

#include <deal.II/lac/solver_bicgstab.h>
#include <deal.II/lac/solver_cg.h>
#include <deal.II/numerics/solution_transfer.h>

#include <fstream>

//! Wrapper intended for use w/ CG algorithms
template <class Number>
class MatrixSum {

public:
  typedef Number value_type;
  typedef std::vector<const dealii::SparseMatrix<Number>*> Matrices;
  MatrixSum(Matrices&& m, std::vector<Number> weights)
    : matrices_(m) {
    assert(m.size() > 0);
    assert(m.size() == weights.size());
    sum_.reinit(matrices_[0]->get_sparsity_pattern());
    sum_.copy_from(*matrices_[0]);
    sum_ *= weights[0];
    for (size_t i = 1; i < m.size(); ++i) {
      sum_.add(weights[i], *matrices_[1]);
    }
  }

  dealii::SparseMatrix<Number>& sum() { return sum_; }

  const dealii::SparseMatrix<Number>& sum() const { return sum_; }

private:
  const Matrices matrices_;
  dealii::SparseMatrix<Number> sum_;
};

ElasticityExample::ElasticityExample(int refine_steps)
  : dof_handler_(triangulation_)
  , fe_(FE_Q<dim>(1), dim) {
  GridGenerator::hyper_cube(triangulation_, -1, 1);
  refine_global(refine_steps);
}

ElasticityExample::~ElasticityExample() { dof_handler_.clear(); }

void ElasticityExample::setup_system() {
  dof_handler_.clear();
  dof_handler_.distribute_dofs(fe_);
  sparsity_pattern_.reinit(dof_handler_.n_dofs(), dof_handler_.n_dofs(), dof_handler_.max_couplings_between_dofs());
  DoFTools::make_sparsity_pattern(dof_handler_, sparsity_pattern_);

  sparsity_pattern_.compress();

  h1_matrix_.reinit(sparsity_pattern_);
  lambda_system_matrix_.reinit(sparsity_pattern_);
  mu_system_matrix_.reinit(sparsity_pattern_);
  system_rhs_.reinit(dof_handler_.n_dofs());
  tmp_data_.reinit(dof_handler_.n_dofs());
}

void ElasticityExample::assemble_h1() {
  QGauss<dim> quadrature_formula(2);

  FEValues<dim> fe_values(fe_, quadrature_formula,
                          update_values | update_gradients | update_quadrature_points | update_JxW_values);

  const unsigned int dofs_per_cell = fe_.dofs_per_cell;
  const unsigned int n_q_points = quadrature_formula.size();

  FullMatrix<Number> h1_cell_matrix(dofs_per_cell, dofs_per_cell);
  std::vector<types::global_dof_index> local_dof_indices(dofs_per_cell);

  // Now we can begin with the loop over all cells:
  typename DoFHandler<dim>::active_cell_iterator cell = dof_handler_.begin_active(), endc = dof_handler_.end();
  for (; cell != endc; ++cell) {
    h1_cell_matrix = 0;
    fe_values.reinit(cell);

    for (unsigned int i = 0; i < dofs_per_cell; ++i) {
      const unsigned int component_i = fe_.system_to_component_index(i).first;

      for (unsigned int j = 0; j < dofs_per_cell; ++j) {
        const unsigned int component_j = fe_.system_to_component_index(j).first;

        for (unsigned int q_point = 0; q_point < n_q_points; ++q_point) {
          h1_cell_matrix(i, j) += fe_values.shape_grad(i, q_point)[component_i] *
                                  fe_values.shape_grad(j, q_point)[component_j] * fe_values.JxW(q_point);
        }
      }
    }

    cell->get_dof_indices(local_dof_indices);
    for (unsigned int i = 0; i < dofs_per_cell; ++i) {
      for (unsigned int j = 0; j < dofs_per_cell; ++j) {
        h1_matrix_.add(local_dof_indices[i], local_dof_indices[j], h1_cell_matrix(i, j));
      }
    }
  }

  std::map<types::global_dof_index, Number> boundary_values;
  VectorTools::interpolate_boundary_values(dof_handler_, 0, ZeroFunction<dim>(dim), boundary_values);
  MatrixTools::apply_boundary_values(boundary_values, h1_matrix_, tmp_data_, system_rhs_);
}

void ElasticityExample::assemble_system() {
  QGauss<dim> quadrature_formula(2);

  FEValues<dim> fe_values(fe_, quadrature_formula,
                          update_values | update_gradients | update_quadrature_points | update_JxW_values);

  const unsigned int dofs_per_cell = fe_.dofs_per_cell;
  const unsigned int n_q_points = quadrature_formula.size();

  FullMatrix<Number> lambda_cell_matrix(dofs_per_cell, dofs_per_cell), mu_cell_matrix(dofs_per_cell, dofs_per_cell);
  Vector<Number> cell_rhs(dofs_per_cell);

  std::vector<types::global_dof_index> local_dof_indices(dofs_per_cell);

  // As was shown in previous examples as well, we need a place where to
  // store the values of the coefficients at all the quadrature points on a
  // cell. In the present situation, we have two coefficients, lambda and
  // mu.
  std::vector<Number> lambda_values(n_q_points);
  std::vector<Number> mu_values(n_q_points);

  // Well, we could as well have omitted the above two arrays since we will
  // use constant coefficients for both lambda and mu, which can be declared
  // like this. They both represent functions always returning the constant
  // value 1.0. Although we could omit the respective factors in the
  // assemblage of the matrix, we use them here for purpose of
  // demonstration.
  Functions::ConstantFunction<dim> lambda(1), mu(1);

  // Then again, we need to have the same for the right hand side. This is
  // exactly as before in previous examples. However, we now have a
  // vector-valued right hand side, which is why the data type of the
  // rhs_values array is changed. We initialize it by
  // n_q_points elements, each of which is a
  // Vector@<double@> with dim elements.
  RightHandSide<dim> right_hand_side;
  std::vector<Vector<Number>> rhs_values(n_q_points, Vector<Number>(dim));

  // Now we can begin with the loop over all cells:
  typename DoFHandler<dim>::active_cell_iterator cell = dof_handler_.begin_active(), endc = dof_handler_.end();
  for (; cell != endc; ++cell) {
    lambda_cell_matrix = 0;
    mu_cell_matrix = 0;
    cell_rhs = 0;

    fe_values.reinit(cell);

    // Next we get the values of the coefficients at the quadrature
    // points. Likewise for the right hand side:
    lambda.value_list(fe_values.get_quadrature_points(), lambda_values);
    mu.value_list(fe_values.get_quadrature_points(), mu_values);

    right_hand_side.vector_value_list(fe_values.get_quadrature_points(), rhs_values);

    // Then assemble the entries of the local stiffness matrix and right
    // hand side vector. This follows almost one-to-one the pattern
    // described in the introduction of this example.  One of the few
    // comments in place is that we can compute the number
    // comp(i), i.e. the index of the only nonzero vector
    // component of shape function i using the
    // fe.system_to_component_index(i).first function call
    // below.
    //
    // (By accessing the first variable of the return value
    // of the system_to_component_index function, you might
    // already have guessed that there is more in it. In fact, the
    // function returns a std::pair@<unsigned int, unsigned
    // int@>, of which the first element is comp(i)
    // and the second is the value base(i) also noted in the
    // introduction, i.e.  the index of this shape function within all the
    // shape functions that are nonzero in this component,
    // i.e. base(i) in the diction of the introduction. This
    // is not a number that we are usually interested in, however.)
    //
    // With this knowledge, we can assemble the local matrix
    // contributions:
    for (unsigned int i = 0; i < dofs_per_cell; ++i) {
      const unsigned int component_i = fe_.system_to_component_index(i).first;

      for (unsigned int j = 0; j < dofs_per_cell; ++j) {
        const unsigned int component_j = fe_.system_to_component_index(j).first;

        for (unsigned int q_point = 0; q_point < n_q_points; ++q_point) {
          lambda_cell_matrix(i, j) +=
              // The first term is (lambda d_i u_i, d_j v_j) + (mu d_i
              // u_j, d_j v_i).  Note that
              // shape_grad(i,q_point) returns the
              // gradient of the only nonzero component of the i-th
              // shape function at quadrature point q_point. The
              // component comp(i) of the gradient, which
              // is the derivative of this only nonzero vector
              // component of the i-th shape function with respect to
              // the comp(i)th coordinate is accessed by the appended
              // brackets.
              fe_values.shape_grad(i, q_point)[component_i] * fe_values.shape_grad(j, q_point)[component_j] *
              lambda_values[q_point] * fe_values.JxW(q_point);
          mu_cell_matrix(i, j) +=
              (fe_values.shape_grad(i, q_point)[component_j] * fe_values.shape_grad(j, q_point)[component_i] *
                   mu_values[q_point] +
               // The second term is (mu nabla u_i, nabla v_j).  We
               // need not access a specific component of the
               // gradient, since we only have to compute the scalar
               // product of the two gradients, of which an
               // overloaded version of the operator* takes care, as
               // in previous examples.
               //
               // Note that by using the ?: operator, we only do this
               // if comp(i) equals comp(j), otherwise a zero is
               // added (which will be optimized away by the
               // compiler).
               ((component_i == component_j)
                    ? (fe_values.shape_grad(i, q_point) * fe_values.shape_grad(j, q_point) * mu_values[q_point])
                    : 0)) *
              fe_values.JxW(q_point);
        }
      }
    }

    // Assembling the right hand side is also just as discussed in the
    // introduction:
    for (unsigned int i = 0; i < dofs_per_cell; ++i) {
      const unsigned int component_i = fe_.system_to_component_index(i).first;

      for (unsigned int q_point = 0; q_point < n_q_points; ++q_point)
        cell_rhs(i) += fe_values.shape_value(i, q_point) * rhs_values[q_point](component_i) * fe_values.JxW(q_point);
    }

    // The transfer from local degrees of freedom into the global matrix
    // and right hand side vector does not depend on the equation under
    // consideration, and is thus the same as in all previous
    // examples. The same holds for the elimination of hanging nodes from
    // the matrix and right hand side, once we are done with assembling
    // the entire linear system:
    cell->get_dof_indices(local_dof_indices);
    for (unsigned int i = 0; i < dofs_per_cell; ++i) {
      for (unsigned int j = 0; j < dofs_per_cell; ++j) {
        lambda_system_matrix_.add(local_dof_indices[i], local_dof_indices[j], lambda_cell_matrix(i, j));
        mu_system_matrix_.add(local_dof_indices[i], local_dof_indices[j], mu_cell_matrix(i, j));
      }

      system_rhs_(local_dof_indices[i]) += cell_rhs(i);
    }
  }
  // The interpolation of the boundary values needs a small modification:
  // since the solution function is vector-valued, so need to be the
  // boundary values. The ZeroFunction constructor accepts a
  // parameter that tells it that it shall represent a vector valued,
  // constant zero function with that many components. By default, this
  // parameter is equal to one, in which case the ZeroFunction
  // object would represent a scalar function. Since the solution vector has
  // dim components, we need to pass dim as number
  // of components to the zero function as well.
  std::map<types::global_dof_index, Number> boundary_values;
  VectorTools::interpolate_boundary_values(dof_handler_, 0, ZeroFunction<dim>(dim), boundary_values);
  MatrixTools::apply_boundary_values(boundary_values, mu_system_matrix_, tmp_data_, system_rhs_);
  MatrixTools::apply_boundary_values(boundary_values, lambda_system_matrix_, tmp_data_, system_rhs_);
}

void ElasticityExample::_solve(Parameter param, VectorType& solution) {
  SolverControl solver_control(20000, 1e-12);
  SolverCG<> cg(solver_control);

  deallog << "Solving for " << dof_handler_.n_dofs() << " unknowns" << std::endl;
  Number lambda(param["lambda"][0]), mu(param["mu"][0]);
  MatrixSum<Number> msum({&lambda_system_matrix_, &mu_system_matrix_}, {lambda, mu});

  auto& sum = msum.sum();
  PreconditionSSOR<> preconditioner;
  preconditioner.initialize(sum, 1.2);
  cg.solve(sum, solution, system_rhs_, preconditioner);
}

void ElasticityExample::visualize(const VectorType& solution, std::string filename) const {
  std::ofstream output(filename);

  DataOut<dim> data_out;
  data_out.attach_dof_handler(dof_handler_);

  // As said above, we need a different name for each component of the
  // solution function. To pass one name for each component, a vector of
  // strings is used. Since the number of components is the same as the
  // number of dimensions we are working in, the following
  // switch statement is used.
  //
  // We note that some graphics programs have restriction as to what
  // characters are allowed in the names of variables. The library therefore
  // supports only the minimal subset of these characters that is supported
  // by all programs. Basically, these are letters, numbers, underscores,
  // and some other characters, but in particular no whitespace and
  // minus/hyphen. The library will throw an exception otherwise, at least
  // if in debug mode.
  //
  // After listing the 1d, 2d, and 3d case, it is good style to let the
  // program die if we run upon a case which we did not consider. Remember
  // that the Assert macro generates an exception if the
  // condition in the first parameter is not satisfied. Of course, the
  // condition false can never be satisfied, so the program
  // will always abort whenever it gets to the default statement:
  std::vector<std::string> solution_names;
  switch (dim) {
    case 1:
      solution_names.push_back("displacement");
      break;
    case 2:
      solution_names.push_back("x_displacement");
      solution_names.push_back("y_displacement");
      break;
    case 3:
      solution_names.push_back("x_displacement");
      solution_names.push_back("y_displacement");
      solution_names.push_back("z_displacement");
      break;
    default:
      throw ExcNotImplemented();
  }

  // After setting up the names for the different components of the solution
  // vector, we can add the solution vector to the list of data vectors
  // scheduled for output. Note that the following function takes a vector
  // of strings as second argument, whereas the one which we have used in
  // all previous examples accepted a string there. In fact, the latter
  // function is only a shortcut for the function which we call here: it
  // puts the single string that is passed to it into a vector of strings
  // with only one element and forwards that to the other function.
  data_out.add_data_vector(solution, solution_names);
  data_out.build_patches();
  data_out.write_vtk(output);
}

void ElasticityExample::refine_global(int refine_steps) {
  triangulation_.refine_global(refine_steps);
  setup_system();
  assemble_h1();
  assemble_system();
}

ElasticityExample::VectorType ElasticityExample::solve(const ElasticityExample::Parameter& param) {
  _solve(param, tmp_data_);
  return VectorType(tmp_data_);
}

const dealii::SparseMatrix<ElasticityExample::Number>& ElasticityExample::lambda_mat() const {
  return lambda_system_matrix_;
}

const dealii::SparseMatrix<ElasticityExample::Number>& ElasticityExample::mu_mat() const { return mu_system_matrix_; }

const dealii::SparseMatrix<ElasticityExample::Number>& ElasticityExample::h1_mat() const { return h1_matrix_; }

const dealii::Vector<ElasticityExample::Number>& ElasticityExample::rhs() const { return system_rhs_; }

ElasticityExample::Number ElasticityExample::h1_0_semi_norm(const Vector<Number>& v) const {
  return std::sqrt(h1_matrix_.matrix_norm_square(v));
}

ElasticityExample::Number ElasticityExample::energy_norm(const Vector<ElasticityExample::Number>& v) const {
  return std::sqrt(mu_system_matrix_.matrix_norm_square(v));
}

ElasticityExample::VectorType ElasticityExample::transfer_to(int refine_steps, const ElasticityExample::VectorType& v) {
  triangulation_.prepare_coarsening_and_refinement();
  dealii::SolutionTransfer<dim, VectorType> trans(dof_handler_);
  trans.prepare_for_pure_refinement();
  refine_global(refine_steps);

  VectorType ret(dof_handler_.n_dofs());
  trans.refine_interpolate(v, ret);
  return ret;
}


#include <cmath>
#include <boost/math/special_functions/digamma.hpp>
#include <boost/math/special_functions/fpclassify.hpp>

#include "constants.hpp"
#include "fastmath.hpp"
#include "logger.hpp"
#include "nlopt/nlopt.h"
#include "shredder.hpp"

double shredder_opt_objective(unsigned int _n, const double* _x,
                              double* _grad, void* data)
{
    UNUSED(_n);

    Shredder* sampler = reinterpret_cast<Shredder*>(data);
    if (_grad) {
        double fx = sampler->f(_x[0], _grad[0]);
        _grad[0] = std::max<double>(std::min<double>(_grad[0], 1e4), -1e4);
        return fx;
    }
    else {
        double d;
        return sampler->f(_x[0], d);
    }
}

static void assert_finite(double x)
{
    if (!boost::math::isfinite(x)) {
        Logger::abort("%f found where finite value expected.", x);
    }
}

Shredder::Shredder(double lower_limit, double upper_limit, double tolerance)
    : lower_limit(lower_limit)
    , upper_limit(upper_limit)
    , tolerance(tolerance)
    , opt(NULL)
{
    opt = nlopt_create(NLOPT_LD_SLSQP, 1);
    nlopt_set_lower_bounds(opt, &lower_limit);
    nlopt_set_upper_bounds(opt, &upper_limit);
    nlopt_set_max_objective(opt, shredder_opt_objective,
            reinterpret_cast<void*>(this));
    nlopt_set_maxeval(opt, 20);

    nlopt_set_ftol_abs(opt, 1e-7);
    nlopt_set_xtol_abs(opt, &tolerance);
}

Shredder::~Shredder()
{
    nlopt_destroy(opt);
}

void Shredder::set_tolerance(double tolerance)
{
    this->tolerance = tolerance;
    nlopt_set_xtol_abs(opt, &this->tolerance);
}

double Shredder::sample(rng_t& rng, double x0)
{
    double d0;
    double lp0 = f(x0, d0);

    assert_finite(lp0);

    double slice_height = fastlog(
            std::max<double>(constants::zero_eps, random_uniform_01(rng))) + lp0;
    assert_finite(slice_height);

    x_min = find_slice_edge(x0, slice_height, lp0, d0, -1);
    x_max = find_slice_edge(x0, slice_height, lp0, d0,  1);

    double x = (x_max + x_min) / 2;
    while (x_max - x_min > tolerance) {
        x = x_min + (x_max - x_min) * random_uniform_01(rng);
        double d;
        double lp = f(x, d);

        if (lp >= slice_height) break;
        else if (x > x0) x_max = x;
        else             x_min = x;
    }

    return x;
}

double Shredder::optimize(double x0)
{
    // bounds may have changed
    nlopt_set_lower_bounds(opt, &lower_limit);
    nlopt_set_upper_bounds(opt, &upper_limit);

    x0 = std::max<double>(std::min<double>(x0, upper_limit), lower_limit);
    //x0 = (lower_limit + upper_limit) / 2;
    double maxf;
    nlopt_result result = nlopt_optimize(opt, &x0, &maxf);

    if (result < 0 && result != NLOPT_ROUNDOFF_LIMITED) {
        Logger::warn("Optimization failed with code %d", (int) result);
    }

    return std::max<double>(std::min<double>(x0, upper_limit), lower_limit);
}

double Shredder::find_slice_edge(double x0, double slice_height,
                                 double lp0, double d0, int direction)
{
    const double lp_eps = 1e-2;
    const double d_eps  = 1e-3;

    // if newton method iterations are not making progress, resort to bisection
    size_t newton_count = 0;

    double lp = lp0 - slice_height;
    double d = d0;
    double x = x0;
    double x_bound_lower, x_bound_upper;

    if (direction < 0) {
        x_bound_lower = lower_limit;
        x_bound_upper = x0;
        double fx = f(lower_limit, d);
        if (boost::math::isfinite(fx) && fx >= slice_height) {
            return lower_limit;
        }
    }
    else {
        x_bound_lower = x0;
        x_bound_upper = upper_limit;
        double fx = f(upper_limit, d);
        if (boost::math::isfinite(fx) && fx >= slice_height) {
            return upper_limit;
        }
    }

    while (fabs(lp) > lp_eps && fabs(x_bound_upper - x_bound_lower) > tolerance) {
        double x1 = x - lp / d;
        if (isnan(d) || d == 0.0 || fabs(d) < d_eps || !boost::math::isfinite(x1)) {
            x1 = (x_bound_lower + x_bound_upper) / 2;
        }

        // if we are very close to the boundry, and this iteration moves us past
        // the boundry, just give up.
        if (direction < 0 && fabs(x - lower_limit) <= tolerance && (x1 < x || lp > 0.0)) break;
        if (direction > 0 && fabs(x - upper_limit) <= tolerance && (x1 > x || lp > 0.0)) break;

        // if we are moving in the wrong direction (i.e. toward the other root),
        // use bisection to correct course.
        if (direction < 0) {
            if (lp > 0) x_bound_upper = x;
            else        x_bound_lower = x;
        }
        else {
            if (lp > 0) x_bound_lower = x;
            else        x_bound_upper = x;
        }

        bool bisect = newton_count >= constants::max_newton_steps ||
            x1 < x_bound_lower + tolerance || x1 > x_bound_upper - tolerance;

        // try using the gradient
        if (!bisect) {
            x = x1;
            lp = f(x, d) - slice_height;
            bisect = !boost::math::isfinite(lp) || !boost::math::isfinite(d);
        }

        // resort to binary search if we seem not to be making progress
        if (bisect) {
            size_t iteration_count = 0;
            while (true) {
                x = (x_bound_lower + x_bound_upper) / 2;
                lp = f(x, d) - slice_height;

                //if (boost::math::isinf(lp) || boost::math::isinf(d)) {
                if (!boost::math::isfinite(lp)) {
                    if (direction < 0) x_bound_lower = x;
                    else               x_bound_upper = x;
                }
                else break;

                if (++iteration_count > 50) {
                    Logger::abort("Slice sampler edge finding is not making progress.");
                }
            }
    }
    else ++newton_count;

        assert_finite(lp);
    }

    assert_finite(x);

    return x;
}

static double sq(double x)
{
    return x * x;
}

static double cb(double x)
{
    return x * x * x;
}

static double lbeta(double x, double y)
{
    return lgamma(x) + lgamma(y) - lgamma(x + y);
}

#if 0
double PoissonLogPdf::f(float lambda, unsigned int k)
{
    return k * fastlog(lambda) - lgammaf(k + 1) - lambda;
}

double PoissonLogPdf::df_dlambda(float lambda, unsigned int k)
{
    return (float) k / lambda - 1;
}
#endif

static const double NEG_LOG_2_PI_DIV_2 = -log(2 * M_PI)/2;

double NormalLogPdf::f(double mu, double sigma, const double* xs, size_t n)
{
    double part1 = n * (NEG_LOG_2_PI_DIV_2 - fastlog(sigma));
    double part2 = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part2 += sq(xs[i] - mu) / (2 * sq(sigma));
    }

    return part1 - part2;
}

float NormalLogPdf::f(float mu, float sigma, const float* xs, size_t n)
{
    float part1 = n * (NEG_LOG_2_PI_DIV_2 - fastlog(sigma));
    float part2 = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part2 += sq(xs[i] - mu) / (2 * sq(sigma));
    }

    return part1 - part2;
}

double NormalLogPdf::f(double mu, double sigma, const double x)
{
    double part1 = NEG_LOG_2_PI_DIV_2 - fastlog(sigma);
    double part2 = sq(x - mu) / (2 * sq(sigma));
    return part1 - part2;
}

double NormalLogPdf::df_dx(double mu, double sigma, const double x)
{
    return (mu - x) / sq(sigma);
}

double NormalLogPdf::df_dx(double mu, double sigma, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += mu - xs[i];
    }
    return part / sq(sigma);
}

double NormalLogPdf::df_dmu(double mu, double sigma, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += xs[i] - mu;
    }
    return part / sq(sigma);
}

float NormalLogPdf::df_dmu(float mu, float sigma, const float* xs, size_t n)
{
    float part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += xs[i] - mu;
    }
    return part / sq(sigma);
}

double NormalLogPdf::df_dsigma(double mu, double sigma, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += sq(xs[i] - mu);
    }

    return part / cb(sigma) - n/sigma;
}

float NormalLogPdf::df_dsigma(float mu, float sigma, const float* xs, size_t n)
{
    float part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += sq(xs[i] - mu);
    }

    return part / cb(sigma) - n/sigma;
}

double LogNormalLogPdf::f(double mu, double sigma, const double* xs, size_t n)
{
    double part1 = n * (NEG_LOG_2_PI_DIV_2 - fastlog(sigma));
    double part2 = 0.0;
    for (size_t i = 0; i < n; ++i) {
        double logx = fastlog(xs[i]);
        part2 += sq(logx - mu) / (2 * sq(sigma)) + logx;
    }

    return part1 - part2;
}

double LogNormalLogPdf::df_dx(double mu, double sigma, double x)
{
    return (mu - fastlog(x)) / (x * sq(sigma)) - 1.0 / x;
}

double LogNormalLogPdf::df_dmu(double mu, double sigma, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += fastlog(xs[i]) - mu;
    }
    return part / sq(sigma);
}

double LogNormalLogPdf::df_dsigma(double mu, double sigma, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += sq(fastlog(xs[i]) - mu);
    }

    return part / cb(sigma) - n/sigma;
}

#if 0
double StudentsTLogPdf::f(double nu, double mu, double sigma, const double* xs, size_t n)
{
    double part1 =
        n * (lgamma((nu + 1) / 2) - lgamma(nu / 2) - fastlog(sqrt(nu * M_PI) * sigma));

    double part2 = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part2 += log1p(sq((xs[i] - mu) / sigma) / nu);
    }

    return part1 - ((nu + 1) / 2) * part2;
}

float StudentsTLogPdf::f(float nu, float mu, float sigma, const float* xs, size_t n)
{
    float part1 =
        n * (lgammaf((nu + 1) / 2) - lgamma(nu / 2) - fastlog(sqrt(nu * M_PI) * sigma));

    float part2 = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part2 += log1p(sq((xs[i] - mu) / sigma) / nu);
    }

    return part1 - ((nu + 1) / 2) * part2;
}

double StudentsTLogPdf::df_dx(double nu, double mu, double sigma, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += (2 * (xs[i] - mu) / sq(sigma) / nu) / (1 + sq((xs[i] - mu) / sigma) / nu);
    }

    return -((nu + 1) / 2) * part;
}

double StudentsTLogPdf::df_dmu(double nu, double mu, double sigma, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += (2 * (xs[i] - mu) / sq(sigma) / nu) / (1 + sq((xs[i] - mu) / sigma) / nu);
    }

    return ((nu + 1) / 2) * part;
}

float StudentsTLogPdf::df_dmu(float nu, float mu, float sigma, const float* xs, size_t n)
{
    float part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += (2 * (xs[i] - mu) / sq(sigma) / nu) / (1 + sq((xs[i] - mu) / sigma) / nu);
    }

    return ((nu + 1) / 2) * part;
}

double StudentsTLogPdf::df_dsigma(double nu, double mu, double sigma, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += (2 * sq((xs[i] - mu) / sigma) / (nu * sigma)) /
                    (1 + sq((xs[i] - mu) / sigma) / nu);
    }

    return ((nu + 1) / 2) * part - n / sigma;
}
#endif

#if 0
double GammaLogPdf::f(double alpha, double beta, const double* xs, size_t n)
{
    double part1 = 0.0, part2 = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part1 += fastlog(xs[i]);
        part2 += xs[i];
    }

    return
        n * (alpha * fastlog(beta) - lgamma(alpha)) +
        (alpha - 1) * part1 -
        beta * part2;
}

float GammaLogPdf::f(float alpha, float beta, const float* xs, size_t n)
{
    float part1 = 0.0, part2 = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part1 += fastlog(xs[i]);
        part2 += xs[i];
    }

    return
        n * (alpha * fastlog(beta) - lgammaf(alpha)) +
        (alpha - 1) * part1 -
        beta * part2;
}

double GammaLogPdf::df_dx(double alpha, double beta, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += (alpha - 1) / xs[i];
    }

    return part - n * beta;
}

double GammaLogPdf::df_dalpha(double alpha, double beta, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += fastlog(xs[i]);
    }

    return n * (fastlog(beta) - boost::math::digamma(alpha)) + part;
}

double GammaLogPdf::df_dbeta(double alpha, double beta, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += xs[i];
    }

    return (double) n * (alpha / beta) - part;
}

float GammaLogPdf::df_dbeta(float alpha, float beta, const float* xs, size_t n)
{
    float part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += xs[i];
    }

    return (float) n * (alpha / beta) - part;
}
#endif

#if 0
double AltGammaLogPdf::f(double mean, double shape, const double* xs, size_t n)
{
    double scale =  mean / shape;
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += (shape - 1.0) * fastlog(xs[i]) - xs[i] / scale;
    }

    double ans = -(double) n * (lgamma(shape) + shape * fastlog(scale)) + part;
    assert_finite(ans);
    return ans;
}

float AltGammaLogPdf::f(float mean, float shape, const float* xs, size_t n)
{
    float scale =  mean / shape;
    float part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += (shape - 1.0) * fastlog(xs[i]) - xs[i] / scale;
    }

    float ans = -(float) n * (lgammaf(shape) + shape * fastlog(scale)) + part;
    assert_finite(ans);
    return ans;
}

double AltGammaLogPdf::f(double mean, double shape, double x)
{
    double scale =  mean / shape;
    double part = (shape - 1.0) * fastlog(x) - x / scale;
    return -(lgamma(shape) + shape * fastlog(scale)) + part;
}

float AltGammaLogPdf::f(float mean, float shape, float x)
{
    float scale =  mean / shape;
    float part = (shape - 1.0) * fastlog(x) - x / scale;
    return -(lgammaf(shape) + shape * fastlog(scale)) + part;
}

double AltGammaLogPdf::df_dx(double mean, double shape, const double* xs, size_t n)
{
    double scale = mean / shape;
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += (shape - 1.0) / xs[i];
    }

    return part - (double) n / scale;
}

double AltGammaLogPdf::df_dx(double mean, double shape, double x)
{
    double scale = mean / shape;
    double part = (shape - 1.0) / x;
    return part - 1.0 / scale;
}

double AltGammaLogPdf::df_dmean(double mean, double shape, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += xs[i];
    }
    part *= shape / sq(mean);
    return part - (double) n * shape / mean;
}

float AltGammaLogPdf::df_dmean(float mean, float shape, const float* xs, size_t n)
{
    float part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += xs[i];
    }
    part *= shape / sq(mean);
    return part - (float) n * shape / mean;
}

double AltGammaLogPdf::df_dshape(double mean, double shape, const double* xs, size_t n)
{
    double scale = mean / shape;
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += fastlog(xs[i]) - xs[i] / mean;
    }

    return (double) n * (-boost::math::digamma(shape) + fastlog(scale) * (mean/sq(shape))) + part;
}

float AltGammaLogPdf::df_dshape(float mean, float shape, const float* xs, size_t n)
{
    float scale = mean / shape;
    float part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += fastlog(xs[i]) - xs[i] / mean;
    }

    return (float) n * (-boost::math::digamma(shape) + fastlog(scale) * (mean/sq(shape))) + part;
}

float AltGammaLogPdf::df_dshape(float mean, float shape, float x)
{
    float scale = mean / shape;
    float part = fastlog(x) - x / mean;
    return (-boost::math::digamma(shape) + fastlog(scale) * (mean/sq(shape))) + part;
}
#endif

double InvGammaLogPdf::f(double alpha, double beta, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += (alpha + 1) * fastlog(xs[i]) + beta / xs[i];
    }

    return n * (alpha * fastlog(beta) - lgamma(alpha)) - part;
}

double InvGammaLogPdf::df_dx(double alpha, double beta, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += beta / sq(xs[i]) - (alpha + 1) / xs[i];
    }

    return part;
}

double InvGammaLogPdf::df_dalpha(double alpha, double beta, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += fastlog(xs[i]);
    }

    return n * (fastlog(beta) - boost::math::digamma(alpha)) - part;
}

double InvGammaLogPdf::df_dbeta(double alpha, double beta, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += 1 / xs[i];
    }

    return n * (alpha / beta) - part;
}

double SqInvGammaLogPdf::f(double alpha, double beta, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        double x = xs[i] * xs[i];
        part += (alpha + 1) * fastlog(x) + beta / x;
    }

    return n * (alpha * fastlog(beta) - lgamma(alpha)) - part;
}

double SqInvGammaLogPdf::df_dx(double alpha, double beta, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        part += 2 * beta / cb(xs[i]) - (2 * alpha + 2) / xs[i];
    }

    return part;
}

double SqInvGammaLogPdf::df_dalpha(double alpha, double beta, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        double x = xs[i] * xs[i];
        part += fastlog(x);
    }

    return n * (fastlog(beta) - boost::math::digamma(alpha)) - part;
}

double SqInvGammaLogPdf::df_dbeta(double alpha, double beta, const double* xs, size_t n)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        double x = xs[i] * xs[i];
        part += 1 / x;
    }

    return n * (alpha / beta) - part;
}

double BetaLogPdf::f(double alpha, double beta, double x)
{
    return (alpha - 1) * fastlog(x) + (beta - 1) * fastlog(1 - x) - lbeta(alpha, beta);
}

double BetaLogPdf::df_dx(double alpha, double beta, double x)
{
    return (alpha - 1) / x - (beta - 1) / (1 - x);
}

double BetaLogPdf::df_dgamma(double gamma, double c, double x)
{
    return c * (fastlog(x / (1 - x)) -
                boost::math::digamma(gamma * c) +
                boost::math::digamma((1 - gamma) * c));
}

double DirichletLogPdf::f(double alpha,
                          const boost::numeric::ublas::matrix<double>* mean,
                          const boost::numeric::ublas::matrix<double>* data,
                          size_t n, size_t m)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        for (size_t j = 0; j < m; ++j) {
            double am = alpha * (*mean)(i, j);
            part += (am - 1) * fastlog((*data)(i, j)) - lgamma(am);
        }
    }

    return n * lgamma(alpha) + part;
}

double DirichletLogPdf::df_dalpha(double alpha,
                                  const boost::numeric::ublas::matrix<double>* mean,
                                  const boost::numeric::ublas::matrix<double>* data,
                                  size_t n, size_t m)
{
    double part = 0.0;
    for (size_t i = 0; i < n; ++i) {
        for (size_t j = 0; j < m; ++j) {
            part += (*mean)(i, j) * (fastlog((*data)(i, j)) -
                                     boost::math::digamma(alpha * (*mean)(i, j)));
        }
    }

    return n * boost::math::digamma(alpha) + part;
}

double LogisticNormalLogPdf::f(double mu, double sigma, double x)
{
    return -fastlog(sigma) - fastlog(sqrt(2*M_PI)) -
           sq(fastlog(x / (1 - x)) - mu) / (2 * sq(sigma)) -
           fastlog(x) - fastlog(1-x);
}

double LogisticNormalLogPdf::df_dx(double mu, double sigma, double x)
{
    double y = fastlog(x / (1 - x));
    return (1/(1-x)) - (1/x) - (mu - y) / (sq(sigma) * (x - 1) * x);
}


// Copyright (c) 2022 CNES
//
// All rights reserved. Use of this source code is governed by a
// BSD-style license that can be found in the LICENSE file.
#pragma once
#include <Eigen/Core>
#include <memory>

#include "pyinterp/detail/math.hpp"
#include "pyinterp/eigen.hpp"

namespace pyinterp::detail::math {

/// Set of coordinates used for interpolation
class CoordsXY {
 public:
  /// Default constructor
  CoordsXY() = delete;

  /// Creates a new instance
  CoordsXY(const Eigen::Index x_size, const Eigen::Index y_size)
      : x_(new Eigen::VectorXd), y_(new Eigen::VectorXd) {
    auto nx = x_size << 1U;
    auto ny = y_size << 1U;
    x_->resize(nx);
    y_->resize(ny);
  }

  /// Creates a new instance from existing coordinates
  CoordsXY(std::shared_ptr<Eigen::VectorXd> x,
           std::shared_ptr<Eigen::VectorXd> y)
      : x_(std::move(x)), y_(std::move(y)) {}

  /// Default destructor
  virtual ~CoordsXY() = default;

  /// Copy constructor
  ///
  /// @param rhs right value
  CoordsXY(const CoordsXY &rhs) = default;

  /// Move constructor
  ///
  /// @param rhs right value
  CoordsXY(CoordsXY &&rhs) noexcept = default;

  /// Copy assignment operator
  ///
  /// @param rhs right value
  auto operator=(const CoordsXY &rhs) -> CoordsXY & = default;

  /// Move assignment operator
  ///
  /// @param rhs right value
  auto operator=(CoordsXY &&rhs) noexcept -> CoordsXY & = default;

  /// Get the half size of the window in abscissa.
  [[nodiscard]] inline auto nx() const noexcept -> Eigen::Index {
    return x_->size() >> 1;
  }

  /// Get the half size of the window in ordinate.
  [[nodiscard]] inline auto ny() const noexcept -> Eigen::Index {
    return y_->size() >> 1;
  }

  /// Get x-coordinates
  constexpr auto x() noexcept -> std::shared_ptr<Eigen::VectorXd> & {
    return x_;
  }

  /// Get x-coordinates
  [[nodiscard]] constexpr auto x() const noexcept
      -> const std::shared_ptr<Eigen::VectorXd> & {
    return x_;
  }

  /// Get y-coordinates
  constexpr auto y() noexcept -> std::shared_ptr<Eigen::VectorXd> & {
    return y_;
  }

  /// Get y-coordinates
  [[nodiscard]] constexpr auto y() const noexcept
      -> const std::shared_ptr<Eigen::VectorXd> & {
    return y_;
  }

  /// Get the ith x-axis.
  [[nodiscard]] inline auto x(const Eigen::Index ix) const -> double {
    return (*x_)(ix);
  }

  /// Get the ith y-axis.
  [[nodiscard]] inline auto y(const Eigen::Index jx) const -> double {
    return (*y_)(jx);
  }

  /// Set the ith x-axis.
  inline auto x(const Eigen::Index ix) -> double & { return (*x_)(ix); }

  /// Get the ith y-axis.
  inline auto y(const Eigen::Index jx) -> double & { return (*y_)(jx); }

  /// Normalizes the angle with respect to the first value of the X axis of this
  /// array.
  [[nodiscard]] inline auto normalize_angle(const double xi) const -> double {
    return math::normalize_angle(xi, (*x_)(0), 360.0);
  }

 private:
  std::shared_ptr<Eigen::VectorXd> x_{};
  std::shared_ptr<Eigen::VectorXd> y_{};
};

/// Set of coordinates/values used for interpolation
///  * q11 = (x1, y1)
///  * q12 = (x1, y2)
///  * .../...
///  * q1n = (x1, yn)
///  * q21 = (x2, y1)
///  * q22 = (x2, y2).
///  * .../...
///  * q2n = (x2, yn)
///  * .../...
///  * qnn = (xn, yn)
///
/// @code
/// Array2D({{x1, x2, ..., xn}, {y1, y2, ..., yn}},
///         {q11, q12, ..., q21, q22, ...., qnn})
/// @endcode
class Frame2D : public CoordsXY {
 public:
  /// Default constructor
  Frame2D() = delete;

  /// Creates a new Array
  Frame2D(const Eigen::Index x_size, const Eigen::Index y_size)
      : CoordsXY(x_size, y_size), q_(new Eigen::MatrixXd) {
    q_->resize(x()->size(), y()->size());
  }

  /// Creates a new Array from existing coordinates/values
  Frame2D(std::shared_ptr<Eigen::VectorXd> x,
          std::shared_ptr<Eigen::VectorXd> y,
          std::shared_ptr<Eigen::MatrixXd> q)
      : CoordsXY(std::move(x), std::move(y)), q_(std::move(q)) {}

  /// Default destructor
  ~Frame2D() override = default;

  /// Copy constructor
  ///
  /// @param rhs right value
  Frame2D(const Frame2D &rhs) = default;

  /// Move constructor
  ///
  /// @param rhs right value
  Frame2D(Frame2D &&rhs) noexcept = default;

  /// Copy assignment operator
  ///
  /// @param rhs right value
  auto operator=(const Frame2D &rhs) -> Frame2D & = default;

  /// Move assignment operator
  ///
  /// @param rhs right value
  auto operator=(Frame2D &&rhs) noexcept -> Frame2D & = default;

  /// Get the values from the array for all x and y coordinates.
  constexpr auto q() noexcept -> std::shared_ptr<Eigen::MatrixXd> & {
    return q_;
  }

  /// Get the values from the array for all x and y coordinates.
  [[nodiscard]] constexpr auto q() const noexcept
      -> const std::shared_ptr<Eigen::MatrixXd> & {
    return q_;
  }

  /// Get the value at coordinate (ix, jx).
  [[nodiscard]] inline auto q(const Eigen::Index ix,
                              const Eigen::Index jx) const -> double {
    return (*q_)(ix, jx);
  }

  /// Get the value at coordinate (ix, jx).
  inline auto q(const Eigen::Index ix, const Eigen::Index jx) -> double & {
    return (*q_)(ix, jx);
  }

  /// Returns true if this instance does not contains at least one Not A Number
  /// (NaN).
  [[nodiscard]] inline auto is_valid() const -> bool { return !q_->hasNaN(); }

 private:
  std::shared_ptr<Eigen::MatrixXd> q_{};
};

/// Set of coordinates/values used for 3D-interpolation
///
/// @tparam Z-Axis type
template <typename T>
class Frame3D : public CoordsXY {
 public:
  /// Default constructor
  Frame3D() = delete;

  /// Creates a new instance
  Frame3D(const Eigen::Index x_size, const Eigen::Index y_size,
          const Eigen::Index z_size)
      : CoordsXY(x_size, y_size), z_() {
    auto nz = z_size << 1U;
    z_.resize(nz);
    q_.resize(nz);

    for (auto iz = 0U; iz < nz; ++iz) {
      q_(iz) = std::make_shared<Eigen::MatrixXd>(x()->size(), y()->size());
    }
  }

  /// Get the set of coordinates/values for the ith z-layer
  [[nodiscard]] auto frame_2d(const Eigen::Index iz) const -> Frame2D {
    return Frame2D(x(), y(), q_(iz));
  }

  /// Default destructor
  ~Frame3D() override = default;

  /// Copy constructor
  ///
  /// @param rhs right value
  Frame3D(const Frame3D &rhs) = default;

  /// Move constructor
  ///
  /// @param rhs right value
  Frame3D(Frame3D &&rhs) noexcept = default;

  /// Copy assignment operator
  ///
  /// @param rhs right value
  auto operator=(const Frame3D &rhs) -> Frame3D & = default;

  /// Move assignment operator
  ///
  /// @param rhs right value
  auto operator=(Frame3D &&rhs) noexcept -> Frame3D & = default;

  /// Get the half size of the window in z.
  [[nodiscard]] inline auto nz() const noexcept -> Eigen::Index {
    return z_.size() >> 1;
  }

  /// Get z-coordinates
  constexpr auto z() noexcept -> Vector<T> & { return z_; }

  /// Get z-coordinates
  [[nodiscard]] constexpr auto z() const noexcept -> const Vector<T> & {
    return z_;
  }

  /// Get the ith z-axis.
  [[nodiscard]] inline auto z(const Eigen::Index ix) const -> T {
    return z_(ix);
  }

  /// Set the ith z-axis.
  inline auto z(const Eigen::Index ix) -> T & { return z_(ix); }

  /// Get the value at coordinate (ix, jx, kx).
  inline auto q(const Eigen::Index ix, const Eigen::Index jx,
                const Eigen::Index kx) -> double & {
    return (*q_(kx))(ix, jx);
  }

  /// Returns true if this instance does not contains at least one Not A Number
  /// (NaN).
  [[nodiscard]] inline auto is_valid() const -> bool {
    for (Eigen::Index kx = 0; kx < q_.size(); ++kx) {
      if ((*q_(kx)).hasNaN()) {
        return false;
      }
    }
    return true;
  }

 private:
  Vector<T> z_;
  Vector<std::shared_ptr<Eigen::MatrixXd>> q_;
};

/// Set of coordinates/values used for 4D-interpolation
///
/// @tparam Z-Axis type
template <typename T>
class Frame4D : public CoordsXY {
 public:
  /// Default constructor
  Frame4D() = delete;

  /// Creates a new instance
  Frame4D(const Eigen::Index x_size, const Eigen::Index y_size,
          const Eigen::Index z_size, const Eigen::Index u_size)
      : CoordsXY(x_size, y_size), z_() {
    auto nz = z_size << 1U;
    auto nu = u_size << 1U;
    z_.resize(nz);
    u_.resize(nu);
    q_.resize(nz, nu);

    for (auto iz = 0U; iz < nz; ++iz) {
      for (auto iu = 0U; iu < nu; ++iu) {
        q_(iz, iu) =
            std::make_shared<Eigen::MatrixXd>(x()->size(), y()->size());
      }
    }
  }

  /// Get the set of coordinates/values for the ith z-layer
  [[nodiscard]] auto frame_2d(const Eigen::Index iz,
                              const Eigen::Index iu) const -> Frame2D {
    return Frame2D(x(), y(), q_(iz, iu));
  }

  /// Default destructor
  ~Frame4D() override = default;

  /// Copy constructor
  ///
  /// @param rhs right value
  Frame4D(const Frame4D &rhs) = default;

  /// Move constructor
  ///
  /// @param rhs right value
  Frame4D(Frame4D &&rhs) noexcept = default;

  /// Copy assignment operator
  ///
  /// @param rhs right value
  auto operator=(const Frame4D &rhs) -> Frame4D & = default;

  /// Move assignment operator
  ///
  /// @param rhs right value
  auto operator=(Frame4D &&rhs) noexcept -> Frame4D & = default;

  /// Get the half size of the window in z.
  [[nodiscard]] inline auto nz() const noexcept -> Eigen::Index {
    return z_.size() >> 1;
  }

  /// Get the half size of the window in u.
  [[nodiscard]] inline auto nu() const noexcept -> Eigen::Index {
    return u_.size() >> 1;
  }

  /// Get z-coordinates
  constexpr auto z() noexcept -> Vector<T> & { return z_; }

  /// Get u-coordinates
  constexpr auto u() noexcept -> Eigen::VectorXd & { return u_; }

  /// Get z-coordinates
  [[nodiscard]] inline auto z() const noexcept -> const Vector<T> & {
    return z_;
  }

  /// Get u-coordinates
  [[nodiscard]] inline auto u() const noexcept -> const Eigen::VectorXd & {
    return u_;
  }

  /// Get the ith z-axis.
  [[nodiscard]] inline auto z(const Eigen::Index ix) const -> T {
    return z_(ix);
  }

  /// Get the ith u-axis.
  [[nodiscard]] inline auto u(const Eigen::Index ix) const -> double {
    return u_(ix);
  }

  /// Set the ith z-axis.
  inline auto z(const Eigen::Index ix) -> T & { return z_(ix); }

  /// Set the ith u-axis.
  inline auto u(const Eigen::Index ix) -> double & { return u_(ix); }

  /// Get the value at coordinate (ix, jx, kx, lx).
  inline auto q(const Eigen::Index ix, const Eigen::Index jx,
                const Eigen::Index kx, const Eigen::Index lx) -> double & {
    return (*q_(kx, lx))(ix, jx);
  }

  /// Returns true if this instance does not contains at least one Not A Number
  /// (NaN).
  [[nodiscard]] inline auto is_valid() const -> bool {
    for (Eigen::Index kx = 0; kx < z_.size(); ++kx) {
      for (Eigen::Index lx = 0; lx < u_.size(); ++lx) {
        if ((*q_(kx, lx)).hasNaN()) {
          return false;
        }
      }
    }
    return true;
  }

 private:
  Vector<T> z_;
  Eigen::VectorXd u_;
  Matrix<std::shared_ptr<Eigen::MatrixXd>> q_;
};

}  // namespace pyinterp::detail::math


#ifndef FLOW_H_
#define FLOW_H_

#include <Eigen/Core>
#include <list>
#include <map>
#include <vector>

#include "qflow/config.hpp"

#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/boykov_kolmogorov_max_flow.hpp>
#include <boost/graph/edmonds_karp_max_flow.hpp>
#include <boost/graph/push_relabel_max_flow.hpp>

#include <lemon/network_simplex.h>
#include <lemon/preflow.h>
#include <lemon/smart_graph.h>

using namespace boost;
using namespace Eigen;

namespace qflow {

class MaxFlowHelper {
   public:
    MaxFlowHelper() {}
    virtual ~MaxFlowHelper(){};
    virtual void resize(int n, int m) = 0;
    virtual void addEdge(int x, int y, int c, int rc, int v, int cost = 1) = 0;
    virtual int compute() = 0;
    virtual void applyTo(std::vector<Vector2i>& edge_diff) = 0;
};

class BoykovMaxFlowHelper : public MaxFlowHelper {
   public:
    typedef int EdgeWeightType;
    typedef adjacency_list_traits<vecS, vecS, directedS> Traits;
    // clang-format off
    typedef adjacency_list < vecS, vecS, directedS,
        property < vertex_name_t, std::string,
        property < vertex_index_t, long,
        property < vertex_color_t, boost::default_color_type,
        property < vertex_distance_t, long,
        property < vertex_predecessor_t, Traits::edge_descriptor > > > > >,

        property < edge_capacity_t, EdgeWeightType,
        property < edge_residual_capacity_t, EdgeWeightType,
        property < edge_reverse_t, Traits::edge_descriptor > > > > Graph;
    // clang-format on

   public:
    BoykovMaxFlowHelper() { rev = get(edge_reverse, g); }
    void resize(int n, int m) {
        vertex_descriptors.resize(n);
        for (int i = 0; i < n; ++i) vertex_descriptors[i] = add_vertex(g);
    }
    int compute() {
        EdgeWeightType flow =
            boykov_kolmogorov_max_flow(g, vertex_descriptors.front(), vertex_descriptors.back());
        return flow;
    }
    void addDirectEdge(Traits::vertex_descriptor& v1, Traits::vertex_descriptor& v2,
                       property_map<Graph, edge_reverse_t>::type& rev, const int capacity,
                       const int inv_capacity, Graph& g, Traits::edge_descriptor& e1,
                       Traits::edge_descriptor& e2) {
        e1 = add_edge(v1, v2, g).first;
        e2 = add_edge(v2, v1, g).first;
        put(edge_capacity, g, e1, capacity);
        put(edge_capacity, g, e2, inv_capacity);

        rev[e1] = e2;
        rev[e2] = e1;
    }
    void addEdge(int x, int y, int c, int rc, int v, int cost = 1) {
        Traits::edge_descriptor e1, e2;
        addDirectEdge(vertex_descriptors[x], vertex_descriptors[y], rev, c, rc, g, e1, e2);
        if (v != -1) {
            edge_to_variables[e1] = std::make_pair(v, -1);
            edge_to_variables[e2] = std::make_pair(v, 1);
        }
    }
    void applyTo(std::vector<Vector2i>& edge_diff) {
        property_map<Graph, edge_capacity_t>::type capacity = get(edge_capacity, g);
        property_map<Graph, edge_residual_capacity_t>::type residual_capacity =
            get(edge_residual_capacity, g);

        graph_traits<Graph>::vertex_iterator u_iter, u_end;
        graph_traits<Graph>::out_edge_iterator ei, e_end;
        for (tie(u_iter, u_end) = vertices(g); u_iter != u_end; ++u_iter)
            for (tie(ei, e_end) = out_edges(*u_iter, g); ei != e_end; ++ei)
                if (capacity[*ei] > 0) {
                    int flow = (capacity[*ei] - residual_capacity[*ei]);
                    if (flow > 0) {
                        auto it = edge_to_variables.find(*ei);
                        if (it != edge_to_variables.end()) {
                            edge_diff[it->second.first / 2][it->second.first % 2] +=
                                it->second.second * flow;
                        }
                    }
                }
    }

   private:
    Graph g;
    property_map<Graph, edge_reverse_t>::type rev;
    std::vector<Traits::vertex_descriptor> vertex_descriptors;
    std::map<Traits::edge_descriptor, std::pair<int, int>> edge_to_variables;
};

class NetworkSimplexFlowHelper : public MaxFlowHelper {
   public:
    using Weight = int;
    using Capacity = int;
    using Graph = lemon::SmartDigraph;
    using Node = Graph::Node;
    using Arc = Graph::Arc;
    template <typename ValueType>
    using ArcMap = lemon::SmartDigraph::ArcMap<ValueType>;
    using Preflow = lemon::Preflow<lemon::SmartDigraph, ArcMap<Capacity>>;
    using NetworkSimplex = lemon::NetworkSimplex<lemon::SmartDigraph, Capacity, Weight>;

   public:
    NetworkSimplexFlowHelper() : cost(graph), capacity(graph), flow(graph), variable(graph) {}
    ~NetworkSimplexFlowHelper(){};
    void resize(int n, int m) {
        nodes.reserve(n);
        for (int i = 0; i < n; ++i) nodes.push_back(graph.addNode());
    }
    void addEdge(int x, int y, int c, int rc, int v, int cst = 1) {
        assert(x >= 0);
        assert(v >= -1);
        if (c) {
            auto e1 = graph.addArc(nodes[x], nodes[y]);
            cost[e1] = cst;
            capacity[e1] = c;
            variable[e1] = std::make_pair(v, 1);
        }

        if (rc) {
            auto e2 = graph.addArc(nodes[y], nodes[x]);
            cost[e2] = cst;
            capacity[e2] = rc;
            variable[e2] = std::make_pair(v, -1);
        }
    }
    int compute() {
        Preflow pf(graph, capacity, nodes.front(), nodes.back());
        NetworkSimplex ns(graph);

        // Run preflow to find maximum flow
        lprintf("push-relabel flow... ");
        pf.runMinCut();
        int maxflow = pf.flowValue();

        // Run network simplex to find minimum cost maximum flow
        ns.costMap(cost).upperMap(capacity).stSupply(nodes.front(), nodes.back(), maxflow);
        auto status = ns.run();
        switch (status) {
            case NetworkSimplex::OPTIMAL:
                ns.flowMap(flow);
                break;
            case NetworkSimplex::INFEASIBLE:
                lputs("NetworkSimplex::INFEASIBLE");
                assert(0);
                break;
            default:
                lputs("Unknown: NetworkSimplex::Default");
                assert(0);
                break;
        }

        return maxflow;
    }
    void applyTo(std::vector<Vector2i>& edge_diff) {
        for (Graph::ArcIt e(graph); e != lemon::INVALID; ++e) {
            int var = variable[e].first;
            if (var == -1) continue;
            int sgn = variable[e].second;
            edge_diff[var / 2][var % 2] -= sgn * flow[e];
        }
    }

   private:
    Graph graph;
    ArcMap<Weight> cost;
    ArcMap<Capacity> capacity;
    ArcMap<Capacity> flow;
    ArcMap<std::pair<int, int>> variable;
    std::vector<Node> nodes;
    std::vector<Arc> edges;
};

#ifdef WITH_GUROBI

#include <gurobi_c++.h>

class GurobiFlowHelper : public MaxFlowHelper {
   public:
    GurobiFlowHelper() {}
    virtual ~GurobiFlowHelper(){};
    virtual void resize(int n, int m) {
        nodes.resize(n * 2);
        edges.resize(m);
    }
    virtual void addEdge(int x, int y, int c, int rc, int v, int cost = 1) {
        nodes[x * 2 + 0].push_back(vars.size());
        nodes[y * 2 + 1].push_back(vars.size());
        vars.push_back(model.addVar(0, c, 0, GRB_INTEGER));
        edges.push_back(std::make_pair(v, 1));

        nodes[y * 2 + 0].push_back(vars.size());
        nodes[x * 2 + 1].push_back(vars.size());
        vars.push_back(model.addVar(0, rc, 0, GRB_INTEGER));
        edges.push_back(std::make_pair(v, -1));
    }
    virtual int compute() {
        std::cerr << "compute" << std::endl;
        int ns = nodes.size() / 2;

        int flow;
        for (int i = 1; i < ns - 1; ++i) {
            GRBLinExpr cons = 0;
            for (auto n : nodes[2 * i + 0]) cons += vars[n];
            for (auto n : nodes[2 * i + 1]) cons -= vars[n];
            model.addConstr(cons == 0);
        }

        // first pass, maximum flow
        GRBLinExpr outbound = 0;
        {
            lprintf("first pass\n");
            for (auto& n : nodes[0]) outbound += vars[n];
            for (auto& n : nodes[1]) outbound -= vars[n];
            model.setObjective(outbound, GRB_MAXIMIZE);
            model.optimize();

            flow = (int)model.get(GRB_DoubleAttr_ObjVal);
            lprintf("Gurobi result: %d\n", flow);
        }

        // second pass, minimum cost flow
        {
            lprintf("second pass\n");
            model.addConstr(outbound == flow);
            GRBLinExpr cost = 0;
            for (auto& v : vars) cost += v;
            model.setObjective(cost, GRB_MINIMIZE);
            model.optimize();

            double optimal_cost = (int)model.get(GRB_DoubleAttr_ObjVal);
            lprintf("Gurobi result: %.3f\n", optimal_cost);
        }
        return flow;
    }
    virtual void applyTo(std::vector<Vector2i>& edge_diff) { assert(0); };

   private:
    GRBEnv env = GRBEnv();
    GRBModel model = GRBModel(env);
    std::vector<GRBVar> vars;
    std::vector<std::pair<int, int>> edges;
    std::vector<std::vector<int>> nodes;
};

#endif

class ECMaxFlowHelper : public MaxFlowHelper {
   public:
    struct FlowInfo {
        int id;
        int capacity, flow;
        int v, d;
        FlowInfo* rev;
    };
    struct SearchInfo {
        SearchInfo(int _id, int _prev_id, FlowInfo* _info)
            : id(_id), prev_id(_prev_id), info(_info) {}
        int id;
        int prev_id;
        FlowInfo* info;
    };
    ECMaxFlowHelper() { num = 0; }
    int num;
    std::vector<FlowInfo*> variable_to_edge;
    void resize(int n, int m) {
        graph.resize(n);
        variable_to_edge.resize(m, 0);
        num = n;
    }
    void addEdge(int x, int y, int c, int rc, int v, int cost = 0) {
        FlowInfo flow;
        flow.id = y;
        flow.capacity = c;
        flow.flow = 0;
        flow.v = v;
        flow.d = -1;
        graph[x].push_back(flow);
        auto& f1 = graph[x].back();
        flow.id = x;
        flow.capacity = rc;
        flow.flow = 0;
        flow.v = v;
        flow.d = 1;
        graph[y].push_back(flow);
        auto& f2 = graph[y].back();
        f2.rev = &f1;
        f1.rev = &f2;
    }
    int compute() {
        int total_flow = 0;
        int count = 0;
        while (true) {
            count += 1;
            std::vector<int> vhash(num, 0);
            std::vector<SearchInfo> q;
            q.push_back(SearchInfo(0, -1, 0));
            vhash[0] = 1;
            int q_front = 0;
            bool found = false;
            while (q_front < q.size()) {
                int vert = q[q_front].id;
                for (auto& l : graph[vert]) {
                    if (vhash[l.id] || l.capacity <= l.flow) continue;
                    q.push_back(SearchInfo(l.id, q_front, &l));
                    vhash[l.id] = 1;
                    if (l.id == num - 1) {
                        found = true;
                        break;
                    }
                }
                if (found) break;
                q_front += 1;
            }
            if (q_front == q.size()) break;
            int loc = q.size() - 1;
            while (q[loc].prev_id != -1) {
                q[loc].info->flow += 1;
                q[loc].info->rev->flow -= 1;
                loc = q[loc].prev_id;
                // int prev_v = q[loc].id;
                // applyFlow(prev_v, current_v, 1);
                // applyFlow(current_v, prev_v, -1);
            }
            total_flow += 1;
        }
        return total_flow;
    }
    void applyTo(std::vector<Vector2i>& edge_diff) {
        for (int i = 0; i < graph.size(); ++i) {
            for (auto& flow : graph[i]) {
                if (flow.flow > 0 && flow.v != -1) {
                    if (flow.flow > 0) {
                        edge_diff[flow.v / 2][flow.v % 2] += flow.d * flow.flow;
                        if (abs(edge_diff[flow.v / 2][flow.v % 2]) > 2) {
                        }
                    }
                }
            }
        }
    }
    void applyFlow(int v1, int v2, int flow) {
        for (auto& it : graph[v1]) {
            if (it.id == v2) {
                it.flow += flow;
                break;
            }
        }
    }
    std::vector<std::list<FlowInfo>> graph;
};

} // namespace qflow

#endif