// This file is part of the dune-hdd project:
//   http://users.dune-project.org/projects/dune-hdd
// Copyright holders: Felix Schindler
// License: BSD 2-Clause License (http://opensource.org/licenses/BSD-2-Clause)

#ifndef DUNE_HDD_LINEARELLIPTIC_ESTIMATORS_SWIPDG_HH
#define DUNE_HDD_LINEARELLIPTIC_ESTIMATORS_SWIPDG_HH

#include <map>
#include <memory>
#include <string>
#include <vector>

#include <boost/numeric/conversion/cast.hpp>

#if HAVE_ALUGRID
# include <dune/grid/alugrid.hh>
#endif

#include <dune/stuff/common/memory.hh>
#include <dune/stuff/common/tmp-storage.hh>
#include <dune/stuff/grid/walker.hh>
#include <dune/stuff/grid/walker/functors.hh>
#include <dune/stuff/playground/functions/ESV2007.hh>

#include <dune/pymor/common/exceptions.hh>
#include <dune/pymor/parameters/base.hh>

#include <dune/gdt/discretefunction/default.hh>
#include <dune/gdt/localevaluation/elliptic.hh>
#include <dune/gdt/localevaluation/product.hh>
#include <dune/gdt/localoperator/codim0.hh>
#include <dune/gdt/operators/oswaldinterpolation.hh>
#include <dune/gdt/operators/projections.hh>
#include <dune/gdt/playground/localevaluation/ESV2007.hh>
#include <dune/gdt/playground/operators/fluxreconstruction.hh>
#include <dune/gdt/spaces/fv/default.hh>
#include <dune/gdt/spaces/rt/pdelab.hh>

namespace Dune {
namespace HDD {
namespace LinearElliptic {
namespace Estimators {
namespace internal {
namespace SWIPDG {

static const size_t over_integrate = 2;

class LocalNonconformityESV2007Base
{
public:
  static std::string id() { return "eta_NC_ESV2007"; }
};

template< class SpaceType, class VectorType, class ProblemType, class GridType >
class LocalNonconformityESV2007
  : public LocalNonconformityESV2007Base
{
public:
  static const bool available = false;
};

#if HAVE_ALUGRID

/**
 *  \brief computes the local nonconformity estimator as defined in ESV2007
 */
template< class SpaceType, class VectorType, class ProblemType >
class LocalNonconformityESV2007< SpaceType, VectorType, ProblemType, ALUGrid< 2, 2, simplex, conforming > >
  : public LocalNonconformityESV2007Base
  , public Stuff::Grid::Functor::Codim0< typename SpaceType::GridViewType >
{
  typedef LocalNonconformityESV2007< SpaceType, VectorType, ProblemType, ALUGrid< 2, 2, simplex, conforming > >
    ThisType;
  typedef Stuff::Grid::Functor::Codim0< typename SpaceType::GridViewType > FunctorBaseType;
public:
  static const bool available = true;

  typedef std::map< std::string, Pymor::Parameter > ParametersMapType;

  typedef typename FunctorBaseType::GridViewType GridViewType;
  typedef typename FunctorBaseType::EntityType   EntityType;

  typedef typename ProblemType::RangeFieldType   RangeFieldType;

private:
  typedef GDT::ConstDiscreteFunction< SpaceType, VectorType > ConstDiscreteFunctionType;
  typedef GDT::DiscreteFunction< SpaceType, VectorType > DiscreteFunctionType;
  typedef typename ConstDiscreteFunctionType::DifferenceType DifferenceType;

  typedef typename ProblemType::DiffusionFactorType::NonparametricType DiffusionFactorType;
  typedef typename ProblemType::DiffusionTensorType::NonparametricType DiffusionTensorType;

  typedef GDT::LocalOperator::Codim0Integral< GDT::LocalEvaluation::Elliptic< DiffusionFactorType,
                                                                              DiffusionTensorType > > LocalOperatorType;
  typedef Stuff::Common::TmpMatricesStorage< RangeFieldType > TmpStorageProviderType;

  static const ProblemType& assert_problem(const ProblemType& problem, const Pymor::Parameter& mu_bar)
  {
    if (mu_bar.type() != problem.parameter_type())
      DUNE_THROW(Pymor::Exceptions::wrong_parameter_type,
                 "Given mu_bar is of type " << mu_bar.type() << " and should be of type " << problem.parameter_type()
                 << "!");
    if (problem.diffusion_tensor()->parametric())
      DUNE_THROW(NotImplemented, "Not implemented for parametric diffusion_tensor!");
    return problem;
  } // ... assert_problem(...)

public:
  static RangeFieldType estimate(const SpaceType& space,
                                 const VectorType& vector,
                                 const ProblemType& problem,
                                 const ParametersMapType parameters = ParametersMapType())
  {
    if (problem.diffusion_factor()->parametric() && parameters.find("mu_bar") == parameters.end())
      DUNE_THROW(Stuff::Exceptions::wrong_input_given, "Given parameters are missing 'mu_bar'!");
    const Pymor::Parameter mu_bar = problem.parametric() ? parameters.at("mu_bar") : Pymor::Parameter();
    ThisType estimator(space, vector, problem, mu_bar);
    Stuff::Grid::Walker< GridViewType > grid_walker(space.grid_view());
    grid_walker.add(estimator);
    grid_walker.walk();
    return std::sqrt(estimator.result_);
  } // ... estimate(...)

  LocalNonconformityESV2007(const SpaceType& space,
                            const VectorType& vector,
                            const ProblemType& problem,
                            const Pymor::Parameter mu_bar = Pymor::Parameter())
    : space_(space)
    , vector_(vector)
    , problem_(assert_problem(problem, mu_bar))
    , problem_mu_bar_(problem_.with_mu(mu_bar))
    , discrete_solution_(space_, vector_)
    , oswald_interpolation_(space_)
    , difference_(Stuff::Common::make_unique< DifferenceType >(discrete_solution_ - oswald_interpolation_))
    , local_operator_(over_integrate,
                      *problem_mu_bar_->diffusion_factor()->affine_part(),
                      *problem_.diffusion_tensor()->affine_part())
    , tmp_local_matrices_({1, local_operator_.numTmpObjectsRequired()}, 1, 1)
    , prepared_(false)
    , result_(0.0)
  {}

  virtual void prepare()
  {
    if (!prepared_) {
      const GDT::Operators::OswaldInterpolation< GridViewType > oswald_interpolation_operator(space_.grid_view());
      oswald_interpolation_operator.apply(discrete_solution_, oswald_interpolation_);
      result_ = 0.0;
      prepared_ = true;
    }
  } // ... prepare(...)

  RangeFieldType compute_locally(const EntityType& entity)
  {
    const auto local_difference = difference_->local_function(entity);
    local_operator_.apply(*local_difference,
                          *local_difference,
                          tmp_local_matrices_.matrices()[0][0],
                          tmp_local_matrices_.matrices()[1]);
    assert(tmp_local_matrices_.matrices()[0][0].rows() >= 1);
    assert(tmp_local_matrices_.matrices()[0][0].cols() >= 1);
    return tmp_local_matrices_.matrices()[0][0][0][0];
  } // ... compute_locally(...)

  virtual void apply_local(const EntityType &entity)
  {
    result_ += compute_locally(entity);
  }

private:
  const SpaceType& space_;
  const VectorType& vector_;
  const ProblemType& problem_;
  const std::shared_ptr< const typename ProblemType::NonparametricType > problem_mu_bar_;
  const ConstDiscreteFunctionType discrete_solution_;
  DiscreteFunctionType oswald_interpolation_;
  std::unique_ptr< const DifferenceType > difference_;
  const LocalOperatorType local_operator_;
  TmpStorageProviderType tmp_local_matrices_;
  bool prepared_;
public:
  RangeFieldType result_;
}; // class LocalNonconformityESV2007< ..., ALUGrid< 2, 2, simplex, conforming >, ... >

#endif // HAVE_ALUGRID

class LocalResidualESV2007Base
{
public:
  static std::string id() { return "eta_R_ESV2007"; }
};

template< class SpaceType, class VectorType, class ProblemType, class GridType >
class LocalResidualESV2007
  : public LocalResidualESV2007Base
{
public:
  static const bool available = false;
};

#if HAVE_ALUGRID

/**
 *  \brief computes the local residual estimator as defined in ESV2007
 */
template< class SpaceType, class VectorType, class ProblemType >
class LocalResidualESV2007< SpaceType, VectorType, ProblemType, ALUGrid< 2, 2, simplex, conforming > >
  : public LocalResidualESV2007Base
  , public Stuff::Grid::Functor::Codim0< typename SpaceType::GridViewType >
{
  typedef LocalResidualESV2007< SpaceType, VectorType, ProblemType, ALUGrid< 2, 2, simplex, conforming > >
      ThisType;
  typedef Stuff::Grid::Functor::Codim0< typename SpaceType::GridViewType > FunctorBaseType;
public:
  static const bool available = true;

  typedef typename FunctorBaseType::GridViewType GridViewType;
  typedef typename FunctorBaseType::EntityType   EntityType;

  typedef typename ProblemType::RangeFieldType RangeFieldType;

private:
  typedef GDT::Spaces::FV::Default< GridViewType, RangeFieldType, 1, 1 > P0SpaceType;
  typedef GDT::DiscreteFunction< P0SpaceType, VectorType > DiscreteFunctionType;
  typedef typename DiscreteFunctionType::DifferenceType DifferenceType;

  typedef typename ProblemType::DiffusionFactorType::NonparametricType DiffusionFactorType;
  typedef typename ProblemType::DiffusionTensorType::NonparametricType DiffusionTensorType;

  typedef typename Stuff::Functions::ESV2007::Cutoff< DiffusionFactorType, DiffusionTensorType > CutoffFunctionType;
  typedef GDT::LocalOperator::Codim0Integral< GDT::LocalEvaluation::Product< CutoffFunctionType > > LocalOperatorType;
  typedef DSC::TmpMatricesStorage< RangeFieldType > TmpStorageProviderType;

  static const ProblemType& assert_problem(const ProblemType& problem)
  {
    if (problem.parametric())
      DUNE_THROW(NotImplemented, "Not implemented yet for parametric problems!");
    assert(problem.diffusion_factor()->has_affine_part());
    assert(problem.diffusion_tensor()->has_affine_part());
    assert(problem.force()->has_affine_part());
    return problem;
  } // ... assert_problem(...)

public:
  static RangeFieldType estimate(const SpaceType& space, const VectorType& /*vector*/, const ProblemType& problem)
  {
    ThisType estimator(space, problem);
    Stuff::Grid::Walker< GridViewType > grid_walker(space.grid_view());
    grid_walker.add(estimator);
    grid_walker.walk();
    return std::sqrt(estimator.result_);
  } // ... estimate(...)

  LocalResidualESV2007(const SpaceType& space, const ProblemType& problem)
    : space_(space)
    , problem_(assert_problem(problem))
    , p0_space_(space_.grid_view())
    , p0_force_(p0_space_)
    , difference_(Stuff::Common::make_unique< DifferenceType >(*problem_.force()->affine_part() - p0_force_))
    , cutoff_function_(*problem_.diffusion_factor()->affine_part(),
                       *problem_.diffusion_tensor()->affine_part())
    , local_operator_(over_integrate, cutoff_function_)
    , tmp_local_matrices_({1, local_operator_.numTmpObjectsRequired()}, 1, 1)
    , prepared_(false)
    , result_(0.0)
  {}

  virtual void prepare()
  {
    if (!prepared_) {
      const GDT::Operators::Projection< GridViewType > projection_operator(space_.grid_view(), over_integrate);
      projection_operator.apply(*problem_.force()->affine_part(), p0_force_);
      result_ = 0.0;
      prepared_ = true;
    }
  } // ... prepare(...)

  RangeFieldType compute_locally(const EntityType& entity)
  {
    const auto local_difference = difference_->local_function(entity);
    local_operator_.apply(*local_difference,
                          *local_difference,
                          tmp_local_matrices_.matrices()[0][0],
                          tmp_local_matrices_.matrices()[1]);
    assert(tmp_local_matrices_.matrices()[0][0].rows() >= 1);
    assert(tmp_local_matrices_.matrices()[0][0].cols() >= 1);
    return tmp_local_matrices_.matrices()[0][0][0][0];
  } // ... compute_locally(...)

  virtual void apply_local(const EntityType &entity)
  {
    result_ += compute_locally(entity);
  }

private:
  const SpaceType& space_;
  const ProblemType& problem_;
  const P0SpaceType p0_space_;
  DiscreteFunctionType p0_force_;
  std::unique_ptr< const DifferenceType > difference_;
  const CutoffFunctionType cutoff_function_;
  const LocalOperatorType local_operator_;
  TmpStorageProviderType tmp_local_matrices_;
  bool prepared_;
public:
  RangeFieldType result_;
}; // class LocalResidualESV2007< ..., ALUGrid< 2, 2, simplex, conforming >, ... >

#endif // HAVE_ALUGRID

class LocalResidualESV2007StarBase
{
public:
  static std::string id() { return "eta_R_ESV2007_*"; }
};

template< class SpaceType, class VectorType, class ProblemType, class GridType >
class LocalResidualESV2007Star
  : public LocalResidualESV2007StarBase
{
public:
  static const bool available = false;
};

#if HAVE_ALUGRID

/**
 *  \brief computes the local residual estimator as defined in ESV2007
 */
template< class SpaceType, class VectorType, class ProblemType >
class LocalResidualESV2007Star< SpaceType, VectorType, ProblemType, ALUGrid< 2, 2, simplex, conforming > >
  : public LocalResidualESV2007StarBase
  , public Stuff::Grid::Functor::Codim0< typename SpaceType::GridViewType >
{
  typedef LocalResidualESV2007Star< SpaceType, VectorType, ProblemType, ALUGrid< 2, 2, simplex, conforming > >
      ThisType;
  typedef Stuff::Grid::Functor::Codim0< typename SpaceType::GridViewType > FunctorBaseType;
public:
  static const bool available = true;

  typedef std::map< std::string, Pymor::Parameter > ParametersMapType;

  typedef typename FunctorBaseType::GridViewType GridViewType;
  typedef typename FunctorBaseType::EntityType   EntityType;

  typedef typename ProblemType::RangeFieldType RangeFieldType;

  static const unsigned int dimDomain = SpaceType::dimDomain;

private:
  typedef GDT::ConstDiscreteFunction< SpaceType, VectorType > ConstDiscreteFunctionType;
  typedef GDT::Spaces::RT::PdelabBased< GridViewType, 0, RangeFieldType, dimDomain > RTN0SpaceType;
  typedef GDT::DiscreteFunction< RTN0SpaceType, VectorType > RTN0DiscreteFunctionType;
  typedef typename RTN0DiscreteFunctionType::DivergenceType DivergenceType;
  typedef typename DivergenceType::DifferenceType DifferenceType;

  typedef typename ProblemType::DiffusionFactorType::NonparametricType DiffusionFactorType;
  typedef typename ProblemType::DiffusionTensorType::NonparametricType DiffusionTensorType;

  typedef typename Stuff::Functions::ESV2007::Cutoff< DiffusionFactorType, DiffusionTensorType > CutoffFunctionType;
  typedef GDT::LocalOperator::Codim0Integral< GDT::LocalEvaluation::Product< CutoffFunctionType > > LocalOperatorType;
  typedef DSC::TmpMatricesStorage< RangeFieldType > TmpStorageProviderType;

  static const ProblemType& assert_problem(const ProblemType& problem, const Pymor::Parameter& mu)
  {
    if (mu.type() != problem.parameter_type())
      DUNE_THROW(Pymor::Exceptions::wrong_parameter_type,
                 "Given mu is of type " << mu.type() << " and should be of type " << problem.parameter_type()
                 << "!");
    if (problem.diffusion_tensor()->parametric())
      DUNE_THROW(NotImplemented, "Not implemented for parametric diffusion_tensor!");
    if (problem.force()->parametric())
      DUNE_THROW(NotImplemented, "Not implemented for parametric force!");
    return problem;
  } // ... assert_problem(...)

public:
  static RangeFieldType estimate(const SpaceType& space,
                                 const VectorType& vector,
                                 const ProblemType& problem,
                                 const ParametersMapType parameters = ParametersMapType())
  {
    if (problem.diffusion_factor()->parametric() && parameters.find("mu") == parameters.end())
      DUNE_THROW(Stuff::Exceptions::wrong_input_given, "Given parameters are missing 'mu'!");
    const Pymor::Parameter mu = problem.parametric() ? parameters.at("mu") : Pymor::Parameter();
    ThisType estimator(space, vector, problem, mu);
    Stuff::Grid::Walker< GridViewType > grid_walker(space.grid_view());
    grid_walker.add(estimator);
    grid_walker.walk();
    return std::sqrt(estimator.result_);
  } // ... estimate(...)

  LocalResidualESV2007Star(const SpaceType& space,
                           const VectorType& vector,
                           const ProblemType& problem,
                           const Pymor::Parameter mu = Pymor::Parameter())
    : space_(space)
    , vector_(vector)
    , problem_(assert_problem(problem, mu))
    , problem_mu_(problem_.with_mu(mu))
    , discrete_solution_(space_, vector_)
    , rtn0_space_(space_.grid_view())
    , diffusive_flux_(rtn0_space_)
    , divergence_(diffusive_flux_.divergence())
    , difference_(*problem_.force()->affine_part() - divergence_)
    , cutoff_function_(*problem_.diffusion_factor()->affine_part(),
                       *problem_.diffusion_tensor()->affine_part())
    , local_operator_(over_integrate, cutoff_function_)
    , tmp_local_matrices_({1, local_operator_.numTmpObjectsRequired()}, 1, 1)
    , prepared_(false)
    , result_(0.0)
  {}

  virtual ~LocalResidualESV2007Star() = default;

  virtual void prepare()
  {
    if (!prepared_) {
      const GDT::Operators::DiffusiveFluxReconstruction< GridViewType, DiffusionFactorType, DiffusionTensorType >
        diffusive_flux_reconstruction(space_.grid_view(),
                                      *problem_mu_->diffusion_factor()->affine_part(),
                                      *problem_.diffusion_tensor()->affine_part(),
                                      over_integrate);
      diffusive_flux_reconstruction.apply(discrete_solution_, diffusive_flux_);
      result_ = 0.0;
      prepared_ = true;
    }
  } // ... prepare(...)

  RangeFieldType compute_locally(const EntityType& entity)
  {
    const auto local_difference = difference_.local_function(entity);
    local_operator_.apply(*local_difference,
                          *local_difference,
                          tmp_local_matrices_.matrices()[0][0],
                          tmp_local_matrices_.matrices()[1]);
    assert(tmp_local_matrices_.matrices()[0][0].rows() >= 1);
    assert(tmp_local_matrices_.matrices()[0][0].cols() >= 1);
    return tmp_local_matrices_.matrices()[0][0][0][0];
  } // ... compute_locally(...)

  virtual void apply_local(const EntityType &entity)
  {
    result_ += compute_locally(entity);
  }

private:
  const SpaceType& space_;
  const VectorType& vector_;
  const ProblemType& problem_;
  const std::shared_ptr< typename ProblemType::NonparametricType > problem_mu_;
  const ConstDiscreteFunctionType discrete_solution_;
  const RTN0SpaceType rtn0_space_;
  RTN0DiscreteFunctionType diffusive_flux_;
  const DivergenceType divergence_;
  const DifferenceType difference_;
  const CutoffFunctionType cutoff_function_;
  const LocalOperatorType local_operator_;
  TmpStorageProviderType tmp_local_matrices_;
  bool prepared_;
public:
  RangeFieldType result_;
}; // class LocalResidualESV2007Star< ..., ALUGrid< 2, 2, simplex, conforming >, ... >

#endif // HAVE_ALUGRID

class LocalDiffusiveFluxESV2007Base
{
public:
  static std::string id() { return "eta_DF_ESV2007"; }
};

template< class SpaceType, class VectorType, class ProblemType, class GridType >
class LocalDiffusiveFluxESV2007
  : public LocalDiffusiveFluxESV2007Base
{
public:
  static const bool available = false;
};

#if HAVE_ALUGRID

/**
 *  \brief computes the local diffusive flux estimator as defined in ESV2007
 */
template< class SpaceType, class VectorType, class ProblemType >
class LocalDiffusiveFluxESV2007< SpaceType, VectorType, ProblemType, ALUGrid< 2, 2, simplex, conforming > >
  : public LocalDiffusiveFluxESV2007Base
  , public Stuff::Grid::Functor::Codim0< typename SpaceType::GridViewType >
{
  typedef LocalDiffusiveFluxESV2007< SpaceType, VectorType, ProblemType, ALUGrid< 2, 2, simplex, conforming > >
      ThisType;
  typedef Stuff::Grid::Functor::Codim0< typename SpaceType::GridViewType > FunctorBaseType;
public:
  static const bool available = true;

  typedef std::map< std::string, Pymor::Parameter > ParametersMapType;

  typedef typename FunctorBaseType::GridViewType GridViewType;
  typedef typename FunctorBaseType::EntityType   EntityType;

  typedef typename ProblemType::RangeFieldType RangeFieldType;

  static const unsigned int dimDomain = SpaceType::dimDomain;

private:
  typedef GDT::ConstDiscreteFunction< SpaceType, VectorType > ConstDiscreteFunctionType;
  typedef GDT::Spaces::RT::PdelabBased< GridViewType, 0, RangeFieldType, dimDomain > RTN0SpaceType;
  typedef GDT::DiscreteFunction< RTN0SpaceType, VectorType > RTN0DiscreteFunctionType;

  typedef typename ProblemType::DiffusionFactorType::NonparametricType DiffusionFactorType;
  typedef typename ProblemType::DiffusionTensorType::NonparametricType DiffusionTensorType;

  typedef GDT::LocalOperator::Codim0Integral<
      GDT::LocalEvaluation::ESV2007::DiffusiveFluxEstimate< DiffusionFactorType,
                                                            RTN0DiscreteFunctionType,
                                                            DiffusionTensorType > > LocalOperatorType;
  typedef DSC::TmpMatricesStorage< RangeFieldType > TmpStorageProviderType;

  static const ProblemType& assert_problem(const ProblemType& problem,
                                           const Pymor::Parameter& mu,
                                           const Pymor::Parameter& mu_hat)
  {
    if (mu.type() != problem.parameter_type())
      DUNE_THROW(Pymor::Exceptions::wrong_parameter_type,
                 "Given mu is of type " << mu.type() << " and should be of type " << problem.parameter_type()
                 << "!");
    if (mu_hat.type() != problem.parameter_type())
      DUNE_THROW(Pymor::Exceptions::wrong_parameter_type,
                 "Given mu_hat is of type " << mu_hat.type() << " and should be of type " << problem.parameter_type()
                 << "!");
    if (problem.diffusion_tensor()->parametric())
      DUNE_THROW(NotImplemented, "Not implemented for parametric diffusion_tensor!");
    return problem;
  } // ... assert_problem(...)

public:
  static RangeFieldType estimate(const SpaceType& space,
                                 const VectorType& vector,
                                 const ProblemType& problem,
                                 const ParametersMapType parameters = ParametersMapType())
  {
    if (problem.diffusion_factor()->parametric() && parameters.find("mu") == parameters.end())
      DUNE_THROW(Stuff::Exceptions::wrong_input_given, "Given parameters are missing 'mu'!");
    if (problem.diffusion_factor()->parametric() && parameters.find("mu_hat") == parameters.end())
      DUNE_THROW(Stuff::Exceptions::wrong_input_given, "Given parameters are missing 'mu_hat'!");
    const Pymor::Parameter mu =     problem.parametric() ? parameters.at("mu")     : Pymor::Parameter();
    const Pymor::Parameter mu_hat = problem.parametric() ? parameters.at("mu_hat") : Pymor::Parameter();
    ThisType estimator(space, vector, problem, mu, mu_hat);
    Stuff::Grid::Walker< GridViewType > grid_walker(space.grid_view());
    grid_walker.add(estimator);
    grid_walker.walk();
    return std::sqrt(estimator.result_);
  } // ... estimate(...)

  LocalDiffusiveFluxESV2007(const SpaceType& space,
                            const VectorType& vector,
                            const ProblemType& problem,
                            const Pymor::Parameter mu = Pymor::Parameter(),
                            const Pymor::Parameter mu_hat = Pymor::Parameter())
    : space_(space)
    , vector_(vector)
    , problem_(assert_problem(problem, mu, mu_hat))
    , problem_mu_(problem_.with_mu(mu))
    , problem_mu_hat_(problem.with_mu(mu_hat))
    , discrete_solution_(space_, vector_)
    , rtn0_space_(space.grid_view())
    , diffusive_flux_(rtn0_space_)
    , local_operator_(over_integrate,
                      *problem_mu_hat_->diffusion_factor()->affine_part(),
                      *problem_.diffusion_tensor()->affine_part(),
                      diffusive_flux_)
    , tmp_local_matrices_({1, local_operator_.numTmpObjectsRequired()}, 1, 1)
    , prepared_(false)
    , result_(0.0)
  {}

  virtual void prepare()
  {
    if (!prepared_) {
      const GDT::Operators::DiffusiveFluxReconstruction< GridViewType, DiffusionFactorType, DiffusionTensorType >
        diffusive_flux_reconstruction(space_.grid_view(),
                                      *problem_mu_->diffusion_factor()->affine_part(),
                                      *problem_.diffusion_tensor()->affine_part(),
                                      over_integrate);
      diffusive_flux_reconstruction.apply(discrete_solution_, diffusive_flux_);
      result_ = 0.0;
      prepared_ = true;
    }
  } // ... prepare(...)

  RangeFieldType compute_locally(const EntityType& entity)
  {
    const auto local_discrete_solution = discrete_solution_.local_function(entity);
    local_operator_.apply(*local_discrete_solution,
                          *local_discrete_solution,
                          tmp_local_matrices_.matrices()[0][0],
                          tmp_local_matrices_.matrices()[1]);
    assert(tmp_local_matrices_.matrices()[0][0].rows() >= 1);
    assert(tmp_local_matrices_.matrices()[0][0].cols() >= 1);
    return tmp_local_matrices_.matrices()[0][0][0][0];
  } // ... compute_locally(...)

  virtual void apply_local(const EntityType &entity)
  {
    result_ += compute_locally(entity);
  }

private:
  const SpaceType& space_;
  const VectorType& vector_;
  const ProblemType& problem_;
  const std::shared_ptr< const typename ProblemType::NonparametricType > problem_mu_;
  const std::shared_ptr< const typename ProblemType::NonparametricType > problem_mu_hat_;
  const ConstDiscreteFunctionType discrete_solution_;
  const RTN0SpaceType rtn0_space_;
  RTN0DiscreteFunctionType diffusive_flux_;
  const LocalOperatorType local_operator_;
  TmpStorageProviderType tmp_local_matrices_;
  bool prepared_;
public:
  RangeFieldType result_;
}; // class LocalDiffusiveFluxESV2007< ..., ALUGrid< 2, 2, simplex, conforming >, ... >

#endif // HAVE_ALUGRID

class ESV2007Base
{
public:
  static std::string id()
  {
    return "eta_ESV2007";
  }
};

template< class SpaceType, class VectorType, class ProblemType, class GridType >
class ESV2007
  : public ESV2007Base
{
public:
  static const bool available = false;
};

#if HAVE_ALUGRID

template< class SpaceType, class VectorType, class ProblemType >
class ESV2007< SpaceType, VectorType, ProblemType, ALUGrid< 2, 2, simplex, conforming > >
  : public ESV2007Base
{
  typedef ALUGrid< 2, 2, simplex, conforming > GridType;
public:
  static const bool available = true;

  typedef typename ProblemType::RangeFieldType RangeFieldType;

  static RangeFieldType estimate(const SpaceType& space, const VectorType& vector, const ProblemType& problem)
  {
    LocalNonconformityESV2007< SpaceType, VectorType, ProblemType, GridType > eta_nc(space, vector, problem);
    LocalResidualESV2007< SpaceType, VectorType, ProblemType, GridType >      eta_r(space, problem);
    LocalDiffusiveFluxESV2007< SpaceType, VectorType, ProblemType, GridType > eta_df(space, vector, problem);
    eta_nc.prepare();
    eta_r.prepare();
    eta_df.prepare();

    RangeFieldType eta_squared(0.0);

    const auto& grid_view = space.grid_view();
    const auto entity_it_end = grid_view.template end< 0 >();
    for (auto entity_it = grid_view.template begin< 0 >(); entity_it != entity_it_end; ++entity_it) {
      const auto& entity = *entity_it;
      eta_squared += eta_nc.compute_locally(entity)
                   + std::pow(std::sqrt(eta_r.compute_locally(entity)) + std::sqrt(eta_df.compute_locally(entity)), 2);
    }
    return std::sqrt(eta_squared);
  } // ... estimate(...)

  static Stuff::LA::CommonDenseVector< RangeFieldType > estimate_local(const SpaceType& space,
                                                                       const VectorType& vector,
                                                                       const ProblemType& problem)
  {
    LocalNonconformityESV2007< SpaceType, VectorType, ProblemType, GridType > eta_nc(space, vector, problem);
    LocalResidualESV2007< SpaceType, VectorType, ProblemType, GridType >      eta_r(space, problem);
    LocalDiffusiveFluxESV2007< SpaceType, VectorType, ProblemType, GridType > eta_df(space, vector, problem);
    eta_nc.prepare();
    eta_r.prepare();
    eta_df.prepare();

    const auto& grid_view = space.grid_view();
    Stuff::LA::CommonDenseVector< RangeFieldType >
        local_indicators(boost::numeric_cast< size_t >(grid_view.indexSet().size(0)), 0.0);
    RangeFieldType eta_squared = 0.0;

    const auto entity_it_end = grid_view.template end< 0 >();
    for (auto entity_it = grid_view.template begin< 0 >(); entity_it != entity_it_end; ++entity_it) {
      const auto& entity = *entity_it;
      const auto index = grid_view.indexSet().index(entity);
      const RangeFieldType eta_t_squared
          = eta_nc.compute_locally(entity)
            + std::pow(std::sqrt(eta_r.compute_locally(entity)) + std::sqrt(eta_df.compute_locally(entity)), 2);
      local_indicators[index] = eta_t_squared;
      eta_squared += eta_t_squared;
    }
    for (auto& element : local_indicators)
      element /= eta_squared;
    return local_indicators;
  } // ... estimate_local(...)
}; // class ESV2007< ..., ALUGrid< 2, 2, simplex, conforming >, ... >

#endif // HAVE_ALUGRID

class ESV2007AlternativeSummationBase
{
public:
  static std::string id()
  {
    return "eta_ESV2007_alt";
  }
};

template< class SpaceType, class VectorType, class ProblemType, class GridType >
class ESV2007AlternativeSummation
  : public ESV2007AlternativeSummationBase
{
public:
  static const bool available = false;
};

#if HAVE_ALUGRID

template< class SpaceType, class VectorType, class ProblemType >
class ESV2007AlternativeSummation< SpaceType, VectorType, ProblemType, ALUGrid< 2, 2, simplex, conforming > >
  : public ESV2007AlternativeSummationBase
{
  typedef ALUGrid< 2, 2, simplex, conforming > GridType;
public:
  static const bool available = true;

  typedef typename ProblemType::RangeFieldType RangeFieldType;

  static RangeFieldType estimate(const SpaceType& space, const VectorType& vector, const ProblemType& problem)
  {
    LocalNonconformityESV2007< SpaceType, VectorType, ProblemType, GridType > eta_nc(space, vector, problem);
    LocalResidualESV2007< SpaceType, VectorType, ProblemType, GridType >      eta_r(space, problem);
    LocalDiffusiveFluxESV2007< SpaceType, VectorType, ProblemType, GridType > eta_df(space, vector, problem);
    eta_nc.prepare();
    eta_r.prepare();
    eta_df.prepare();

    RangeFieldType eta_nc_squared(0.0);
    RangeFieldType eta_r_squared(0.0);
    RangeFieldType eta_df_squared(0.0);

    const auto& grid_view = space.grid_view();
    const auto entity_it_end = grid_view.template end< 0 >();
    for (auto entity_it = grid_view.template begin< 0 >(); entity_it != entity_it_end; ++entity_it) {
      const auto& entity = *entity_it;
      eta_nc_squared += eta_nc.compute_locally(entity);
      eta_r_squared += eta_r.compute_locally(entity);
      eta_df_squared += eta_df.compute_locally(entity);
    }
    return std::sqrt(eta_nc_squared) + std::sqrt(eta_r_squared) + std::sqrt(eta_df_squared);
  } // ... estimate(...)

  static Stuff::LA::CommonDenseVector< RangeFieldType > estimate_local(const SpaceType& space,
                                                                       const VectorType& vector,
                                                                       const ProblemType& problem)
  {
    LocalNonconformityESV2007< SpaceType, VectorType, ProblemType, GridType > eta_nc(space, vector, problem);
    LocalResidualESV2007< SpaceType, VectorType, ProblemType, GridType >      eta_r(space, problem);
    LocalDiffusiveFluxESV2007< SpaceType, VectorType, ProblemType, GridType > eta_df(space, vector, problem);
    eta_nc.prepare();
    eta_r.prepare();
    eta_df.prepare();

    const auto grid_view = space.grid_view();
    Stuff::LA::CommonDenseVector< RangeFieldType >
        local_indicators(boost::numeric_cast< size_t >(grid_view.indexSet().size(0)), 0.0);
    RangeFieldType eta_nc_squared(0.0);
    RangeFieldType eta_r_squared(0.0);
    RangeFieldType eta_df_squared(0.0);

    const auto entity_it_end = grid_view.template end< 0 >();
    for (auto entity_it = grid_view.template begin< 0 >(); entity_it != entity_it_end; ++entity_it) {
      const auto& entity = *entity_it;
      const auto index = grid_view.indexSet().index(entity);
      const RangeFieldType eta_nc_t_squared = eta_nc.compute_locally(entity);
      const RangeFieldType eta_r_t_squared = eta_r.compute_locally(entity);
      const RangeFieldType eta_df_t_squared = eta_df.compute_locally(entity);
      eta_nc_squared += eta_nc_t_squared;
      eta_r_squared += eta_r_t_squared;
      eta_df_squared += eta_df_t_squared;
      local_indicators[index] = 3.0 *(eta_nc_t_squared + eta_r_t_squared + eta_df_t_squared);
    }
    const RangeFieldType eta_squared
        = std::pow(std::sqrt(eta_nc_squared) + std::sqrt(eta_r_squared) + std::sqrt(eta_df_squared), 2);
    for (auto& element : local_indicators)
      element /= eta_squared;
    return local_indicators;
  } // ... estimate_local(...)
}; // class ESV2007AlternativeSummation< ..., ALUGrid< 2, 2, simplex, conforming >, ... >

#endif // HAVE_ALUGRID

} // namespace SWIPDG
} // namespace internal

template< class SpaceType, class VectorType, class ProblemType, class GridType >
class SWIPDG
{
public:
  typedef typename ProblemType::RangeFieldType RangeFieldType;

private:
  template< class IndividualEstimator, bool available = false >
  class Caller
  {
  public:
    static std::vector< std::string > append(std::vector< std::string > in)
    {
      return in;
    }

    static bool equals(const std::string& /*type*/)
    {
      return false;
    }

    static RangeFieldType estimate(const SpaceType& /*space*/,
                                   const VectorType& /*vector*/,
                                   const ProblemType& /*problem*/)
    {
      DUNE_THROW(Stuff::Exceptions::internal_error, "This should not happen!");
      return RangeFieldType(0);
    }

    static Stuff::LA::CommonDenseVector< RangeFieldType > estimate_local(const SpaceType& /*space*/,
                                                                         const VectorType& /*vector*/,
                                                                         const ProblemType& /*problem*/)
    {
      DUNE_THROW(Stuff::Exceptions::internal_error, "This should not happen!");
      return Stuff::LA::CommonDenseVector< RangeFieldType >();
    }
  }; // class Caller

  template< class IndividualEstimator >
  class Caller< IndividualEstimator, true >
  {
  public:
    static std::vector< std::string > append(std::vector< std::string > in)
    {
      in.push_back(IndividualEstimator::id());
      return in;
    }

    static bool equals(const std::string& type)
    {
      return IndividualEstimator::id() == type;
    }

    static RangeFieldType estimate(const SpaceType& space, const VectorType& vector, const ProblemType& problem)
    {
      return IndividualEstimator::estimate(space, vector, problem);
    }

    static Stuff::LA::CommonDenseVector< RangeFieldType > estimate_local(const SpaceType& space,
                                                                         const VectorType& vector,
                                                                         const ProblemType& problem)
    {
      return IndividualEstimator::estimate_local(space, vector, problem);
    }
  }; // class Caller< ..., true >

  template< class IndividualEstimator >
  static std::vector< std::string > call_append(std::vector< std::string > in)
  {
    return Caller< IndividualEstimator, IndividualEstimator::available >::append(in);
  }

  template< class IndividualEstimator >
  static bool call_equals(const std::string& type)
  {
    return Caller< IndividualEstimator, IndividualEstimator::available >::equals(type);
  }

  template< class IndividualEstimator >
  static RangeFieldType call_estimate(const SpaceType& space, const VectorType& vector, const ProblemType& problem)
  {
    return Caller< IndividualEstimator, IndividualEstimator::available >::estimate(space, vector, problem);
  }

  template< class IndividualEstimator >
  static Stuff::LA::CommonDenseVector< RangeFieldType > call_estimate_local(const SpaceType& space,
                                                                            const VectorType& vector,
                                                                            const ProblemType& problem)
  {
    return Caller< IndividualEstimator, IndividualEstimator::available >::estimate_local(space, vector, problem);
  }

  typedef internal::SWIPDG::LocalNonconformityESV2007
      < SpaceType, VectorType, ProblemType, GridType >              LocalNonconformityESV2007Type;
  typedef internal::SWIPDG::LocalResidualESV2007
      < SpaceType, VectorType, ProblemType, GridType >              LocalResidualESV2007Type;
  typedef internal::SWIPDG::LocalResidualESV2007Star
      < SpaceType, VectorType, ProblemType, GridType >              LocalResidualESV2007StarType;
  typedef internal::SWIPDG::LocalDiffusiveFluxESV2007
      < SpaceType, VectorType, ProblemType, GridType >              LocalDiffusiveFluxESV2007Type;
  typedef internal::SWIPDG::ESV2007
      < SpaceType, VectorType, ProblemType, GridType >              ESV2007Type;
  typedef internal::SWIPDG::ESV2007AlternativeSummation
      < SpaceType, VectorType, ProblemType, GridType >              ESV2007AlternativeSummationType;

public:
  static std::vector< std::string > available()
  {
    std::vector< std::string > tmp;
    tmp = call_append< LocalNonconformityESV2007Type >(tmp);
    tmp = call_append< LocalResidualESV2007Type >(tmp);
    tmp = call_append< LocalResidualESV2007StarType >(tmp);
    tmp = call_append< LocalDiffusiveFluxESV2007Type >(tmp);
    tmp = call_append< ESV2007Type >(tmp);
    tmp = call_append< ESV2007AlternativeSummationType >(tmp);
    return tmp;
  } // ... available(...)

  static std::vector< std::string > available_local()
  {
    std::vector< std::string > tmp;
    tmp = call_append< ESV2007Type >(tmp);
    tmp = call_append< ESV2007AlternativeSummationType >(tmp);
    return tmp;
  } // ... available_local(...)

  static RangeFieldType estimate(const SpaceType& space,
                                 const VectorType& vector,
                                 const ProblemType& problem,
                                 const std::string type)
  {
    if (call_equals< LocalNonconformityESV2007Type >(type))
      return call_estimate< LocalNonconformityESV2007Type >(space, vector, problem);
    else if (call_equals< LocalResidualESV2007Type >(type))
      return call_estimate< LocalResidualESV2007Type >(space, vector, problem);
    else if (call_equals< LocalResidualESV2007StarType >(type))
      return call_estimate< LocalResidualESV2007StarType >(space, vector, problem);
    else if (call_equals< LocalDiffusiveFluxESV2007Type >(type))
      return call_estimate< LocalDiffusiveFluxESV2007Type >(space, vector, problem);
    else if (call_equals< ESV2007Type >(type))
      return call_estimate< ESV2007Type >(space, vector, problem);
    else if (call_equals< ESV2007AlternativeSummationType >(type))
      return call_estimate< ESV2007AlternativeSummationType >(space, vector, problem);
    else
      DUNE_THROW(Stuff::Exceptions::you_are_using_this_wrong,
                 "Requested type '" << type << "' is not one of available()!");
  } // ... estimate(...)

  static Stuff::LA::CommonDenseVector< RangeFieldType > estimate_local(const SpaceType& space,
                                                                       const VectorType& vector,
                                                                       const ProblemType& problem,
                                                                       const std::string type)
  {
    if (call_equals< ESV2007Type >(type))
      return call_estimate_local< ESV2007Type >(space, vector, problem);
    else if (call_equals< ESV2007AlternativeSummationType >(type))
      return call_estimate_local< ESV2007AlternativeSummationType >(space, vector, problem);
    else
      DUNE_THROW(Stuff::Exceptions::you_are_using_this_wrong,
                 "Requested type '" << type << "' is not one of available_local()!");
  } // ... estimate_local(...)
}; // class SWIPDG

} // namespace Discretizations
} // namespace Estimators
} // namespace HDD
} // namespace Dune

#endif // DUNE_HDD_LINEARELLIPTIC_ESTIMATORS_SWIPDG_HH


#include "GraphSampler.h"
#include <random>
#include <cassert>
#include <iostream>
#include <cmath>
#include <boost/math/distributions/hypergeometric.hpp>
#include <boost/math/distributions/hypergeometric.hpp>
#include <boost/math/policies/policy.hpp>

using namespace std;

GraphSampler::GraphSampler(TriangleCounter* counter) : counter_(counter) {
	if (counter_){
		counter_->clear();
	}
}

GraphSampler::~GraphSampler(){}

FixedPSampler::FixedPSampler(double p, bool use_sample_and_hold, TriangleCounter* counter)
	: GraphSampler(counter), p_(p), use_sample_and_hold_(use_sample_and_hold){
}

FixedPSampler::~FixedPSampler(){}

void FixedPSampler::exec_operation(const EdgeUpdate& update){

	counter_->new_update(update);

	if (update.is_add){
		double u_rand = (double)rand() / ((double)RAND_MAX+1.0);

		if(use_sample_and_hold_){ // always count first
			counter_->add_triangles(update.node_u, update.node_v, 1.0); //Weight not used
		}

		if (u_rand < p_){
			// Undirected graph
			if(!use_sample_and_hold_){ // count only if sampled
				counter_->add_triangles(update.node_u, update.node_v, 1.0); //Weight not used
			}
			counter_->add_edge_sample(update.node_u, update.node_v);

		}
	} else { // Remove are always executed (if not present no effect)
		assert(!use_sample_and_hold_); // Not supported.
		bool succeed = counter_->remove_edge_sample(update.node_u, update.node_v);
		// Here !use_sample_and_hold_
		if (succeed){ // it was present
			counter_->remove_triangles(update.node_u, update.node_v, 1.0); //Weight not used
		}
	}
}

double FixedPSampler::get_triangle_est(){
	if (!use_sample_and_hold_){
		return (double)counter_->triangles()*pow(1.0/p_,3);
	} else {
		return (double)counter_->triangles()*pow(1.0/p_,2);
	}
}

double FixedPSampler::get_triangle_est_local(int node){
	assert(counter_->is_local());
	if (!use_sample_and_hold_){
		return (double)counter_->triangles_local(node)*pow(1.0/p_,3);
	} else {
		return (double)counter_->triangles_local(node)*pow(1.0/p_,2);
	}
}

//RESERVOIR SAMPLER

ReservoirSampler::ReservoirSampler(size_t reservoir_size, bool use_sample_and_hold, TriangleCounter* counter)
	: GraphSampler(counter), reservoir_size_(reservoir_size), use_sample_and_hold_(use_sample_and_hold){
		reservoir_.reserve(reservoir_size);
}

ReservoirSampler::~ReservoirSampler(){}

void ReservoirSampler::add_reservoir(const pair<int,int> edge){
	//cout<<"ADD RES"<<edge.first<<" "<<edge.second<<endl;
	assert(reservoir_map_.find(edge) == reservoir_map_.end());
	reservoir_.push_back(edge);
	reservoir_map_.insert(make_pair(edge, reservoir_.size()-1));
	counter_->add_edge_sample(edge.first, edge.second);
	if (!use_sample_and_hold_){
		counter_->add_triangles(edge.first, edge.second, 1.0); //Weight not used
	}
}

void ReservoirSampler::delete_reservoir(const pair<int,int> edge){
	if (reservoir_map_.find(edge) == reservoir_map_.end()){
		return;
	}
	int pos = reservoir_map_.at(edge);
	reservoir_map_.erase(edge);
	if (pos < reservoir_.size() -1){ // not the last item
		pair<int, int> last_edge = reservoir_.back();
		reservoir_[pos] = last_edge;
		reservoir_map_[last_edge] = pos;
	}
	reservoir_.pop_back();
	bool succ = counter_->remove_edge_sample(edge.first, edge.second);
	if (!use_sample_and_hold_){
		counter_->remove_triangles(edge.first, edge.second, 1.0); //Weight not used
	}

	assert(succ);
	assert(reservoir_.size()<=reservoir_size_);
	assert(reservoir_.size()==reservoir_map_.size());
}

void ReservoirSampler::exec_operation(const EdgeUpdate& update){
	assert(update.is_add); //only add supported
	assert(update.node_u != update.node_v);
	int max_n = (int)max(update.node_u, update.node_v);
	int min_n = (int)min(update.node_u, update.node_v);

	pair<int,int> edge = make_pair(min_n, max_n);

	counter_->new_update(update);

	//if(update.is_add){
	assert(reservoir_map_.find(edge) == reservoir_map_.end());

	// Prob of sampling at this step if using sample and hold
  double p = 0;
  unsigned long long int t = counter_->edges_present_original();

  if (t >= 2 && use_sample_and_hold_){
    p = ((double)reservoir_size_/t)*((double)(reservoir_size_-1)/(t-1));
    p = min(p, 1.0);
  }

	if (use_sample_and_hold_){ // Always count and rever remove.
		counter_->add_triangles(update.node_u, update.node_v, (double)1.0/p);
	}

	// Enough space
	if (reservoir_.size() < reservoir_size_){
		add_reservoir(edge);
	} else {
		double u_rand = (double)rand() / ((double)RAND_MAX+1.0);
		double thres = ((double)reservoir_size_)/counter_->edges_present_original();
		if (u_rand < thres){
			// Doing the exchange
			int rand_pos = rand()%reservoir_size_;
			const pair<int,int>& to_remove = reservoir_[rand_pos];
			delete_reservoir(to_remove);
			add_reservoir(edge);
		}
	}
	//} else { // is remove
	//	delete_reservoir(edge);
	//}
}

 double ReservoirSampler::get_triangle_est(){
	if (!use_sample_and_hold_){

    double p = 0;
    unsigned long long int t = counter_->edges_present_original();

    if (t >= 3){
      p = ((double)reservoir_size_/t)*((double)(reservoir_size_-1)/(t-1))
          *((double)(reservoir_size_-2)/(t-2));
      p = min(p, 1.0);
      return (double)counter_->triangles()/p;

    } else {
      return 0; // no triangles possible;
    }

	} else {
		return counter_->triangles_weight();
	}
}

 double ReservoirSampler::get_triangle_est_local(int node){
	assert(counter_->is_local());

	if (!use_sample_and_hold_){
    unsigned long long int t = counter_->edges_present_original();
    if (t >= 3){
      double p = ((double)reservoir_size_/t)*((double)(reservoir_size_-1)/(t-1))
          *((double)(reservoir_size_-2)/(t-2));
      p = min(p, 1.0);
      return (double)counter_->triangles_local(node)/p;

    } else {
      return 0; // no triangles possible;
    }
	} else {
		return counter_->triangles_weight_local(node);
	}
}

// ****************************************
// Reservoir add & remove

ReservoirAddRemSampler::ReservoirAddRemSampler(size_t reservoir_size, TriangleCounter* counter)
	: GraphSampler(counter), reservoir_size_(reservoir_size), d_i_(0), d_o_(0){
		reservoir_.reserve(reservoir_size);
}

ReservoirAddRemSampler::~ReservoirAddRemSampler(){}

void ReservoirAddRemSampler::add_reservoir(const pair<int,int> edge){
	//cout<<"ADD RES"<<edge.first<<" "<<edge.second<<endl;
  int before_size = reservoir_.size();

	assert(reservoir_map_.find(edge) == reservoir_map_.end());
	reservoir_.push_back(edge);
	reservoir_map_.insert(make_pair(edge, reservoir_.size()-1));
	bool succ = counter_->add_edge_sample(edge.first, edge.second);

  assert (succ);
  assert(reservoir_.size()<=reservoir_size_);
  assert(reservoir_.size()==reservoir_map_.size());
  assert (reservoir_map_.find(edge) != reservoir_map_.end());
  assert(reservoir_.size()== before_size +1);
}

void ReservoirAddRemSampler::delete_reservoir(const pair<int,int> edge){
  //cout<<"REM RES"<<edge.first<<" "<<edge.second<<endl;

  assert (reservoir_map_.find(edge) != reservoir_map_.end());
  int before_size = reservoir_.size();

	int pos = reservoir_map_.at(edge);
	reservoir_map_.erase(edge);
	if (pos < reservoir_.size() -1){ // not the last item
		pair<int, int> last_edge = reservoir_.back();
		reservoir_[pos] = last_edge;
		reservoir_map_[last_edge] = pos;
	}
	reservoir_.pop_back();
	bool succ = counter_->remove_edge_sample(edge.first, edge.second);
	counter_->remove_triangles(edge.first, edge.second, 1.0); //Weight not used

	assert(succ);
	assert(reservoir_.size()<=reservoir_size_);
	assert(reservoir_.size()==reservoir_map_.size());
  assert (reservoir_map_.find(edge) == reservoir_map_.end());
  assert(reservoir_.size()== before_size -1);
}

void ReservoirAddRemSampler::exec_operation(const EdgeUpdate& update){
	assert(update.node_u != update.node_v);
	int max_n = (int)max(update.node_u, update.node_v);
	int min_n = (int)min(update.node_u, update.node_v);

	pair<int,int> edge = make_pair(min_n, max_n);

	counter_->new_update(update);

	if(update.is_add){
		assert(reservoir_map_.find(edge) == reservoir_map_.end());

    if (d_o_ + d_i_ > 0) { // case d_o + d_i > 0
      double u_rand = (double)rand() / ((double)RAND_MAX+1.0);
			double thres = ((double)d_i_)/(d_i_+d_o_);
      if (u_rand < thres){ //with pro d_i / (d_i + d_o)
        d_i_ --;
        assert(reservoir_.size() < reservoir_size_);

        add_reservoir(edge);
        counter_->add_triangles(edge.first, edge.second, 1.0); //Weight not used
      } else {
        d_o_ --;
      }
    } else if (reservoir_.size() < reservoir_size_){ // enough space and d_i + d_o = 0

			add_reservoir(edge);
      counter_->add_triangles(edge.first, edge.second, 1.0); //Weight not used

		} else { // reservoid full and d_i + d_o = 0
      assert (counter_->edges_present_original()>reservoir_size_);

			double u_rand = (double)rand() / ((double)RAND_MAX+1.0);
			double thres = ((double)reservoir_size_)/(counter_->edges_present_original());
			if (u_rand < thres){
				// Doing the exchange
				int rand_pos = rand()%reservoir_size_;
				pair<int,int> to_remove = reservoir_[rand_pos];

        size_t before_size = reservoir_.size();

				delete_reservoir(to_remove);
				add_reservoir(edge);

        assert(reservoir_.size()<=reservoir_size_);
        assert(reservoir_.size()==reservoir_map_.size());
        assert (reservoir_map_.find(to_remove) == reservoir_map_.end());
        assert (reservoir_map_.find(edge) != reservoir_map_.end());
        assert(reservoir_.size()== before_size);

				counter_->add_triangles(edge.first, edge.second, 1.0); //Weight not used
			}
		}
	} else { // is remove
    assert(counter_->edges_present_original()>=0);

		if (reservoir_map_.find(edge) != reservoir_map_.end()){//Was present
      d_i_ ++; // deletion in sample

			delete_reservoir(edge);
		} else { // the edge deleted was not in the reservoir
      d_o_ ++; // deletion in sample

    }
	}
}

double ReservoirAddRemSampler::prob_sampling_triangle() const{
  if (counter_->edges_present_original() < 3 || reservoir_.size() < 3) {
    return 0;
  }
  unsigned long long int s = counter_->edges_present_original();

  //cout<<"SIZES: "<<reservoir_.size()<< " " <<s<<endl;
  assert ((unsigned long long int)reservoir_.size()<=s);

  double mt = reservoir_.size();

  double p = 0;

  if (reservoir_.size() == s) {
    p = 1;
  } else {
    p = (mt/s)*((mt-1)/(s-1))*((mt-2)/(s-2));
  }

  assert (p<=1);

  unsigned long long int n = min((unsigned long long int)reservoir_size_, s+d_i_+d_o_);

  boost::math::hypergeometric_distribution<double> hyper(s, n, d_i_+d_o_+s);
  double kt = 0;

  for (int i = 0; i<=2 ; i++){
    if (i >= max(0ull, (unsigned long long int)(n) +  - d_i_-d_o_) && i <= min((unsigned long long int)(n), s)){
      kt += boost::math::pdf<double>(hyper, i);
    }
  }

  return p*(1-kt);
}

double ReservoirAddRemSampler::get_triangle_est(){

  double prob = prob_sampling_triangle();
	return counter_->triangles() / prob;
}

double ReservoirAddRemSampler::get_triangle_est_local(int node){
	assert(counter_->is_local());

  double prob = prob_sampling_triangle();
	return counter_->triangles_local(node) / prob;
}

// ALI PINAR Paper

PinarSampler::PinarSampler(size_t edge_res_size, size_t wedge_res_size)
	: GraphSampler(NULL/* no need of TriangleCounter*/), t_(0), tot_wedges_(0), fraction_closed_(0.0), edge_res_size_(edge_res_size), wedge_res_size_(wedge_res_size){
		edge_reservoir_.resize(edge_res_size);
		wedge_reservoir_.resize(wedge_res_size);
		wedge_closed_.resize(wedge_res_size);
}

PinarSampler::~PinarSampler(){}

double PinarSampler::get_triangle_est(){
	return fraction_closed_*t_*t_/
		(static_cast<double>(edge_res_size_)*(edge_res_size_-1))*tot_wedges_;
}

void PinarSampler::exec_operation(const EdgeUpdate& update){
	assert(update.is_add); //only add supported
	assert(update.node_u != update.node_v);
	int max_n = (int)max(update.node_u, update.node_v);
	int min_n = (int)min(update.node_u, update.node_v);

	t_+=1;

	pair<int,int> edge = make_pair(min_n, max_n);

	long long int closed = 0;

	// Check if closing wedges
	for (int i = 0; i<wedge_reservoir_.size(); i++){
		const pair<int,pair<int,int>> &wedge = wedge_reservoir_[i];

		if ((wedge.second.first == max_n && wedge.second.second == min_n)
				||(wedge.second.first == min_n && wedge.second.second == max_n)
		){
				wedge_closed_[i] = true;

		}
		closed += (wedge_closed_[i] ? 1 : 0);
	}

	// Update edge reservoir
	bool updated = false;
	for (int i = 0; i<edge_reservoir_.size(); i++){
		double u_rand = (double)rand() / ((double)RAND_MAX+1.0);
		if (u_rand <= 1.0/t_){
			edge_reservoir_[i] = edge;
			updated = true;
		}
	}

	if(updated){
		// count total num of wedges
		UDynGraph sub_graph;
		for (const auto& e: edge_reservoir_){
			sub_graph.add_edge(e.first, e.second);
		}
		vector<int> nodes;
		sub_graph.nodes(&nodes);

		tot_wedges_ = 0;
		for (const auto& n:nodes){
			int deg = sub_graph.degree(n);
			tot_wedges_+= (deg*(deg-1))/2;
		}

		// get new wedges with this edge...
		vector<pair<int,pair<int,int>>> new_wedges;
		vector<int> neighbors_min;
		sub_graph.neighbors(min_n, &neighbors_min);
		for (const auto & neighbor: neighbors_min){
			// wedge min_n-neighbor, min_n-max_n
			if (neighbor != max_n){
				new_wedges.push_back(make_pair(min_n, make_pair(neighbor, max_n)));
			}
		}
		vector<int> neighbors_max;
		sub_graph.neighbors(max_n, &neighbors_max);
		for (const auto & neighbor: neighbors_max){
			// wedge max_n-neighbor, max_n-min_n
			if (neighbor != min_n){
				new_wedges.push_back(make_pair(max_n, make_pair(neighbor, min_n)));
			}
		}

		for (int i = 0; i<wedge_reservoir_.size(); i++){
			double u_rand = (double)rand() / ((double)RAND_MAX+1.0);
			if (u_rand <= 1.0*new_wedges.size()/tot_wedges_){
				wedge_reservoir_[i] = new_wedges[rand()%new_wedges.size()];
				closed -= (wedge_closed_[i] ? 1 : 0);
				wedge_closed_[i] = false; // This make absolutely no sense but it is done in ali pinar paper So I implemented it as stated. ****
				closed += (wedge_closed_[i] ? 1 : 0);
			}
		}

		fraction_closed_ = 1.0*closed/ wedge_reservoir_.size();
	}

}


//
//  bpp.cpp
//  PhyloAcc
//
//  Created by hzr on 3/8/16.
//  Copyright © 2016 hzr. All rights reserved.
//

#include "bpp.hpp"
#include <armadillo>
#include <sys/types.h>
#include <dirent.h>
#include<queue>

#include <cmath>
#include <cassert>

#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_cdf.h>

#include <fstream>
#include <sstream>
#include <iostream>
#include <iomanip>
#include <algorithm>
#include <ctype.h>

#include "newick.h"
#include "utils.h"
#include "bpp_c.hpp"

using namespace std;
using namespace arma;

// load the phylogenetic tree
void BPP::InitPhyloTree(PhyloTree & tree) //, double indel_pi), double indel, double indel2
{

//    cx_mat bvec;
//    cx_vec aval;
//
//    if(num_base <= 4)  // no indel
//    {
//
//    }else{
//        tree.subs_rate *= (1-indel);
//        mat B = ones<mat>(4,1) * indel;
//        mat C = ones<mat>(1,5) * indel2;
//        //tree.subs_rate.diag() -= indel;
//        tree.subs_rate.insert_cols(4, B);
//        tree.subs_rate.insert_rows(4, C);
//        //tree.subs_rate(4,4) = -0.04;
//        colvec c = sum(tree.subs_rate,1);
//        tree.subs_rate.diag() -=c;
//
////        tree.pi.insert_rows(4,1);
////        tree.pi.head(4) *= 1 - indel_pi;
////        tree.pi[4] = indel_pi;
//
//    }
//
//    eig_gen(aval, bvec, tree.subs_rate);
//    eigenval = conv_to<mat>::from(aval);
//    eigenvec = conv_to<mat>::from(bvec).t();
//    eigeninv = inv(eigenvec);
//    submat = tree.subs_rate;
//
//    //cout <<eigenval;
//    //cout <<"eigenvec: " << eigenvec;
//    //cout <<tree.subs_rate;
//
//    mat a = null(tree.subs_rate.t());
//
//
//    pi = a/accu(a); //tree.pi;
//    //cout <<"pi: " <<  pi.t();
//
//    log_pi = log(pi);

    submat = tree.subs_rate;
    children    = new int[N][2];
    parent      = new int[N];
    distances   = new double[N];

    for(int s=0; s<N; s++)
    {
        distances[s] = tree.distances[s];

    }

    for(int i=0; i<N; i++)
    {
        children[i][0] = -1;
        children[i][1] = -1;
        parent[i] = -1;
    }

    for(int i=0; i<N; i++)
    {
        int p = -1;
        for(int j=0; j<N; j++)
        {
            if (tree.dag[i][j])
            {
                p++;
                children[i][p] = j;
                parent[j] = i;
            }
        }
    }

//    distances[children[N-1][1]] += distances[children[N-1][0]];
//    distances[children[N-1][0]] = 0;
//    moveroot = children[N-1][0];

    //distances[83] += distances[42];  // modify distance for root!
    //distances[42] =0 ;

}

// try to match the phylogenetic profile and tree
void BPP::MatchProfAndTree(PhyloProf & _prof, PhyloTree & _tree)
{

    // try to match the species
    bool success_match = true;
    int S = _tree.S;
    int S2 = _prof.S;
    vector<int> reorder(S);  //each species in the tree where is in prof
    for(int s1=0; s1<S; s1++)
    {
        bool has_same_species = false;
        string sname1 = _tree.species_names[s1];
        for(int s2=0; s2<S2; s2++)
        {
            string sname2 = _prof.species_names[s2];
            //            cout << sname1 << " ? " << sname2 << endl;
            if (sname1 == sname2)
            {
                has_same_species = true;
                reorder[s1] = s2;
                break;
            }
        }
        if (!has_same_species)
        {
            cout << "No matrix species " << _prof.species_names[s1] << " found in tree." << endl;
            success_match = false;
            break;
        }
    }

    if (!success_match) // if cannot match literally
    {
        cout << endl << "The species in phylogenetic profile and tree cannot be matched literally:" << endl;
        cout << "The program will use the default mapping in data:" << endl;
        for(int s=0; s<S; s++)
            cout << "(" << _prof.species_names[s] << "\t=  " << _tree.species_names[s] << ")" << endl;
        cout << endl;
    }
    else                // if successully matched
    {
        cout << "The species in profile and tree match perfectly. Reorder the species in profile matrix by the tree." << endl << endl;
        vector<string> old_X = _prof.X;
        for(int s=0; s<S; s++)
        {
            int reorder_s = reorder[s];
            _prof.X[s] = old_X[reorder_s];
        }
        _prof.species_names = _tree.species_names;
    }

}

void BPP::InitMCMC(int _num_burn, int _num_mcmc, int _num_thin)
{
    // init parameters
    num_burn = _num_burn;

    num_mcmc = _num_mcmc;
    num_thin = _num_thin;

    last_time = time(NULL);

    // init MCMC sampling storage
    Max_Z = vector<vector < vector <int > >> (3,vector < vector <int >>  (C,vector <int > (N,0)));
    cur_Z = vector<vector < vector <int > >> (3,vector < vector <int >>  (C,vector <int > (N,0)));

    log_liks_null = vector <double>(C,0);
    log_liks_Z = vector<vector <double>>(3,vector <double>(C,0));
    log_liks_sgl = vector <double>(C,0);
    log_liks_resZ = vector <double>(C,0);

    log_liks_curZ = vector <double>(C,0);
    log_liks_propZ = vector <double>(C,0);
    MH_ratio_gain = vector <double>(C,0);
    MH_ratio_loss = vector <double>(C,0);

    cur_crate = vector <double>(C,ratio0);
    cur_nrate = vector <double>(C,ratio1);

    cur_lrate = vector <double>(C,ind_lrate);
    cur_lrate2 = vector <double>(C,ind_lrate2);
    cur_grate = vector <double>(C, ind_grate);

}

void BPP::sample_proposal(int iter, double & lrate_prop, double & grate_prop, ofstream & output)
{
    //indel_prop =gsl_ran_gamma(RNG,20 ,indel/20);
    //indel2_prop =gsl_ran_gamma(RNG,10 ,indel2/10);
    lrate_prop =gsl_ran_beta(RNG, ind_lrate * vlr, (1 - ind_lrate) *vlr); // let vlr == vgr!!
    grate_prop =gsl_ran_beta(RNG, ind_grate * vgr, (1 - ind_grate - ind_lrate) *vgr); //gsl_ran_gamma(RNG, vgr, ind_grate/vgr);
    grate_prop = grate_prop * (1 - lrate_prop);

//    submat.submat(0,0,3,3) *= (1 - indel_prop)/(1 - indel);
//    submat.col(4) = ones<mat>(5,1) * indel_prop;
//    submat.row(4) = ones<mat>(1,5) * indel2_prop;
//    colvec c = sum(submat,1);
//    submat.diag() -=c;
//
//    cx_mat bvec;
//    cx_vec aval;
//    eig_gen(aval, bvec, submat);
//    eigenvalprop = conv_to<mat>::from(aval);
//    eigenvecprop = conv_to<mat>::from(bvec).t();
//    eigeninvprop = inv(eigenvecprop);
//
//    mat a = null(submat.t());
//    piprop = a/accu(a);
//    //log_piprop = log(a/accu(a));
//
//    //cout <<"proposed subs matrix: " << submat << endl;
//    //cout <<"piprop: " <<  piprop.t();

    output << iter << "\t"<< nprior_a<< "\t"<< nprior_b <<"\t"<< cprior_a << "\t"<< cprior_b << "\t"<< indel << "\t"<< indel2 << "\t"<<ind_grate<< "\t"<< ind_lrate <<endl;
}

double BPP::log_lik(vector< vector<vec> > & lambda, double _indel, double _indel2, int start1, int end1, vector<unsigned int> & v, double p)
{
    // compute loglik
    double result =0;
    mat x(2,2);

    int rr = *subtree.rbegin();
    // 1. sending the lambda msg from leaves bottom up through the network
    for(vector<int>::iterator it = subtree.begin(); it!=subtree.end(); it++) //int s=S; s<N; s++)
    {
        int s = *it;
        if(s<S) continue;
        int* p = children[s];
        for(int it = start1; it < end1; it++)  lambda[v[it]][s].fill(0);

        for(int cc=0;cc<2;cc++)
        {
            int chi = p[cc];
            assert(chi != -1);
            if(distances[chi]>0 )
            {
                double tt = (1 - exp(-(_indel + _indel2) * distances[chi]))/(_indel + _indel2);
                x.at(1,0) = _indel * tt;
                x.at(0,0) = 1 - x.at(1,0);

                x.at(0,1) = _indel2 * tt;
                x.at(1,1) = 1 - x.at(0,1);

                //cout << x;
                x = log(x);

            }
            else{
                x.fill(-INFINITY); //83, root
                x.diag().fill(0);
            }

            #pragma omp parallel for schedule (guided)
            for(int it = start1; it < end1; it++)  lambda[v[it]][s] +=  BPP::log_multi(x,lambda[v[it]][chi]);
        }

    }

    // 2. processing the distribution of root species

    for(int it = start1; it < end1; it++)
    {

//        if(lambda[v[it]][N-1][1] < -1e3)
//        {
//            cout <<v[it]<<": "<< lambda[v[it]][N-1].t();
//            cout <<children[N-1][0]<<": " << lambda[v[it]][children[N-1][0]].t();
//            cout <<children[N-1][1]<<": " << lambda[v[it]][children[N-1][1]].t();
//
//        }
        lambda[v[it]][rr][0] += log(1-p); //N-1
        lambda[v[it]][rr][1] += log(p) ;
        result += BPP::log_exp_sum(lambda[v[it]][rr]);
    }

    return(result);
}

void BPP::sample_hyperparam(int iter, vector<int> & ids, ofstream & output) // recompute log_TM, double indel_prop, double indel2_prop,
{
// 0702 no sample rates
    //indepent MH to sample hyperparam of rates
    double p=1,r = 1; //hyperparam for shape
    double q=0.1,s = 0.1; // hyperparam for scale

    double vna = 100, vnb = 100, vca = 100, vcb = 100;
    double nprior_a_prop =gsl_ran_gamma(RNG, vna, nprior_a/vna);
    double cprior_a_prop =gsl_ran_gamma(RNG, vca, cprior_a/vca);

    double nprior_b_prop =gsl_ran_gamma(RNG, vnb, nprior_b/vnb);
    double cprior_b_prop =gsl_ran_gamma(RNG, vcb, cprior_b/vcb);

    //MH proposal
    double sum_r = 0;
    double log_prod_r = 0;
    //double var_r = 0;
    //for(vector<double>::iterator it = cur_nrate.begin(); it< cur_nrate.end(); it++)
    for(std::size_t i = 0; i < ids.size(); i++ )  // 0702
    {
        int c = ids[i];
        sum_r += cur_nrate[c];
        //var_r += pow(*it, 2);
        log_prod_r += log(cur_nrate[c]);
    }

    double M_ratio = (nprior_a_prop - 1) * (log(p) + log_prod_r) - (q + sum_r)/nprior_b_prop - (ids.size() + r)*lgamma(nprior_a_prop) - log(nprior_b_prop) * nprior_a_prop * (s + ids.size());
    M_ratio -= (nprior_a - 1) * (log(p) + log_prod_r) - (q + sum_r)/nprior_b - (ids.size() + r)*lgamma(nprior_a) - log(nprior_b) * nprior_a * (s + ids.size());

    double H_ratio = log(gsl_ran_gamma_pdf(nprior_a,vna,nprior_a_prop/vna)) - log(gsl_ran_gamma_pdf(nprior_a_prop,vna,nprior_a/vna)) + log(gsl_ran_gamma_pdf(nprior_b,vnb,nprior_b_prop/vnb)) - log(gsl_ran_gamma_pdf(nprior_b_prop,vnb,nprior_b/vnb));

    cout << "nrate_MH_ratio: " << M_ratio <<", " << H_ratio << ", " << nprior_a << ", " << nprior_a_prop << ", " << nprior_b << ", " << nprior_b_prop << endl;

    if(log(gsl_rng_uniform(RNG)) < M_ratio + H_ratio)
    {
        nprior_a = nprior_a_prop;
        nprior_b = nprior_b_prop;
    }

    sum_r = 0;
    log_prod_r = 0;
    //for(vector<double>::iterator it = cur_crate.begin(); it< cur_crate.end(); it++)
    for(std::size_t i = 0; i < ids.size(); i++ )  // 0702
    {
        int c = ids[i];
        sum_r += cur_crate[c];
        //var_r += pow(*it, 2);
        log_prod_r += log(cur_crate[c]);
    }

    M_ratio = (cprior_a_prop - 1) * (log(p) + log_prod_r) - (q + sum_r)/cprior_b_prop - (ids.size() + r)*lgamma(cprior_a_prop) - log(cprior_b_prop) * cprior_a_prop * (s + ids.size()) - ((cprior_a - 1) * (log(p) + log_prod_r) - (q + sum_r)/cprior_b - (ids.size() + r)*lgamma(cprior_a) - log(cprior_b) * cprior_a * (s + ids.size()));

    H_ratio = log(gsl_ran_gamma_pdf(cprior_a,vca,cprior_a_prop/vca)) - log(gsl_ran_gamma_pdf(cprior_a_prop,vca,cprior_a/vca)) + log(gsl_ran_gamma_pdf(cprior_b,vcb,cprior_b_prop/vcb)) - log(gsl_ran_gamma_pdf(cprior_b_prop,vcb,cprior_b/vcb));

    cout << "crate_MH_ratio: " << M_ratio + H_ratio << ", " << cprior_a << ", " << cprior_a_prop <<", " << cprior_b << ", " << cprior_b_prop << endl;

    if(log(gsl_rng_uniform(RNG)) < M_ratio + H_ratio)
    {
        cprior_a = cprior_a_prop;
        cprior_b = cprior_b_prop;
    }

  // sample hyperparameters of lrate and grate, exponential prior for prior_l_a and prior_l_b, prior_g_a and prior_g_b
    double u = gsl_rng_uniform(RNG)*(1.3 - 0.7) + 0.7; // generate uniform from (1/prop_n, prop_n);
    double prior_l_a_prop = prior_l_a *  u;

    u = gsl_rng_uniform(RNG)*(1.3 - 0.7) + 0.7; // generate uniform from (1/prop_n, prop_n);
    double prior_l_b_prop = prior_l_b *  u;
    double log_p = 0, log_pc = 0;
    for(std::size_t i = 0; i < ids.size(); i++ )  // 0702
    {
        int c = ids[i];
        log_p += log(cur_lrate[c]);
        log_pc += log(1 - cur_lrate[c]);
    }

    M_ratio = (prior_l_a_prop - prior_l_a) * (log_p - 1) + (prior_l_b_prop - prior_l_b) * (log_pc - 1);
    M_ratio += ids.size() * (gsl_sf_lnbeta(prior_l_a, prior_l_b) - gsl_sf_lnbeta(prior_l_a_prop, prior_l_b_prop));

    H_ratio = log(prior_l_a) - log(prior_l_a_prop) + log(prior_l_b) - log(prior_l_b_prop);

    cout << "lrate_MH_ratio: " << M_ratio + H_ratio << ", " << prior_l_a << ", " << prior_l_a_prop <<", " << prior_l_b << ", " << prior_l_b_prop  << endl;

    if(log(gsl_rng_uniform(RNG)) < M_ratio + H_ratio)
    {
        prior_l_a = prior_l_a_prop;
        prior_l_b = prior_l_b_prop;
    }

    // sample grate
    u = gsl_rng_uniform(RNG)*(1.3 - 0.7) + 0.7; // generate uniform from (1/prop_n, prop_n);
    double prior_g_a_prop = prior_g_a *  u;

    u = gsl_rng_uniform(RNG)*(1.3 - 0.7) + 0.7; // generate uniform from (1/prop_n, prop_n);
    double prior_g_b_prop = prior_g_b *  u;
    log_p = 0; log_pc = 0;
    for(std::size_t i = 0; i < ids.size(); i++ )  // 0702
    {
        int c = ids[i];
        log_p += log(cur_grate[c]);
        log_pc += log(1 - cur_grate[c]);
    }

    M_ratio = (prior_g_a_prop - prior_g_a) * (log_p - 1) + (prior_g_b_prop - prior_g_b) * (log_pc - 1);
    M_ratio += ids.size() * (gsl_sf_lnbeta(prior_g_a, prior_g_b) - gsl_sf_lnbeta(prior_g_a_prop, prior_g_b_prop));

    H_ratio = log(prior_g_a) - log(prior_g_a_prop) + log(prior_g_b) - log(prior_g_b_prop);

    cout << "grate_MH_ratio: " << M_ratio + H_ratio << ", " << prior_g_a << ", " << prior_g_a_prop <<", " << prior_g_b << ", " << prior_g_b_prop  << endl;

    if(log(gsl_rng_uniform(RNG)) < M_ratio + H_ratio)
    {
        prior_g_a = prior_g_a_prop;
        prior_g_b = prior_g_b_prop;
    }

     output << iter << "\t"<< nprior_a<< "\t"<< nprior_b <<"\t"<< cprior_a << "\t"<< cprior_b << "\t"<< prior_l_a << "\t"<< prior_l_b << "\t"<< prior_g_a << "\t"<< prior_g_b <<endl;

}

void BPP::getUppertree(int root, vector<int>& child, set<int> & visited_init)  // include root!
{
    for(vector<int>::iterator it = child.begin(); it!=child.end(); it++)
    {
        int p = *it;
       while(p!=root)
       {
           visited_init.insert(p);
           p = parent[p];
       }

    }

    visited_init.insert(root);

}

void BPP::getSubtree(int root, vector<int> & visited_init)  // traverse from root to children, include root
{

    int j = root;

    //cout << nodes_names[j]<<"\t";

    if(children[j][0]!=-1)
    {

        for(int chi=0;chi<2;chi++)
        {
            getSubtree(children[j][chi], visited_init);

        }

    }

    visited_init.push_back(j);

}

void BPP::getSubtree(int root, set<int>& child, vector<int> & visited_init)  // traverse from root, stop at children, 74 & 64; do include 1-S!
{

    int j = root;

    //cout << nodes_names[j]<<"\t";
    if(child.find(j) != child.end())
    {

        visited_init.push_back(j);

        return;
    }

    if(children[j][0]!=-1)
    {

        for(int chi=0;chi<2;chi++)
        {
            getSubtree(children[j][chi], child, visited_init);

        }

      //  visited_init.push_back(j);

    }

     visited_init.push_back(j);

}

void BPP::Output_init(PhyloProf & prof, string output_path, vector<int> & ids){

        string outpath_elem = output_path+ "_elem_lik.txt";
        ofstream out_lik(outpath_elem.c_str());
        out_lik.precision(8);
        out_lik << "No.\tID\tloglik_Null\tloglik_Acc\tloglik_Full\tlogBF1\tlogBF2\tloglik_Max_M0\tloglik_Max_M1\tloglik_Max_M2"<<endl;
        //for(int cc=0; cc<C;cc++)
        for(vector<int>::iterator it = ids.begin(); it !=ids.end(); it++)
        {
            int cc = *it;
            out_lik <<cc << "\t" << prof.element_names[cc] << "\t" << log_liks_null[cc] <<"\t"  <<log_liks_resZ[cc] <<"\t"  <<log_liks_sgl[cc]<< "\t";
            out_lik << log_liks_resZ[cc] -  log_liks_null[cc] << "\t" << log_liks_resZ[cc] -  log_liks_sgl[cc];
            //for(int r=0;r<3;r++)
            out_lik <<"\t" <<log_liks_Z[0][cc] << "\t" <<log_liks_Z[2][cc]<<"\t" <<log_liks_Z[1][cc];
            out_lik << endl;
        }

        out_lik.close();

    ofstream out_z;
    for(int r =0;r<3;r++)
    {
            if(r == 2)
            {
                outpath_elem = output_path+"_M" +to_string(1) + "_elem_Z.txt";
            }else if(r == 1)
            {
                outpath_elem = output_path+"_M" +to_string(2) + "_elem_Z.txt";
            }else{
                outpath_elem = output_path+"_M" +to_string(0) + "_elem_Z.txt";
            }

            out_z.open(outpath_elem.c_str());
            out_z<<"No.";
            for(int s =0 ;s<N;s++){  // header: species name
                out_z<< "\t" << nodes_names[s] ;
            }
            out_z <<endl;

            for(vector<int>::iterator it = ids.begin(); it !=ids.end(); it++)
            {
                int c = *it;
                out_z<<c;
                for(int s=0; s<N;s++)
                    out_z <<"\t" << Max_Z[r][c][s];
                out_z <<endl;
            }
        out_z.close();
    }

}

void BPP::Output_init0(PhyloProf & prof, ofstream& out_lik, vector<int> & ids){

    //for(int cc=0; cc<C;cc++)
    for(vector<int>::iterator it = ids.begin(); it !=ids.end(); it++)
    {
        int cc = *it;
        out_lik <<cc << "\t" << prof.element_names[cc] << "\t"  <<log_liks_sgl[cc]<< "\t"<<log_liks_Z[1][cc];
        out_lik << endl;
    }

    //out_lik.close();

}


/**
 * \file
 * \copyright
 * Copyright (c) 2012-2021, OpenGeoSys Community (http://www.opengeosys.org)
 *            Distributed under a Modified BSD License.
 *              See accompanying file LICENSE.txt or
 *              http://www.opengeosys.org/project/license
 *
 */

#include "LocalLinearLeastSquaresExtrapolator.h"

#include <Eigen/SVD>

#include "BaseLib/Logging.h"
#include "ExtrapolatableElementCollection.h"
#include "MathLib/LinAlg/Eigen/EigenMapTools.h"
#include "MathLib/LinAlg/LinAlg.h"
#include "MathLib/LinAlg/MatrixVectorTraits.h"
#include "NumLib/Assembler/SerialExecutor.h"
#include "NumLib/Function/Interpolation.h"

namespace NumLib
{
LocalLinearLeastSquaresExtrapolator::LocalLinearLeastSquaresExtrapolator(
    NumLib::LocalToGlobalIndexMap const& dof_table)
    : _dof_table_single_component(dof_table)
{
    /* Note in case the following assertion fails:
     * If you copied the extrapolation code, for your processes from
     * somewhere, note that the code from the groundwater flow process might
     * not suit your needs: It is a special case and is therefore most
     * likely too simplistic. You better adapt the extrapolation code from
     * some more advanced process, like the TES process.
     */
    if (dof_table.getNumberOfGlobalComponents() != 1)
    {
        OGS_FATAL(
            "The d.o.f. table passed must be for one variable that has only "
            "one component!");
    }
}

void LocalLinearLeastSquaresExtrapolator::extrapolate(
    const unsigned num_components,
    ExtrapolatableElementCollection const& extrapolatables,
    const double t,
    std::vector<GlobalVector*> const& x,
    std::vector<NumLib::LocalToGlobalIndexMap const*> const& dof_table)
{
    auto const num_nodal_dof_result =
        _dof_table_single_component.dofSizeWithoutGhosts() * num_components;

    std::vector<GlobalIndexType> ghost_indices;
    {  // Create num_components times version of ghost_indices arranged by
       // location. For example for 3 components and ghost_indices {5,6,10} we
       // compute {15, 16, 17,  18, 19, 20,  30, 31, 32}.
        auto const& single_component_ghost_indices =
            _dof_table_single_component.getGhostIndices();
        auto const single_component_ghost_indices_size =
            single_component_ghost_indices.size();
        ghost_indices.reserve(single_component_ghost_indices_size *
                              num_components);
        for (unsigned i = 0; i < single_component_ghost_indices_size; ++i)
        {
            for (unsigned c = 0; c < num_components; ++c)
            {
                ghost_indices.push_back(
                    single_component_ghost_indices[i] * num_components + c);
            }
        }
    }

    if (!_nodal_values ||
#ifdef USE_PETSC
        _nodal_values->getLocalSize() + _nodal_values->getGhostSize()
#else
        _nodal_values->size()
#endif
            != static_cast<GlobalIndexType>(num_nodal_dof_result))
    {
        _nodal_values = MathLib::MatrixVectorTraits<GlobalVector>::newInstance(
            {num_nodal_dof_result, num_nodal_dof_result, &ghost_indices,
             nullptr});
    }
    _nodal_values->setZero();

    // counts the writes to each nodal value, i.e., the summands in order to
    // compute the average afterwards
    auto counts =
        MathLib::MatrixVectorTraits<GlobalVector>::newInstance(*_nodal_values);
    counts->setZero();

    auto const size = extrapolatables.size();
    for (std::size_t i = 0; i < size; ++i)
    {
        extrapolateElement(i, num_components, extrapolatables, t, x, dof_table,
                           *counts);
    }
    MathLib::LinAlg::finalizeAssembly(*_nodal_values);

    MathLib::LinAlg::componentwiseDivide(*_nodal_values, *_nodal_values,
                                         *counts);
}

void LocalLinearLeastSquaresExtrapolator::calculateResiduals(
    const unsigned num_components,
    ExtrapolatableElementCollection const& extrapolatables,
    const double t,
    std::vector<GlobalVector*> const& x,
    std::vector<NumLib::LocalToGlobalIndexMap const*> const& dof_table)
{
    auto const num_element_dof_result = static_cast<GlobalIndexType>(
        _dof_table_single_component.size() * num_components);

    if (!_residuals || _residuals->size() != num_element_dof_result)
    {
#ifndef USE_PETSC
        _residuals.reset(new GlobalVector{num_element_dof_result});
#else
        _residuals.reset(new GlobalVector{num_element_dof_result, false});
#endif
    }

    if (static_cast<std::size_t>(num_element_dof_result) !=
        extrapolatables.size() * num_components)
    {
        OGS_FATAL("mismatch in number of D.o.F.");
    }

    auto const size = extrapolatables.size();
    for (std::size_t i = 0; i < size; ++i)
    {
        calculateResidualElement(i, num_components, extrapolatables, t, x,
                                 dof_table);
    }
    MathLib::LinAlg::finalizeAssembly(*_residuals);
}

void LocalLinearLeastSquaresExtrapolator::extrapolateElement(
    std::size_t const element_index,
    const unsigned num_components,
    ExtrapolatableElementCollection const& extrapolatables,
    const double t,
    std::vector<GlobalVector*> const& x,
    std::vector<NumLib::LocalToGlobalIndexMap const*> const& dof_table,
    GlobalVector& counts)
{
    auto const& integration_point_values =
        extrapolatables.getIntegrationPointValues(
            element_index, t, x, dof_table, _integration_point_values_cache);

    // Empty vector means to ignore the values and not to change the counts.
    if (integration_point_values.empty())
    {
        return;
    }

    auto const& N_0 = extrapolatables.getShapeMatrix(element_index, 0);
    auto const num_nodes = static_cast<unsigned>(N_0.cols());
    auto const num_values =
        static_cast<unsigned>(integration_point_values.size());

    if (num_values % num_components != 0)
    {
        OGS_FATAL(
            "The number of computed integration point values is not divisible "
            "by the number of num_components. Maybe the computed property is "
            "not a {:d}-component vector for each integration point.",
            num_components);
    }

    // number of integration points in the element
    const auto num_int_pts = num_values / num_components;

    if (num_int_pts < num_nodes)
    {
        OGS_FATAL(
            "Least squares is not possible if there are more nodes than "
            "integration points.");
    }

    auto const pair_it_inserted = _qr_decomposition_cache.emplace(
        std::make_pair(num_nodes, num_int_pts), CachedData{});

    auto& cached_data = pair_it_inserted.first->second;
    if (pair_it_inserted.second)
    {
        DBUG("Computing new singular value decomposition");

        // interpolation_matrix * nodal_values = integration_point_values
        // We are going to pseudo-invert this relation now using singular value
        // decomposition.
        auto& interpolation_matrix = cached_data.A;
        interpolation_matrix.resize(num_int_pts, num_nodes);

        interpolation_matrix.row(0) = N_0;
        for (unsigned int_pt = 1; int_pt < num_int_pts; ++int_pt)
        {
            auto const& shp_mat =
                extrapolatables.getShapeMatrix(element_index, int_pt);
            assert(shp_mat.cols() == num_nodes);

            // copy shape matrix to extrapolation matrix row-wise
            interpolation_matrix.row(int_pt) = shp_mat;
        }

        // JacobiSVD is extremely reliable, but fast only for small matrices.
        // But we usually have small matrices and we don't compute very often.
        // Cf.
        // http://eigen.tuxfamily.org/dox/group__TopicLinearAlgebraDecompositions.html
        //
        // Decomposes interpolation_matrix = U S V^T.
        Eigen::JacobiSVD<Eigen::MatrixXd> svd(
            interpolation_matrix, Eigen::ComputeThinU | Eigen::ComputeThinV);

        auto const& S = svd.singularValues();
        auto const& U = svd.matrixU();
        auto const& V = svd.matrixV();

        // Compute and save the pseudo inverse V * S^{-1} * U^T.
        auto const rank = svd.rank();
        assert(rank == num_nodes);

        // cf. http://eigen.tuxfamily.org/dox/JacobiSVD_8h_source.html
        cached_data.A_pinv.noalias() = V.leftCols(rank) *
                                       S.head(rank).asDiagonal().inverse() *
                                       U.leftCols(rank).transpose();
    }
    else if (cached_data.A.row(0) != N_0)
    {
        OGS_FATAL("The cached and the passed shapematrices differ.");
    }

    auto const& global_indices =
        _dof_table_single_component(element_index, 0).rows;

    if (num_components == 1)
    {
        auto const integration_point_values_vec =
            MathLib::toVector(integration_point_values);

        // Apply the pre-computed pseudo-inverse.
        Eigen::VectorXd const nodal_values =
            cached_data.A_pinv * integration_point_values_vec;

        // TODO does that give rise to PETSc problems? E.g., writing to ghost
        // nodes? Furthermore: Is ghost nodes communication necessary for PETSc?
        _nodal_values->add(global_indices, nodal_values);
        counts.add(global_indices,
                   std::vector<double>(global_indices.size(), 1.0));
    }
    else
    {
        auto const integration_point_values_mat = MathLib::toMatrix(
            integration_point_values, num_components, num_int_pts);

        // Apply the pre-computed pseudo-inverse.
        Eigen::MatrixXd const nodal_values =
            cached_data.A_pinv * integration_point_values_mat.transpose();

        std::vector<GlobalIndexType> indices;
        indices.reserve(num_components * global_indices.size());

        // _nodal_values is ordered location-wise
        for (unsigned comp = 0; comp < num_components; ++comp)
        {
            transform(cbegin(global_indices), cend(global_indices),
                      back_inserter(indices),
                      [&](auto const i) { return num_components * i + comp; });
        }

        // Nodal_values are passed as a raw pointer, because PETScVector and
        // EigenVector implementations differ slightly.
        _nodal_values->add(indices, nodal_values.data());
        counts.add(indices, std::vector<double>(indices.size(), 1.0));
    }
}

void LocalLinearLeastSquaresExtrapolator::calculateResidualElement(
    std::size_t const element_index,
    const unsigned num_components,
    ExtrapolatableElementCollection const& extrapolatables,
    const double t,
    std::vector<GlobalVector*> const& x,
    std::vector<NumLib::LocalToGlobalIndexMap const*> const& dof_table)
{
    auto const& int_pt_vals = extrapolatables.getIntegrationPointValues(
        element_index, t, x, dof_table, _integration_point_values_cache);

    auto const num_values = static_cast<unsigned>(int_pt_vals.size());
    if (num_values % num_components != 0)
    {
        OGS_FATAL(
            "The number of computed integration point values is not divisible "
            "by the number of num_components. Maybe the computed property is "
            "not a {:d}-component vector for each integration point.",
            num_components);
    }

    // number of integration points in the element
    const auto num_int_pts = num_values / num_components;

    const auto& global_indices =
        _dof_table_single_component(element_index, 0).rows;
    const auto num_nodes = static_cast<unsigned>(global_indices.size());

    auto const& interpolation_matrix =
        _qr_decomposition_cache.find({num_nodes, num_int_pts})->second.A;

    Eigen::VectorXd nodal_vals_element(num_nodes);
    auto const int_pt_vals_mat =
        MathLib::toMatrix(int_pt_vals, num_components, num_int_pts);

    MathLib::LinAlg::setLocalAccessibleVector(
        *_nodal_values);  // For access in the for-loop.
    for (unsigned comp = 0; comp < num_components; ++comp)
    {
        // filter nodal values of the current element
        for (unsigned i = 0; i < num_nodes; ++i)
        {
            // TODO PETSc negative indices?
            auto const idx = num_components * global_indices[i] + comp;
            nodal_vals_element[i] = _nodal_values->get(idx);
        }

        double const residual = (interpolation_matrix * nodal_vals_element -
                                 int_pt_vals_mat.row(comp).transpose())
                                    .squaredNorm();

        auto const eidx =
            static_cast<GlobalIndexType>(num_components * element_index + comp);
        // The residual is set to the root mean square value.
        auto const root_mean_square = std::sqrt(residual / num_int_pts);
        _residuals->set(eidx, root_mean_square);
    }
}

}  // namespace NumLib


/*
(***********************************************************************)
(*                                                                     *)
(* The ACME project                                                    *)
(*                                                                     *)
(* Copyright (c) 2020-2021, Davide Stocco and Enrico Bertolazzi.       *)
(*                                                                     *)
(* The ACME project and its components are supplied under the terms of *)
(* the open source BSD 2-Clause License. The contents of the ACME      *)
(* project and its components may not be copied or disclosed except in *)
(* accordance with the terms of the BSD 2-Clause License.              *)
(*                                                                     *)
(* URL: https://opensource.org/licenses/BSD-2-Clause                   *)
(*                                                                     *)
(*    Davide Stocco                                                    *)
(*    Department of Industrial Engineering                             *)
(*    University of Trento                                             *)
(*    e-mail: davide.stocco@unitn.it                                   *)
(*                                                                     *)
(*    Enrico Bertolazzi                                                *)
(*    Department of Industrial Engineering                             *)
(*    University of Trento                                             *)
(*    e-mail: enrico.bertolazzi@unitn.it                               *)
(*                                                                     *)
(***********************************************************************)
*/

///
/// file: acme_point.hh
///

#pragma once

#ifndef INCLUDE_ACME_POINT
#define INCLUDE_ACME_POINT

#include <Eigen/Core>
#include <Eigen/Dense>
#include <Eigen/Geometry>

#include "acme.hh"
#include "acme_entity.hh"
#include "acme_math.hh"

namespace acme
{

  /*\
   |               _       _
   |   _ __   ___ (_)_ __ | |_
   |  | '_ \ / _ \| | '_ \| __|
   |  | |_) | (_) | | | | | |_
   |  | .__/ \___/|_|_| |_|\__|
   |  |_|
  \*/

  //! Point class container
  /**
   * Specialization of Eigen::Matrix class
   */
  class point : public vec3, public entity
  {
  public:
    using vec3::Matrix;

    //! Point class destructor
    ~point(void);

    //! Point class constructor
    point(void);

    //! This constructor allows you to construct matrix from Eigen expressions
    template <typename derived>
    point(
      Eigen::MatrixBase<derived> const &other //!< Eigen matrix object
      ) : vec3(other)
    {
    }

    //! This method allows you to assign Eigen expressions to matrix
    template <typename derived>
    point &
    operator=(
      Eigen::MatrixBase<derived> const &other //!< Eigen matrix object
    )
    {
      this->vec3::operator=(other);
      return *this;
    }

    //! Translate point by vector
    void
    translate(
      vec3 const &vector_in //!< Input translation vector
      ) override;

    //! Transform point with affine transformation matrix
    void
    transform(
      affine const &affine_in //!< 4x4 affine transformation matrix
      ) override;

    //! Check if entity is degenerated
    bool
    isDegenerated(
      acme::real tolerance = EPSILON //!< Tolerance
    ) const override;

    //! Return object hierarchical level
    integer
    level(void) const override
    {
      return 1;
    }

    //! Return object type as string
    std::string
    type(void) const override
    {
      return "point";
    }

    //! Check whether the object is no entity
    bool
    isNone(void) const override
    {
      return false;
    }

    //! Check whether the object is a point
    bool
    isPoint(void) const override
    {
      return true;
    }

    //! Check whether the object is a line
    bool
    isLine(void) const override
    {
      return false;
    }

    //! Check whether the object is a ray
    bool
    isRay(void) const override
    {
      return false;
    }

    //! Check whether the object is a plane
    bool
    isPlane(void) const override
    {
      return false;
    }

    //! Check whether the object is a segment
    bool
    isSegment(void) const override
    {
      return false;
    }

    //! Check whether the object is a triangle
    bool
    isTriangle(void) const override
    {
      return false;
    }

    //! Check whether the object is a disk
    bool
    isDisk(void) const override
    {
      return false;
    }

    //! Check whether the object is a ball
    bool
    isBall(void) const override
    {
      return false;
    }

    //! Check whether in the point is clampable
    bool
    isClampable(void) const override
    {
      return true;
    }

    //! Check whether in the point is non-clampable
    bool
    isNonClampable(void) const override
    {
      return false;
    }

    //! Get minumum and maximum values along axes
    bool
    clamp(
      acme::vec3 &min, //!< Input minimum point
      acme::vec3 &max  //!< Input maximum point
    ) const override;

    //! Get minumum and maximum values along axes
    bool
    clamp(
      acme::real &min_x, //!< Input x value of minimum point
      acme::real &min_y, //!< Input y value of minimum point
      acme::real &min_z, //!< Input z value of minimum point
      acme::real &max_x, //!< Input x value of maximum point
      acme::real &max_y, //!< Input y value of maximum point
      acme::real &max_z  //!< Input z value of maximum point
    ) const override;

  }; // class point

  static point const NAN_POINT       = point::Constant(QUIET_NAN); //!< Not-a-Number static const point object
  static point       THROWAWAY_POINT = point(NAN_POINT);           //!< Throwaway static non-const point object

} // namespace acme

#endif

///
/// eof: acme_point.hh
///