/* 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_dc_serialization

#include <boost/test/unit_test.hpp>
#include <core/util/test_macros.hpp>
#include <cstdio>
#include <ml/neural_net/model_spec.hpp>
#include <toolkits/coreml_export/mlmodel_include.hpp>
#include <toolkits/drawing_classifier/drawing_classifier.hpp>

namespace turi {
namespace drawing_classifer {
namespace {

class drawing_classifier_mock
    : public turi::drawing_classifier::drawing_classifier {
 public:
  drawing_classifier_mock() = default;

  drawing_classifier_mock(std::unique_ptr<neural_net::model_spec> ms)
      : drawing_classifier({}, std::move(ms), nullptr, nullptr, nullptr) {}

  std::unique_ptr<neural_net::model_spec> get_model_spec_copy() const {
    return clone_model_spec_for_test();
  }

  std::unique_ptr<neural_net::model_spec> get_model_spec() const {
    return init_model();
  }

};

BOOST_AUTO_TEST_CASE(test_dc_init_model) {
  // states
  constexpr unsigned int num_classes = 10;
  const std::string target = "target";
  const std::vector<std::string> features = {"0", "1"};

  // init_model
  drawing_classifier_mock dc;
  dc.add_or_update_state(
      {{"target", target},
       {"num_classes", num_classes},
       {"features", flex_list(features.begin(), features.end())}});

  auto nn_spec = dc.get_model_spec();

  const CoreML::Specification::NeuralNetwork& nn = nn_spec->get_coreml_spec();
  /*
   * 1 input,
   * 3 (conv layers, relu, maxpool), 1 flatten, 2 dense, 1 softmax
   * in total: 14 layers
   */
  TS_ASSERT_EQUALS(nn.layers_size(), 14);

  const std::string _suffix = "";

  /* layer 0: concat layer */
  {
    auto concat_layer = nn.layers(0);
    TS_ASSERT(concat_layer.has_concat());
    TS_ASSERT(concat_layer.output_size() == 1);
    TS_ASSERT_EQUALS(concat_layer.output(0), "features");
    TS_ASSERT_EQUALS(concat_layer.name(), "features");
    for (int ii = 0; ii < concat_layer.input_size(); ii++) {
      TS_ASSERT_EQUALS(concat_layer.input(ii), features.at(ii));
    }
  }

  {
    const std::map<int, int> layer_num_to_channels = {
        {0, 16}, {1, 32}, {2, 64}};
    for (auto& x : layer_num_to_channels) {
      int layer_index = x.first * 3 + 1;
      const auto& convlayer = nn.layers(layer_index);
      TS_ASSERT(convlayer.has_convolution());
      TS_ASSERT_EQUALS(convlayer.name(),
                       "drawing_conv" + std::to_string(x.first) + _suffix);
      TS_ASSERT_EQUALS(convlayer.convolution().outputchannels(), x.second);
      TS_ASSERT_EQUALS(convlayer.convolution().kernelchannels(),
                       x.first == 0 ? 1 : x.second / 2);
      TS_ASSERT_EQUALS(convlayer.convolution().stride(0), 1);
      TS_ASSERT_EQUALS(convlayer.convolution().stride(1), 1);
      TS_ASSERT_EQUALS(convlayer.convolution().kernelsize(0), 3);
      TS_ASSERT_EQUALS(convlayer.convolution().kernelsize(1), 3);
      TS_ASSERT(convlayer.convolution().has_same());

      const auto& relu_layer = nn.layers(layer_index + 1);
      TS_ASSERT_EQUALS(relu_layer.name(),
                       "drawing_relu" + std::to_string(x.first) + _suffix);
      TS_ASSERT(relu_layer.has_activation());
      TS_ASSERT(relu_layer.activation().has_relu());

      const auto& pool_layer = nn.layers(layer_index + 2);
      TS_ASSERT(pool_layer.has_pooling());
      TS_ASSERT_EQUALS(pool_layer.name(),
                       "drawing_pool" + std::to_string(x.first) + _suffix);
      TS_ASSERT_EQUALS(pool_layer.pooling().kernelsize(0), 2);
      TS_ASSERT_EQUALS(pool_layer.pooling().kernelsize(1), 2);
      TS_ASSERT_EQUALS(pool_layer.pooling().stride(0), 2);
      TS_ASSERT_EQUALS(pool_layer.pooling().stride(1), 2);
      TS_ASSERT(pool_layer.pooling().has_valid());
    }
  }

  unsigned layer_index = 10;
  {
    const auto& flatten_layer = nn.layers(layer_index);
    TS_ASSERT(flatten_layer.has_flatten());
    TS_ASSERT_EQUALS(flatten_layer.name(), "drawing_flatten0" + _suffix);
  }

  layer_index++;
  {
    const auto& dense_layer = nn.layers(layer_index);
    TS_ASSERT(dense_layer.has_innerproduct());
    TS_ASSERT_EQUALS(dense_layer.name(), "drawing_dense0" + _suffix);
    TS_ASSERT_EQUALS(dense_layer.innerproduct().inputchannels(), 64 * 3 * 3);
    TS_ASSERT_EQUALS(dense_layer.innerproduct().outputchannels(), 128);
  }

  layer_index++;
  {
    const auto& dense_layer = nn.layers(layer_index);
    TS_ASSERT_EQUALS(dense_layer.name(), "drawing_dense1" + _suffix);
    TS_ASSERT(dense_layer.has_innerproduct());
    TS_ASSERT_EQUALS(dense_layer.innerproduct().inputchannels(), 128);
    TS_ASSERT_EQUALS(dense_layer.innerproduct().outputchannels(), num_classes);
  }

  layer_index++;
  TS_ASSERT(layer_index == 13);
  {
    const auto& softmax_layer = nn.layers(layer_index);
    TS_ASSERT(softmax_layer.has_softmax());
    TS_ASSERT_EQUALS(softmax_layer.output(0), target + "Probability");
  }
}

BOOST_AUTO_TEST_CASE(test_export_coreml) {
  // minimum startup code
  const std::string target = "target";
  const std::vector<std::string> features = {"0", "1"};
  const std::vector<std::string> labels = {"0", "1"};

  turi::drawing_classifier::drawing_classifier dc;
  dc.add_or_update_state(
      {{"target", target},
       {"num_classes", labels.size()},
       {"classes", flex_list(labels.begin(), labels.end())},
       {"max_iterations", 300},
       {"warm_start", false},
       {"features", flex_list(features.begin(), features.end())}});

  auto ml_model_wrapper = dc.export_to_coreml("", /* debug no throw */ true);
  TS_ASSERT(ml_model_wrapper != nullptr);

  const auto& my_model_spec = ml_model_wrapper->coreml_model()->getProto();
  TS_ASSERT_EQUALS(my_model_spec.specificationversion(), 1);

  auto& my_model_desc = my_model_spec.description();
  TS_ASSERT_EQUALS(my_model_desc.input_size(), features.size());

  // test input image type
  TS_ASSERT_EQUALS(my_model_desc.input(0).name(), "image");
  TS_ASSERT(my_model_desc.input(0).type().has_imagetype());

  auto input_feature_type = my_model_desc.input(0).type().imagetype();
  TS_ASSERT_EQUALS(input_feature_type.colorspace(),
                   CoreML::Specification::ImageFeatureType::GRAYSCALE);
  TS_ASSERT_EQUALS(input_feature_type.width(), 28);
  TS_ASSERT_EQUALS(input_feature_type.height(), 28);

  TS_ASSERT_EQUALS(my_model_desc.output_size(), 2);
  TS_ASSERT_EQUALS(my_model_desc.output(0).name(), target + "Probability");
  TS_ASSERT_EQUALS(my_model_desc.output(1).name(), target);

  TS_ASSERT_EQUALS(my_model_desc.predictedfeaturename(), target);
  TS_ASSERT_EQUALS(my_model_desc.predictedprobabilitiesname(),
                   target + "Probability");
}

BOOST_AUTO_TEST_CASE(test_save_load) {
  auto remove_if_exist = [](const char* fname) {
    std::FILE* fp = nullptr;
    if ((fp = std::fopen(fname, "r")) != NULL) {
      if (std::fclose(fp) != 0) {
        std::stringstream ss;
        ss << "cannot close file: " << fname << std::endl;
        throw(turi::error::io_error(ss.str()));
      }

      if (std::remove(fname) != 0) {
        std::stringstream ss;
        ss << "cannot remove file: " << fname << std::endl;
        throw(turi::error::io_error(ss.str()));
      }
    }
  };

  auto load_save_compare = [&remove_if_exist](
                               drawing_classifier_mock& dc,
                               drawing_classifier_mock& dc_other) {
    constexpr auto my_file = "./test_dc_serialization.cxx.save.txt";
    remove_if_exist(my_file);
    std::stringstream ss;

    // save impl from the first mock
    {
      std::fstream out_file(my_file,
                            std::ios_base::out | std::ios_base::binary);
      turi::oarchive oarch(out_file);
      dc.save_impl(oarch);
      out_file.close();
    }

    // load impl to other dc mock instance
    {
      std::fstream in_file(my_file, std::ios_base::in | std::ios_base::binary);
      turi::iarchive iarch(in_file);
      dc_other.load_version(iarch, dc_other.get_version());
      in_file.close();
    }

    // clean myself
    if (std::remove(my_file)) {
      std::stringstream ss;
      ss << "fail to remove file:" << my_file << std::endl;
      std::perror(ss.str().c_str());
    }

    // compare weights in memory
    auto original_spec = dc.get_model_spec_copy();
    TS_ASSERT(original_spec != nullptr);
    auto original_view = original_spec->export_params_view();

    auto loaded_spec = dc_other.get_model_spec_copy();
    TS_ASSERT(loaded_spec != nullptr);
    auto loaded_view = loaded_spec->export_params_view();

    TS_ASSERT(original_view.size() > 1);
    TS_ASSERT_EQUALS(original_view.size(), loaded_view.size());

    for (const auto& entry : original_view) {
      // std::cout << entry.first << std::endl;
      auto& loaded_weights = loaded_view.at(entry.first);
      auto& original_weights = entry.second;
      TS_ASSERT(loaded_weights.size() > 0);
      TS_ASSERT_EQUALS(original_weights.size(), loaded_weights.size());
      TS_ASSERT_EQUALS(
       std::memcmp(loaded_weights.data(), original_weights.data(),
                       loaded_weights.size()),
          0);
    }
  };

  // states
  constexpr unsigned int num_classes = 10;
  const std::string target = "target";
  const std::vector<std::string> features = {"0", "1"};

  // random init; avoid segfault

  drawing_classifier_mock dummy;
  dummy.add_or_update_state(
      {{"target", target},
       {"num_classes", num_classes},
       {"random_seed", 1},
       {"features", flex_list(features.begin(), features.end())}});

  // model spec should be different since all weights are random
  // generated
  auto spec1 = dummy.get_model_spec();
  TS_ASSERT(spec1 != nullptr);

  dummy.add_or_update_state({{"random_seed", 2}});
  auto spec2 = dummy.get_model_spec();
  TS_ASSERT(spec2 != nullptr);

  auto view1 = spec1->export_params_view();
  auto view2 = spec2->export_params_view();

  bool is_same = true;

  for (const auto& entry : view1) {
    auto& weights2 = view2.at(entry.first);
    auto& weights1 = entry.second;
    TS_ASSERT(weights1.size() > 0);
    TS_ASSERT_EQUALS(weights1.size(), weights2.size());
    if (std::memcmp(weights1.data(), weights2.data(), weights1.size())) {
      is_same = false;
      break;
    }
  }

  // make sure 2 models params view are not same
  // due to the fact that they use different seed!
  TS_ASSERT(!is_same);

  // start from 2 model spec with different weights view
  drawing_classifier_mock dc(std::move(spec1));
  dc.add_or_update_state(
      {{"target", target},
       {"num_classes", num_classes},
       {"features", flex_list(features.begin(), features.end())}});

  // load from a different instance
  drawing_classifier_mock dc_other(std::move(spec2));
  dc_other.add_or_update_state(
      {{"target", target},
       {"num_classes", num_classes},
       {"features", flex_list(features.begin(), features.end())}});

  load_save_compare(dc, dc_other);
}

}  // anonymous namespace
}  // namespace drawing_classifer
}  // namespace turi


// Copyright (c) 2018-2019 The Bitcoin Unlimited developers
// Distributed under the MIT software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include "blockrelay/graphene.h"
#include "blockrelay/graphene_set.h"
#include "bloom.h"
#include "fastfilter.h"
#include "hashwrapper.h"
#include "primitives/block.h"
#include "primitives/transaction.h"
#include "serialize.h"
#include "streams.h"
#include "test/test_bitcoin.h"

#include <boost/test/unit_test.hpp>
#include <cassert>
#include <chrono>
#include <cmath>
#include <iostream>
#include <iomanip>

#define MAX_GRAPHENE_SET_VERSION 4

size_t ProjectedGrapheneSizeBytes(uint64_t version, uint64_t nBlockTxs, uint64_t nExcessTxs, uint64_t nSymDiff, bool computeOptimized=false)
{
    const int SERIALIZATION_OVERHEAD = 11;
    auto fpr = [nExcessTxs](int a) { return a / float(nExcessTxs); };

    CIblt iblt(nSymDiff, CGrapheneSet::GetCIbltVersion(version));
    size_t ibltBytes = ::GetSerializeSize(iblt, SER_NETWORK, PROTOCOL_VERSION) - SERIALIZATION_OVERHEAD;

    size_t filterBytes;
    if (computeOptimized)
    {
        CVariableFastFilter filter(nBlockTxs, fpr(nSymDiff));
        filterBytes = ::GetSerializeSize(filter, SER_NETWORK, PROTOCOL_VERSION) - SERIALIZATION_OVERHEAD;
    }
    else
    {
        CBloomFilter filter(
            nBlockTxs, fpr(nSymDiff), insecure_rand_ctx.rand32(), BLOOM_UPDATE_ALL, true, std::numeric_limits<uint32_t>::max());
        filterBytes = ::GetSerializeSize(filter, SER_NETWORK, PROTOCOL_VERSION) - SERIALIZATION_OVERHEAD;
    }

    return filterBytes + ibltBytes;
}

// Create a deterministic hash by providing an index
uint256 GetHash(unsigned int nIndex)
{
    std::stringstream ss;
    ss << std::setfill('0') << std::setw(32) << nIndex;
    return uint256S(ss.str());
}

BOOST_FIXTURE_TEST_SUITE(graphene_tests, BasicTestingSetup)

BOOST_AUTO_TEST_CASE(graphene_set_encodes_and_decodes)
{
    uint64_t version = MAX_GRAPHENE_SET_VERSION;
    uint256 senderArr[] = {
        SerializeHash(3), SerializeHash(1), SerializeHash(2), SerializeHash(7), SerializeHash(11), SerializeHash(4)};
    std::vector<uint256> senderItems(senderArr, senderArr + sizeof(senderArr) / sizeof(uint256));
    uint256 receiverArr[] = {
        SerializeHash(7), SerializeHash(2), SerializeHash(4), SerializeHash(-1), SerializeHash(1), SerializeHash(11)};
    std::vector<uint256> receiverItems(receiverArr, receiverArr + sizeof(receiverArr) / sizeof(uint256));

    // unordered graphene sets
    {
        CGrapheneSet senderGrapheneSet(6, 6, senderItems, 0, 0, version, 0, false, false, true);
        std::vector<uint64_t> reconciledCheapHashes = senderGrapheneSet.Reconcile(receiverItems);

        std::vector<uint64_t> senderCheapHashes;
        for (uint256 item : senderItems)
            senderCheapHashes.push_back(senderGrapheneSet.GetShortID(item));

        std::sort(senderCheapHashes.begin(), senderCheapHashes.end(), [](uint64_t i1, uint64_t i2) { return i1 < i2; });
        std::sort(reconciledCheapHashes.begin(), reconciledCheapHashes.end(),
            [](uint64_t i1, uint64_t i2) { return i1 < i2; });

        BOOST_CHECK_EQUAL_COLLECTIONS(reconciledCheapHashes.begin(), reconciledCheapHashes.end(),
            senderCheapHashes.begin(), senderCheapHashes.end());
    }

    // ordered graphene sets
    {
        CGrapheneSet senderGrapheneSet(6, 6, senderItems, 0, 0, version, 0, false, true, true);
        std::vector<uint64_t> reconciledCheapHashes = senderGrapheneSet.Reconcile(receiverItems);

        std::vector<uint64_t> senderCheapHashes;
        for (uint256 item : senderItems)
            senderCheapHashes.push_back(senderGrapheneSet.GetShortID(item));

        BOOST_CHECK_EQUAL_COLLECTIONS(reconciledCheapHashes.begin(), reconciledCheapHashes.end(),
            senderCheapHashes.begin(), senderCheapHashes.end());
    }
}

BOOST_AUTO_TEST_CASE(graphene_set_decodes_multiple_sizes)
{
    size_t nItemList[] = {1, 10, 50, 500, 5000, 10000};
    int nNumHashes = 0;
    uint64_t version = MAX_GRAPHENE_SET_VERSION;
    for (size_t nItems : nItemList)
    {
        std::vector<uint256> senderItems;
        std::vector<uint256> baseReceiverItems;
        for (size_t i = 1; i <= nItems; i++)
        {
            nNumHashes++;
            const uint256 &hash = GetHash(nNumHashes);
            senderItems.push_back(hash);
            baseReceiverItems.push_back(hash);
        }

        // Add 10 more items to receiver mempool
        {
            std::vector<uint256> receiverItems = baseReceiverItems;
            for (size_t j = 1; j < 11; j++)
            {
                nNumHashes++;
                receiverItems.push_back(SerializeHash(GetHash(nNumHashes)));
            }

            CGrapheneSet senderGrapheneSet(
                receiverItems.size(), receiverItems.size(), senderItems, 0, 0, version, 0, false, true, true);
            std::vector<uint64_t> reconciledCheapHashes = senderGrapheneSet.Reconcile(receiverItems);

            std::vector<uint64_t> senderCheapHashes;
            for (auto &hash : senderItems)
                senderCheapHashes.push_back(senderGrapheneSet.GetShortID(hash));

            BOOST_CHECK_EQUAL_COLLECTIONS(reconciledCheapHashes.begin(), reconciledCheapHashes.end(),
                senderCheapHashes.begin(), senderCheapHashes.end());
        }

        // Add 100 more items to receiver mempool
        {
            std::vector<uint256> receiverItems = baseReceiverItems;
            for (size_t j = 1; j < 101; j++)
            {
                nNumHashes++;
                receiverItems.push_back(SerializeHash(GetHash(nNumHashes)));
            }

            CGrapheneSet senderGrapheneSet(
                receiverItems.size(), receiverItems.size(), senderItems, 0, 0, version, 0, false, true, true);
            std::vector<uint64_t> reconciledCheapHashes = senderGrapheneSet.Reconcile(receiverItems);

            std::vector<uint64_t> senderCheapHashes;
            for (auto &hash : senderItems)
                senderCheapHashes.push_back(senderGrapheneSet.GetShortID(hash));

            BOOST_CHECK_EQUAL_COLLECTIONS(reconciledCheapHashes.begin(), reconciledCheapHashes.end(),
                senderCheapHashes.begin(), senderCheapHashes.end());
        }
    }
}

BOOST_AUTO_TEST_CASE(graphene_set_finds_brute_force_opt_for_small_blocks)
{
    uint64_t version = MAX_GRAPHENE_SET_VERSION;
    CGrapheneSet grapheneSet(version);

    uint16_t approx_items_thresh = version >= 4 ? APPROX_ITEMS_THRESH_REDUCE_CHECK : APPROX_ITEMS_THRESH;
    int n = (int)std::floor(approx_items_thresh / 2);
    int mu = 100;
    int m = (int)std::floor(n / 8) + mu;

    int best_a = 1;
    size_t best_size = std::numeric_limits<size_t>::max();
    int a = 1;
    for (a = 1; a < m - mu; a++)
    {
        size_t totalBytes = ProjectedGrapheneSizeBytes(version, n, m - mu, a);
        size_t totalBytesOpt = ProjectedGrapheneSizeBytes(version, n, m - mu, a, true);

        BOOST_CHECK_EQUAL(totalBytes, totalBytesOpt);

        if (totalBytes < best_size)
        {
            best_size = totalBytes;
            best_a = a;
        }
    }

    BOOST_CHECK_EQUAL(grapheneSet.OptimalSymDiff(version, n, m, m - mu, 1), best_a);
}

BOOST_AUTO_TEST_CASE(graphene_set_finds_approx_opt_for_large_blocks)
{
    uint64_t version = MAX_GRAPHENE_SET_VERSION;
    uint16_t approx_items_thresh = version >= 4 ? APPROX_ITEMS_THRESH_REDUCE_CHECK : APPROX_ITEMS_THRESH;
    int n = 4 * approx_items_thresh;
    int mu = 1000;
    int m = approx_items_thresh + mu;
    CGrapheneSet grapheneSet(version);
    double optSymDiff = grapheneSet.OptimalSymDiff(version, n, m, m - mu, 0);
    double fpr = CGrapheneSet::BloomFalsePositiveRate(optSymDiff, m - mu);
    uint8_t checksumBits = CGrapheneSet::NChecksumBits(n * CIblt::OptimalOverhead(n), CIblt::OptimalNHash(n), m, fpr, UNCHECKED_ERROR_TOL);
    auto approxSymDiff = [n, checksumBits]() {
        return std::max(
            1.0, std::round(FILTER_CELL_SIZE * n / ((checksumBits + 8 * IBLT_FIXED_CELL_SIZE) * IBLT_DEFAULT_OVERHEAD * LN2SQUARED)));
    };

    BOOST_CHECK_EQUAL(approxSymDiff(), optSymDiff);
}

BOOST_AUTO_TEST_CASE(graphene_set_approx_opt_close_to_optimal)
{
    for (uint64_t version=0; version < MAX_GRAPHENE_SET_VERSION; version++)
    {
        int n = version >= 4 ? APPROX_ITEMS_THRESH_REDUCE_CHECK : APPROX_ITEMS_THRESH;
        int mu = 100;
        int m = (int)std::ceil(n / APPROX_EXCESS_RATE) + mu;
        CGrapheneSet grapheneSet(version);

        float totalBytesApprox = (float)ProjectedGrapheneSizeBytes(version, n, m - mu, CGrapheneSet::ApproxOptimalSymDiff(version, n));
        float totalBytesBrute =
            (float)ProjectedGrapheneSizeBytes(version, n, m - mu, grapheneSet.BruteForceSymDiff(n, m, m - mu, MAX_CHECKSUM_BITS));
        float totalBytesBruteOpt =
            (float)ProjectedGrapheneSizeBytes(version, n, m - mu, grapheneSet.BruteForceSymDiff(n, m, m - mu, MAX_CHECKSUM_BITS), true);

        BOOST_CHECK_CLOSE(totalBytesApprox, totalBytesBrute, 15);
        BOOST_CHECK_CLOSE(totalBytesApprox, totalBytesBruteOpt, 15);
    }
}

BOOST_AUTO_TEST_CASE(graphene_set_decodes_empty_intersection)
{
    uint64_t version = MAX_GRAPHENE_SET_VERSION;
    uint256 senderArr[] = {SerializeHash(-7), SerializeHash(-2), SerializeHash(-4), SerializeHash(-1),
        SerializeHash(-5), SerializeHash(-11), SerializeHash(3), SerializeHash(1), SerializeHash(2), SerializeHash(7),
        SerializeHash(11), SerializeHash(4)};
    std::vector<uint256> senderItems(senderArr, senderArr + sizeof(senderArr) / sizeof(uint256));

    // includes no transactions from block
    uint256 receiverArr[] = {SerializeHash(-7), SerializeHash(-2), SerializeHash(-4), SerializeHash(-1),
        SerializeHash(-5), SerializeHash(-11)};
    std::vector<uint256> receiverItems(receiverArr, receiverArr + sizeof(receiverArr) / sizeof(uint256));

    CGrapheneSet senderGrapheneSet(6, 12, senderItems, 0, 0, version, 0, false, true, true);
    std::vector<uint64_t> reconciledCheapHashes = senderGrapheneSet.Reconcile(receiverItems);

    std::vector<uint64_t> senderCheapHashes;
    for (uint256 item : senderItems)
        senderCheapHashes.push_back(senderGrapheneSet.GetShortID(item));

    BOOST_CHECK_EQUAL_COLLECTIONS(
        reconciledCheapHashes.begin(), reconciledCheapHashes.end(), senderCheapHashes.begin(), senderCheapHashes.end());
}

BOOST_AUTO_TEST_CASE(graphene_set_can_serde)
{
    uint64_t version = MAX_GRAPHENE_SET_VERSION;
    std::vector<uint256> senderItems;
    CDataStream ss(SER_DISK, 0);

    senderItems.push_back(SerializeHash(3));
    CGrapheneSet sentGrapheneSet(1, 1, senderItems, 0, 0, version, 0, false, false, true);
    CGrapheneSet receivedGrapheneSet(version);

    ss << sentGrapheneSet;
    ss >> receivedGrapheneSet;

    BOOST_CHECK_EQUAL(receivedGrapheneSet.Reconcile(senderItems)[0], sentGrapheneSet.GetShortID(senderItems[0]));
}

BOOST_AUTO_TEST_CASE(graphene_set_version_check)
{
    for (uint64_t version = 0; version <= MAX_GRAPHENE_SET_VERSION; version++)
    {
        std::vector<uint256> senderItems;
        CDataStream ss(SER_DISK, 0);

        senderItems.push_back(SerializeHash(1));
        senderItems.push_back(SerializeHash(2));
        senderItems.push_back(SerializeHash(3));
        CGrapheneSet sentGrapheneSet(3, 3, senderItems, 0, 0, version, false, false, true);
        CGrapheneSet receivedGrapheneSet(version);

        ss << sentGrapheneSet;
        ss >> receivedGrapheneSet;
    }
}

BOOST_AUTO_TEST_CASE(item_rank_encodes_and_decodes)
{
    uint64_t itemArr[4] = {1, 20, 500, 7000};
    std::vector<uint64_t> inputItems(itemArr, itemArr + sizeof(itemArr) / sizeof(uint64_t));
    uint16_t nBits = 13;

    std::vector<unsigned char> encoded = CGrapheneSet::EncodeRank(inputItems, nBits);
    std::vector<uint64_t> outputItems = CGrapheneSet::DecodeRank(encoded, inputItems.size(), nBits);

    BOOST_CHECK_EQUAL_COLLECTIONS(outputItems.begin(), outputItems.end(), inputItems.begin(), inputItems.end());
}

BOOST_AUTO_TEST_CASE(compute_optimized_graphene_set_can_serde)
{
    uint64_t version = MAX_GRAPHENE_SET_VERSION;
    std::vector<uint256> senderItems;
    CDataStream ss(SER_DISK, 0);

    senderItems.push_back(SerializeHash(3));
    CGrapheneSet sentGrapheneSet(1, 1, senderItems, 0, 0, version, 0, true, false, true);
    CGrapheneSet receivedGrapheneSet(version, true);

    ss << sentGrapheneSet;
    ss >> receivedGrapheneSet;

    BOOST_CHECK_EQUAL(receivedGrapheneSet.Reconcile(senderItems)[0], sentGrapheneSet.GetShortID(senderItems[0]));
}

BOOST_AUTO_TEST_CASE(graphene_set_cpu_check)
{
    size_t nItems = 10000;
    int nNumHashes = 0;

    std::vector<uint256> senderItems;
    std::vector<uint64_t> senderCheapHashes;
    std::vector<uint256> baseReceiverItems;
    for (size_t i = 1; i <= nItems; i++)
    {
        nNumHashes++;
        const uint256 &hash = GetHash(nNumHashes);
        senderItems.push_back(hash);
        senderCheapHashes.push_back(hash.GetCheapHash());
        baseReceiverItems.push_back(hash);
    }

    // Add 10000 more items to receiver mempool
    std::vector<uint256> receiverItems = baseReceiverItems;
    for (size_t j = 1; j < 10000; j++)
    {
        nNumHashes++;
        receiverItems.push_back(SerializeHash(GetHash(nNumHashes)));
    }

    auto legacyStart = std::chrono::high_resolution_clock::now();
    CGrapheneSet legacyGrapheneSet(
        receiverItems.size(), receiverItems.size(), senderItems, 0, 0, 0, 0, false, false, true);
    legacyGrapheneSet.Reconcile(receiverItems);
    auto legacyFinish = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> legacyElapsed = legacyFinish - legacyStart;
    std::cout << "Legacy elapsed time: " << legacyElapsed.count() << " s\n";

    auto sipStart = std::chrono::high_resolution_clock::now();
    CGrapheneSet sipGrapheneSet(
        receiverItems.size(), receiverItems.size(), senderItems, 0, 0, MAX_GRAPHENE_SET_VERSION, 0, false, false, true);
    sipGrapheneSet.Reconcile(receiverItems);
    auto sipFinish = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> sipElapsed = sipFinish - sipStart;
    std::cout << "Sip elapsed time: " << sipElapsed.count() << " s\n";

    auto fastStart = std::chrono::high_resolution_clock::now();
    CGrapheneSet fastGrapheneSet(
        receiverItems.size(), receiverItems.size(), senderItems, 0, 0, MAX_GRAPHENE_SET_VERSION, 0, true, false, true);
    fastGrapheneSet.Reconcile(receiverItems);
    auto fastFinish = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> fastElapsed = fastFinish - fastStart;
    std::cout << "Fast elapsed time: " << fastElapsed.count() << " s\n";
}

BOOST_AUTO_TEST_CASE(graphene_block_can_serde)
{
    uint64_t version = GRAPHENE_MAX_VERSION_SUPPORTED;
    // regular graphene block
    {
        CBlock block;
        CTransaction tx;
        CDataStream stream(
            ParseHex("01000000010b26e9b7735eb6aabdf358bab62f9816a21ba9ebdb719d5299e88607d722c190000000008b4830450220070aca4"
                     "4506c5cef3a16ed519d7c3c39f8aab192c4e1c90d065f37b8a4af6141022100a8e160b856c2d43d27d8fba71e5aef6405b864"
                     "3ac4cb7cb3c462aced7f14711a0141046d11fee51b0e60666d5049a9101a72741df480b96ee26488a4d3466b95c9a40ac5eee"
                     "f87e10a5cd336c19a84565f80fa6c547957b7700ff4dfbdefe76036c339ffffffff021bff3d11000000001976a91404943fdd"
                     "508053c75000106d3bc6e2754dbcff1988ac2f15de00000000001976a914a266436d2965547608b9e15d9032a7b9d64fa4318"
                     "8ac00000000"),
            SER_DISK, CLIENT_VERSION);
        stream >> tx;
        const CTransactionRef ptx = MakeTransactionRef(tx);
        block.vtx.push_back(ptx);
        CGrapheneBlock senderGrapheneBlock(MakeBlockRef(block), 5, 6, version, false);
        CGrapheneBlock receiverGrapheneBlock(4);
        CDataStream ss(SER_DISK, 0);

        ss << senderGrapheneBlock;
        ss >> receiverGrapheneBlock;
    }

    // compute optimized graphene block
    {
        CBlock block;
        CTransaction tx;
        CDataStream stream(
            ParseHex("01000000010b26e9b7735eb6aabdf358bab62f9816a21ba9ebdb719d5299e88607d722c190000000008b4830450220070aca4"
                     "4506c5cef3a16ed519d7c3c39f8aab192c4e1c90d065f37b8a4af6141022100a8e160b856c2d43d27d8fba71e5aef6405b864"
                     "3ac4cb7cb3c462aced7f14711a0141046d11fee51b0e60666d5049a9101a72741df480b96ee26488a4d3466b95c9a40ac5eee"
                     "f87e10a5cd336c19a84565f80fa6c547957b7700ff4dfbdefe76036c339ffffffff021bff3d11000000001976a91404943fdd"
                     "508053c75000106d3bc6e2754dbcff1988ac2f15de00000000001976a914a266436d2965547608b9e15d9032a7b9d64fa4318"
                     "8ac00000000"),
            SER_DISK, CLIENT_VERSION);
        stream >> tx;
        const CTransactionRef ptx = MakeTransactionRef(tx);
        block.vtx.push_back(ptx);
        CGrapheneBlock senderGrapheneBlock(MakeBlockRef(block), 5, 6, version, true);
        CGrapheneBlock receiverGrapheneBlock(4, true);
        CDataStream ss(SER_DISK, 0);

        ss << senderGrapheneBlock;
        ss >> receiverGrapheneBlock;
    }
}

BOOST_AUTO_TEST_CASE(nchecksumbits_gives_correct_value)
{
    double tol = 1 / std::pow(2, 11);
    uint8_t bits = CGrapheneSet::NChecksumBits(10, 2, 1, 0.5, tol);

    BOOST_CHECK_EQUAL(bits, 11);
}

BOOST_AUTO_TEST_CASE(graphene_failure_recovery_primitives)
{
    size_t nItems = 100;
    int nNumHashes = 0;
    uint64_t grapheneSetVersion = CGrapheneBlock::GetGrapheneSetVersion(GRAPHENE_MAX_VERSION_SUPPORTED);
    uint32_t ibltSalt = 3;

    std::vector<uint256> senderItems;
    std::set<uint64_t> senderCheapHashes;
    std::vector<uint256> baseReceiverItems;
    for (size_t i = 1; i <= nItems; i++)
    {
        nNumHashes++;
        const uint256 &hash = GetHash(nNumHashes);
        senderItems.push_back(hash);
        senderCheapHashes.insert(hash.GetCheapHash());
        baseReceiverItems.push_back(hash);
    }

    // Add 10000 more items to receiver mempool
    std::vector<uint256> receiverItems = baseReceiverItems;
    for (size_t j = 1; j < 10000; j++)
    {
        nNumHashes++;
        receiverItems.push_back(SerializeHash(GetHash(nNumHashes)));
    }

    CGrapheneSet fastGrapheneSet(
        receiverItems.size(), receiverItems.size(), senderItems, 0, 0, MAX_GRAPHENE_SET_VERSION, 0, true, false, true);
    fastGrapheneSet.Reconcile(receiverItems);

    uint64_t nPassedFilterItems = 12;
    uint64_t nReceiverUniverseItems = 200;
    double fSenderBloomFpr = 1.0;

	CVariableFastFilter receiverFilter = fastGrapheneSet.FailureRecoveryFilter(senderItems, nItems, nPassedFilterItems, nReceiverUniverseItems, FAILURE_RECOVERY_SUCCESS_RATE, fSenderBloomFpr, grapheneSetVersion);
    CIblt recoveryIblt = fastGrapheneSet.FailureRecoveryIblt(senderCheapHashes, nItems, nPassedFilterItems, nReceiverUniverseItems, FAILURE_RECOVERY_SUCCESS_RATE, fSenderBloomFpr, grapheneSetVersion, ibltSalt);
}

BOOST_AUTO_TEST_CASE(graphene_failure_recovery_params_sanity)
{
    // exmple from https://github.com/bissias/graphene-experiments/blob/master/jupyter/graphene_v2_param_estimates.ipynb
    // variable names are from graphene paper https://people.cs.umass.edu/~gbiss/graphene.sigcomm.pdf
    uint64_t z = 10;
    uint64_t n = 10;
    uint64_t m = 20;
    double f_S = 0.1;
    double beta = 0.9;
    uint64_t x_star_desired = 4;
    uint64_t y_star_desired = 5;

    uint64_t x_star_actual = CGrapheneSet::LowerBoundTruePositives(n, z, m, f_S, beta);
    uint64_t y_star_actual = CGrapheneSet::UpperBoundFalsePositives(n, z, m, f_S, beta);

    BOOST_CHECK_EQUAL(x_star_desired, x_star_actual);
    BOOST_CHECK_EQUAL(y_star_desired, y_star_actual);
}

BOOST_AUTO_TEST_SUITE_END()


#define BOOST_TEST_MODULE "C++ Unit Tests for metaLBM"
#include <boost/test/included/unit_test.hpp>
#include <boost/test/unit_test.hpp>
#include <boost/test/floating_point_comparison.hpp>
namespace tt = boost::test_tools;

#include <iostream>

#define NPROCS 1
#include "Lattice.h"
typedef double valueType;
typedef lbm::Lattice<valueType, lbm::LatticeType::D1Q3> L;
#include "Collision.h"

using namespace lbm;

BOOST_AUTO_TEST_SUITE(TestCollision)

BOOST_AUTO_TEST_CASE(TestBGK) {
  constexpr CollisionType collisionType = CollisionType::BGK;
  valueType f[100]
    Collision<valueType, collisionType> collision(0.6, f);

  BOOST_TEST(collision.getAlpha() == (valueType) 2,
             tt::tolerance(1e-15));
}

BOOST_AUTO_TEST_SUITE_END()


#include "assign_bilinear_bases_dist.hxx"
#include "set_dual_objective_b.hxx"
#include "../../SDP.hxx"

#include <boost/filesystem.hpp>

void read_blocks(const boost::filesystem::path &sdp_directory, SDP &sdp);
void read_objectives(const boost::filesystem::path &sdp_directory,
                     const El::Grid &grid, El::BigFloat &objective_const,
                     El::DistMatrix<El::BigFloat> &dual_objective_b);
void read_bilinear_bases(
  const boost::filesystem::path &sdp_directory, const Block_Info &block_info,
  const El::Grid &grid,
  std::vector<El::Matrix<El::BigFloat>> &bilinear_bases_local,
  std::vector<El::DistMatrix<El::BigFloat>> &bilinear_bases_dist);

void read_primal_objective_c(const boost::filesystem::path &sdp_directory,
                             const std::vector<size_t> &block_indices,
                             const El::Grid &grid,
                             Block_Vector &primal_objective_c);
void read_free_var_matrix(const boost::filesystem::path &sdp_directory,
                          const std::vector<size_t> &block_indices,
                          const El::Grid &grid, Block_Matrix &free_var_matrix);

SDP::SDP(const boost::filesystem::path &sdp_directory,
         const Block_Info &block_info, const El::Grid &grid)
{
  read_objectives(sdp_directory, grid, objective_const, dual_objective_b);
  read_bilinear_bases(sdp_directory, block_info, grid, bilinear_bases_local,
                      bilinear_bases_dist);
  read_primal_objective_c(sdp_directory, block_info.block_indices, grid,
                          primal_objective_c);
  read_free_var_matrix(sdp_directory, block_info.block_indices, grid,
                       free_var_matrix);
}

SDP::SDP(const El::BigFloat &objective_const_input,
         const std::vector<El::BigFloat> &dual_objective_b_input,
         const std::vector<std::vector<El::BigFloat>> &primal_objective_c_input,
         const std::vector<El::Matrix<El::BigFloat>> &free_var_input,
         const Block_Info &block_info, const El::Grid &grid)
    : yp_to_y(dual_objective_b_input.size(), dual_objective_b_input.size(),
              grid),
      objective_const(objective_const_input)
{
  set_dual_objective_b(dual_objective_b_input, grid, dual_objective_b);

  auto &block_indices(block_info.block_indices);
  bilinear_bases_local.resize(2 * block_indices.size());

  for(size_t block(0); block != block_indices.size(); ++block)
    {
      bilinear_bases_local[2 * block].Resize(1, 1);
      bilinear_bases_local[2 * block + 1].Resize(0, 1);
      bilinear_bases_local[2 * block](0, 0) = 1;
    }
  assign_bilinear_bases_dist(bilinear_bases_local, grid, bilinear_bases_dist);

  for(size_t block(0); block != block_indices.size(); ++block)
    {
      primal_objective_c.blocks.emplace_back(
        primal_objective_c_input.at(block_indices.at(block)).size(), 1, grid);
      // TODO: copied from sdp_solve/SDP/SDP/read_primal_objective_c.cxx
      auto &primal_block(primal_objective_c.blocks.back());
      const size_t local_height(primal_block.LocalHeight());
      if(primal_block.GlobalCol(0) == 0)
        {
          for(size_t row = 0; row < local_height; ++row)
            {
              const size_t global_row(primal_block.GlobalRow(row));
              primal_block.SetLocal(
                row, 0,
                primal_objective_c_input.at(block_indices.at(block))
                  .at(global_row));
            }
        }
    }

  const int64_t B_Height([&]() {
    int64_t result(0);
    for(auto &block : primal_objective_c_input)
      {
        result += block.size();
      }
    return result;
  }());
  const int64_t B_Width(dual_objective_b.Height());
  El::DistMatrix<El::BigFloat> B(B_Height, B_Width, grid);

  int64_t row_block(0);
  for(size_t block(0); block != block_indices.size(); ++block)
    {
      const int64_t block_height(
        free_var_input.at(block_indices.at(block)).Height()),
        block_width(free_var_input.at(block_indices.at(block)).Width());
      for(int64_t row_offset(0); row_offset != block_height; ++row_offset)
        {
          const size_t row(row_block + row_offset);
          for(int64_t column(0); column != block_width; ++column)
            {
              if(B.IsLocal(row, column))
                {
                  B.SetLocal(B.LocalRow(row), B.LocalCol(column),
                             free_var_input.at(block_indices.at(block))(
                               row_offset, column));
                }
            }
        }
      row_block += block_height;
    }

  El::DistMatrix<El::BigFloat> U(grid), s(grid), V(grid);
  El::SVD(B, U, s, V);

  for(int64_t row(0); row < s.LocalHeight(); ++row)
    {
      for(int64_t column(0); column < s.LocalWidth(); ++column)
        {
          s.SetLocal(row, column, 1 / s.GetLocal(row, column));
        }
    }
  El::Copy(V, yp_to_y);
  El::DiagonalScale(El::LeftOrRight::RIGHT, El::Orientation::NORMAL, s,
                    yp_to_y);

  El::DistMatrix<El::BigFloat> b_new(s.Height(), s.Width(), grid);
  El::Copy(dual_objective_b, b_new);
  El::Gemv(El::Orientation::TRANSPOSE, El::BigFloat(1.0), yp_to_y, b_new,
           El::BigFloat(0.0), dual_objective_b);

  free_var_matrix.blocks.reserve(block_indices.size());
  int64_t global_row(0);
  for(size_t block(0); block != block_indices.size(); ++block)
    {
      free_var_matrix.blocks.emplace_back(
        free_var_input.at(block_indices.at(block)).Height(),
        free_var_input.at(block_indices.at(block)).Width(), grid);
      auto &free_var_block(free_var_matrix.blocks.back());
      for(int64_t row(0); row != free_var_block.Height(); ++row)
        {
          for(int64_t column(0); column != free_var_block.Width(); ++column)
            {
              if(free_var_block.IsLocal(row, column))
                {
                  free_var_block.SetLocal(free_var_block.LocalRow(row),
                                          free_var_block.LocalCol(column),
                                          U.Get(global_row, column));
                }
            }
          ++global_row;
        }
    }
}


#include <blitzml/base/solver.h>
#include <blitzml/base/vector_util.h>
#include <blitzml/base/math_util.h>

#include <iostream>
#include <limits>
#include <algorithm>
#include <utility>
#include <string.h>
#include <cmath>

using std::vector;
using std::swap;
using std::max;
using std::min;

namespace BlitzML {

void Solver::solve(const Dataset *data, const Parameters *params,
                   value_t *result, char *status_message,
                   const char *log_directory) {
  solver_timer.reset();
  solver_logger.set_suppress_warnings(suppress_warnings());
  solver_logger.set_log_directory(log_directory);
  solver_logger.throttle_logging_with_interval(max_time() / 100);

  this->data = data;
  this->params = params;

  num_components = initial_problem_size();

  ws.set_max_size(num_components);
  ws.clear();

  initialize_priority_vectors();
  initialize(result);
  scheduler.reset();
  scheduler.set_strong_convexity_constant(strong_convexity_constant());

  obj_vals.set_dual_obj(compute_dual_obj());
  obj_vals.set_primal_obj_y(1e30);

  iteration_number = 0;
  while (++iteration_number) {
    Timer overhead_timer;
    update_alpha();
    update_y();
    update_obj_vals();
    log_info(should_force_log());
    record_subproblem_progress();
    if (use_screening()) {
      screen_components();
    }
    scheduler.record_overhead_time(overhead_timer.elapsed_time());

    setup_subproblem();

    subproblem_solve_timer.reset();
    solve_subproblem();
    subproblem_solve_timer.pause_timing();

    Solver::ConvergenceStatus convergence_status = get_convergence_status();
    if (convergence_status != NOT_CONVERGED) {
      ++iteration_number;
      log_info(true);
      set_status_message(status_message, convergence_status);
      fill_result(result);
      solver_logger.close_files();
      return;
    }
  }
}

void Solver::initialize_priority_vectors() {
  priority_values.resize(num_components);
}

void Solver::update_alpha() {
  if (!use_working_sets() && iteration_number > 1) {
    alpha = 1.;
  } else {
    alpha = compute_alpha();
  }
  debug("alpha is %0.15f", alpha);
}

void Solver::update_obj_vals() {
  duality_gap_last_iteration = obj_vals.duality_gap();
  dual_objective_last_iteration = obj_vals.dual_obj();

  value_t primal_obj_y;
  if (alpha == 1. && iteration_number > 1) {
    primal_obj_y = obj_vals.primal_obj_z();
  } else {
    primal_obj_y = compute_primal_obj_y();
  }
  obj_vals.set_primal_obj_y(primal_obj_y);
}

void Solver::record_subproblem_progress() {
  if (iteration_number == 1) {
    return;
  }

  value_t subproblem_time = subproblem_solve_timer.elapsed_time();
  value_t subproblem_duality_gap = obj_vals.primal_obj_x() - obj_vals.dual_obj();
  scheduler.record_subproblem_progress(obj_vals.duality_gap(),
                                       subproblem_duality_gap,
                                       subproblem_time,
                                       alpha);
}

value_t Solver::strong_convexity_constant() const {
  return 1;
}

void Solver::setup_subproblem() {
  if (use_working_sets()) {
    set_subproblem_parameters();
    record_subproblem_size();
  } else {
    setup_subproblem_no_working_sets();
  }
}

void Solver::setup_subproblem_no_working_sets() {
  if (iteration_number == 1) {
    ws.clear();
    for (index_t i = 0; i < num_components; ++i) {
      ws.add_index(i);
    }
  }
  subproblem_params.epsilon = 0.01;
  subproblem_params.max_iterations = 5;
  subproblem_params.min_time = 0.;
  subproblem_params.time_limit = max_time();
}

void Solver::update_priority_values() {
  int total = 0;
  for (index_t i = 0; i < num_components; ++i) {
    unsigned short capsule_ind = compute_priority_value(i);
    priority_values[i] = capsule_ind;
    if (capsule_ind < capsule_candidates.size()) {
      ++total;
      problem_size_diffs[capsule_ind] += size_of_component(i);
    }
  }
}

void Solver::set_initial_subproblem_state(SubproblemState& initial_state) const {
  initial_state.dual_obj = obj_vals.dual_obj();
  initial_state.x = x;
}

bool Solver::check_sufficient_dual_progress(
        const SubproblemState& initial_state, value_t epsilon) const {
  value_t dual_obj_progress = obj_vals.dual_obj() - initial_state.dual_obj;
  value_t norm_diff_sq = norm_diff_sq_z_initial_x(initial_state);
  value_t min_dual_progress = 0.5 * strong_convexity_constant() * (1 - epsilon) * norm_diff_sq;
  return (dual_obj_progress >= min_dual_progress);
}

bool Solver::should_force_log() {
  if (iteration_number == 1) {
    return true;
  }
  if (obj_vals.duality_gap() < 0.2 * last_log_gap) {
    return true;
  }
  return false;
}

void Solver::log_info(bool force_log) {
  solver_timer.pause_timing();

  double elapsed_time = solver_timer.elapsed_time();
  if (verbose()) {
    print("Time: %5e, dual objective: %7e, duality gap: %7e",
           elapsed_time, obj_vals.dual_obj(), obj_vals.duality_gap());
  }

  bool logged = solver_logger.log_new_point(elapsed_time, obj_vals.dual_obj(), force_log);
  if (logged) {
    solver_logger.log_value<value_t>("primal_obj_y", obj_vals.primal_obj_y());
    solver_logger.log_value<unsigned>("iteration", iteration_number);
    solver_logger.log_value<value_t>("alpha", alpha);
    solver_logger.log_value<value_t>("xi", subproblem_params.xi);
    solver_logger.log_value<value_t>("epsilon", subproblem_params.epsilon);
    solver_logger.log_value<size_t>("working_set_size", ws.size());
    solver_logger.log_value<index_t>("num_not_eliminated", num_components);
    if (log_vectors()) {
      solver_logger.log_vector<value_t>("y", y);
      solver_logger.log_vector<value_t>("x", x);
    }
    log_variables(solver_logger);
    last_log_gap = obj_vals.duality_gap();
  }

  solver_timer.continue_timing();
}

Solver::ConvergenceStatus Solver::get_convergence_status() {
  double elapsed_time = solver_timer.elapsed_time();
  if (obj_vals.duality_gap() == 0. && (elapsed_time > min_time())) {
    return REACHED_STOPPING_TOLERANCE;
  } else if (obj_vals.dual_obj() != 0.
             && obj_vals.duality_gap() / fabs(obj_vals.dual_obj()) < tolerance()
             && elapsed_time > min_time()) {
    return REACHED_STOPPING_TOLERANCE;
  }
  if (elapsed_time > max_time()) {
    return EXCEEDED_TIME_LIMIT;
  }
  if (iteration_number > max_iterations()) {
    return EXCEEDED_MAX_ITERATIONS;
  }
  if (iteration_number > 1
      && obj_vals.dual_obj() <= dual_objective_last_iteration
      && obj_vals.duality_gap() >= duality_gap_last_iteration) {
    return REACHED_MACHINE_PRECISION;
  }

  return NOT_CONVERGED;
}

void Solver::set_status_message(char* status_message,
                                Solver::ConvergenceStatus status) {
  switch (status) {
    case REACHED_STOPPING_TOLERANCE:
      strcpy(status_message, "Reached stopping tolerance");
      break;
    case EXCEEDED_TIME_LIMIT:
      strcpy(status_message, "Exceeded time limit");
      break;
    case EXCEEDED_MAX_ITERATIONS:
      strcpy(status_message, "Exceeded iterations limit");
      break;
    case REACHED_MACHINE_PRECISION:
      strcpy(status_message, "Reached machine precision");
      break;
    default:
      strcpy(status_message, "Convergence status unknown");
      break;
  }
}

value_t Solver::compute_d_sq() const {
  return l2_norm_diff_sq(x, y);
}

void Solver::set_subproblem_parameters() {
  scheduler.set_current_duality_gap(obj_vals.duality_gap());
  scheduler.set_current_dsq(compute_d_sq());

  scheduler.capsule_candidates(capsule_candidates);
  size_t num_capsule_candidates = capsule_candidates.size();

  problem_size_diffs.assign(num_capsule_candidates, 0.);
  update_priority_values();

  vector<value_t> problem_size_candidates(num_capsule_candidates, 0);
  problem_size_candidates[0] = problem_size_diffs[0];
  for (size_t ind = 1; ind < num_capsule_candidates; ++ind) {
    problem_size_candidates[ind] = problem_size_candidates[ind-1] + problem_size_diffs[ind];
  }

  double overhead_time_est = scheduler.overhead_time_estimate();
  subproblem_params.max_iterations = 1000;
  subproblem_params.min_time = 0.1 * overhead_time_est;

  value_t best_epsilon;
  double best_time_limit;
  unsigned short best_capsule_ind;
  if (iteration_number > 1) {
    scheduler.set_problem_sizes(problem_size_candidates);
    scheduler.optimize_xi_and_epsilon(best_capsule_ind, best_epsilon, best_time_limit);
  } else {
    best_epsilon = 0.5;
    best_time_limit = 2 * overhead_time_est;
    best_capsule_ind = 0;
    while ((capsule_candidates[best_capsule_ind].xi < 1.0) &&
           (problem_size_candidates[best_capsule_ind] < problem_size_candidates.back())) {
      ++best_capsule_ind;
    }
    subproblem_params.max_iterations = 1;
  }
  subproblem_params.epsilon = best_epsilon;

  double max_time_limit = overhead_time_est * 15;
  double min_time_limit = overhead_time_est * 0.5;
  subproblem_params.time_limit = max(min(best_time_limit, max_time_limit), min_time_limit);
  debug("time ratio is %f", subproblem_params.time_limit / overhead_time_est);
  debug("epsilon is %f", subproblem_params.epsilon);

  subproblem_params.xi = capsule_candidates[best_capsule_ind].xi;

  ws.clear();
  for (index_t i = 0; i < num_components; ++i) {
    if (priority_values[i] <= best_capsule_ind) {
      ws.add_index(i);
    }
  }
  debug("size working set is %d", ws.size());
}

void Solver::record_subproblem_size() {
  size_t size_subproblem = 0;
  for (const_index_itr i = ws.begin_sorted(); i != ws.end_sorted(); ++i) {
    size_subproblem += size_of_component(*i);
  }
  scheduler.record_subproblem_size(subproblem_params.xi, static_cast<value_t>(size_subproblem));
}

} // namespace BlitzML


//==================================================================================================
/**
  Copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
**/
//==================================================================================================
#ifndef BOOST_SIMD_ARCH_X86_AVX2_SIMD_FUNCTION_MULTIPLIES_HPP_INCLUDED
#define BOOST_SIMD_ARCH_X86_AVX2_SIMD_FUNCTION_MULTIPLIES_HPP_INCLUDED

#include <boost/simd/detail/overload.hpp>
#include <boost/simd/constant/constant.hpp>

namespace boost { namespace simd { namespace ext
{
   namespace bd = boost::dispatch;
   namespace bs = boost::simd;

   BOOST_DISPATCH_OVERLOAD( multiplies_
                          , (typename A0)
                          , bs::avx2_
                          , bs::pack_<bd::ints8_<A0>, bs::avx_>
                          , bs::pack_<bd::ints8_<A0>, bs::avx_>
                          )
   {
      BOOST_FORCEINLINE A0 operator()( const A0& a0, const A0& a1) const BOOST_NOEXCEPT
      {
        auto const mask = Constant<bd::upgrade_t<A0,signed>,0x00FF>();
        auto const hmul = _mm256_mullo_epi16( _mm256_srli_epi16(a0.storage(), 8)
                                            , _mm256_srli_epi16(a1.storage(), 8)
                                            );

        return _mm256_or_si256( _mm256_and_si256(mask, _mm256_mullo_epi16(a0,a1))
                              , _mm256_slli_epi16( _mm256_and_si256(mask, hmul), 8)
                              );
      }
   };

   BOOST_DISPATCH_OVERLOAD( multiplies_
                          , (typename A0)
                          , bs::avx2_
                          , bs::pack_<bd::ints16_<A0>, bs::avx_>
                          , bs::pack_<bd::ints16_<A0>, bs::avx_>
                          )
   {
      BOOST_FORCEINLINE A0 operator()( const A0& a0, const A0& a1) const BOOST_NOEXCEPT
      {
        return _mm256_mullo_epi16(a0, a1);
      }
   };

   BOOST_DISPATCH_OVERLOAD( multiplies_
                          , (typename A0)
                          , bs::avx2_
                          , bs::pack_<bd::ints32_<A0>, bs::avx_>
                          , bs::pack_<bd::ints32_<A0>, bs::avx_>
                          )
   {
      BOOST_FORCEINLINE A0 operator()( const A0& a0, const A0& a1) const BOOST_NOEXCEPT
      {
        return _mm256_mullo_epi32(a0, a1);
      }
   };
} } }

#endif


#include <Configuration.h> // Package Path
#include <avatar_locomanipulation/enable_pinocchio_with_hpp_fcl.h> // Enable HPP FCL
// Multibody
#include "pinocchio/multibody/model.hpp"
#include "pinocchio/multibody/data.hpp"
#include "pinocchio/multibody/geometry.hpp"
// Algorithms
#include "pinocchio/algorithm/kinematics.hpp"
#include "pinocchio/algorithm/geometry.hpp"
#include "pinocchio/algorithm/model.hpp"
#include "pinocchio/algorithm/jacobian.hpp"
#include "pinocchio/algorithm/frames.hpp" // Jacobian Frame Computation
// Parsers
#include "pinocchio/parsers/urdf.hpp"
#include "pinocchio/parsers/srdf.hpp"
// Spatial
#include "pinocchio/spatial/fcl-pinocchio-conversions.hpp"

// Standard
#include <math.h>
#include <map>
#include <iostream>
#include <boost/shared_ptr.hpp>

int main(int argc, char ** argv){
	// models and geomModels
	pinocchio::Model robot_model, cart_model, appended_model;
  pinocchio::GeometryModel robot_geomModel, cart_geomModel, appended_geomModel;
  // data and geomdata
  std::unique_ptr<pinocchio::Data> robot_data, cart_data, appended_data;
  std::unique_ptr<pinocchio::GeometryData> robot_geomData, cart_geomData, appended_geomData;

  // Define the Robot--------
	std::string filename = THIS_PACKAGE_PATH"models/valkyrie_simplified_collisions.urdf";
  std::string srdf_filename = THIS_PACKAGE_PATH"models/valkyrie_disable_collisions.srdf";
  std::string meshDir  = THIS_PACKAGE_PATH"../val_model/";
  // build model and geomModel
  pinocchio::urdf::buildModel(filename, pinocchio::JointModelFreeFlyer(),robot_model);
  pinocchio::urdf::buildGeom(robot_model, filename, pinocchio::COLLISION, robot_geomModel, meshDir );
  // Add and remove collision pairs
  robot_geomModel.addAllCollisionPairs();
  pinocchio::srdf::removeCollisionPairs(robot_model, robot_geomModel, srdf_filename, false);
  // Define the data and geomData
  robot_data = std::unique_ptr<pinocchio::Data>(new pinocchio::Data(robot_model));
	robot_geomData = std::unique_ptr<pinocchio::GeometryData>(new pinocchio::GeometryData(robot_geomModel));
	// Define the configuration of Robot
  Eigen::VectorXd q_robot;
  q_robot = Eigen::VectorXd::Zero(robot_model.nq);

  double theta = 0;//M_PI/4.0;
  Eigen::AngleAxis<double> aa(theta, Eigen::Vector3d(0.0, 0.0, 1.0));
  Eigen::Quaternion<double> init_quat(1.0, 0.0, 0.0, 0.0); //Initialized to remember the w component comes first
  init_quat = aa;

  q_robot[3] = init_quat.x(); q_robot[4] = init_quat.y(); q_robot[5] = init_quat.z(); q_robot[6] = init_quat.w(); // Set up the quaternion in q

  q_robot[2] = 1.0; // set z value to 1.0, this is the pelvis location

  q_robot[7 + robot_model.getJointId("leftHipPitch") - 2] = -0.3;
  q_robot[7 + robot_model.getJointId("rightHipPitch") - 2] = -0.3;
  q_robot[7 + robot_model.getJointId("leftKneePitch") - 2] = 0.6;
  q_robot[7 + robot_model.getJointId("rightKneePitch") - 2] = 0.6;
  q_robot[7 + robot_model.getJointId("leftAnklePitch") - 2] = -0.3;
  q_robot[7 + robot_model.getJointId("rightAnklePitch") - 2] = -0.3;

  q_robot[7 + robot_model.getJointId("rightShoulderPitch") - 2] = -0.2;
  q_robot[7 + robot_model.getJointId("rightShoulderRoll") - 2] = 1.1;
  q_robot[7 + robot_model.getJointId("rightElbowPitch") - 2] = 0.4;
  q_robot[7 + robot_model.getJointId("rightForearmYaw") - 2] = 1.5;

  q_robot[7 + robot_model.getJointId("leftShoulderPitch") - 2] = -0.2;
  q_robot[7 + robot_model.getJointId("leftShoulderRoll") - 2] = -1.1;
  q_robot[7 + robot_model.getJointId("leftElbowPitch") - 2] = -0.4;
  q_robot[7 + robot_model.getJointId("leftForearmYaw") - 2] = 1.5;

	// Define the Cart--------
	filename = THIS_PACKAGE_PATH"models/test_cart.urdf";
  meshDir  = THIS_PACKAGE_PATH"models/cart/";
  // build model and geomModel
  pinocchio::urdf::buildModel(filename, pinocchio::JointModelFreeFlyer(),cart_model);
  pinocchio::urdf::buildGeom(cart_model, filename, pinocchio::COLLISION, cart_geomModel, meshDir );

  //------ Suggested Fix
  // Apply a prefix
  std::string prefix ("cart/");
  for (pinocchio::JointIndex i = 1; i < cart_model.joints.size(); ++i) {
    cart_model.names[i] = prefix + cart_model.names[i];
  }
  for (pinocchio::FrameIndex i = 0; i < cart_model.frames.size(); ++i) {
    ::pinocchio::Frame& f = cart_model.frames[i];
    f.name = prefix + f.name;
  }
  //----------------------
  // Add and remove collision pairs
  cart_geomModel.addAllCollisionPairs();
  // Define the data and geomData
  cart_data = std::unique_ptr<pinocchio::Data>(new pinocchio::Data(cart_model));
	cart_geomData = std::unique_ptr<pinocchio::GeometryData>(new pinocchio::GeometryData(cart_geomModel));
	// Define the configuration of the cart
  Eigen::VectorXd q_cart;
  q_cart = Eigen::VectorXd::Zero(cart_model.nq);
  q_cart[0] = -0.05;  q_cart[1] = 0.0;  q_cart[2] = 0.0;
  // q_cart[0] = -0.06;  q_cart[1] = -0.0085;  q_cart[2] = -0.04;
  double theta1 = 0;//M_PI/4.0;
  Eigen::AngleAxis<double> bb(theta1, Eigen::Vector3d(0.0, 0.0, 1.0)); // yaw pi/4 to the left
  Eigen::Quaternion<double> quat_init; quat_init =  bb;
  q_cart[3] = quat_init.x();// 0.0;
  q_cart[4] = quat_init.y(); //0.0;
  q_cart[5] = quat_init.z(); //sin(theta/2.0);
  q_cart[6] = quat_init.w(); //cos(theta/2.0);

  // Perform initial forward kinematics
  pinocchio::forwardKinematics(robot_model, *robot_data, q_robot);
  // Compute Joint Jacobians
  pinocchio::computeJointJacobians(robot_model,*robot_data, q_robot);
  // Update Frame Placements
  pinocchio::updateFramePlacements(robot_model, *robot_data);
  // Perform initial forward kinematics
  pinocchio::forwardKinematics(cart_model, *cart_data, q_cart);
  // Compute Joint Jacobians
  pinocchio::computeJointJacobians(cart_model,*cart_data, q_cart);
  // Update Frame Placements
  pinocchio::updateFramePlacements(cart_model, *cart_data);

  // Append the object onto the robot, and fill appended RobotModel
  pinocchio::appendModel(robot_model, cart_model, robot_geomModel, cart_geomModel, 0, pinocchio::SE3::Identity(), appended_model, appended_geomModel);

  appended_data = std::unique_ptr<pinocchio::Data>(new pinocchio::Data(appended_model));
  appended_geomData = std::unique_ptr<pinocchio::GeometryData>(new pinocchio::GeometryData(appended_geomModel));

  Eigen::VectorXd q_appended;
  q_appended = Eigen::VectorXd::Zero(appended_model.nq);

  for(int i=0; i<cart_model.nq; ++i){
  	q_appended[i] = q_cart[i];
  }
  int i = cart_model.nq;
  for(int j=0; j<robot_model.nq; ++j){
  	q_appended[i] = q_robot[j];
  	++i;
  }

  // Perform initial forward kinematics
  pinocchio::forwardKinematics(appended_model, *appended_data, q_appended);
  // Compute Joint Jacobians
  pinocchio::computeJointJacobians(appended_model,*appended_data, q_appended);
  // Update Frame Placements
  pinocchio::updateFramePlacements(appended_model, *appended_data);
  // Update geometry Placements
  pinocchio::updateGeometryPlacements(appended_model, *appended_data, appended_geomModel, *appended_geomData, q_appended);

  std::cout << "appended_model: \n" << appended_model << std::endl;
  std::cout << "appended_geomModel: \n" << appended_geomModel << std::endl;
  // List Operational Space Frames
  for (int k=0 ; k<appended_model.frames.size() ; ++k){
    std::cout << "frame:" << k << " " << appended_model.frames[k].name << " : " << appended_data->oMf[k].translation().transpose() << std::endl;
  }
}


// Distributed under the MIT License.
// See LICENSE.txt for details.

#pragma once

#include <array>
#include <boost/rational.hpp>
#include <cstddef>
#include <memory>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#include "DataStructures/Tensor/TypeAliases.hpp"

/// \cond
template <size_t VolumeDim, typename TargetFrame>
class Block;
template <size_t VolumeDim>
class BlockNeighbor;
template <typename SourceFrame, typename TargetFrame, size_t Dim>
class CoordinateMapBase;
class DataVector;
template <size_t VolumeDim>
class Direction;
template <size_t VolumeDim, typename TargetFrame>
class Domain;
template <size_t VolumeDim>
class ElementId;
/// \endcond

// Test that the Blocks in the Domain are constructed correctly.
template <size_t VolumeDim, typename Fr = Frame::Inertial>
void test_domain_construction(
    const Domain<VolumeDim, Fr>& domain,
    const std::vector<
        std::unordered_map<Direction<VolumeDim>, BlockNeighbor<VolumeDim>>>&
        expected_block_neighbors,
    const std::vector<std::unordered_set<Direction<VolumeDim>>>&
        expected_external_boundaries,
    const std::vector<
        std::unique_ptr<CoordinateMapBase<Frame::Logical, Fr, VolumeDim>>>&
        expected_maps) noexcept;

// Test that two neighboring Blocks abut each other.
template <size_t VolumeDim, typename Fr = Frame::Inertial>
void test_physical_separation(
    const std::vector<Block<VolumeDim, Fr>>& blocks) noexcept;

// Fraction of the logical volume of a block covered by an element
// The sum of this over all the elements of a block should be one
template <size_t VolumeDim>
boost::rational<size_t> fraction_of_block_volume(
    const ElementId<VolumeDim>& element_id) noexcept;

// Test that the Elements of the initial domain are connected and cover the
// computational domain, as well as that neighboring Elements  are at the same
// refinement level.
template <size_t VolumeDim, typename Fr = Frame::Inertial>
void test_initial_domain(const Domain<VolumeDim, Fr>& domain,
                         const std::vector<std::array<size_t, VolumeDim>>&
                             initial_refinement_levels) noexcept;

// Euclidean basis vector along the given `Direction` and in the given `Frame`.
template <size_t SpatialDim, typename Fr = Frame::Inertial>
tnsr::i<DataVector, SpatialDim, Fr> euclidean_basis_vector(
    const Direction<SpatialDim>& direction,
    const DataVector& used_for_size) noexcept;


#include <boost/geometry/strategies/cartesian/side_of_intersection.hpp>


// Copyright 2005 The Trustees of Indiana University.

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

//  Authors: Douglas Gregor
//           Andrew Lumsdaine
#include "graph_types.hpp"

namespace boost {
  // Work around an annoying qualified "boost::get" call.
  // using boost::graph::python::get;
}

#include <boost/graph/astar_search.hpp>
#include <boost/graph/iteration_macros.hpp>

namespace boost { namespace graph { namespace python {

#define BGL_PYTHON_VISITOR maybe_astar_visitor
#define BGL_PYTHON_EVENTS_HEADER <boost/graph/python/astar_events.hpp>
#include <boost/graph/python/visitor.hpp>
#undef BGL_PYTHON_EVENTS_HEADER
#undef BGL_PYTHON_VISITOR

using boost::python::object;

// Function object adaptor for the A* heuristic
template<typename Graph>
struct astar_heuristic_adaptor
{
  typedef typename graph_traits<Graph>::vertex_descriptor Vertex;

  astar_heuristic_adaptor(object h) : h(h) { }

  float operator()(Vertex v)
  {
    return boost::python::extract<float>(h(v));
  }

 private:
  object h;
};

template<typename Graph>
void
astar_search
  (const Graph& g,
   typename graph_traits<Graph>::vertex_descriptor s,
   object in_heuristic,
   object in_visitor,
   vector_property_map<
      typename graph_traits<Graph>::vertex_descriptor,
      typename property_map<Graph, vertex_index_t>::const_type>* in_predecessor,
   vector_property_map<
      float,
      typename property_map<Graph, vertex_index_t>::const_type>* in_cost,
   vector_property_map<
      float,
      typename property_map<Graph, vertex_index_t>::const_type>* in_distance,
   vector_property_map<
      float,
      typename property_map<Graph, edge_index_t>::const_type>* in_weight,
   vector_property_map<
      default_color_type,
      typename property_map<Graph, vertex_index_t>::const_type>* in_color)
{
  typedef typename graph_traits<Graph>::vertex_descriptor vertex_descriptor;
  typedef typename property_map<Graph, vertex_index_t>::const_type
    VertexIndexMap;
  typedef typename property_map<Graph, edge_index_t>::const_type
    EdgeIndexMap;
  typedef vector_property_map<vertex_descriptor, VertexIndexMap> PredecessorMap;
  typedef vector_property_map<float, VertexIndexMap> CostMap;
  typedef vector_property_map<float, VertexIndexMap> DistanceMap;
  typedef vector_property_map<float, EdgeIndexMap> WeightMap;
  typedef vector_property_map<default_color_type, VertexIndexMap> ColorMap;

  PredecessorMap predecessor(in_predecessor? *in_predecessor
                             /*default*/: PredecessorMap(num_vertices(g),
                                                         get(vertex_index, g)));
  CostMap cost(in_cost? *in_cost
               /*default*/: CostMap(num_vertices(g), get(vertex_index, g)));
  DistanceMap distance(in_distance? *in_distance
                       /*default*/: DistanceMap(num_vertices(g),
                                                get(vertex_index, g)));
  WeightMap weight(in_weight? *in_weight
                   /*default*/: WeightMap(num_edges(g), get(edge_index, g)));
  ColorMap color(in_color? *in_color
                 /*default*/: ColorMap(num_vertices(g), get(vertex_index, g)));

  if (!in_weight)
    BGL_FORALL_EDGES_T(e, g, Graph)
      put(weight, e, 1.0);

  if (in_visitor != object()) {
    typedef queue<vertex_descriptor> Queue;
    maybe_astar_visitor vis(in_visitor);
    boost::astar_search(g, s, astar_heuristic_adaptor<Graph>(in_heuristic),
                        predecessor_map(predecessor).
                        rank_map(cost).
                        distance_map(distance).
                        weight_map(weight).
                        color_map(color).
                        visitor(vis));
  } else {
    boost::astar_search(g, s, astar_heuristic_adaptor<Graph>(in_heuristic),
                        predecessor_map(predecessor).
                        rank_map(cost).
                        distance_map(distance).
                        weight_map(weight).
                        color_map(color));
  }
}

void export_astar_search()
{
  using boost::python::arg;
  using boost::python::def;
  using boost::python::object;

#define UNDIRECTED_GRAPH(Name,Type)                                     \
  {                                                                     \
    typedef property_map< Type , vertex_index_t>::const_type            \
      VertexIndexMap;                                                   \
    typedef property_map< Type , edge_index_t>::const_type              \
      EdgeIndexMap;                                                     \
    typedef graph_traits< Type >::vertex_descriptor vertex_descriptor;  \
    typedef vector_property_map<vertex_descriptor, VertexIndexMap>      \
      PredecessorMap;                                                   \
    typedef vector_property_map<float, VertexIndexMap> CostMap;         \
    typedef vector_property_map<float, VertexIndexMap> DistanceMap;     \
    typedef vector_property_map<float, EdgeIndexMap> WeightMap;         \
    typedef vector_property_map<default_color_type, VertexIndexMap>     \
      VertexColorMap;                                                   \
                                                                        \
    def("astar_search",                                                 \
        &boost::graph::python::astar_search< Type >,                    \
        (arg("graph"),                                                  \
         arg("root_vertex"),                                            \
         arg("heuristic"),                                              \
         arg("visitor"),                                     \
         arg("predecessor_map") = static_cast<PredecessorMap*>(0),      \
         arg("cost_map") = static_cast<CostMap*>(0),                    \
         arg("distance_map") = static_cast<DistanceMap*>(0),            \
         arg("weight_map") = static_cast<WeightMap*>(0),                \
         arg("color_map") = static_cast<VertexColorMap*>(0)));          \
  }
#include "graphs.hpp"
}

} } } // end namespace boost::graph::python


/*******************************************************************************
 *
 * Data structures for the symbolic manipulation of linear constraints.
 *
 * Author: Arnaud J. Venet (arnaud.j.venet@nasa.gov)
 * Contributor: Jorge A. Navas (jorge.navas@sri.com)
 *
 * Notices:
 *
 * Copyright (c) 2011 United States Government as represented by the
 * Administrator of the National Aeronautics and Space Administration.
 * All Rights Reserved.
 *
 * Disclaimers:
 *
 * No Warranty: THE SUBJECT SOFTWARE IS PROVIDED "AS IS" WITHOUT ANY WARRANTY OF
 * ANY KIND, EITHER EXPRESSED, IMPLIED, OR STATUTORY, INCLUDING, BUT NOT LIMITED
 * TO, ANY WARRANTY THAT THE SUBJECT SOFTWARE WILL CONFORM TO SPECIFICATIONS,
 * ANY IMPLIED WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE,
 * OR FREEDOM FROM INFRINGEMENT, ANY WARRANTY THAT THE SUBJECT SOFTWARE WILL BE
 * ERROR FREE, OR ANY WARRANTY THAT DOCUMENTATION, IF PROVIDED, WILL CONFORM TO
 * THE SUBJECT SOFTWARE. THIS AGREEMENT DOES NOT, IN ANY MANNER, CONSTITUTE AN
 * ENDORSEMENT BY GOVERNMENT AGENCY OR ANY PRIOR RECIPIENT OF ANY RESULTS,
 * RESULTING DESIGNS, HARDWARE, SOFTWARE PRODUCTS OR ANY OTHER APPLICATIONS
 * RESULTING FROM USE OF THE SUBJECT SOFTWARE.  FURTHER, GOVERNMENT AGENCY
 * DISCLAIMS ALL WARRANTIES AND LIABILITIES REGARDING THIRD-PARTY SOFTWARE,
 * IF PRESENT IN THE ORIGINAL SOFTWARE, AND DISTRIBUTES IT "AS IS."
 *
 * Waiver and Indemnity:  RECIPIENT AGREES TO WAIVE ANY AND ALL CLAIMS AGAINST
 * THE UNITED STATES GOVERNMENT, ITS CONTRACTORS AND SUBCONTRACTORS, AS WELL
 * AS ANY PRIOR RECIPIENT.  IF RECIPIENT'S USE OF THE SUBJECT SOFTWARE RESULTS
 * IN ANY LIABILITIES, DEMANDS, DAMAGES, EXPENSES OR LOSSES ARISING FROM SUCH
 * USE, INCLUDING ANY DAMAGES FROM PRODUCTS BASED ON, OR RESULTING FROM,
 * RECIPIENT'S USE OF THE SUBJECT SOFTWARE, RECIPIENT SHALL INDEMNIFY AND HOLD
 * HARMLESS THE UNITED STATES GOVERNMENT, ITS CONTRACTORS AND SUBCONTRACTORS,
 * AS WELL AS ANY PRIOR RECIPIENT, TO THE EXTENT PERMITTED BY LAW.
 * RECIPIENT'S SOLE REMEDY FOR ANY SUCH MATTER SHALL BE THE IMMEDIATE,
 * UNILATERAL TERMINATION OF THIS AGREEMENT.
 *
 ******************************************************************************/

#pragma once

#include <memory>
#include <optional>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#include <boost/container/flat_map.hpp>
#include <boost/functional/hash.hpp>

#include "crab/patricia_trees.hpp"
#include "crab/types.hpp"

namespace crab {

class linear_expression_t final {

  public:
    using component_t = std::pair<number_t, variable_t>;
    using variable_set_t = patricia_tree_set<variable_t>;

  private:
    using map_t = boost::container::flat_map<variable_t, number_t>;
    using map_ptr = std::shared_ptr<map_t>;
    using pair_t = typename map_t::value_type;

    map_ptr _map;
    number_t _cst;

    linear_expression_t(map_ptr map, number_t cst) : _map(std::move(map)), _cst(std::move(cst)) {}

    linear_expression_t(const map_t& map, number_t cst) : _map(std::make_shared<map_t>()), _cst(std::move(cst)) {
        *this->_map = map;
    }

    void add(variable_t x, const number_t& n) {
        typename map_t::iterator it = this->_map->find(x);
        if (it != this->_map->end()) {
            number_t r = it->second + n;
            if (r == 0) {
                this->_map->erase(it);
            } else {
                it->second = r;
            }
        } else {
            if (n != 0) {
                this->_map->insert(pair_t(x, n));
            }
        }
    }

  public:
    using iterator = typename map_t::iterator;
    using const_iterator = typename map_t::const_iterator;

    linear_expression_t() : _map(std::make_shared<map_t>()), _cst(0) {}

    linear_expression_t(linear_expression_t&& other) = default;
    linear_expression_t(const linear_expression_t& other) = default;

    explicit linear_expression_t(number_t n) : _map(std::make_shared<map_t>()), _cst(std::move(n)) {}

    linear_expression_t(signed long long int n) : _map(std::make_shared<map_t>()), _cst(number_t(n)) {}

    linear_expression_t(variable_t x) : _map(std::make_shared<map_t>()), _cst(0) {
        this->_map->insert(pair_t(x, number_t(1)));
    }

    linear_expression_t(const number_t& n, variable_t x) : _map(std::make_shared<map_t>()), _cst(0) {
        this->_map->insert(pair_t(x, n));
    }

    linear_expression_t& operator=(const linear_expression_t& e) {
        if (this != &e) {
            this->_map = e._map;
            this->_cst = e._cst;
        }
        return *this;
    }

    const_iterator begin() const { return this->_map->begin(); }

    const_iterator end() const { return this->_map->end(); }

    iterator begin() { return this->_map->begin(); }

    iterator end() { return this->_map->end(); }

    size_t hash() const {
        size_t res = 0;
        for (const auto& p : *this) {
            boost::hash_combine(res, p);
        }
        boost::hash_combine(res, _cst);
        return res;
    }

    // syntactic equality
    bool equal(const linear_expression_t& o) const {
        if (is_constant()) {
            if (!o.is_constant()) {
                return false;
            } else {
                return (constant() == o.constant());
            }
        } else {
            if (constant() != o.constant()) {
                return false;
            }

            if (size() != o.size()) {
                return false;
            } else {
                for (const_iterator it = begin(), jt = o.begin(), et = end(); it != et; ++it, ++jt) {
                    if (((*it).first != (*jt).first) || ((*it).second != (*jt).second)) {
                        return false;
                    }
                }
                return true;
            }
        }
    }

    bool is_constant() const { return (this->_map->empty()); }

    number_t constant() const { return this->_cst; }

    std::size_t size() const { return this->_map->size(); }

    number_t operator[](variable_t x) const {
        typename map_t::const_iterator it = this->_map->find(x);
        if (it != this->_map->end()) {
            return it->second;
        } else {
            return 0;
        }
    }

    template <typename RenamingMap>
    linear_expression_t rename(const RenamingMap& map) const {
        number_t cst(this->_cst);
        linear_expression_t new_exp(cst);
        for (auto v : this->variables()) {
            auto const it = map.find(v);
            if (it != map.end()) {
                variable_t v_out((*it).second);
                new_exp = new_exp + linear_expression_t(this->operator[](v), v_out);
            } else {
                new_exp = new_exp + linear_expression_t(this->operator[](v), v);
            }
        }
        return new_exp;
    }

    linear_expression_t operator+(number_t n) const {
        linear_expression_t r(this->_map, this->_cst + std::move(n));
        return r;
    }

    linear_expression_t operator+(int n) const { return this->operator+(number_t(n)); }

    linear_expression_t operator+(variable_t x) const {
        linear_expression_t r(*this->_map, this->_cst);
        r.add(x, number_t(1));
        return r;
    }

    linear_expression_t operator+(const linear_expression_t& e) const {
        linear_expression_t r(*this->_map, this->_cst + e._cst);
        for (typename map_t::const_iterator it = e._map->begin(); it != e._map->end(); ++it) {
            r.add(it->first, it->second);
        }
        return r;
    }

    linear_expression_t operator-(const number_t& n) const { return this->operator+(-n); }

    linear_expression_t operator-(int n) const { return this->operator+(-number_t(n)); }

    linear_expression_t operator-(variable_t x) const {
        linear_expression_t r(*this->_map, this->_cst);
        r.add(x, number_t(-1));
        return r;
    }

    linear_expression_t operator-() const { return this->operator*(number_t(-1)); }

    linear_expression_t operator-(const linear_expression_t& e) const {
        linear_expression_t r(*this->_map, this->_cst - e._cst);
        for (typename map_t::const_iterator it = e._map->begin(); it != e._map->end(); ++it) {
            r.add(it->first, -it->second);
        }
        return r;
    }

    linear_expression_t operator*(const number_t& n) const {
        if (n == 0) {
            return linear_expression_t();
        } else {
            map_ptr map = std::make_shared<map_t>();
            for (typename map_t::const_iterator it = this->_map->begin(); it != this->_map->end(); ++it) {
                number_t c = n * it->second;
                if (c != 0) {
                    map->insert(pair_t(it->first, c));
                }
            }
            return linear_expression_t(map, n * this->_cst);
        }
    }

    linear_expression_t operator*(int n) const { return operator*(number_t(n)); }

    variable_set_t variables() const {
        variable_set_t variables;
        for (const auto& v_c : *this) {
            variables += v_c.first;
        }
        return variables;
    }

    void write(std::ostream& o) const {
        bool start = true;
        for (auto [v, n] : *this) {
            if (n > 0 && !start) {
                o << "+";
            }
            if (n == -1) {
                o << "-";
            } else if (n != 1) {
                o << n << "*";
            }
            o << v;
            start = false;
        }
        if (this->_cst > 0 && !this->_map->empty()) {
            o << "+";
        }
        if (this->_cst != 0 || this->_map->empty()) {
            o << this->_cst;
        }
    }

    // for dgb
    void dump() { write(std::cout); }

}; // class linear_expression_t

inline std::ostream& operator<<(std::ostream& o, const linear_expression_t& e) {
    e.write(o);
    return o;
}

inline std::size_t hash_value(const linear_expression_t& e) { return e.hash(); }

struct linear_expression_hasher_t {
    size_t operator()(const linear_expression_t& e) const { return e.hash(); }
};

struct linear_expression_equal_t {
    bool operator()(const linear_expression_t& e1, const linear_expression_t& e2) const { return e1.equal(e2); }
};

using linear_expression_unordered_set =
    std::unordered_set<linear_expression_t, linear_expression_hasher_t, linear_expression_equal_t>;

template <typename Value>
using linear_expression_unordered_map =
    std::unordered_map<linear_expression_t, Value, linear_expression_hasher_t, linear_expression_equal_t>;

class linear_constraint_t final {

  public:
    using variable_set_t = patricia_tree_set<variable_t>;
    using constraint_kind_t = enum { EQUALITY, DISEQUATION, INEQUALITY, STRICT_INEQUALITY };
    using iterator = typename linear_expression_t::iterator;
    using const_iterator = typename linear_expression_t::const_iterator;

  private:
    constraint_kind_t _kind;
    linear_expression_t _expr;
    // This flag has meaning only if _kind == INEQUALITY or STRICT_INEQUALITY.
    // If true the inequality is signed otherwise unsigned.
    // By default all constraints are signed.
    bool _signedness;

  public:
    linear_constraint_t() : _kind(EQUALITY), _signedness(true) {}
    linear_constraint_t(linear_constraint_t&& other) = default;
    linear_constraint_t(const linear_constraint_t& other) = default;

    linear_constraint_t(linear_expression_t expr, constraint_kind_t kind)
        : _kind(kind), _expr(std::move(expr)), _signedness(true) {}

    linear_constraint_t(linear_expression_t expr, constraint_kind_t kind, bool signedness)
        : _kind(kind), _expr(std::move(expr)), _signedness(signedness) {
        if (_kind != INEQUALITY && _kind != STRICT_INEQUALITY) {
            CRAB_ERROR("Only inequalities can have signedness information");
        }
    }

    static linear_constraint_t get_true() {
        linear_constraint_t res(linear_expression_t(number_t(0)), EQUALITY);
        return res;
    }

    static linear_constraint_t get_false() {
        linear_constraint_t res(linear_expression_t(number_t(0)), DISEQUATION);
        return res;
    }

    bool is_tautology() const {
        switch (this->_kind) {
        case DISEQUATION: return (this->_expr.is_constant() && this->_expr.constant() != 0);
        case EQUALITY: return (this->_expr.is_constant() && this->_expr.constant() == 0);
        case INEQUALITY: return (this->_expr.is_constant() && this->_expr.constant() <= 0);
        case STRICT_INEQUALITY: return (this->_expr.is_constant() && this->_expr.constant() < 0);
        default: CRAB_ERROR("Unreachable");
        }
    }

    bool is_contradiction() const {
        switch (this->_kind) {
        case DISEQUATION: return (this->_expr.is_constant() && this->_expr.constant() == 0);
        case EQUALITY: return (this->_expr.is_constant() && this->_expr.constant() != 0);
        case INEQUALITY: return (this->_expr.is_constant() && this->_expr.constant() > 0);
        case STRICT_INEQUALITY: return (this->_expr.is_constant() && this->_expr.constant() >= 0);
        default: CRAB_ERROR("Unreachable");
        }
    }

    bool is_inequality() const { return (this->_kind == INEQUALITY); }

    bool is_strict_inequality() const { return (this->_kind == STRICT_INEQUALITY); }

    bool is_equality() const { return (this->_kind == EQUALITY); }

    bool is_disequation() const { return (this->_kind == DISEQUATION); }

    const linear_expression_t& expression() const { return this->_expr; }

    constraint_kind_t kind() const { return this->_kind; }

    bool is_signed() const {
        if (_kind != INEQUALITY && _kind != STRICT_INEQUALITY) {
            CRAB_WARN("Only inequalities have signedness");
        }
        return _signedness;
    }

    bool is_unsigned() const { return (!is_signed()); }

    const_iterator begin() const { return this->_expr.begin(); }

    const_iterator end() const { return this->_expr.end(); }

    iterator begin() { return this->_expr.begin(); }

    iterator end() { return this->_expr.end(); }

    std::size_t size() const { return this->_expr.size(); }

    // syntactic equality
    bool equal(const linear_constraint_t& o) const {
        return (_kind == o._kind && _signedness == o._signedness && _expr.equal(o._expr));
    }

    size_t hash() const {
        size_t res = 0;
        boost::hash_combine(res, _expr);
        boost::hash_combine(res, _kind);
        if (_kind == INEQUALITY || _kind == STRICT_INEQUALITY) {
            boost::hash_combine(res, _signedness);
        }
        return res;
    }

    index_t index() const {
        // XXX: to store linear constraints in patricia trees
        return (index_t)hash();
    }

    number_t operator[](variable_t x) const { return this->_expr.operator[](x); }

    linear_constraint_t negate() const {

        if (is_tautology()) {
            return get_false();
        } else if (is_contradiction()) {
            return get_true();
        } else {
            switch (kind()) {
            case INEQUALITY: {
                // try to take advantage if we use z_number_t.
                // negate(e <= 0) = e >= 1
                return linear_constraint_t(-(expression() - 1), INEQUALITY, is_signed());
            }
            case STRICT_INEQUALITY: {
                // negate(x + y < 0)  <-->  x + y >= 0 <--> -x -y <= 0
                linear_expression_t e = -this->_expr;
                return linear_constraint_t(e, INEQUALITY, is_signed());
            }
            case EQUALITY: return linear_constraint_t(this->_expr, DISEQUATION);
            case DISEQUATION: return linear_constraint_t(this->_expr, EQUALITY);
            default: CRAB_ERROR("Cannot negate linear constraint");
            }
        }
    }

    template <typename RenamingMap>
    linear_constraint_t rename(const RenamingMap& map) const {
        linear_expression_t e = this->_expr.rename(map);
        return linear_constraint_t(e, this->_kind, is_signed());
    }

    void write(std::ostream& o) const {
        if (this->is_contradiction()) {
            o << "false";
        } else if (this->is_tautology()) {
            o << "true";
        } else {
            linear_expression_t e = this->_expr - this->_expr.constant();
            o << e;
            switch (this->_kind) {
            case INEQUALITY: {
                if (is_signed()) {
                    o << " <= ";
                } else {
                    o << " <=_u ";
                }
                break;
            }
            case STRICT_INEQUALITY: {
                if (is_signed()) {
                    o << " < ";
                } else {
                    o << " <_u ";
                }
                break;
            }
            case EQUALITY: {
                o << " = ";
                break;
            }
            case DISEQUATION: {
                o << " != ";
                break;
            }
            }
            number_t c = -this->_expr.constant();
            o << c;
        }
    }

    // for dgb
    void dump() { write(std::cout); }

}; // class linear_constraint_t

inline std::ostream& operator<<(std::ostream& o, const linear_constraint_t& c) {
    c.write(o);
    return o;
}

inline std::size_t hash_value(const linear_constraint_t& e) { return e.hash(); }

inline linear_expression_t var_sub(variable_t x, const number_t& n) { return linear_expression_t(x).operator-(n); }
inline linear_expression_t var_sub(variable_t x, variable_t y) { return linear_expression_t(x).operator-(y); }
inline linear_expression_t var_add(variable_t x, number_t n) { return linear_expression_t(x).operator+(std::move(n)); }
inline linear_expression_t var_add(variable_t x, variable_t y) { return linear_expression_t(x).operator+(y); }
inline linear_expression_t var_mul(const number_t& n, variable_t x) { return linear_expression_t(n, x); }
} // namespace crab


/********************************************************************************
Copyright (c) 2016, TRACLabs, Inc.
All rights reserved.

Redistribution and use in source and binary forms, with or without modification,
 are permitted provided that the following conditions are met:

    1. Redistributions of source code must retain the above copyright notice,
       this list of conditions and the following disclaimer.

    2. Redistributions in binary form must reproduce the above copyright notice,
       this list of conditions and the following disclaimer in the documentation
       and/or other materials provided with the distribution.

    3. Neither the name of the copyright holder nor the names of its contributors
       may be used to endorse or promote products derived from this software
       without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
OF THE POSSIBILITY OF SUCH DAMAGE.
********************************************************************************/

#include <boost/date_time.hpp>
#include <trac_ik/trac_ik.hpp>
#include <ros/ros.h>
#include <kdl/chainiksolverpos_nr_jl.hpp>

double fRand(double min, double max)
{
  double f = (double)rand() / RAND_MAX;
  return min + f * (max - min);
}

void test(ros::NodeHandle& nh, double num_samples, std::string chain_start, std::string chain_end, double timeout, std::string urdf_param)
{

  double eps = 1e-5;

  // This constructor parses the URDF loaded in rosparm urdf_param into the
  // needed KDL structures.  We then pull these out to compare against the KDL
  // IK solver.
  TRAC_IK::TRAC_IK tracik_solver(chain_start, chain_end, urdf_param, timeout, eps);

  KDL::Chain chain;
  KDL::JntArray ll, ul; //lower joint limits, upper joint limits

  bool valid = tracik_solver.getKDLChain(chain);

  if (!valid) {
    ROS_ERROR("There was no valid KDL chain found");
    return;
  }

  valid = tracik_solver.getKDLLimits(ll,ul);

  if (!valid) {
    ROS_ERROR("There were no valid KDL joint limits found");
    return;
  }

  assert(chain.getNrOfJoints() == ll.data.size());
  assert(chain.getNrOfJoints() == ul.data.size());

  ROS_INFO ("Using %d joints",chain.getNrOfJoints());

  // Set up KDL IK
  KDL::ChainFkSolverPos_recursive fk_solver(chain); // Forward kin. solver
  KDL::ChainIkSolverVel_pinv vik_solver(chain); // PseudoInverse vel solver
  KDL::ChainIkSolverPos_NR_JL kdl_solver(chain,ll,ul,fk_solver, vik_solver, 1, eps); // Joint Limit Solver
  // 1 iteration per solve (will wrap in timed loop to compare with TRAC-IK)

  // Create Nominal chain configuration midway between all joint limits
  KDL::JntArray nominal(chain.getNrOfJoints());

  for (uint j=0; j<nominal.data.size(); j++) {
    nominal(j) = (ll(j)+ul(j))/2.0;
  }

  // Create desired number of valid, random joint configurations
  std::vector<KDL::JntArray> JointList;
  KDL::JntArray q(chain.getNrOfJoints());

  for (uint i=0; i < num_samples; i++) {
    for (uint j=0; j<ll.data.size(); j++) {
      q(j)=fRand(ll(j), ul(j));
    }
    JointList.push_back(q);
  }

  boost::posix_time::ptime start_time;
  boost::posix_time::time_duration diff;

  KDL::JntArray result;
  KDL::Frame end_effector_pose;
  int rc;

  double total_time=0;
  uint success=0;

  ROS_INFO_STREAM("*** Testing KDL with "<<num_samples<<" random samples");

  for (uint i=0; i < num_samples; i++) {
    fk_solver.JntToCart(JointList[i],end_effector_pose);
    double elapsed = 0;
    result=nominal; // start with nominal
    start_time = boost::posix_time::microsec_clock::local_time();
    do {
      q=result; // when iterating start with last solution
      rc=kdl_solver.CartToJnt(q,end_effector_pose,result);
      diff = boost::posix_time::microsec_clock::local_time() - start_time;
      elapsed = diff.total_nanoseconds() / 1e9;
    } while (rc < 0 && elapsed < timeout);
    total_time+=elapsed;
    if (rc>=0)
      success++;

    if (int((double)i/num_samples*100)%10 == 0)
      ROS_INFO_STREAM_THROTTLE(1,int((i)/num_samples*100)<<"\% done");
  }

  ROS_INFO_STREAM("KDL found "<<success<<" solutions ("<<100.0*success/num_samples<<"\%) with an average of "<<total_time/num_samples<<" secs per sample");

  total_time=0;
  success=0;

  ROS_INFO_STREAM("*** Testing TRAC-IK with "<<num_samples<<" random samples");

  for (uint i=0; i < num_samples; i++) {
    fk_solver.JntToCart(JointList[i],end_effector_pose);
    double elapsed = 0;
    start_time = boost::posix_time::microsec_clock::local_time();
    rc=tracik_solver.CartToJnt(nominal,end_effector_pose,result);
    diff = boost::posix_time::microsec_clock::local_time() - start_time;
    elapsed = diff.total_nanoseconds() / 1e9;
    total_time+=elapsed;
    if (rc>=0)
      success++;

    if (int((double)i/num_samples*100)%10 == 0)
      ROS_INFO_STREAM_THROTTLE(1,int((i)/num_samples*100)<<"\% done");
  }

  ROS_INFO_STREAM("TRAC-IK found "<<success<<" solutions ("<<100.0*success/num_samples<<"\%) with an average of "<<total_time/num_samples<<" secs per sample");

}

int main(int argc, char** argv)
{
  srand(1);
  ros::init(argc, argv, "ik_tests");
  ros::NodeHandle nh("~");

  int num_samples;
  std::string chain_start, chain_end, urdf_param;
  double timeout;

  nh.param("num_samples", num_samples, 1000);
  nh.param("chain_start", chain_start, std::string(""));
  nh.param("chain_end", chain_end, std::string(""));

  if (chain_start=="" || chain_end=="") {
    ROS_FATAL("Missing chain info in launch file");
    exit (-1);
  }

  nh.param("timeout", timeout, 0.005);
  nh.param("urdf_param", urdf_param, std::string("/robot_description"));

  if (num_samples < 1)
    num_samples = 1;

  test(nh, num_samples, chain_start, chain_end, timeout, urdf_param);

  // Useful when you make a script that loops over multiple launch files that test different robot chains
  // std::vector<char *> commandVector;
  // commandVector.push_back((char*)"killall");
  // commandVector.push_back((char*)"-9");
  // commandVector.push_back((char*)"roslaunch");
  // commandVector.push_back(NULL);

  // char **command = &commandVector[0];
  // execvp(command[0],command);

  return 0;
}


#include <iostream>
#include <vector>

#include <boost/lexical_cast.hpp>

#include "util.hpp"
#include "test_bp_vector_common.hpp"

#include "bp_vector.hpp"
#include "mapper.hpp"

#include "perftest_common.hpp"

// this generic trait enables easy comparisons with other BP
// implementations

struct succinct_bp_vector_traits
{
    typedef succinct::bit_vector_builder builder_type;
    typedef succinct::bp_vector bp_vector_type;

    static inline void build(builder_type& builder, bp_vector_type& bp)
    {
	bp_vector_type(&builder, true, false).swap(bp);
    }

    static inline std::string log_header()
    {
	return std::string("SUCCINCT");
    }

    static inline double bits_per_bp(bp_vector_type& vec)
    {
	return double(succinct::mapper::size_of(vec))
            * 8.0 / double(vec.size());
    }

};

template <typename BpVector>
double time_visit(BpVector const& bp, size_t sample_size = 1000000)
{
    std::vector<char> random_bits;
    for (size_t i = 0; i < sample_size; ++i) {
        random_bits.push_back(rand() > (RAND_MAX / 2));
    }

    volatile size_t foo = 0; // to prevent the compiler to optimize away the loop

    size_t find_close_performed = 0;
    size_t steps_done = 0;
    double elapsed;
    SUCCINCT_TIMEIT(elapsed) {
        while (steps_done < sample_size) {
            size_t cur_node = 1; // root

            while (bp[cur_node] && steps_done < sample_size) {
                if (random_bits[steps_done++]) {
                    size_t next_node = bp.find_close(cur_node);
		    cur_node = next_node + 1;
                    find_close_performed += 1;
                } else {
                    cur_node += 1;
                }
            }
            foo = cur_node;
        }
    }

    (void)foo; // silence warning
    return elapsed / double(find_close_performed);
}

template <typename BpVectorTraits>
void build_random_binary_tree(typename BpVectorTraits::bp_vector_type& bp, size_t size)
{
    typename BpVectorTraits::builder_type builder;
    succinct::random_binary_tree(builder, size);
    BpVectorTraits::build(builder, bp);
}

template <typename BpVectorTraits>
void bp_benchmark(size_t runs)
{
    srand(42); // make everything deterministic
    static const size_t sample_size = 10000000;

    std::cout << BpVectorTraits::log_header() << std::endl;
    std::cout << "log_height" "\t" "find_close_us" "\t" "bits_per_bp" << std::endl;

    for (size_t ln = 10; ln <= 28; ln += 2) {
	size_t n = 1 << ln;
	double elapsed = 0;
	double bits_per_bp = 0;
	for (size_t run = 0; run < runs; ++run) {
	    typename BpVectorTraits::bp_vector_type bp;
	    build_random_binary_tree<BpVectorTraits>(bp, n);
	    elapsed += time_visit(bp, sample_size);
	    bits_per_bp += BpVectorTraits::bits_per_bp(bp);
	}
	std::cout << ln
                  << "\t" << elapsed / double(runs)
                  << "\t" << bits_per_bp / double(runs)
                  << std::endl;
    }
}

int main(int argc, char** argv)
{
    size_t runs;

    if (argc == 2) {
        runs = boost::lexical_cast<size_t>(argv[1]);
    } else {
        runs = 1;
    }

    bp_benchmark<succinct_bp_vector_traits>(runs);
}


#include "graph_cut.h"

#include <unordered_map>

#include <boost/graph/stoer_wagner_min_cut.hpp>
#include <boost/property_map/property_map.hpp>
#include <boost/typeof/typeof.hpp>

extern "C" {
    #include "graclus/metisLib/metis.h"
}

#include <glog/logging.h>

namespace GraphSfM {
namespace graph {

// Wrapper class for weighted, undirected Graclus graph.
class GraclusGraph {
public:
    GraclusGraph(const std::vector<std::pair<int, int>>& edges,
                 const std::vector<int>& weights)
    {
        CHECK_EQ(edges.size(), weights.size());

        std::unordered_map<int, std::vector<std::pair<int, int>>> adjacency_list;
        for (size_t i = 0; i < edges.size(); ++i) {
            const auto& edge = edges[i];
            const auto weight = weights[i];
            const int vertex_idx1 = GetVertexIdx(edge.first);
            const int vertex_idx2 = GetVertexIdx(edge.second);
            adjacency_list[vertex_idx1].emplace_back(vertex_idx2, weight);
            adjacency_list[vertex_idx2].emplace_back(vertex_idx1, weight);
        }

        xadj_.reserve(vertex_id_to_idx_.size() + 1);
        adjncy_.reserve(2 * edges.size());
        adjwgt_.reserve(2 * edges.size());

        idxtype edge_idx = 0;
        for (size_t i = 0; i < vertex_id_to_idx_.size(); ++i) {
            xadj_.push_back(edge_idx);

            if (adjacency_list.count(i) == 0) { continue; }

            for (const auto& edge : adjacency_list[i]) {
                edge_idx += 1;
                adjncy_.push_back(edge.first);
                adjwgt_.push_back(edge.second);
            }
        }

        xadj_.push_back(edge_idx);

        CHECK_EQ(edge_idx, 2 * edges.size());
        CHECK_EQ(xadj_.size(), vertex_id_to_idx_.size() + 1);
        CHECK_EQ(adjncy_.size(), 2 * edges.size());
        CHECK_EQ(adjwgt_.size(), 2 * edges.size());

        data.gdata = data.rdata = nullptr;

        data.nvtxs = vertex_id_to_idx_.size();
        data.nedges = 2 * edges.size();
        data.mincut = data.minvol = -1;

        data.xadj = xadj_.data();
        data.adjncy = adjncy_.data();

        data.vwgt = nullptr;
        data.adjwgt = adjwgt_.data();

        data.adjwgtsum = nullptr;
        data.label = nullptr;
        data.cmap = nullptr;

        data.where = data.pwgts = nullptr;
        data.id = data.ed = nullptr;
        data.bndptr = data.bndind = nullptr;
        data.rinfo = nullptr;
        data.vrinfo = nullptr;
        data.nrinfo = nullptr;

        data.ncon = 1;
        data.nvwgt = nullptr;
        data.npwgts = nullptr;

        data.vsize = nullptr;

        data.coarser = data.finer = nullptr;
    }

    int GetVertexIdx(const int id) {
        const auto it = vertex_id_to_idx_.find(id);
        if (it == vertex_id_to_idx_.end()) {
            const int idx = vertex_id_to_idx_.size();
            vertex_id_to_idx_.emplace(id, idx);
            vertex_idx_to_id_.emplace(idx, id);
            return idx;
        } else {
            return it->second;
        }
    }

    int GetVertexId(const int idx) { return vertex_idx_to_id_.at(idx); }

    GraphType data;

private:
    std::unordered_map<int, int> vertex_id_to_idx_;
    std::unordered_map<int, int> vertex_idx_to_id_;
    std::vector<idxtype> xadj_;
    std::vector<idxtype> adjncy_;
    std::vector<idxtype> adjwgt_;
};

void ComputeMinGraphCutStoerWagner(
    const std::vector<std::pair<int, int>>& edges,
    const std::vector<int>& weights, int* cut_weight,
    std::vector<char>* cut_labels)
{
    CHECK_EQ(edges.size(), weights.size());
    CHECK_GE(edges.size(), 2);

    typedef boost::property<boost::edge_weight_t, int> edge_weight_t;
    typedef boost::adjacency_list<boost::vecS, boost::vecS, boost::undirectedS,
                                boost::no_property, edge_weight_t>
        undirected_graph_t;

    int max_vertex_index = 0;
    for (const auto& edge : edges) {
        CHECK_GE(edge.first, 0);
        CHECK_GE(edge.second, 0);
        max_vertex_index = std::max(max_vertex_index, edge.first);
        max_vertex_index = std::max(max_vertex_index, edge.second);
    }

    const undirected_graph_t graph(edges.begin(), edges.end(), weights.begin(),
                                   max_vertex_index + 1, edges.size());

    const auto parities = boost::make_one_bit_color_map(
        boost::num_vertices(graph), boost::get(boost::vertex_index, graph));

    *cut_weight = boost::stoer_wagner_min_cut(graph, boost::get(boost::edge_weight, graph),
                                  boost::parity_map(parities));

    cut_labels->resize(boost::num_vertices(graph));
    for (size_t i = 0; i < boost::num_vertices(graph); ++i) {
        (*cut_labels)[i] = boost::get(parities, i);
    }
}

std::unordered_map<int, int> ComputeNormalizedMinGraphCut(
    const std::vector<std::pair<int, int>>& edges,
    const std::vector<int>& weights, const int num_parts)
{
    GraclusGraph graph(edges, weights);

    const int levels =
        amax((graph.data.nvtxs) / (40 * log2_metis(num_parts)), 20 * (num_parts));

    std::vector<idxtype> cut_labels(graph.data.nvtxs);

    int options[11];
    options[0] = 0;
    int wgtflag = 1;
    int numflag = 0;
    int chain_length = 0;
    int edgecut;
    int var_num_parts = num_parts;

    MLKKM_PartGraphKway(&graph.data.nvtxs, graph.data.xadj, graph.data.adjncy,
                        graph.data.vwgt, graph.data.adjwgt, &wgtflag, &numflag,
                        &var_num_parts, &chain_length, options, &edgecut,
                        cut_labels.data(), levels);

    float lbvec[MAXNCON];
    ComputePartitionBalance(&graph.data, num_parts, cut_labels.data(), lbvec);
    ComputeNCut(&graph.data, &cut_labels[0], num_parts);
    std::unordered_map<int, int> labels;
    for (size_t idx = 0; idx < cut_labels.size(); ++idx) {
        labels.emplace(graph.GetVertexId(idx), cut_labels[idx]);
    }

    return labels;
}

}  // namespace graph
}  // namespace GraphSfM