/****************************************************************************
 * Copyright (c) 2017-2021 by the ArborX authors                            *
 * All rights reserved.                                                     *
 *                                                                          *
 * This file is part of the ArborX library. ArborX is                       *
 * distributed under a BSD 3-clause license. For the licensing terms see    *
 * the LICENSE file in the top-level directory.                             *
 *                                                                          *
 * SPDX-License-Identifier: BSD-3-Clause                                    *
 ****************************************************************************/

#include "ArborX_EnableDeviceTypes.hpp" // ARBORX_DEVICE_TYPES
#include "ArborX_EnableViewComparison.hpp"
#include <ArborX_DetailsUtils.hpp>
#include <ArborX_Exception.hpp>

#include <Kokkos_Core.hpp>

#include <boost/test/unit_test.hpp>

#include <algorithm>
#include <numeric>
#include <vector>

#define BOOST_TEST_MODULE StandardAlgorithms

namespace tt = boost::test_tools;

BOOST_AUTO_TEST_CASE_TEMPLATE(iota, DeviceType, ARBORX_DEVICE_TYPES)
{
  using ExecutionSpace = typename DeviceType::execution_space;
  ExecutionSpace space{};
  int const n = 10;
  double const val = 3.;
  Kokkos::View<double *, DeviceType> v("v", n);
  ArborX::iota(space, v, val);
  std::vector<double> v_ref(n);
  std::iota(v_ref.begin(), v_ref.end(), val);
  auto v_host = Kokkos::create_mirror_view(v);
  Kokkos::deep_copy(v_host, v);
  BOOST_TEST(v_ref == v_host, tt::per_element());

  Kokkos::View<int[3], DeviceType> w("w");
  ArborX::iota(space, w);
  std::vector<int> w_ref = {0, 1, 2};
  auto w_host = Kokkos::create_mirror_view(w);
  Kokkos::deep_copy(w_host, w);
  BOOST_TEST(w_ref == w_host, tt::per_element());
}

BOOST_AUTO_TEST_CASE_TEMPLATE(prefix_sum, DeviceType, ARBORX_DEVICE_TYPES)
{
  using ExecutionSpace = typename DeviceType::execution_space;
  ExecutionSpace space{};
  int const n = 10;
  Kokkos::View<int *, DeviceType> x("x", n);
  std::vector<int> x_ref(n, 1);
  x_ref.back() = 0;
  auto x_host = Kokkos::create_mirror_view(x);
  for (int i = 0; i < n; ++i)
    x_host(i) = x_ref[i];
  Kokkos::deep_copy(x, x_host);
  Kokkos::View<int *, DeviceType> y("y", n);
  ArborX::exclusivePrefixSum(space, x, y);
  std::vector<int> y_ref(n);
  std::iota(y_ref.begin(), y_ref.end(), 0);
  auto y_host = Kokkos::create_mirror_view(y);
  Kokkos::deep_copy(y_host, y);
  Kokkos::deep_copy(x_host, x);
  BOOST_TEST(y_host == y_ref, tt::per_element());
  BOOST_TEST(x_host == x_ref, tt::per_element());
  // in-place
  ArborX::exclusivePrefixSum(space, x, x);
  Kokkos::deep_copy(x_host, x);
  BOOST_TEST(x_host == y_ref, tt::per_element());
  int const m = 11;
  BOOST_TEST(n != m);
  Kokkos::View<int *, DeviceType> z("z", m);
  BOOST_CHECK_THROW(ArborX::exclusivePrefixSum(space, x, z),
                    ArborX::SearchException);
  Kokkos::View<double[3], DeviceType> v("v");
  auto v_host = Kokkos::create_mirror_view(v);
  v_host(0) = 1.;
  v_host(1) = 1.;
  v_host(2) = 0.;
  Kokkos::deep_copy(v, v_host);
  ArborX::exclusivePrefixSum(space, v);
  Kokkos::deep_copy(v_host, v);
  std::vector<double> v_ref = {0., 1., 2.};
  BOOST_TEST(v_host == v_ref, tt::per_element());
  Kokkos::View<double *, DeviceType> w("w", 4);
  BOOST_CHECK_THROW(ArborX::exclusivePrefixSum(space, v, w),
                    ArborX::SearchException);
  v_host(0) = 1.;
  v_host(1) = 0.;
  v_host(2) = 0.;
  Kokkos::deep_copy(v, v_host);
  Kokkos::resize(w, 3);
  ArborX::exclusivePrefixSum(space, v, w);
  auto w_host = Kokkos::create_mirror_view(w);
  Kokkos::deep_copy(w_host, w);
  std::vector<double> w_ref = {0., 1., 1.};
  BOOST_TEST(w_host == w_ref, tt::per_element());
}

BOOST_AUTO_TEST_CASE_TEMPLATE(last_element, DeviceType, ARBORX_DEVICE_TYPES)
{
  Kokkos::View<int *, DeviceType> v("v", 2);
  auto v_host = Kokkos::create_mirror_view(v);
  v_host(0) = 33;
  v_host(1) = 24;
  Kokkos::deep_copy(v, v_host);
  BOOST_TEST(ArborX::lastElement(v) == 24);
  Kokkos::View<int *, DeviceType> w("w", 0);
  BOOST_CHECK_THROW(ArborX::lastElement(w), ArborX::SearchException);
  Kokkos::View<double[1], DeviceType> u("u");
  Kokkos::deep_copy(u, 3.14);
  BOOST_TEST(ArborX::lastElement(u) == 3.14);
}

BOOST_AUTO_TEST_CASE_TEMPLATE(minmax, DeviceType, ARBORX_DEVICE_TYPES)
{
  using ExecutionSpace = typename DeviceType::execution_space;
  ExecutionSpace space{};

  Kokkos::View<double[4], DeviceType> v("v");
  auto v_host = Kokkos::create_mirror_view(v);
  v_host(0) = 3.14;
  v_host(1) = 1.41;
  v_host(2) = 2.71;
  v_host(3) = 1.62;
  Kokkos::deep_copy(v, v_host);

  auto const result_float = ArborX::minMax(space, v);
  BOOST_TEST(std::get<0>(result_float) == 1.41);
  BOOST_TEST(std::get<1>(result_float) == 3.14);
  Kokkos::View<int *, DeviceType> w("w", 0);
  BOOST_CHECK_THROW(ArborX::minMax(space, w), ArborX::SearchException);
  Kokkos::resize(w, 1);
  Kokkos::deep_copy(w, 255);
  auto const result_int = ArborX::minMax(space, w);
  BOOST_TEST(std::get<0>(result_int) == 255);
  BOOST_TEST(std::get<1>(result_int) == 255);

  // testing use case in ORNL-CEES/DataTransferKit#336
  Kokkos::View<int[2][3], DeviceType> u("u");
  auto u_host = Kokkos::create_mirror_view(u);
  u_host(0, 0) = 1; // x
  u_host(0, 1) = 2; // y
  u_host(0, 2) = 3; // z
  u_host(1, 0) = 4; // x
  u_host(1, 1) = 5; // y
  u_host(1, 2) = 6; // Z
  Kokkos::deep_copy(u, u_host);
#if 0
    // FIXME might be an issue with CUDA
    auto const minmax_x =
        ArborX::minMax( Kokkos::subview( u, Kokkos::ALL, 0 ) );
    BOOST_TEST( std::get<0>( minmax_x ) == 1 );
    BOOST_TEST( std::get<1>( minmax_x ) == 4 );
    auto const minmax_y =
        ArborX::minMax( Kokkos::subview( u, Kokkos::ALL, 1 ) );
    BOOST_TEST( std::get<0>( minmax_y ) == 2 );
    BOOST_TEST( std::get<1>( minmax_y ) == 5 );
#endif
}

BOOST_AUTO_TEST_CASE_TEMPLATE(accumulate, DeviceType, ARBORX_DEVICE_TYPES)
{
  using ExecutionSpace = typename DeviceType::execution_space;
  ExecutionSpace space{};

  Kokkos::View<int[6], DeviceType> v("v");
  Kokkos::deep_copy(v, 5);
  BOOST_TEST(ArborX::accumulate(space, v, 3) == 33);

  Kokkos::View<int *, DeviceType> w("w", 5);
  ArborX::iota(space, w, 2);
  BOOST_TEST(ArborX::accumulate(space, w, 4) == 24);
}

BOOST_AUTO_TEST_CASE_TEMPLATE(adjacent_difference, DeviceType,
                              ARBORX_DEVICE_TYPES)
{
  using ExecutionSpace = typename DeviceType::execution_space;
  ExecutionSpace space{};

  Kokkos::View<int[5], DeviceType> v("v");
  auto v_host = Kokkos::create_mirror_view(v);
  v_host(0) = 2;
  v_host(1) = 4;
  v_host(2) = 6;
  v_host(3) = 8;
  v_host(4) = 10;
  Kokkos::deep_copy(v, v_host);
  // In-place operation is not allowed
  BOOST_CHECK_THROW(ArborX::adjacentDifference(space, v, v),
                    ArborX::SearchException);
  auto w = Kokkos::create_mirror(DeviceType(), v);
  BOOST_CHECK_NO_THROW(ArborX::adjacentDifference(space, v, w));
  auto w_host = Kokkos::create_mirror_view(w);
  Kokkos::deep_copy(w_host, w);
  std::vector<int> w_ref(5, 2);
  BOOST_TEST(w_host == w_ref, tt::per_element());

  Kokkos::View<float *, DeviceType> x("x", 10);
  Kokkos::deep_copy(x, 3.14);
  BOOST_CHECK_THROW(ArborX::adjacentDifference(space, x, x),
                    ArborX::SearchException);
  Kokkos::View<float[10], DeviceType> y("y");
  BOOST_CHECK_NO_THROW(ArborX::adjacentDifference(space, x, y));
  std::vector<float> y_ref(10);
  y_ref[0] = 3.14;
  auto y_host = Kokkos::create_mirror_view(y);
  Kokkos::deep_copy(y_host, y);
  BOOST_TEST(y_host == y_ref, tt::per_element());

  Kokkos::resize(x, 5);
  BOOST_CHECK_THROW(ArborX::adjacentDifference(space, y, x),
                    ArborX::SearchException);
}

BOOST_AUTO_TEST_CASE_TEMPLATE(min_and_max, DeviceType, ARBORX_DEVICE_TYPES)
{
  using ExecutionSpace = typename DeviceType::execution_space;
  ExecutionSpace space{};

  Kokkos::View<int[4], DeviceType> v("v");
  ArborX::iota(space, v);
  BOOST_TEST(ArborX::min(space, v) == 0);
  BOOST_TEST(ArborX::max(space, v) == 3);

  Kokkos::View<int *, DeviceType> w("w", 7);
  ArborX::iota(space, w, 2);
  BOOST_TEST(ArborX::min(space, w) == 2);
  BOOST_TEST(ArborX::max(space, w) == 8);
}


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

#ifndef MPO19_CLUSTERING_22072020
#define MPO19_CLUSTERING_22072020

#include <mpo19/typedefs.hpp>
#include <mpo19/cluster.hpp>
#include <boost/range/algorithm.hpp>

namespace mpo19{
template<typename Linkage>
struct hierarchical_clustering : public Linkage {

    using dendrogram_t  = std::vector<cluster>;

    using Linkage::Linkage;

    dendrogram_t operator()(size_t t)
    {
        using std::swap;

        size_t n = Linkage::get_input_size();

        assert(t < n && "t must be strictly smaller than n");

        dendrogram_t dendrogram;

        //Initialize dendrogram
        for(index_t i = 0; i != index_t(n); ++i) {
            dendrogram.emplace_back(cluster::leaf_t(i));
        }

        for(size_t l = 0; l < n - t; ++l) {
            //get indices of clusters to be merged
            auto p = get_merge_indices();

            size_t i = std::get<0>(p);
            size_t j = std::get<1>(p);

            assert(i < j && "By convention i should be smaller than j");
            assert(0 <= i && i < n && "i should be between 0 and n");
            assert(0 <= j && j < n && "j should be between 0 and n");

            //merge the two clusters at position i and j
            merge(dendrogram[i], dendrogram[j], l + 1);
        }

        assert(dendrogram.size() == n && "dendrogram should be of size n here");
        //Erase all clusters that only refer to other clusters through a pointer
        dendrogram.erase(
        boost::remove_if(dendrogram, [](cluster const& c) {
            return !c.is_cluster();
        })
        , dendrogram.end());

        assert(dendrogram.size() == t && "dendrogram should be of size t here");

        return dendrogram;
    }

    dendrogram_t partial_evaluation(size_t t, dendrogram_t dendrogram = dendrogram_t{})
    {
        using std::swap;
        using boost::range::count_if;

        size_t n = Linkage::get_input_size();

        assert(t < n && "t must be strictly smaller than n");

        size_t start;
        if(dendrogram.empty()) {
            //Initialize dendrogram
            for(index_t i = 0; i != index_t(n); ++i) {
                dendrogram.emplace_back(cluster::leaf_t(i));
            }
            start = 0;
        }
        else {
            assert(dendrogram.size() == n);
            start = n - count_if(dendrogram, [](cluster const& c){
                return c.is_cluster();
            });
        }

        for(size_t l = start; l < n - t; ++l) {
            //get indices of clusters to be merged
            auto p = get_merge_indices();

            size_t i = std::get<0>(p);
            size_t j = std::get<1>(p);

            assert(i < j && "By convention i should be smaller than j");
            assert(0 <= i && i < n && "i should be between 0 and n");
            assert(0 <= j && j < n && "j should be between 0 and n");

            //merge the two clusters at position i and j
            merge(dendrogram[i], dendrogram[j], l + 1);
        }

        assert(dendrogram.size() == n && "dendrogram should be of size n here");

        return dendrogram;
    }

private:
    using Linkage::get_merge_indices;
};

}
#endif //MPO19_CLUSTERING_22072020


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

#include "Tracker.h"

#include "math/numeric/svd.h"
#include "util/log.h"

#include <boost/scoped_array.hpp>

#include <cmath>

#define LOG_MODULE ::logging::LogModuleEnum::COMPLEXITY

namespace complexity {

namespace {

} // anonymous namespace

void TrackWithEqn::addMeasurement(long double size, long double timeVal, long double weight) {
  unsigned int nTerms = eqn_.getNTerms();

  if (measurements_.size() > nTerms) {
    MeasurementSet::iterator it = measurements_.begin();
    if (it->weight < weight) {
      if (measurements_.size() > 2*nTerms) {
        measurements_.erase(it);
      }
      measurements_.insert(Measurement(size, timeVal, weight));
      dirty_ = true;
    }
  } else {
    measurements_.insert(Measurement(size, timeVal, weight));
    dirty_ = true;
  }
}

const Equation& TrackWithEqn::getEquation() const {
  if (dirty_) { // If new measurement exist

    // E x C = T, solve for C

    int nMeasurements = measurements_.size();
    int nTerms = eqn_.getNTerms();
    if (nTerms > nMeasurements) {  // Not enough measurements
      return eqn_;
    }

    // Allocate memory for E and T all in one shot
    boost::scoped_array<long double> Etmp(
      new long double[nMeasurements * (nTerms + 2)]);
    long double** E = reinterpret_cast<long double**>(Etmp.get());
    for (int i = 0; i < nMeasurements; ++i) {
      E[i] = Etmp.get() + nMeasurements + i * nTerms;
    }
    // Last nMeasurements elements
    long double* T = Etmp.get() + nMeasurements * (nTerms + 1);

    // Fill in E and T matrices from measurement map
    MeasurementSet::const_iterator it = measurements_.begin();
    for (int i = 0; it != measurements_.end(); ++it, ++i) {
      eqn_.evaluateTermsAt(it->size, E[i]);
      T[i] = it->value;
    }

    // Solve for coefficients using singular value decomposition
    // Note that logically we consider this object constant
    // we are just refining equation coefficients in light of new measurements
    TrackWithEqn* this_ = const_cast<TrackWithEqn*>(this);
    math::numeric::solveViaSVD(
        E, (this_->eqn_).getCoeffs(), T, nMeasurements, nTerms);

    this_->dirty_ = false;
  }
  return eqn_;
}

long double TrackWithEqn::getResidual() const {
  const Equation& eqn = this->getEquation();
  MeasurementSet::const_iterator it = measurements_.begin();
  long double residual = 0.0;
  for (; it != measurements_.end(); ++it) {
    long double diff = eqn.evaluateAt(it->size) - it->value;
    residual += diff * diff;
  }
  return std::sqrt(residual);
}

}  // namespace complexity


/* Copyright © 2019 Apple Inc. All rights reserved.
 *
 * Use of this source code is governed by a BSD-3-clause license that can
 * be found in the LICENSE.txt file or at
 * https://opensource.org/licenses/BSD-3-Clause
 */

#define BOOST_TEST_MODULE test_parameter_sampler

#include <toolkits/object_detection/one_shot_object_detection/util/parameter_sampler.hpp>

#include <boost/test/unit_test.hpp>
#include <core/util/test_macros.hpp>

#include <unordered_map>

static constexpr int kMaxDimension = 6000;
static constexpr int kNumRows = 1000;
static constexpr int kBaseDimensionSize = 2500;

namespace turi {
namespace one_shot_object_detection {
namespace {

ParameterSampler create_parameter_sampler(std::default_random_engine &engine) {
  std::uniform_int_distribution<int> dist(0, kMaxDimension);

  size_t starter_width = dist(engine);
  size_t starter_height = dist(engine);
  size_t dx = 0;
  size_t dy = 0;
  ParameterSampler sampler(starter_width, starter_height, dx, dy);
  return sampler;
}

std::unordered_map<std::string, std::vector<double> >
constuct_angle_samples_map(ParameterSampler &sampler, int seed) {
  std::unordered_map<std::string, std::vector<double> > angle_samples_map;
  std::vector<double> thetas;
  std::vector<double> phis;
  std::vector<double> gammas;
  for (int i = 0; i < kNumRows; i++) {
    size_t background_width = kBaseDimensionSize + i;
    size_t background_height = kBaseDimensionSize - i;
    sampler.sample(background_width, background_height, seed, i);
    double theta = sampler.get_theta();
    double phi = sampler.get_phi();
    double gamma = sampler.get_gamma();
    thetas.push_back(theta);
    phis.push_back(phi);
    gammas.push_back(gamma);
  }
  angle_samples_map["theta"] = thetas;
  angle_samples_map["phi"] = phis;
  angle_samples_map["gamma"] = gammas;
  return angle_samples_map;
}

int count_in_range(std::vector<double> samples, double lower, double upper) {
  int count = 0;
  for (double sample : samples) {
    if (sample >= lower && sample <= upper) {
      count += 1;
    }
  }
  return count;
}

bool angles_match_distribution(const std::vector<double> &angles,
                               const std::vector<double> &angle_means,
                               double angle_stdev) {
  // for every mean, assert that the count of angles in
  // [mean-stdev, mean+stdev] is greater than the count of angles in
  // [mean-2*stdev, mean-stdev] U [mean+stdev, mean+2*stdev]
  int first_stdev_count;
  int second_stdev_count;
  for (double mean : angle_means) {
    first_stdev_count =
        count_in_range(angles, mean - angle_stdev, mean + angle_stdev);
    second_stdev_count =
        (count_in_range(angles, mean - 2 * angle_stdev, mean - angle_stdev) +
         count_in_range(angles, mean + angle_stdev, mean + 2 * angle_stdev));
    TS_ASSERT_LESS_THAN_EQUALS(second_stdev_count, first_stdev_count);
  }
  return true;
}

BOOST_AUTO_TEST_CASE(test_parameter_distributions) {
  std::default_random_engine engine;
  int seed = 500;
  engine.seed(seed);
  ParameterSampler sampler = create_parameter_sampler(engine);
  std::unordered_map<std::string, std::vector<double> > all_angles =
      constuct_angle_samples_map(sampler, seed);
  std::vector<double> thetas = all_angles["theta"];
  std::vector<double> phis = all_angles["phi"];
  std::vector<double> gammas = all_angles["gamma"];
  TS_ASSERT(angles_match_distribution(thetas, sampler.get_theta_means(),
                                      sampler.get_theta_stdev()));
  TS_ASSERT(angles_match_distribution(phis, sampler.get_phi_means(),
                                      sampler.get_phi_stdev()));
  TS_ASSERT(angles_match_distribution(gammas, sampler.get_gamma_means(),
                                      sampler.get_gamma_stdev()));
}

}  // namespace
}  // namespace one_shot_object_detection
}  // namespace turi


// Copyright 2018 Hans Dembinski
//
// Distributed under the Boost Software License, version 1.0.
// (See accompanying file LICENSE_1_0.txt
// or copy at http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_HISTOGRAM_ACCUMULATORS_SUM_HPP
#define BOOST_HISTOGRAM_ACCUMULATORS_SUM_HPP

#include <boost/histogram/fwd.hpp>
#include <cmath>
#include <type_traits>

namespace boost {
namespace histogram {
namespace accumulators {

/**
  Uses Neumaier algorithm to compute accurate sums.

  The algorithm uses memory for two floats and is three to
  five times slower compared to a simple floating point
  number used to accumulate a sum, but the relative error
  of the sum is at the level of the machine precision,
  independent of the number of samples.

  A. Neumaier, Zeitschrift fuer Angewandte Mathematik
  und Mechanik 54 (1974) 39-51.
*/
template <typename RealType>
class sum {
public:
  sum() = default;

  /// Initialize sum to value
  explicit sum(const RealType& value) noexcept : large_(value) {}

  /// Set sum to value
  sum& operator=(const RealType& value) noexcept {
    large_ = value;
    small_ = 0;
    return *this;
  }

  /// Increment sum by one
  sum& operator++() { return operator+=(1); }

  /// Increment sum by value
  sum& operator+=(const RealType& value) {
    auto temp = large_ + value; // prevent optimization
    if (std::abs(large_) >= std::abs(value))
      small_ += (large_ - temp) + value;
    else
      small_ += (value - temp) + large_;
    large_ = temp;
    return *this;
  }

  /// Scale by value
  sum& operator*=(const RealType& value) {
    large_ *= value;
    small_ *= value;
    return *this;
  }

  template <class T>
  bool operator==(const sum<T>& rhs) const noexcept {
    return large_ == rhs.large_ && small_ == rhs.small_;
  }

  template <class T>
  bool operator!=(const T& rhs) const noexcept {
    return !operator==(rhs);
  }

  /// Return large part of the sum.
  const RealType& large() const { return large_; }

  /// Return small part of the sum.
  const RealType& small() const { return small_; }

  // allow implicit conversion to RealType
  operator RealType() const { return large_ + small_; }

  template <class Archive>
  void serialize(Archive&, unsigned /* version */);

private:
  RealType large_ = RealType();
  RealType small_ = RealType();
};

} // namespace accumulators
} // namespace histogram
} // namespace boost

#ifndef BOOST_HISTOGRAM_DOXYGEN_INVOKED
namespace std {
template <class T, class U>
struct common_type<boost::histogram::accumulators::sum<T>,
                   boost::histogram::accumulators::sum<U>> {
  using type = boost::histogram::accumulators::sum<common_type_t<T, U>>;
};
} // namespace std
#endif

#endif


// Software License for MTL
//
// Copyright (c) 2007 The Trustees of Indiana University.
//               2008 Dresden University of Technology and the Trustees of Indiana University.
//               2010 SimuNova UG (haftungsbeschränkt), www.simunova.com.
// All rights reserved.
// Authors: Peter Gottschling and Andrew Lumsdaine
//
// This file is part of the Matrix Template Library
//
// See also license.mtl.txt in the distribution.

#include <iostream>
// #include <boost/test/minimal.hpp>

#include <boost/numeric/mtl/operation/update.hpp>
#include <boost/numeric/mtl/matrix/compressed2D.hpp>
#include <boost/numeric/mtl/matrix/parameter.hpp>
#include <boost/random.hpp>
#include <boost/timer.hpp>

using namespace mtl;
using namespace std;
using namespace boost;

template <typename Generator>
double generator_overhead(int s, Generator gen)
{
    boost::timer t;
    while (! gen.finished() ) {
	int      r, c;
	double   v;
	gen(r, c, v);
    }
    return t.elapsed();
}

template <typename Generator>
double insert_petsc(int s, Generator gen, double overhead)
{
    return 0.0;
}

template <typename Generator>
double insert_mtl2(int s, Generator gen, double overhead)
{
    return 0.0;
}

template <typename Generator>
double insert_mtl4(int s, Generator gen, double overhead)
{
    typedef typename Generator::value_type                                             value_type;
    typedef mat::parameters<row_major, mtl::index::c_index, non_fixed::dimensions>   parameters;
    typedef compressed2D<value_type, parameters>                                       matrix_type;
    matrix_type   matrix(non_fixed::dimensions(s, s));

    boost::timer t;
    compressed2D_inserter<value_type, parameters>  inserter(matrix, 7);
    while (! gen.finished() ) {
	int      r, c;
	double   v;
	gen(r, c, v);
	inserter(r, c) << v;
	// cout << "A[" << r << ", " << c << "] = " << v << '\n';
    }
    return t.elapsed();
}

// produce a matrix with s nonzeros, randomly distributed
template <typename Size, typename Value, typename RandomGen,
	  typename RowDistribution = uniform_int<Size>,
	  typename ColumnDistribution = RowDistribution>
struct random_generator
{
    typedef Value value_type;

    explicit random_generator( Size s, RandomGen mygen = RandomGen() ) :
	mygen( mygen ), nnz( 5 * s ),
	pick_row( RowDistribution(0,s-1) ), pick_col( ColumnDistribution(0,s-1) ) {}

    bool finished()
    {
	return !(bool)nnz;
    }

    void operator()( Size& i, Size& j, value_type& v )
    {
	i = pick_row( mygen );
	j = pick_col( mygen );
	v = 1.;
	nnz--;
    }

    RandomGen          mygen;
    Size               nnz;
    RowDistribution    pick_row;
    ColumnDistribution pick_col;
};

// produce a matrix with s nonzeros, randomly distributed
template <typename Size, typename Value, typename RandomGen,
	  typename ColumnDistribution = uniform_int<Size> >
struct random_column_generator
{
    typedef Value value_type;

    explicit random_column_generator( Size s, RandomGen mygen = RandomGen() ) :
	mygen( mygen ), s( s ), nnz( 5 * s ),
	pick_col( ColumnDistribution(0,s-1) ) {}

    bool finished()
    {
	return !(bool)nnz;
    }

    void operator()( Size& i, Size& j, value_type& v )
    {
	i = s - (nnz+4) / 5;
	// if (i < 0 || i >= s) cout << "i = " << i << ", s = " << s << ", nnz = " << nnz << '\n';
	assert(i >= 0 && i < s);
	j = pick_col( mygen );
	assert(j >= 0 && j < s);
	v = 1.;
	nnz--;
    }

    RandomGen          mygen;
    Size               s, nnz;
    ColumnDistribution pick_col;
};

template <typename Size, typename Value>
struct poisson_generator
{
    typedef Value  value_type;

    explicit poisson_generator(int s) : s(s), count(0), row(0), offset(2) {     // s must be 2^k 100
        if (s == 9) {                                      // only to test a 9x9 matrix
	    d1= d2= 3;                                     // only to test a 9x9 matrix
	} else {
	    assert(s % 100 == 0);
	    d1= 100, d2= s/d1;
	    for (; d2 > d1; d1<<= 1) d2= s / d1;
	}
	nnz= 5 * s - 2 * d1 - 2 * d2;
    }

    bool finished() {
	return count >= nnz;
    }

    bool valid_offset() {
	switch (offset) {
	    case 0: return row >= d2; // northern nb
	    case 1: return row % d2;  // western nb
	    case 2: return true;
	    case 3: return (row + 1) % d2; // eastern nb
	    case 4: return row + d2 < s;   // southern nb
	}
	assert(true); // shouldn't be reached
	return false;
    }

    void next_offset() {
	offset++;
	if (offset == 5) offset= 0, row++;
    }

    void next_nnz() {
	do
	    {next_offset(); }
	while (! valid_offset());
	count++;
    }

    void operator() (Size& r, Size& c, Value& v) {
	r= row;
	switch (offset) {
	    case 0: c= r - d2; v= -1; break;
	    case 1: c= r - 1; v= -1; break;
	    case 2: c= r; v= 4; break;
	    case 3: c= r + 1; v= -1; break;
	    case 4: c= r + d2; v= -1;
	}
	next_nnz();
    }

    int s, d1, d2, count,
	nnz,                   // number of non-zeros in matrix
	row, offset;           // current row and which entry in the row
};

template <typename Generator>
void run(int max_size)
{
    for (int s= 10000; s < max_size; s*= 2) {
	cout << "Size is " << s << '\n';

	Generator gen0(s);
	double overhead= generator_overhead(s, gen0);
	cout << "Generator overhead: " << overhead << "s\n";

	Generator gen1(s);
	insert_petsc(s, gen1, overhead);

	Generator gen2(s);
	insert_mtl2(s, gen2, overhead);

	Generator gen3(s);
	double mtl4t= insert_mtl4(s, gen3, overhead);
	cout << "Inserting into MTL 4: " << mtl4t << "s\n";
    }
}

void check_dims()
{
    for (int s= 100;  s < 1000000; s*= 2) {
	int d1= 10, d2= s/d1;
	for (; d2 > d1; d1<<= 1)
	    d2= s / d1;
	std::cout << d1 << " * " << d2 << " = " << s << '\n';
    }
}

int main(int argc, char* argv[])
{
    // check_dims();
    poisson_generator<int, double> poisson_9(9);
    insert_mtl4(9, poisson_9, 0.0);

    minstd_rand gen;
    random_generator<int, double, minstd_rand> random_9(9);
    insert_mtl4(9, random_9, 0.0);

    if( argc > 1 ) {
	run<poisson_generator<int, double> > (atoi(argv[1]));
	run<random_generator<int, double, minstd_rand> > (atoi(argv[1]));
	run<random_column_generator<int, double, minstd_rand> > (atoi(argv[1]));
    }

    return 0;
}


#pragma once

#include <functional>
#include <set>
#include <unordered_set>
#include <boost/container/flat_set.hpp>
#include <string>
#include <cmath>        // std::abs

#include "nifty/tools/changable_priority_queue.hxx"
#include "nifty/graph/edge_contraction_graph.hxx"
#include "nifty/graph/agglo/cluster_policies/cluster_policies_common.hxx"
#include <iostream>
#include <algorithm>    // std::max
#include <xtensor/xview.hpp>
using namespace xt::placeholders;

namespace nifty{
    namespace graph{
        namespace agglo{

// UCM: ultra contour map (enable an edge union-find datastructure)
// UPDATE_RULE: the type of linkage criteria implemented in ./details/merge_rules.hxx

            template<
                    class GRAPH, class UPDATE_RULE, bool ENABLE_UCM
            >
            class GaspClusterPolicy{
                typedef GaspClusterPolicy<
                        GRAPH, UPDATE_RULE,  ENABLE_UCM
                > SelfType;

            private:
                typedef typename GRAPH:: template EdgeMap<uint8_t> UInt8EdgeMap;
                typedef typename GRAPH:: template EdgeMap<float> FloatEdgeMap;
                typedef typename GRAPH:: template NodeMap<float> FloatNodeMap;

                // Flat sets for storing cannot-link constraints:
                typedef boost::container::flat_set<uint64_t> SetType;
                typedef typename GRAPH:: template NodeMap<SetType > NonLinkConstraints;

                typedef UPDATE_RULE UpdateRuleType;
            public:
                typedef typename UpdateRuleType::SettingsType UpdateRuleSettingsType;

                // input types
                typedef GRAPH                                       GraphType;
                typedef FloatEdgeMap                                EdgePrioType;
                typedef FloatEdgeMap                                EdgeSizesType;
                typedef FloatNodeMap                                NodeSizesType;

                struct SettingsType{
                    UpdateRuleSettingsType updateRule;
                    uint64_t numberOfNodesStop{1};
                    double sizeRegularizer{0.};
                    bool addNonLinkConstraints{false};
                    bool mergeConstrainedEdgesAtTheEnd{false};
                    bool collectStats{false};
                };

                enum class EdgeStates : uint8_t {
                    LOCAL = 0, // Edge represent e.g. two direct neighbors that can be contracted
                    LIFTED = 1, // Lifted edges cannot be contracted until it is merged with some local one
                };

                typedef EdgeContractionGraph<GraphType, SelfType>   EdgeContractionGraphType;

                friend class EdgeContractionGraph<GraphType, SelfType, ENABLE_UCM> ;
            private:

                // internal types
                typedef nifty::tools::ChangeablePriorityQueue< float , std::greater<float> > QueueType;

            public:

                template<class SIGNED_WEIGHTS, class IS_LOCAL_EDGE, class EDGE_SIZES, class NODE_SIZES>
                GaspClusterPolicy(const GraphType &,
                                  const SIGNED_WEIGHTS & ,
                                  const IS_LOCAL_EDGE &,
                                  const EDGE_SIZES & ,
                                  const NODE_SIZES & ,
                                  const SettingsType & settings = SettingsType());

                std::pair<uint64_t, double> edgeToContractNext() const;
                bool isDone();

                // callback called by edge contraction graph

                EdgeContractionGraphType & edgeContractionGraph();

            private:
                double pqMergePrio(const uint64_t edge) const;
                double computeWeight(const uint64_t edge) const;

            public:
                // callbacks called by edge contraction graph
                void contractEdge(const uint64_t edgeToContract);
                void mergeNodes(const uint64_t aliveNode, const uint64_t deadNode);
                void mergeEdges(const uint64_t aliveEdge, const uint64_t deadEdge);
                void contractEdgeDone(const uint64_t edgeToContract);

                bool isEdgeConstrained(const uint64_t edge){
                    const auto uv = edgeContractionGraph_.uv(edge);
                    const auto u = uv.first;
                    const auto v = uv.second;
                    const auto & setU  = nonLinkConstraints_[u];
                    const auto & setV  = nonLinkConstraints_[v];
                    // This find operation gets expensive (log(N)), so we look into the smaller set:
                    // NIFTY_CHECK((setU.find(v)!=setU.end()) == (setV.find(u)!=setV.end()),"");
                    if (setU.size() < setV.size()) {
                        return setU.find(v)!=setU.end();
                    } else {
                        return setV.find(u)!=setV.end();
                    }
                }

                bool isMergeAllowed(const uint64_t edge) const{
                    // Here we do not care about the fact that an edge is lifted or not.
                    // We just look at the priority

                    return accumulated_weights_[edge] > 0.;
                }

                double edgeCostInPQ(const uint64_t edge) const{
                    const auto priority = accumulated_weights_[edge];
                    if (settings_.addNonLinkConstraints) {
                        return std::abs(priority);
                    } else {
                        return priority;
                    }
                }

                void addNonLinkConstraint(const uint64_t edge){
                    const auto reprEdge = edgeContractionGraph_.findRepresentativeEdge(edge);
                    NIFTY_ASSERT(accumulated_weights_[reprEdge]<=0.);
                    const auto uv = edgeContractionGraph_.uv(reprEdge);
                    const auto u = uv.first;
                    const auto v = uv.second;
                    // This insert operation, first performs a search (logN) and then insert (if not found)accumulated_weights_[reprEdge]
                    nonLinkConstraints_[uv.first].insert(uv.second);
                    nonLinkConstraints_[uv.second].insert(uv.first);
                }

                auto exportFinalNodeDataOriginalGraph(){
                    // Export node data of the ORIGINAL graph:
                    NIFTY_ASSERT(settings_.collectStats);
                    typename xt::xtensor<float, 2>::shape_type retshape;
                    retshape[0] = graph_.nodeIdUpperBound()+1;
                    retshape[1] = 4;
                    xt::xtensor<float, 2> out(retshape);

                    graph_.forEachNode([&](const uint64_t node) {
                        out(node, 0) = maxNodeSize_per_iter_[node];
                        out(node, 1) = maxCostInPQ_per_iter_[node];
                        out(node, 2) = meanNodeSize_per_iter_[node];
                        out(node, 3) = variance_[node];
                    });
                    return out;
                }

                auto exportFinalEdgeDataContractedGraph(){
                    NIFTY_ASSERT(settings_.collectStats);

                    // Export edge data of the final contracted graph:
                    xt::xtensor<float, 2> out = xt::zeros<float>({uint64_t(edgeContractionGraph().numberOfEdges()), uint64_t(4)});
                    xt::xtensor<bool, 1> notVisitedEdges = xt::ones<bool>({graph_.edgeIdUpperBound()+1});
                    uint64_t edge_counter = 0;
                    graph_.forEachEdge([&](const uint64_t edge) {
                        const auto cEdge = edgeContractionGraph_.findRepresentativeEdge(edge);
                        const auto uv = edgeContractionGraph_.uv(cEdge);
                        const auto u = edgeContractionGraph_.findRepresentativeNode(uv.first);
                        const auto v = edgeContractionGraph_.findRepresentativeNode(uv.second);
                        if (u != v && notVisitedEdges(cEdge)) {
                            out(edge_counter, 0) = u;
                            out(edge_counter, 1) = v;
                            out(edge_counter, 2) = accumulated_weights_[cEdge];
                            out(edge_counter, 3) = accumulated_weights_.weight(cEdge);
                            edge_counter++;
                            notVisitedEdges(cEdge) = false;
                        };
                    });

                    // -------
                    // Export data merges/constraints:
                    // -------

                    // First, we collect information/statistics about constraints:
                    //  - sum of constrained edge sizes
                    //  - sum of constrained edge weights
                    //  - max constraint
                    //  - min constraint
                    //  - (collect more stats about size of the edges...?)
                    xt::xtensor<float, 2> constraintStatistics = xt::zeros<float>({uint64_t(graph_.edgeIdUpperBound()+1), uint64_t(6)});
                    graph_.forEachEdge([&](const uint64_t edge){
                        if (constraintsStats_(edge, 0) == 1.) {
                            // Find the ID of the final boundary (could be now merged or not):
                            const auto cEdge = edgeContractionGraph_.findRepresentativeEdge(edge);
                            constraintStatistics(cEdge, 0) += constraintsStats_(edge,2); // Add constraint size
                            constraintStatistics(cEdge, 1) += constraintsStats_(edge,1)*constraintsStats_(edge,2); // Increase sum
                            // The weight of the constrained edge will be always negative, so we keep it in mind while saving the min/max
                            constraintStatistics(cEdge, 2) = std::max(constraintStatistics(cEdge, 2),-constraintsStats_(edge,1));
                            constraintStatistics(cEdge, 3) = std::min(constraintStatistics(cEdge, 3),constraintsStats_(edge,1));
                            // Export stats in case the onstrained edge became positive at some point:
                            constraintStatistics(cEdge, 4) += constraintsStats_(edge,4); // Add constraint size
                            constraintStatistics(cEdge, 5) += constraintsStats_(edge,3)*constraintsStats_(edge,4); // Increase sum
                        }

                    });

                    // Then, we collect information about merged edges and remaining edges in the final graph:
                    //  - was the edge merged or not (1 if yes, 0 if not)
                    //  - edge weight (either before to be merged or in the final graph)
                    //  - edge size (either before to be merged or in the final graph)
                    xt::xtensor<float, 2> mergeStatistics = xt::zeros<float>({uint64_t(graph_.edgeIdUpperBound()+1), uint64_t(4)});
                    graph_.forEachEdge([&](const uint64_t edge){
                        const auto cEdge = edgeContractionGraph_.findRepresentativeEdge(edge);

                        // First, we fill the remaining parts of the constraintStatistics array:
                        constraintStatistics(edge, 0) = constraintStatistics(cEdge, 0);
                        constraintStatistics(edge, 1) = constraintStatistics(cEdge, 1);
                        constraintStatistics(edge, 2) = constraintStatistics(cEdge, 2);
                        constraintStatistics(edge, 3) = constraintStatistics(cEdge, 3);
                        constraintStatistics(edge, 4) = constraintStatistics(cEdge, 4);
                        constraintStatistics(edge, 5) = constraintStatistics(cEdge, 5);

                        // Find the ID of the final boundary (could be now merged or not):
                        const auto uv = edgeContractionGraph_.uv(cEdge);
                        const auto u = edgeContractionGraph_.findRepresentativeNode(uv.first);
                        const auto v = edgeContractionGraph_.findRepresentativeNode(uv.second);
                        if (u != v) {
                            // In this case, the edge was not merged:
                            mergeStatistics(edge, 0) = 0.;
                        } else {
                            // In this case, the edge was merged:
                            mergeStatistics(edge, 0) = 1.;
                        }
                        // Now save the information about the boundary:
                        mergeStatistics(edge, 1) = accumulated_weights_[cEdge];
                        mergeStatistics(edge, 2) = accumulated_weights_.weight(cEdge);
                        mergeStatistics(edge, 3) = mergeStats_(cEdge);
                    });

                    return std::make_tuple(out, constraintStatistics, mergeStatistics);
                }

                auto exportAction(){
                    NIFTY_ASSERT(settings_.collectStats);
                    // Exports actions done
                    // xt::xtensor<float, 2> out = xt::zeros<float>({uint64_t(nb_performed_contractions_+100), uint64_t(5)});
                    xt::xtensor<uint64_t , 2> out = xt::view(actionStats_, xt::all(), xt::all());
                    return out;
                }

            private:

                // INPUT
                const GraphType &   graph_;

                NonLinkConstraints nonLinkConstraints_;

                UPDATE_RULE accumulated_weights_;
                uint64_t nb_performed_contractions_;

                // Stats:
                uint64_t nb_edges_popped_;
                xt::xtensor<uint64_t , 2> actionStats_;
                xt::xtensor<float , 2> constraintsStats_;
                xt::xtensor<float , 1> mergeStats_;

                // State of the edges (LOCAL or LIFTED)
                typename GRAPH:: template EdgeMap<EdgeStates>  edgeState_;

                SettingsType        settings_;

                // Contracted graph:
                EdgeContractionGraphType edgeContractionGraph_;

                // Priority queue:
                QueueType pq_;

                uint64_t edgeToContractNext_;
                double   edgeToContractNextMergePrio_;

                NodeSizesType nodeSizes_;
                NodeSizesType maxNodeSize_per_iter_;
                NodeSizesType meanNodeSize_per_iter_;
                NodeSizesType variance_;
                NodeSizesType maxCostInPQ_per_iter_;
                uint64_t max_node_size_;
                uint64_t sum_node_size_;
                uint64_t quadratic_sum_node_size_;
            };

            template<class GRAPH, class UPDATE_RULE, bool ENABLE_UCM>
            template<class SIGNED_WEIGHTS, class IS_LOCAL_EDGE,class EDGE_SIZES,class NODE_SIZES>
            inline GaspClusterPolicy<GRAPH, UPDATE_RULE, ENABLE_UCM>::
            GaspClusterPolicy(
                    const GraphType & graph,
                    const SIGNED_WEIGHTS & signedWeights,
                    const IS_LOCAL_EDGE & isLocalEdge,
                    const EDGE_SIZES      & edgeSizes,
                    const NODE_SIZES      & nodeSizes,
                    const SettingsType & settings
            )
                    :   graph_(graph),
                        nonLinkConstraints_(graph),
                        accumulated_weights_(graph, signedWeights, edgeSizes, settings.updateRule),
                        edgeState_(graph),
                        nodeSizes_(graph),
                        pq_(graph.edgeIdUpperBound()+1),
                        settings_(settings),
                        edgeContractionGraph_(graph, *this),
                        maxNodeSize_per_iter_(graph),
                        maxCostInPQ_per_iter_(graph),
                        variance_(graph),
                        meanNodeSize_per_iter_(graph),
                        max_node_size_(0),
                        sum_node_size_(0),
                        quadratic_sum_node_size_(0),
                        nb_performed_contractions_(0),
                        nb_edges_popped_(0)
            {
                if (settings_.collectStats) {
                    // TODO: update shape actionStats
                    actionStats_ = xt::zeros<uint64_t>({uint64_t(graph_.numberOfNodes()), uint64_t(5)});
                    constraintsStats_ = xt::zeros<float>({uint64_t(graph_.numberOfEdges()), uint64_t(5)});
                    mergeStats_ = xt::zeros<float>({uint64_t(graph_.numberOfEdges())});
                }
                graph_.forEachNode([&](const uint64_t node) {
                    nodeSizes_[node] = nodeSizes[node];
                    sum_node_size_ += nodeSizes[node];
                    quadratic_sum_node_size_ += nodeSizes[node] * nodeSizes[node];
                    if (nodeSizes[node] > max_node_size_)
                        max_node_size_ = uint8_t(nodeSizes[node]);
                });

                graph_.forEachEdge([&](const uint64_t edge){
                    const auto loc = isLocalEdge[edge];
                    edgeState_[edge] = (loc == 1 ? EdgeStates::LOCAL : EdgeStates::LIFTED);
                    pq_.push(edge, this->computeWeight(edge));
                });
            }

            template<class GRAPH, class UPDATE_RULE, bool ENABLE_UCM>
            inline std::pair<uint64_t, double>
            GaspClusterPolicy<GRAPH, UPDATE_RULE, ENABLE_UCM>::
            edgeToContractNext() const {
                return std::pair<uint64_t, double>(edgeToContractNext_,edgeToContractNextMergePrio_) ;
            }

            template<class GRAPH, class UPDATE_RULE, bool ENABLE_UCM>
            inline bool
            GaspClusterPolicy<GRAPH, UPDATE_RULE, ENABLE_UCM>::isDone(
            ){
                while(true) {
                    while(!pq_.empty() && !isNegativeInf(pq_.topPriority()) && edgeContractionGraph_.numberOfNodes() > settings_.numberOfNodesStop){
                        // Here we already know that the edge is not lifted
                        // (Otherwise we would have inf cost in PQ)
                        const auto nextActioneEdge = pq_.top();

                        if (settings_.collectStats) {
                            // TODO: we can have more than nb_graph_edges popped from PQ, so old code could give seg fault
                            // nb_edges_popped_++;
                            // actionStats_(nb_edges_popped, 0) = graph_.uv(nextActioneEdge).first;
                            // actionStats_(nb_edges_popped, 1) = graph_.uv(nextActioneEdge).second;
                            // actionStats_(nb_edges_popped, 2) = nextActioneEdge;
                            // actionStats_(nb_edges_popped, 3) = pq_.topPriority();
                        }

                        // Check if some early constraints were enforced:
                        if (settings_.addNonLinkConstraints) {
                            if(this->isEdgeConstrained(nextActioneEdge)) {
                                if (settings_.collectStats) {
                                    // Check if the constrained edge is now positive and in case remember about it:
                                    if (this->isMergeAllowed(nextActioneEdge)) {
                                        const auto reprEdge = edgeContractionGraph_.findRepresentativeEdge(
                                                nextActioneEdge);
                                        constraintsStats_(reprEdge, 3) = accumulated_weights_[reprEdge];
                                        constraintsStats_(reprEdge, 4) = accumulated_weights_.weight(reprEdge);
                                    }
                                    // Remember that an edge was constrained at this iteration:
                                    // actionStats_(nb_edges_popped_, 4) = 1;
                                }
                                pq_.push(nextActioneEdge, -1.0*std::numeric_limits<double>::infinity());
                                continue;
                            }
                        }

                        // Here we check if we are allowed to make the merge:
                        if(this->isMergeAllowed(nextActioneEdge)){
                            edgeToContractNext_ = nextActioneEdge;
                            edgeToContractNextMergePrio_ = pq_.topPriority();
                            if (settings_.collectStats) {
                                // actionStats_(nb_performed_contractions_, 0) = graph_.uv(nextActioneEdge).first;
                                // actionStats_(nb_performed_contractions_, 1) = graph_.uv(nextActioneEdge).second;
                                // actionStats_(nb_performed_contractions_, 2) = nextActioneEdge;
                                // actionStats_(nb_performed_contractions_, 3) = pq_.topPriority();
                                // actionStats_(nb_performed_contractions_, 4) = 1;
                            }
                            return false;
                        }
                        else{
                            if (! settings_.addNonLinkConstraints) {
                                // In this case we know that we reached priority zero, so we can already stop
                                return true;
                            }
                            this->addNonLinkConstraint(nextActioneEdge);
                            pq_.push(nextActioneEdge, -1.0*std::numeric_limits<double>::infinity());
                            if (settings_.collectStats) {
                                // actionStats_(nb_edges_popped_, 4) = 2;
                            }

                            // Remember about weight and size of the constrained edge:
                            if (settings_.collectStats) {
                                const auto reprEdge = edgeContractionGraph_.findRepresentativeEdge(nextActioneEdge);
                                constraintsStats_(reprEdge, 0) = 1.;
                                constraintsStats_(reprEdge, 1) = accumulated_weights_[reprEdge]; // Remember the edge weight
                                constraintsStats_(reprEdge, 2) = accumulated_weights_.weight(reprEdge); // Remember the edge size
                            }
                        }
                    }
                    if (settings_.addNonLinkConstraints && settings_.mergeConstrainedEdgesAtTheEnd) {
                        // We push again all values to PQ and merge what is left and positive:
                        settings_.addNonLinkConstraints = false;
                        graph_.forEachEdge([&](const uint64_t e){
                            const auto cEdge = edgeContractionGraph_.findRepresentativeEdge(e);
                            const auto uv = edgeContractionGraph_.uv(cEdge);
                            const auto u = edgeContractionGraph_.findRepresentativeNode(uv.first);
                            const auto v = edgeContractionGraph_.findRepresentativeNode(uv.second);
                            if (u != v) {
                                pq_.push(cEdge, this->computeWeight(cEdge));
                            }
                        });
                        continue;
                    }
                    return true;
                }
            }

            template<class GRAPH, class UPDATE_RULE, bool ENABLE_UCM>
            inline double
            GaspClusterPolicy<GRAPH, UPDATE_RULE, ENABLE_UCM>::
            pqMergePrio(
                    const uint64_t edge
            ) const {
                const auto s = edgeState_[edge];
                double costInPQ;
                if(s == EdgeStates::LOCAL){
                    costInPQ = this->edgeCostInPQ(edge);
                }
                else{
                    // In this case the edge is lifted, so we need to be careful.
                    // It can be inserted in the PQ to constrain, but not to merge.
                    // REMARK: The second condition  "isMergeAllowed" is actually not necessary, because it is anyway checked
                    // again in isDone before to actually contract the edge...
                    if (!settings_.addNonLinkConstraints || this->isMergeAllowed(edge))
                        costInPQ = -1.0*std::numeric_limits<double>::infinity();
                    else {
                        costInPQ = this->edgeCostInPQ(edge);
                    }
                }
                return costInPQ;
            }

            template<class GRAPH, class UPDATE_RULE, bool ENABLE_UCM>
            inline void
            GaspClusterPolicy<GRAPH, UPDATE_RULE, ENABLE_UCM>::
            contractEdge(
                    const uint64_t edgeToContract
            ){
                // Remember about the highest cost in PQ:
                if (settings_.collectStats) {
                    mergeStats_(edgeToContract) = nb_performed_contractions_;
                    actionStats_(nb_performed_contractions_, 0) = edgeContractionGraph_.uv(edgeToContract).first;
                    actionStats_(nb_performed_contractions_, 1) = edgeContractionGraph_.uv(edgeToContract).second;
                    actionStats_(nb_performed_contractions_, 2) = edgeToContract;
                    actionStats_(nb_performed_contractions_, 3) = 1;
                    // actionStats_(nb_performed_contractions_, 3) = pq_.topPriority();

                }
                maxCostInPQ_per_iter_[nb_performed_contractions_] = this->edgeCostInPQ(edgeToContract);
                pq_.deleteItem(edgeToContract);
            }

            template<class GRAPH, class UPDATE_RULE, bool ENABLE_UCM>
            inline typename GaspClusterPolicy<GRAPH, UPDATE_RULE, ENABLE_UCM>::EdgeContractionGraphType &
            GaspClusterPolicy<GRAPH, UPDATE_RULE, ENABLE_UCM>::
            edgeContractionGraph(){
                return edgeContractionGraph_;
            }

            template<class GRAPH, class UPDATE_RULE, bool ENABLE_UCM>
            inline void
            GaspClusterPolicy<GRAPH, UPDATE_RULE, ENABLE_UCM>::
            mergeNodes(
                    const uint64_t aliveNode,
                    const uint64_t deadNode
            ){
                // Save data about max_node_size
                maxNodeSize_per_iter_[nb_performed_contractions_] = max_node_size_;
                const auto remaining_nodes = edgeContractionGraph_.numberOfNodes();
                meanNodeSize_per_iter_[nb_performed_contractions_] = float(sum_node_size_) / float(remaining_nodes);
                variance_[nb_performed_contractions_] =  float(quadratic_sum_node_size_) / float(remaining_nodes) - std::pow(meanNodeSize_per_iter_[nb_performed_contractions_], 2);

                quadratic_sum_node_size_ += std::pow(nodeSizes_[deadNode] + nodeSizes_[aliveNode], 2) - std::pow(nodeSizes_[aliveNode], 2) - std::pow(nodeSizes_[deadNode], 2);

                nodeSizes_[aliveNode] += nodeSizes_[deadNode];
                if (nodeSizes_[aliveNode] > max_node_size_)
                    max_node_size_ = uint64_t(nodeSizes_[aliveNode]);

                if (settings_.addNonLinkConstraints) {
                    auto  & aliveNodeNlc = nonLinkConstraints_[aliveNode];
                    const auto & deadNodeNlc = nonLinkConstraints_[deadNode];
                    aliveNodeNlc.insert(deadNodeNlc.begin(), deadNodeNlc.end());

                    for(const auto v : deadNodeNlc){
                        auto & nlc = nonLinkConstraints_[v];

                        // best way to change values in set...
                        nlc.erase(deadNode);
                        nlc.insert(aliveNode);
                    }

                    aliveNodeNlc.erase(deadNode);
                }

            }

            template<class GRAPH, class UPDATE_RULE, bool ENABLE_UCM>
            inline void
            GaspClusterPolicy<GRAPH, UPDATE_RULE, ENABLE_UCM>::
            mergeEdges(
                    const uint64_t aliveEdge,
                    const uint64_t deadEdge
            ){

                NIFTY_ASSERT_OP(aliveEdge,!=,deadEdge);
                NIFTY_ASSERT(pq_.contains(aliveEdge));
                NIFTY_ASSERT(pq_.contains(deadEdge));

                pq_.deleteItem(deadEdge);

                // update priority:
                accumulated_weights_.merge(aliveEdge, deadEdge);

                // update state
                // const auto oldStateAlive = edgeState_[aliveEdge];
                auto & sa = edgeState_[aliveEdge];
                const auto  sd = edgeState_[deadEdge];
                if(
                        sa == EdgeStates::LOCAL || sd == EdgeStates::LOCAL
                        ){
                    sa = EdgeStates::LOCAL;
                }
                else if( sa == EdgeStates::LIFTED  && sd == EdgeStates::LIFTED )
                {
                    sa = EdgeStates::LIFTED;
                }

                const auto sr = settings_.sizeRegularizer;
                bool add_to_PQ = true;
                if (settings_.addNonLinkConstraints) {
                    if (this->isEdgeConstrained(aliveEdge)) {
                        // If the edge was already constrained, it does not make sense to add it again to PQ
                        add_to_PQ = false;
                    }
                }
                if (sr > 0.00001)
                    add_to_PQ = false;
                if (add_to_PQ)
                    pq_.push(aliveEdge, this->computeWeight(aliveEdge));

            }

            template<class GRAPH, class UPDATE_RULE, bool ENABLE_UCM>
            inline void
            GaspClusterPolicy<GRAPH, UPDATE_RULE, ENABLE_UCM>::
            contractEdgeDone(
                    const uint64_t edgeToContract
            ){
                // HERE WE UPDATE the PQ when a SizeReg is used:
                const auto sr = settings_.sizeRegularizer;
                if (sr > 0.000001) {
                    const auto u = edgeContractionGraph_.nodeOfDeadEdge(edgeToContract);
                    for(auto adj : edgeContractionGraph_.adjacency(u)){
                        const auto edge = adj.edge();
                        pq_.push(edge, computeWeight(edge));
                    }
                }

                nb_performed_contractions_++;

            }

            template<class GRAPH, class UPDATE_RULE, bool ENABLE_UCM>
            inline double
            GaspClusterPolicy<GRAPH, UPDATE_RULE,  ENABLE_UCM>::
            computeWeight(
                    const uint64_t edge
            ) const {
                const auto fromEdge = this->pqMergePrio(edge); // This -inf if the edge was lifted
                const auto sr = settings_.sizeRegularizer;

                if (sr > 0.000001 && !isNegativeInf(fromEdge))
                {
                    const auto uv = edgeContractionGraph_.uv(edge);
                    const auto sizeU = nodeSizes_[uv.first];
                    const auto sizeV = nodeSizes_[uv.second];
                    const auto sFac = 2.0 / ( 1.0/std::pow(sizeU,sr) + 1.0/std::pow(sizeV,sr) );
                    return fromEdge * (1. / sFac);
                } else {
                    return fromEdge;
                }
            }

        } // namespace agglo
    } // namespace nifty::graph
} // namespace nifty