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

#include <chain.h>
#include <chainparams.h>
#include <validation.h>
#include <pow.h>
#include <random.h>
#include <util.h>
#include <nacapow/nacapow.h>
#include <test/test_bitcoin.h>

#include <boost/test/unit_test.hpp>

BOOST_FIXTURE_TEST_SUITE(pow_tests, BasicTestingSetup)

/* Test calculation of next difficulty target with no constraints applying */

// What to implement:
// 1. Generate 2016 blocks
// 2. Then, lets use. Case A:  c(seq(0.3 * 60, 0.4 * 60, length=2013), 0.499, 0.5 * 60, 0.5 * 60) for the nTime of blocks
// The target would be c(seq(5284900, 6284900, length=2013), 3284901, 3284900, 3284900)
// Case A, the found target value should be 3284901

BOOST_AUTO_TEST_CASE(get_next_work_case_a)
{
  const auto chainParams = CreateChainParams(CBaseChainParams::MAIN);
  uint64_t nTimes[2016];
  uint32_t targets[2016];
  double delta_time = (0.4 - 0.3) * 60. / 2012.;
  double delta_target = (6284900. - 5284900.) / 2012.;
  unsigned int counter = 0;
  for (int i=0; i < 2016; i++)
  {
    if (i < 2012)
    {
      nTimes[i] = (uint64_t)(delta_time * i + (0.3 * 60)) ;
      targets[i] = delta_target * i + 5284900;
    }
    else if (i == 2013)
    {
      nTimes[i] = (uint64_t)(0.499 * 60);
      targets[i] = 3284901;
    }
    else
    {
      nTimes[i] = (uint64_t)(0.5 * 60);
      targets[i] = 3284900;
    }
  }
  // Initialize the chain with a Genesis block
  CBlockIndex* genesis = new CBlockIndex();
  genesis->pprev = nullptr;
  genesis->nBits = (uint32_t)526378596;
  genesis->nTime = 1566247964 + nTimes[0];
  chainActive.SetTip(genesis);
  // arith_uint256 biggestTarget = arith_uint256(
      // "0x005FE664FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF");
  while (chainActive.Tip()->nHeight < 2015) {
    CBlockIndex* prev = chainActive.Tip();
    CBlockIndex* next = new CBlockIndex();
    NacaOutput output;
    output.subhash = targets[counter];
    for (unsigned int i = 0; i < 7; i++)
    {
      output.xcomp[i] = 1.5 + i;
      //output.xcomp[i] = 0.;
    }
    unsigned char* output_ch = (unsigned char*)(&output);
    std::vector<unsigned char> vec(output_ch, output_ch + sizeof(output));
    uint256 uint_out = uint256(vec);
    next->pprev = prev;
    next->pprev->phashBlock = &uint_out;
    next->nHeight = prev->nHeight + 1;
    next->nBits = 526378596;
    next->nTime = prev->nTime + nTimes[counter];
    //next->BuildSkip();
    chainActive.SetTip(next);
    counter += 1;
  }
  BOOST_CHECK_EQUAL(chainActive.Tip()->nHeight, 2015);
  unsigned int calculatedNBits;
  arith_uint256 computedTarget;
  CBlockIndex* pindexLast = chainActive.Tip();
  calculatedNBits = CalculateNextWorkRequired(pindexLast, genesis->GetBlockTime(), chainParams->GetConsensus());
  printf("nBits adjusted: %u\n", calculatedNBits);
  //computedTarget.SetCompact(calculatedNBits);
  //printf("nBits original value: %u\n", genesis->nBits);
  //printf("nBits adjusted value AFTER: %u\n", calculatedNBits);
  //printf("Corresponing target is: %s\n", computedTarget.ToString().c_str());
  // Here check that the new nBits value computed has decreased (make it easier to generated a block)
  BOOST_CHECK(calculatedNBits < genesis->nBits);
}

bool TimeComp(const uint64_t& i, const uint64_t& j) { return(i < j);}
BOOST_AUTO_TEST_CASE(get_next_work_case_b)
{
// For case B:
// Then, lets use. Case B:  c(0.566 * 60, 0.567 * 60, 0.569 * 60, seq(0.6 * 60, 0.8 * 60, length=2013)) for the nTime of blocks
// The target would be c(seq(1284900, 2284900, length=2013), 3284901, 3284902, 3284900)
// For case B, the target value should be 3284900 * 1.2 (nBits with GetCompact() to be checked)
// 3. Each block have the genesis nBits value
  const auto chainParams = CreateChainParams(CBaseChainParams::MAIN);
  uint64_t nTimes[2016];
  uint32_t targets[2016];
  double delta_time = (0.8 - 0.6) * 60. / 2012.;
  double delta_target = (2284900. - 1284900.) / 2012.;
  unsigned int counter = 0;
  for (int i=0; i < 2016; i++)
  {
    if (i < 2012)
    {
      nTimes[i] = (uint64_t)(delta_time * i + (0.6 * 60)) ;
      targets[i] = delta_target * i + 1284900;
    }
    else if (i == 2013)
    {
      nTimes[i] = (uint64_t)(0.566 * 60);
      targets[i] = 3284901;
    }
    else if (i == 2014)
    {
      nTimes[i] = (uint64_t)(0.567 * 60);
      targets[i] = 3284902;
    } else
    {
      nTimes[i] = (uint64_t)(0.569 * 60);
      targets[i] = 3284900;
    }
  }
  // Reordering nTimes
  std::vector<uint64_t>nTimesVec(nTimes, std::end(nTimes));
  //printf("nTimes generated - First: %u - Last: %u\n", nTimesVec[0], nTimesVec[10]);
  std::sort(nTimesVec.begin(), nTimesVec.end(), TimeComp);

  // Initialize the chain with a Genesis block
  CBlockIndex* genesis = new CBlockIndex();
  genesis->pprev = nullptr;
  genesis->nBits = (uint32_t)523378598;
  genesis->nTime = 1566247964 + nTimesVec[0];
  chainActive.SetTip(genesis);
  while (chainActive.Tip()->nHeight < 2015) {
    CBlockIndex* prev = chainActive.Tip();
    CBlockIndex* next = new CBlockIndex();
    NacaOutput output;
    output.subhash = targets[counter];
    for (unsigned int i = 0; i < 7; i++)
    {
      output.xcomp[i] = 1.5 + i;
    }
    unsigned char* output_ch = (unsigned char*)(&output);
    std::vector<unsigned char> vec(output_ch, output_ch + sizeof(output));
    uint256 uint_out = uint256(vec);
    next->pprev = prev;
    next->pprev->phashBlock = &uint_out;
    next->nHeight = prev->nHeight + 1;
    next->nBits = 523378598;
    next->nTime = prev->nTime + nTimesVec[counter];
    //next->BuildSkip();
    chainActive.SetTip(next);
    counter += 1;
  }
  BOOST_CHECK_EQUAL(chainActive.Tip()->nHeight, 2015);
  unsigned int calculatedNBits;
  arith_uint256 computedTarget;
  CBlockIndex* pindexLast = chainActive.Tip();
  calculatedNBits = CalculateNextWorkRequired(pindexLast, genesis->GetBlockTime(), chainParams->GetConsensus());
  printf("nBits adjusted value BEFORE: %u\n", calculatedNBits);
  //computedTarget.SetCompact(calculatedNBits);
  //printf("nBits original value: %u\n", genesis->nBits);
  //printf("nBits adjusted value AFTER: %u\n", calculatedNBits);
  // Here check that the new nBits value computed has increased (make it harder to generated a block)
  BOOST_CHECK(calculatedNBits > genesis->nBits);
}

BOOST_AUTO_TEST_CASE(GetBlockProofEquivalentTime_test)
{
    const auto chainParams = CreateChainParams(CBaseChainParams::MAIN);
    std::vector<CBlockIndex> blocks(10000);
    for (int i = 0; i < 10000; i++) {
        blocks[i].pprev = i ? &blocks[i - 1] : nullptr;
        blocks[i].nHeight = i;
        blocks[i].nTime = 1269211443 + i * chainParams->GetConsensus().nPowTargetSpacing;
        blocks[i].nBits = 0x207fffff; /* target 0x7fffff000... */
        blocks[i].nChainWork = i ? blocks[i - 1].nChainWork + GetBlockProof(blocks[i - 1]) : arith_uint256(0);
    }

    for (int j = 0; j < 1000; j++) {
        CBlockIndex *p1 = &blocks[InsecureRandRange(10000)];
        CBlockIndex *p2 = &blocks[InsecureRandRange(10000)];
        CBlockIndex *p3 = &blocks[InsecureRandRange(10000)];

        int64_t tdiff = GetBlockProofEquivalentTime(*p1, *p2, *p3, chainParams->GetConsensus());
        BOOST_CHECK_EQUAL(tdiff, p1->GetBlockTime() - p2->GetBlockTime());
    }
}

BOOST_AUTO_TEST_SUITE_END()


// Copyright (c) by respective owners including Yahoo!, Microsoft, and
// individual contributors. All rights reserved. Released under a BSD (revised)
// license as described in the file LICENSE.

#include "reductions/cb/details/large_action_space.h"
#include "test_common.h"
#include "vw/core/qr_decomposition.h"
#include "vw/core/rand48.h"
#include "vw/core/rand_state.h"
#include "vw/core/reductions/cb/cb_explore_adf_common.h"
#include "vw/core/reductions/cb/cb_explore_adf_large_action_space.h"
#include "vw/core/vw.h"

#include <boost/test/unit_test.hpp>

using internal_action_space =
    VW::cb_explore_adf::cb_explore_adf_base<VW::cb_explore_adf::cb_explore_adf_large_action_space>;

BOOST_AUTO_TEST_CASE(creation_of_the_og_A_matrix)
{
  auto d = 2;
  auto& vw = *VW::initialize("--cb_explore_adf --large_action_space --full_predictions --max_actions " +
          std::to_string(d) + " --quiet --random_seed 5",
      nullptr, false, nullptr, nullptr);

  {
    VW::multi_ex examples;

    examples.push_back(VW::read_example(vw, "0:1.0:0.5 | 1 2 3"));

    vw.learn(examples);
    vw.finish_example(examples);
  }

  std::vector<std::string> e_r;
  vw.l->get_enabled_reductions(e_r);
  if (std::find(e_r.begin(), e_r.end(), "cb_explore_adf_large_action_space") == e_r.end())
  { BOOST_FAIL("cb_explore_adf_large_action_space not found in enabled reductions"); }

  VW::LEARNER::multi_learner* learner =
      as_multiline(vw.l->get_learner_by_name_prefix("cb_explore_adf_large_action_space"));

  auto action_space = (internal_action_space*)learner->get_internal_type_erased_data_pointer_test_use_only();

  BOOST_CHECK_EQUAL(action_space != nullptr, true);

  {
    VW::multi_ex examples;

    examples.push_back(VW::read_example(vw, "0:1.0:0.5 | 1 2 3"));

    vw.predict(examples);

    action_space->explore._generate_A(examples);

    auto num_actions = examples.size();
    // gather the feature indexes

    for (size_t action_index = 0; action_index < num_actions; action_index++)
    {
      auto* ex = examples[action_index];
      // test sanity - test assumes no shared features
      BOOST_CHECK_EQUAL(!CB::ec_is_example_header(*ex), true);
      for (auto ns : ex->indices)
      {
        for (auto ft_index : ex->feature_space[ns].indices)
        {
          BOOST_CHECK_EQUAL(
              action_space->explore._A.coeffRef(action_index, (ft_index & vw.weights.dense_weights.mask())),
              vw.weights.dense_weights[ft_index]);
        }
      }
    }
    vw.finish_example(examples);
  }
  VW::finish(vw);
}

BOOST_AUTO_TEST_CASE(check_interactions_on_Y)
{
  auto d = 2;
  std::vector<std::pair<VW::workspace*, bool>> vws;
  auto* vw_no_interactions = VW::initialize("--cb_explore_adf --large_action_space --full_predictions --max_actions " +
          std::to_string(d) + " --quiet --random_seed 5",
      nullptr, false, nullptr, nullptr);

  vws.push_back({vw_no_interactions, false});

  auto* vw_yes_interactions = VW::initialize("--cb_explore_adf --large_action_space --full_predictions --max_actions " +
          std::to_string(d) + " --quiet --random_seed 5 -q ::",
      nullptr, false, nullptr, nullptr);

  vws.push_back({vw_yes_interactions, true});

  size_t interactions_rows = 0;
  size_t non_interactions_rows = 0;

  for (auto& vw_pair : vws)
  {
    auto& vw = *std::get<0>(vw_pair);
    bool interactions = std::get<1>(vw_pair);

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "0:1.0:0.5 |f 1 2 3"));
      examples.push_back(VW::read_example(vw, "|f a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "|f a_4 a_5 a_6"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "|f 1 2 3"));
      examples.push_back(VW::read_example(vw, "0:0.9:0.5 |f a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "|f a_4 a_5 a_6"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "|f 1 2 3"));
      examples.push_back(VW::read_example(vw, "|f a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "0:0.8:0.5 |f a_4 a_5 a_6"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    std::vector<std::string> e_r;
    vw.l->get_enabled_reductions(e_r);
    if (std::find(e_r.begin(), e_r.end(), "cb_explore_adf_large_action_space") == e_r.end())
    { BOOST_FAIL("cb_explore_adf_large_action_space not found in enabled reductions"); }

    VW::LEARNER::multi_learner* learner =
        as_multiline(vw.l->get_learner_by_name_prefix("cb_explore_adf_large_action_space"));

    auto action_space = (internal_action_space*)learner->get_internal_type_erased_data_pointer_test_use_only();

    BOOST_CHECK_EQUAL(action_space != nullptr, true);

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "|f 1 2 3"));
      examples.push_back(VW::read_example(vw, "|f a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "|f a_4 a_5 a_6"));

      vw.predict(examples);

      std::set<uint64_t> non_zero_rows;
      for (int k = 0; k < action_space->explore.Y.outerSize(); ++k)
      {
        for (Eigen::SparseMatrix<float>::InnerIterator it(action_space->explore.Y, k); it; ++it)
        { non_zero_rows.emplace(it.row()); }
      }

      if (!interactions) { non_interactions_rows = non_zero_rows.size(); }
      if (interactions) { interactions_rows = non_zero_rows.size(); }
      vw.finish_example(examples);
    }
    VW::finish(vw);
  }
  BOOST_CHECK_GT(interactions_rows, non_interactions_rows);
}

BOOST_AUTO_TEST_CASE(check_interactions_on_B)
{
  auto d = 2;
  std::vector<std::pair<VW::workspace*, bool>> vws;
  auto* vw_no_interactions = VW::initialize("--cb_explore_adf --large_action_space --full_predictions --max_actions " +
          std::to_string(d) + " --quiet --random_seed 5",
      nullptr, false, nullptr, nullptr);

  vws.push_back({vw_no_interactions, false});

  auto* vw_yes_interactions = VW::initialize("--cb_explore_adf --large_action_space --full_predictions --max_actions " +
          std::to_string(d) + " --quiet --random_seed 5 -q ::",
      nullptr, false, nullptr, nullptr);

  vws.push_back({vw_yes_interactions, true});

  Eigen::MatrixXf B_non_interactions;
  Eigen::MatrixXf B_interactions;

  for (auto& vw_pair : vws)
  {
    auto& vw = *std::get<0>(vw_pair);
    bool interactions = std::get<1>(vw_pair);

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "0:1.0:0.5 |f 1 2 3"));
      examples.push_back(VW::read_example(vw, "|f a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "|f a_4 a_5 a_6"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "|f 1 2 3"));
      examples.push_back(VW::read_example(vw, "0:0.9:0.5 |f a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "|f a_4 a_5 a_6"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "|f 1 2 3"));
      examples.push_back(VW::read_example(vw, "|f a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "0:0.8:0.5 |f a_4 a_5 a_6"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    std::vector<std::string> e_r;
    vw.l->get_enabled_reductions(e_r);
    if (std::find(e_r.begin(), e_r.end(), "cb_explore_adf_large_action_space") == e_r.end())
    { BOOST_FAIL("cb_explore_adf_large_action_space not found in enabled reductions"); }

    VW::LEARNER::multi_learner* learner =
        as_multiline(vw.l->get_learner_by_name_prefix("cb_explore_adf_large_action_space"));

    auto action_space = (internal_action_space*)learner->get_internal_type_erased_data_pointer_test_use_only();

    BOOST_CHECK_EQUAL(action_space != nullptr, true);

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "|f 1 2 3"));
      examples.push_back(VW::read_example(vw, "|f a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "|f a_4 a_5 a_6"));

      vw.predict(examples);

      if (!interactions) { B_non_interactions = action_space->explore.B; }
      if (interactions) { B_interactions = action_space->explore.B; }
      vw.finish_example(examples);
    }
    VW::finish(vw);
  }
  BOOST_CHECK_EQUAL(B_interactions.isApprox(B_non_interactions), false);
}

BOOST_AUTO_TEST_CASE(check_At_times_Omega_is_Y)
{
  auto d = 2;
  std::vector<std::pair<VW::workspace*, bool>> vws;
  auto* vw_epsilon = VW::initialize("--cb_explore_adf --large_action_space --full_predictions --max_actions " +
          std::to_string(d) + " --quiet --random_seed 5",
      nullptr, false, nullptr, nullptr);

  vws.push_back({vw_epsilon, false});

  auto* vw_squarecb =
      VW::initialize("--cb_explore_adf --squarecb --large_action_space --full_predictions --max_actions " +
              std::to_string(d) + " --quiet --random_seed 5",
          nullptr, false, nullptr, nullptr);

  vws.push_back({vw_squarecb, true});

  for (auto& vw_pair : vws)
  {
    auto& vw = *std::get<0>(vw_pair);
    auto apply_diag_M = std::get<1>(vw_pair);

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "0:1.0:0.5 | 1 2 3"));
      examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "| a_4 a_5 a_6"));
      examples.push_back(VW::read_example(vw, "| a_7 a_8 a_9"));
      examples.push_back(VW::read_example(vw, "| a_10 a_11 a_12"));
      examples.push_back(VW::read_example(vw, "| a_13 a_14 a_15"));
      examples.push_back(VW::read_example(vw, "| a_16 a_17 a_18"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "| 1 2 3"));
      examples.push_back(VW::read_example(vw, "0:0.9:0.5 | a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "| a_4 a_5 a_6"));
      examples.push_back(VW::read_example(vw, "| a_7 a_8 a_9"));
      examples.push_back(VW::read_example(vw, "| a_10 a_11 a_12"));
      examples.push_back(VW::read_example(vw, "| a_13 a_14 a_15"));
      examples.push_back(VW::read_example(vw, "| a_16 a_17 a_18"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "| 1 2 3"));
      examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "0:0.8:0.5 | a_4 a_5 a_6"));
      examples.push_back(VW::read_example(vw, "| a_7 a_8 a_9"));
      examples.push_back(VW::read_example(vw, "| a_10 a_11 a_12"));
      examples.push_back(VW::read_example(vw, "| a_13 a_14 a_15"));
      examples.push_back(VW::read_example(vw, "| a_16 a_17 a_18"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "| 1 2 3"));
      examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "| a_4 a_5 a_6"));
      examples.push_back(VW::read_example(vw, "0:0.7:0.5 | a_7 a_8 a_9"));
      examples.push_back(VW::read_example(vw, "| a_10 a_11 a_12"));
      examples.push_back(VW::read_example(vw, "| a_13 a_14 a_15"));
      examples.push_back(VW::read_example(vw, "| a_16 a_17 a_18"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "| 1 2 3"));
      examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "| a_4 a_5 a_6"));
      examples.push_back(VW::read_example(vw, "| a_7 a_8 a_9"));
      examples.push_back(VW::read_example(vw, "0:0.6:0.5 | a_10 a_11 a_12"));
      examples.push_back(VW::read_example(vw, "| a_13 a_14 a_15"));
      examples.push_back(VW::read_example(vw, "| a_16 a_17 a_18"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "| 1 2 3"));
      examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "| a_4 a_5 a_6"));
      examples.push_back(VW::read_example(vw, "| a_7 a_8 a_9"));
      examples.push_back(VW::read_example(vw, "| a_10 a_11 a_12"));
      examples.push_back(VW::read_example(vw, "0:0.5:0.5 | a_13 a_14 a_15"));
      examples.push_back(VW::read_example(vw, "| a_16 a_17 a_18"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "| 1 2 3"));
      examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "| a_4 a_5 a_6"));
      examples.push_back(VW::read_example(vw, "| a_7 a_8 a_9"));
      examples.push_back(VW::read_example(vw, "| a_10 a_11 a_12"));
      examples.push_back(VW::read_example(vw, "| a_13 a_14 a_15"));
      examples.push_back(VW::read_example(vw, "0:0.4:0.5 | a_16 a_17 a_18"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    std::vector<std::string> e_r;
    vw.l->get_enabled_reductions(e_r);
    if (std::find(e_r.begin(), e_r.end(), "cb_explore_adf_large_action_space") == e_r.end())
    { BOOST_FAIL("cb_explore_adf_large_action_space not found in enabled reductions"); }

    VW::LEARNER::multi_learner* learner =
        as_multiline(vw.l->get_learner_by_name_prefix("cb_explore_adf_large_action_space"));

    auto action_space = (internal_action_space*)learner->get_internal_type_erased_data_pointer_test_use_only();

    BOOST_CHECK_EQUAL(action_space != nullptr, true);

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "| 1 2 3"));
      examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "| a_4 a_5 a_6"));
      examples.push_back(VW::read_example(vw, "| a_7 a_8 a_9"));
      examples.push_back(VW::read_example(vw, "| a_10 a_11 a_12"));
      examples.push_back(VW::read_example(vw, "| a_13 a_14 a_15"));
      examples.push_back(VW::read_example(vw, "| a_16 a_17 a_18"));

      vw.predict(examples);

      action_space->explore.calculate_shrink_factor(examples[0]->pred.a_s);
      action_space->explore._generate_A(examples);
      action_space->explore.generate_Y(examples);

      uint64_t num_actions = examples[0]->pred.a_s.size();

      // Generate Omega
      std::vector<Eigen::Triplet<float>> omega_triplets;
      uint64_t max_ft_index = 0;
      uint64_t seed = vw.get_random_state()->get_current_state() * 10.f;

      for (uint64_t action_index = 0; action_index < num_actions; action_index++)
      {
        auto* ex = examples[action_index];
        // test sanity - test assumes no shared features
        BOOST_CHECK_EQUAL(!CB::ec_is_example_header(*ex), true);
        for (auto ns : ex->indices)
        {
          for (uint64_t col = 0; col < d; col++)
          {
            auto combined_index = action_index + col + seed;
            auto mm = merand48_boxmuller(combined_index);
            omega_triplets.push_back(Eigen::Triplet<float>(action_index, col, mm));
          }
        }
      }

      Eigen::SparseMatrix<float> Omega(num_actions, d);
      Omega.setFromTriplets(omega_triplets.begin(), omega_triplets.end(), [](const float& a, const float& b) {
        assert(a == b);
        return b;
      });

      Eigen::SparseMatrix<float> diag_M(num_actions, num_actions);

      if (apply_diag_M)
      {
        for (Eigen::Index i = 0; i < action_space->explore.shrink_factors.size(); i++)
        { diag_M.coeffRef(i, i) = action_space->explore.shrink_factors[i]; }
      }
      else
      {
        diag_M.setIdentity();
      }

      Eigen::SparseMatrix<float> Yd(action_space->explore.Y.rows(), d);

      Yd = action_space->explore._A.transpose() * diag_M * Omega;
      // Orthonormalize Yd
      VW::gram_schmidt(Yd);
      BOOST_CHECK_EQUAL(Yd.isApprox(action_space->explore.Y), true);

      vw.finish_example(examples);
    }
    VW::finish(vw);
  }
}

BOOST_AUTO_TEST_CASE(check_A_times_Y_is_B)
{
  auto d = 2;
  std::vector<std::pair<VW::workspace*, bool>> vws;
  auto* vw_epsilon = VW::initialize("--cb_explore_adf --large_action_space --full_predictions --max_actions " +
          std::to_string(d) + " --quiet --random_seed 5",
      nullptr, false, nullptr, nullptr);

  vws.push_back({vw_epsilon, false});

  auto* vw_squarecb =
      VW::initialize("--cb_explore_adf --squarecb --large_action_space --full_predictions --max_actions " +
              std::to_string(d) + " --quiet --random_seed 5",
          nullptr, false, nullptr, nullptr);

  vws.push_back({vw_squarecb, true});

  for (auto& vw_pair : vws)
  {
    auto& vw = *std::get<0>(vw_pair);
    auto apply_diag_M = std::get<1>(vw_pair);

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "0:1.0:0.5 | 1 2 3"));
      examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    std::vector<std::string> e_r;
    vw.l->get_enabled_reductions(e_r);
    if (std::find(e_r.begin(), e_r.end(), "cb_explore_adf_large_action_space") == e_r.end())
    { BOOST_FAIL("cb_explore_adf_large_action_space not found in enabled reductions"); }

    VW::LEARNER::multi_learner* learner =
        as_multiline(vw.l->get_learner_by_name_prefix("cb_explore_adf_large_action_space"));

    auto action_space = (internal_action_space*)learner->get_internal_type_erased_data_pointer_test_use_only();

    BOOST_CHECK_EQUAL(action_space != nullptr, true);

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "0:1.0:0.5 | 1 2 3"));
      examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));

      vw.predict(examples);

      action_space->explore._generate_A(examples);
      action_space->explore.calculate_shrink_factor(examples[0]->pred.a_s);
      action_space->explore.generate_Y(examples);
      action_space->explore.generate_B(examples);

      auto num_actions = examples[0]->pred.a_s.size();
      Eigen::SparseMatrix<float> diag_M(num_actions, num_actions);

      if (apply_diag_M)
      {
        for (Eigen::Index i = 0; i < action_space->explore.shrink_factors.size(); i++)
        { diag_M.coeffRef(i, i) = action_space->explore.shrink_factors[i]; }
      }
      else
      {
        diag_M.setIdentity();
      }

      Eigen::MatrixXf B = diag_M * action_space->explore._A * action_space->explore.Y;
      BOOST_CHECK_EQUAL(B.isApprox(action_space->explore.B), true);

      vw.finish_example(examples);
    }
    VW::finish(vw);
  }
}

BOOST_AUTO_TEST_CASE(check_B_times_P_is_Z)
{
  auto d = 2;

  std::vector<std::pair<VW::workspace*, bool>> vws;
  auto* vw_epsilon = VW::initialize("--cb_explore_adf --large_action_space --full_predictions --max_actions " +
          std::to_string(d) + " --quiet --random_seed 5",
      nullptr, false, nullptr, nullptr);

  vws.push_back({vw_epsilon, false});

  auto* vw_squarecb =
      VW::initialize("--cb_explore_adf --squarecb --large_action_space --full_predictions --max_actions " +
              std::to_string(d) + " --quiet --random_seed 5",
          nullptr, false, nullptr, nullptr);

  vws.push_back({vw_squarecb, true});

  for (auto& vw_pair : vws)
  {
    auto& vw = *std::get<0>(vw_pair);

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "0:1.0:0.5 | 1 2 3"));
      examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    std::vector<std::string> e_r;
    vw.l->get_enabled_reductions(e_r);
    if (std::find(e_r.begin(), e_r.end(), "cb_explore_adf_large_action_space") == e_r.end())
    { BOOST_FAIL("cb_explore_adf_large_action_space not found in enabled reductions"); }

    VW::LEARNER::multi_learner* learner =
        as_multiline(vw.l->get_learner_by_name_prefix("cb_explore_adf_large_action_space"));

    auto action_space = (internal_action_space*)learner->get_internal_type_erased_data_pointer_test_use_only();

    BOOST_CHECK_EQUAL(action_space != nullptr, true);

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "0:1.0:0.5 | 1 2 3"));
      examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));

      vw.predict(examples);

      action_space->explore.calculate_shrink_factor(examples[0]->pred.a_s);
      action_space->explore._generate_A(examples);
      action_space->explore.generate_Y(examples);
      action_space->explore.generate_B(examples);
      action_space->explore.generate_Z(examples);

      Eigen::MatrixXf P(d, d);

      uint64_t seed = vw.get_random_state()->get_current_state() * 10.f;

      for (size_t row = 0; row < d; row++)
      {
        for (size_t col = 0; col < d; col++)
        {
          auto combined_index = row + col + seed;
          auto mm = merand48_boxmuller(combined_index);
          P(row, col) = mm;
        }
      }

      Eigen::MatrixXf Zp = action_space->explore.B * P;
      VW::gram_schmidt(Zp);
      BOOST_CHECK_EQUAL(Zp.isApprox(action_space->explore.Z), true);
      vw.finish_example(examples);
    }

    VW::finish(vw);
  }
}

BOOST_AUTO_TEST_CASE(check_final_truncated_SVD_validity)
{
  auto d = 3;

  std::vector<std::pair<VW::workspace*, bool>> vws;
  auto* vw_wo_interactions = VW::initialize("--cb_explore_adf --large_action_space --full_predictions --max_actions " +
          std::to_string(d) + " --quiet --random_seed 5",
      nullptr, false, nullptr, nullptr);

  vws.push_back({vw_wo_interactions, false});

  auto* vw_w_interactions = VW::initialize("--cb_explore_adf --large_action_space --full_predictions --max_actions " +
          std::to_string(d) + " --quiet --random_seed 5 -q ::",
      nullptr, false, nullptr, nullptr);

  vws.push_back({vw_w_interactions, false});

  auto* vw_wo_interactions_sq =
      VW::initialize("--cb_explore_adf --squarecb --large_action_space --full_predictions --max_actions " +
              std::to_string(d) + " --quiet --random_seed 5",
          nullptr, false, nullptr, nullptr);

  vws.push_back({vw_wo_interactions_sq, true});

  auto* vw_w_interactions_sq =
      VW::initialize("--cb_explore_adf --squarecb --large_action_space --full_predictions --max_actions " +
              std::to_string(d) + " --quiet --random_seed 5 -q ::",
          nullptr, false, nullptr, nullptr);

  vws.push_back({vw_w_interactions_sq, true});

  auto* vw_w_interactions_sq_sparse_weights = VW::initialize(
      "--cb_explore_adf --squarecb --sparse_weights --large_action_space --full_predictions --max_actions " +
          std::to_string(d) + " --quiet --random_seed 5 -q ::",
      nullptr, false, nullptr, nullptr);

  vws.push_back({vw_w_interactions_sq_sparse_weights, true});

  for (auto& vw_pair : vws)
  {
    auto& vw = *std::get<0>(vw_pair);
    auto apply_diag_M = std::get<1>(vw_pair);
    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "0:1.0:0.5 |f 1 2 3"));
      examples.push_back(VW::read_example(vw, "|f a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "|f a_4 a_5 a_6"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "|f 1 2 3"));
      examples.push_back(VW::read_example(vw, "0:1.0:0.5 |f a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "|f a_4 a_5 a_6"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "|f 1 2 3"));
      examples.push_back(VW::read_example(vw, "|f a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "0:1.0:0.5 |f a_4 a_5 a_6"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    std::vector<std::string> e_r;
    vw.l->get_enabled_reductions(e_r);
    if (std::find(e_r.begin(), e_r.end(), "cb_explore_adf_large_action_space") == e_r.end())
    { BOOST_FAIL("cb_explore_adf_large_action_space not found in enabled reductions"); }

    VW::LEARNER::multi_learner* learner =
        as_multiline(vw.l->get_learner_by_name_prefix("cb_explore_adf_large_action_space"));

    auto action_space = (internal_action_space*)learner->get_internal_type_erased_data_pointer_test_use_only();

    BOOST_CHECK_EQUAL(action_space != nullptr, true);

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "|f 1 2 3"));
      examples.push_back(VW::read_example(vw, "|f a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "|f a_4 a_5 a_6"));
      action_space->explore._populate_all_SVD_components();

      vw.predict(examples);

      action_space->explore.calculate_shrink_factor(examples[0]->pred.a_s);
      action_space->explore.randomized_SVD(examples);

      action_space->explore._generate_A(examples);
      {
        Eigen::FullPivLU<Eigen::MatrixXf> lu_decomp(action_space->explore._A);
        auto rank = lu_decomp.rank();
        // for test set actual rank of A
        action_space->explore._set_rank(rank);
        // should have a rank larger than 1 for the test
        BOOST_CHECK_GT(rank, 1);
      }

      vw.predict(examples);

      action_space->explore.calculate_shrink_factor(examples[0]->pred.a_s);
      action_space->explore.randomized_SVD(examples);

      auto num_actions = examples.size();

      // U dimensions should be K x d
      BOOST_CHECK_EQUAL(action_space->explore.U.rows(), num_actions);
      BOOST_CHECK_EQUAL(action_space->explore.U.cols(), d);

      // truncated randomized SVD reconstruction
      for (int i = 0; i < action_space->explore.U.cols(); ++i)
      {
        BOOST_CHECK_SMALL(1.f - action_space->explore.U.col(i).norm(), FLOAT_TOL);
        for (int j = 0; j < i; ++j)
        { BOOST_CHECK_SMALL(action_space->explore.U.col(i).dot(action_space->explore.U.col(j)), FLOAT_TOL); }
      }

      for (int i = 0; i < action_space->explore._V.cols(); ++i)
      {
        BOOST_CHECK_SMALL(1.f - action_space->explore._V.col(i).norm(), FLOAT_TOL);
        for (int j = 0; j < i; ++j)
        { BOOST_CHECK_SMALL(action_space->explore._V.col(i).dot(action_space->explore._V.col(j)), FLOAT_TOL); }
      }

      Eigen::SparseMatrix<float> diag_M(num_actions, num_actions);

      if (apply_diag_M)
      {
        for (Eigen::Index i = 0; i < action_space->explore.shrink_factors.size(); i++)
        { diag_M.coeffRef(i, i) = action_space->explore.shrink_factors[i]; }
      }
      else
      {
        diag_M.setIdentity();
      }

      BOOST_CHECK_SMALL(
          ((diag_M * action_space->explore._A) -
              action_space->explore.U * action_space->explore._S.asDiagonal() * action_space->explore._V.transpose())
              .norm(),
          FLOAT_TOL);

      // compare singular values with actual SVD singular values
      Eigen::MatrixXf A_dense = diag_M * action_space->explore._A;
      Eigen::JacobiSVD<Eigen::MatrixXf> svd(A_dense, Eigen::ComputeThinU | Eigen::ComputeThinV);
      Eigen::VectorXf S = svd.singularValues();

      for (size_t i = 0; i < action_space->explore._S.rows(); i++)
      { BOOST_CHECK_SMALL(S(i) - action_space->explore._S(i), FLOAT_TOL); }

      vw.finish_example(examples);
    }
    VW::finish(vw);
  }
}

BOOST_AUTO_TEST_CASE(check_shrink_factor)
{
  auto d = 2;
  std::vector<std::pair<VW::workspace*, bool>> vws;
  auto* vw_epsilon = VW::initialize("--cb_explore_adf --large_action_space --full_predictions --max_actions " +
          std::to_string(d) + " --quiet --random_seed 5",
      nullptr, false, nullptr, nullptr);

  vws.push_back({vw_epsilon, false});

  auto* vw_squarecb =
      VW::initialize("--cb_explore_adf --squarecb --large_action_space --full_predictions --max_actions " +
              std::to_string(d) + " --quiet --random_seed 5",
          nullptr, false, nullptr, nullptr);

  vws.push_back({vw_squarecb, true});

  for (auto& vw_pair : vws)
  {
    auto& vw = *std::get<0>(vw_pair);
    auto apply_diag_M = std::get<1>(vw_pair);

    std::vector<std::string> e_r;
    vw.l->get_enabled_reductions(e_r);
    if (std::find(e_r.begin(), e_r.end(), "cb_explore_adf_large_action_space") == e_r.end())
    { BOOST_FAIL("cb_explore_adf_large_action_space not found in enabled reductions"); }

    VW::LEARNER::multi_learner* learner =
        as_multiline(vw.l->get_learner_by_name_prefix("cb_explore_adf_large_action_space"));

    auto action_space = (internal_action_space*)learner->get_internal_type_erased_data_pointer_test_use_only();

    BOOST_CHECK_EQUAL(action_space != nullptr, true);

    VW::multi_ex examples;

    examples.push_back(VW::read_example(vw, "| 1 2 3"));
    examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));
    examples.push_back(VW::read_example(vw, "| a_4 a_5 a_6"));
    examples.push_back(VW::read_example(vw, "| a_7 a_8 a_9"));
    examples.push_back(VW::read_example(vw, "| a_10 a_11 a_12"));
    examples.push_back(VW::read_example(vw, "| a_13 a_14 a_15"));
    examples.push_back(VW::read_example(vw, "| a_16 a_17 a_18"));

    vw.predict(examples);

    auto num_actions = examples[0]->pred.a_s.size();

    BOOST_CHECK_EQUAL(num_actions, 7);

    action_space->explore.calculate_shrink_factor(examples[0]->pred.a_s);

    Eigen::SparseMatrix<float> diag_M(num_actions, num_actions);
    Eigen::SparseMatrix<float> identity_diag_M(num_actions, num_actions);
    identity_diag_M.setIdentity();

    for (Eigen::Index i = 0; i < action_space->explore.shrink_factors.size(); i++)
    { diag_M.coeffRef(i, i) = action_space->explore.shrink_factors[i]; }

    if (apply_diag_M) { BOOST_CHECK_EQUAL(diag_M.isApprox(identity_diag_M), false); }
    else
    {
      BOOST_CHECK_EQUAL(diag_M.isApprox(identity_diag_M), true);
    }

    vw.finish_example(examples);
    VW::finish(vw);
  }
}

BOOST_AUTO_TEST_CASE(check_finding_max_volume)
{
  auto d = 3;
  auto& vw = *VW::initialize("--cb_explore_adf --large_action_space --full_predictions --max_actions " +
          std::to_string(d) + " --quiet --random_seed 5",
      nullptr, false, nullptr, nullptr);
  VW::cb_explore_adf::cb_explore_adf_large_action_space largecb(
      /*d=*/0, /*gamma_scale=*/1.f, /*gamma_exponent=*/0.f, /*c=*/2, false, &vw);
  largecb.U = Eigen::MatrixXf{{1, 2, 3}, {4, 5, 6}, {7, 8, 9}, {0, 0, 0}, {7, 5, 3}, {6, 4, 8}};
  Eigen::MatrixXf X{{1, 2, 3}, {3, 2, 1}, {2, 1, 3}};

  auto p = largecb.find_max_volume(0, X);
  BOOST_CHECK_SMALL(p.first - 30, FLOAT_TOL);
  BOOST_CHECK_EQUAL(p.second, 2);

  p = largecb.find_max_volume(1, X);
  BOOST_CHECK_SMALL(p.first - 27, FLOAT_TOL);
  BOOST_CHECK_EQUAL(p.second, 4);

  p = largecb.find_max_volume(2, X);
  BOOST_CHECK_SMALL(p.first - 24, FLOAT_TOL);
  BOOST_CHECK_EQUAL(p.second, 5);

  VW::finish(vw);
}

BOOST_AUTO_TEST_CASE(check_spanner_results_squarecb)
{
  auto d = 2;
  std::vector<std::pair<VW::workspace*, bool>> vws;

  auto* vw_full_preds =
      VW::initialize("--cb_explore_adf --squarecb --large_action_space --full_predictions --max_actions " +
              std::to_string(d) + " --quiet --random_seed 5",
          nullptr, false, nullptr, nullptr);

  vws.push_back({vw_full_preds, true});

  auto* vw_sparse_preds = VW::initialize("--cb_explore_adf --squarecb --large_action_space --max_actions " +
          std::to_string(d) + " --quiet --random_seed 5",
      nullptr, false, nullptr, nullptr);

  vws.push_back({vw_sparse_preds, false});

  for (auto& vw_pair : vws)
  {
    auto& vw = *std::get<0>(vw_pair);
    auto full_preds = std::get<1>(vw_pair);

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "0:1.0:0.5 | 1 2 3"));
      examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "| a_4 a_5 a_6"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "| 1 2 3"));
      examples.push_back(VW::read_example(vw, "0:1.0:0.5 | a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "| a_4 a_5 a_6"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "| 1 2 3"));
      examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "0:1.0:0.5 | a_4 a_5 a_6"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    std::vector<std::string> e_r;
    vw.l->get_enabled_reductions(e_r);
    if (std::find(e_r.begin(), e_r.end(), "cb_explore_adf_large_action_space") == e_r.end())
    { BOOST_FAIL("cb_explore_adf_large_action_space not found in enabled reductions"); }

    VW::LEARNER::multi_learner* learner =
        as_multiline(vw.l->get_learner_by_name_prefix("cb_explore_adf_large_action_space"));

    auto action_space = (internal_action_space*)learner->get_internal_type_erased_data_pointer_test_use_only();
    BOOST_CHECK_EQUAL(action_space != nullptr, true);

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "| 1 2 3"));
      examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "| a_4 a_5 a_6"));

      vw.predict(examples);

      const auto num_actions = examples.size();
      const auto& preds = examples[0]->pred.a_s;

      // Only d actions have non-zero scores.
      if (full_preds) { BOOST_CHECK_EQUAL(preds.size(), num_actions); }
      else
      {
        BOOST_CHECK_EQUAL(preds.size(), d);
      }
      BOOST_CHECK_SMALL(preds[0].score - 0.697270989f, FLOAT_TOL);
      BOOST_CHECK_EQUAL(preds[0].action, 1);

      BOOST_CHECK_SMALL(preds[1].score - 0.30272904f, FLOAT_TOL);
      BOOST_CHECK_EQUAL(preds[1].action, 2);

      if (full_preds)
      {
        BOOST_CHECK_SMALL(preds[2].score, FLOAT_TOL);
        BOOST_CHECK_EQUAL(preds[2].action, 0);
      }

      vw.finish_example(examples);
    }
    VW::finish(vw);
  }
}

BOOST_AUTO_TEST_CASE(check_spanner_results_epsilon_greedy)
{
  auto d = 2;
  float epsilon = 0.2f;

  std::vector<std::pair<VW::workspace*, bool>> vws;

  auto* vw_full_preds = VW::initialize("--cb_explore_adf --epsilon " + std::to_string(epsilon) +
          " --large_action_space --full_predictions --max_actions " + std::to_string(d) + " --quiet --random_seed 5",
      nullptr, false, nullptr, nullptr);

  vws.push_back({vw_full_preds, true});

  auto* vw_sparse_preds = VW::initialize("--cb_explore_adf --epsilon " + std::to_string(epsilon) +
          " --large_action_space --max_actions " + std::to_string(d) + " --quiet --random_seed 5",
      nullptr, false, nullptr, nullptr);

  vws.push_back({vw_sparse_preds, false});

  for (auto& vw_pair : vws)
  {
    auto& vw = *std::get<0>(vw_pair);
    auto full_preds = std::get<1>(vw_pair);

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "0:1.0:0.5 | 1 2 3"));
      examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "| a_4 a_5 a_6"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "| 1 2 3"));
      examples.push_back(VW::read_example(vw, "0:1.0:0.5 | a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "| a_4 a_5 a_6"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "| 1 2 3"));
      examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "0:1.0:0.5 | a_4 a_5 a_6"));

      vw.learn(examples);
      vw.finish_example(examples);
    }

    std::vector<std::string> e_r;
    vw.l->get_enabled_reductions(e_r);
    if (std::find(e_r.begin(), e_r.end(), "cb_explore_adf_large_action_space") == e_r.end())
    { BOOST_FAIL("cb_explore_adf_large_action_space not found in enabled reductions"); }

    VW::LEARNER::multi_learner* learner =
        as_multiline(vw.l->get_learner_by_name_prefix("cb_explore_adf_large_action_space"));

    auto action_space = (internal_action_space*)learner->get_internal_type_erased_data_pointer_test_use_only();
    BOOST_CHECK_EQUAL(action_space != nullptr, true);

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "| 1 2 3"));
      examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "| a_4 a_5 a_6"));

      vw.predict(examples);

      const auto num_actions = examples.size();
      const auto& preds = examples[0]->pred.a_s;
      // Only d actions have non-zero scores.
      if (full_preds) { BOOST_CHECK_EQUAL(preds.size(), num_actions); }
      else
      {
        BOOST_CHECK_EQUAL(preds.size(), d);
      }

      size_t num_actions_non_zeroed = d;
      float epsilon_ur = epsilon / num_actions_non_zeroed;
      BOOST_CHECK_SMALL(preds[0].score - (epsilon_ur + (1.f - epsilon)), FLOAT_TOL);
      BOOST_CHECK_EQUAL(preds[0].action, 2);

      BOOST_CHECK_SMALL(preds[1].score - epsilon_ur, FLOAT_TOL);
      BOOST_CHECK_EQUAL(preds[1].action, 0);

      if (full_preds)
      {
        BOOST_CHECK_SMALL(preds[2].score, FLOAT_TOL);
        BOOST_CHECK_EQUAL(preds[2].action, 1);
      }

      vw.finish_example(examples);
    }
    VW::finish(vw);
  }
}

BOOST_AUTO_TEST_CASE(check_uniform_probabilities_before_learning)
{
  auto d = 2;
  std::vector<std::pair<VW::workspace*, bool>> vws;
  auto* vw_epsilon = VW::initialize("--cb_explore_adf --large_action_space --full_predictions --max_actions " +
          std::to_string(d) + " --quiet --random_seed 5",
      nullptr, false, nullptr, nullptr);

  vws.push_back({vw_epsilon, false});

  auto* vw_squarecb =
      VW::initialize("--cb_explore_adf --squarecb --large_action_space --full_predictions --max_actions " +
              std::to_string(d) + " --quiet --random_seed 5",
          nullptr, false, nullptr, nullptr);

  vws.push_back({vw_squarecb, true});

  for (auto& vw_pair : vws)
  {
    auto& vw = *std::get<0>(vw_pair);
    auto apply_diag_M = std::get<1>(vw_pair);

    VW::LEARNER::multi_learner* learner =
        as_multiline(vw.l->get_learner_by_name_prefix("cb_explore_adf_large_action_space"));

    {
      VW::multi_ex examples;

      examples.push_back(VW::read_example(vw, "| 1 2 3"));
      examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));
      examples.push_back(VW::read_example(vw, "| a_4 a_5 a_6"));

      learner->predict(examples);

      const auto num_actions = examples.size();
      const auto& preds = examples[0]->pred.a_s;
      BOOST_CHECK_EQUAL(preds.size(), num_actions);
      for (const auto& pred : preds) { BOOST_CHECK_SMALL(pred.score - (1.f / 3.f), FLOAT_TOL); }

      vw.finish_example(examples);
    }
    VW::finish(vw);
  }
}

BOOST_AUTO_TEST_CASE(check_probabilities_when_d_is_larger)
{
  auto d = 3;
  auto& vw = *VW::initialize("--cb_explore_adf --large_action_space --full_predictions --max_actions " +
          std::to_string(d) + " --quiet --random_seed 5",
      nullptr, false, nullptr, nullptr);

  {
    VW::multi_ex examples;

    examples.push_back(VW::read_example(vw, "0:1.0:0.5 | 1 2 3"));
    examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));
    examples.push_back(VW::read_example(vw, "| a_4 a_5 a_6"));

    vw.learn(examples);
    vw.finish_example(examples);
  }

  {
    VW::multi_ex examples;

    examples.push_back(VW::read_example(vw, "| 1 2 3"));
    examples.push_back(VW::read_example(vw, "0:1.0:0.5 | a_1 a_2 a_3"));
    examples.push_back(VW::read_example(vw, "| a_4 a_5 a_6"));

    vw.learn(examples);
    vw.finish_example(examples);
  }

  {
    VW::multi_ex examples;

    examples.push_back(VW::read_example(vw, "| 1 2 3"));
    examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));
    examples.push_back(VW::read_example(vw, "0:1.0:0.5 | a_4 a_5 a_6"));

    vw.learn(examples);
    vw.finish_example(examples);
  }

  std::vector<std::string> e_r;
  vw.l->get_enabled_reductions(e_r);
  if (std::find(e_r.begin(), e_r.end(), "cb_explore_adf_large_action_space") == e_r.end())
  { BOOST_FAIL("cb_explore_adf_large_action_space not found in enabled reductions"); }

  VW::LEARNER::multi_learner* learner =
      as_multiline(vw.l->get_learner_by_name_prefix("cb_explore_adf_large_action_space"));

  auto action_space = (internal_action_space*)learner->get_internal_type_erased_data_pointer_test_use_only();
  BOOST_CHECK_EQUAL(action_space != nullptr, true);

  {
    VW::multi_ex examples;

    examples.push_back(VW::read_example(vw, "| 1 2 3"));
    examples.push_back(VW::read_example(vw, "| a_1 a_2 a_3"));
    examples.push_back(VW::read_example(vw, "| a_4 a_5 a_6"));

    vw.predict(examples);

    const auto num_actions = examples.size();
    const auto& preds = examples[0]->pred.a_s;
    BOOST_CHECK_EQUAL(preds.size(), num_actions);
    BOOST_CHECK_SMALL(preds[0].score - 0.966666639f, FLOAT_TOL);
    BOOST_CHECK_SMALL(preds[1].score - 0.0166666675f, FLOAT_TOL);
    BOOST_CHECK_SMALL(preds[2].score - 0.0166666675f, FLOAT_TOL);

    vw.finish_example(examples);
  }
  VW::finish(vw);
}


#include "align.h"
#include <pcl/surface/mls.h>
#include <pcl/io/pcd_io.h>
#include <pcl/visualization/pcl_visualizer.h>
#include <boost/thread/thread.hpp>
#include <pcl/console/parse.h>

bool FINE_FLAG = 1;             /*Fine registration*/
bool GLOBAL_FLAG = 0;           /*Global alignment*/
bool DENOISING = 0;             /*Denoise - outlier rejection*/
bool DEPTH_AVG = 0;             /*Down-sampling (recommand apply it before mls)*/
bool MLS = 0;                   /*Denoise - moving least square*/
double MIN_DISTANCE = 0.0035f;  /*Registration threshold*/
double DOWN_SAMPLE = 0.002f;    /*Downsampling voxel size*/
bool LOOP_CLOSURE = 1;          /*Loop closure*/
int GITER=3;                    /*Global alginment iterations*/
double HESSIAN = 200;           /*Hessian threshold for sift*/

// Since the images are cropped, the cropped location is also provided
// If no crop, set the parameters below as zeros
int TOP_LEFT1 = -1;             /*Crop location - x*/
int TOP_LEFT2 = -1;             /*Crop location - y*/

// Global var for kinect_v1
/*-----------------------------------------------------------------*/
double FOCAL = 570.3;
double MM_PER_M = 1000;
int WIDTH = 640;
int HEIGHT = 480;
double DEPTH_THRESHOLD = 1.2;   /*The largest depth we capture*/
/*-----------------------------------------------------------------*/

pcl::visualization::PCLVisualizer *p;

// print sturctions
#define PRINT_TRANSFORM_PAIR 0
#define PRINT_TRANSFORM 1

// saving structions
#define SAVE_EACH_PCD 0

void
showScene (pcl::visualization::PCLVisualizer& viewer, Eigen::Matrix4f m, PointCloud output);

/* Registration: combine the initial and fine registration together
 * @params
        src_points: source point cloud
        tgt_points: target point cloud
        src_key: the 3D SURF point in the source pcd
        tgt_key: the 3D SURF point in the target pcd
        correspondences: pointer of correspondences of the 2D surf key points
        initial_matrix:initial transform
 * @output
        output point cloud
        initial: flag to indicate whether to apply initial reg.
        threshold: threshold for ICP
 */

Eigen::Matrix4f twostepAlign (const PointCloud &src_points,
    const PointCloud &tgt_points, const PointCloud &src_key,
    const PointCloud &tgt_key, pcl::CorrespondencesPtr correspondences,
    const Eigen::Matrix4f initial_matrix, PointCloud &output, bool initial,
    double threshold)
{
  Eigen::Matrix4f matrix1;  // initial
  Eigen::Matrix4f matrix2;  //fine*initial
  Eigen::Matrix4f matrix_final;

  // smart pointer no need to delete it at the end
  PointCloudPtr src (new PointCloud (src_points));
  PointCloudPtr tgt (new PointCloud (tgt_points));

  bool flag = 1;

  // initial alignment
  if (initial && flag == 1)
  {
	  if (correspondences->size () < 3)
	  {
	    printf (
	        "The correspondence points are less than 3. The RANSAC based initial alignment will be swithched off.  \n");
	    printf ("Please relax the threshold for SURF feature mtaching.  \n");
	    flag = 0;
	  }
	  else
		  matrix1 = computeInitialAlignment (src_key, tgt_key, correspondences);
  }

  // Switch off the initial alignment
  if (!initial || flag == 0)
  {
    matrix1 << 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1;
  }
  PointCloudPtr points_transformed (new PointCloud);
  pcl::transformPointCloud (*src, *points_transformed, matrix1);

  // fine registration
  if (FINE_FLAG)
  {
    matrix2 = fineAlign (points_transformed, tgt, output, threshold,initial) * matrix1;  //pairwise alignment matrix
  }

  else  // Switch off the fine alignment
  {
    matrix2 = matrix1;
    output = *points_transformed;
  }

  // initial_matrix: the initial alignment refered to the first frame
  // output is the transformed point cloud refered to the firs frame
  matrix_final = initial_matrix * matrix2;
  pcl::transformPointCloud (output, output, initial_matrix);

  // print the transform for debugging
  if (PRINT_TRANSFORM_PAIR)
  {
    cout << "Final Tranform: " << endl;
    for (int k = 0; k < 3; k++)
    {
      for (int j = 0; j < 4; j++)
      {
        cout << matrix_final (k, j) << " ";
      }
      cout << endl;
    }
  }

  // return the final ransfrom relative to the first frame
  return matrix_final;
}

/* perform the pairwise alignment for a bunch of data (global alignment)
 * @params
        data: a bunch of point clouds for pairwise registration
        correspondencesdata: the correpondence information for pairwise frames that is pre-processed
        key3data: the 3D feature points in each frame
        output: the output pcd combining all the registered pcds
        out_data: a bunch of output point clouds for each frame
        matrix_buffer: a bunch a matrix storing the transform
 */

void pairwiseAlign (
    vector<PointCloud, Eigen::aligned_allocator<PointCloud> > data,
    vector<pcl::CorrespondencesPtr> correspondencesdata,
    vector<PointCloud, Eigen::aligned_allocator<PointCloud> > key3data,
    PointCloud &output,
    vector<PointCloud, Eigen::aligned_allocator<PointCloud> > &out_data,
    vector<Eigen::Matrix4f, Eigen::aligned_allocator<Eigen::Matrix4f> > & matrix_buffer,
    bool initial, double threshold)
{
  cout << "pariwise registration ..." << endl;

  Eigen::Matrix4f matrix_initial;
  Eigen::Matrix4f matrix;
  matrix_initial << 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1;

  output = data[0];
  if (initial == 1)
  {
    p->addPointCloud (data[0].makeShared (),"3");
    p->spinOnce ();
  }

  matrix_buffer.push_back (matrix_initial);

  if (SAVE_EACH_PCD)
  {
    pcl::io::savePCDFile ("pcd/p1.pcd", data[0], true);
  }

  out_data.push_back (data[0]);

  for (unsigned int i = 1; i < data.size (); i++)
  {

    PointCloudPtr temp (new PointCloud ());
    if (initial)
      cout << "# " << i << "->" << i + 1 << ":" << endl;
    matrix = twostepAlign (data[i], data[i - 1], key3data[i], key3data[i - 1],
        correspondencesdata[i - 1], matrix_initial, *temp, initial, threshold);

    if (SAVE_EACH_PCD)
    {
      string savename;
      stringstream strStream;
      strStream << i + 1;
      savename = "pcd/p" + strStream.str () + ".pcd";
      pcl::io::savePCDFile (savename.c_str (), *temp, true);
    }

    out_data.push_back (*temp);

    // for global align
    if (matrix_buffer.size () < i + 1)
      matrix_buffer.push_back (matrix);
    else
      matrix_buffer[i] = matrix * matrix_buffer[i];

    matrix_initial = matrix;

    output = output + *temp;
    if (initial==1)
      showScene (*p, matrix_buffer[i], output);
  }

}

int main (int argc, char** argv)
{
  double f;
  double vd;
  double r;
  int iter;
  double h;

  // argiment parsing
  if (pcl::console::parse_argument (argc, argv, "-r", r) >= 0)
  {
    MIN_DISTANCE = r;
    cout << "adjust the registration threshold to " << r << endl;
  }

  if (pcl::console::parse_argument (argc, argv, "-giter", iter) >= 0)
  {
    GITER = iter;
    cout << "adjust global alignment iterations to " << iter <<" iterations"<< endl;
  }

  if (pcl::console::parse_argument (argc, argv, "-hessian", h) >= 0)
  {
    HESSIAN=h;
    cout << "adjust SIFT matching threshold to " << h << endl;
  }

  if (pcl::console::find_argument (argc, argv, "-h") >= 0 || argc == 1)
  {
    printUsage (argv[0]);
    return 0;
  }

  if (pcl::console::find_argument  (argc, argv, "-g") >= 0)
  {
    GLOBAL_FLAG = 1;
    cout << "Global Reistration is ON" << endl;
  }

  if (pcl::console::find_argument (argc, argv, "-lcoff") >= 0)
   {
    LOOP_CLOSURE = 0;
    cout << "Loop CLosure Detection is OFF" << endl;
   }

  if (pcl::console::find_argument (argc, argv, "-m") >= 0)
  {
    MLS = 1;
    cout << "Moving least square smoothing is ON" << endl;
  }

  if (pcl::console::find_argument (argc, argv, "-dn") >= 0)
  {
    DENOISING= 1;
    cout << "De-noising (statistic) is ON" << endl;
  }

  if (pcl::console::parse_argument (argc, argv, "-vd", vd) >= 0)
  {
    DEPTH_AVG = 1;
    cout << "Down-sampling is ON" << endl;
    if (vd>0.00001)
    {
      DOWN_SAMPLE = vd;
      cout<<"Down sample threshold is changed to "<<DOWN_SAMPLE<<endl;
    }

  }

  vector<PointCloud, Eigen::aligned_allocator<PointCloud> > data;
  vector<PointCloud, Eigen::aligned_allocator<PointCloud> > key3data;
  vector<vector<KeyPoint> > key2data;
  vector<pcl::CorrespondencesPtr> correspondencesdata;

  Mat previousrgb;
  vector<KeyPoint> previouskey2d;

  // read the rgb-d data one by one
  ifstream index;
  string root_dir = argv[1];
  index.open (string(root_dir + "/input.txt").c_str());

  int i = 0;
  while (!index.eof ())
  {
    string file;
    index >> file;
    if (file.empty ())
      break;

    file = root_dir + "/" + file;

    Mat rgbimg;
    Mat depimg;
    Mat mask;

    // read data
    readImg (file, "_rgb.png", rgbimg);
    readImg (file, "_depth.png", depimg);
    readImg (file, "_mask.png", mask);

    if (rgbimg.empty () || depimg.empty () || mask.empty ())
    {
      printf (
          "Error: The RGB-D file does not exit. Please note that the right format should be: \n");
      printf (
          "X_rgb.png, X_depth.png, X_mask.png with the same size. ('X' is your input) \n");
      return (0);
    }

    if (rgbimg.size () != depimg.size () || rgbimg.size () != mask.size ())
    {
      printf (
          "Error: The RGB-D files are not consistent. The rgb, depth and mask images should have the same size. \n");
      return (0);
    }

    cout << "loading:  " << file << endl;

    // check the image size and the offest coordinate
    int x1, y1;
    if (rgbimg.rows == 480 && rgbimg.cols == 640)
    {
      // if the image is full size (480*640), no offset coordinate
      x1 = 0;
      y1 = 0;
    }
    else
    {
      // if not, read the offset file
      ifstream loc;
      string location1 = file + "_loc.txt";
      loc.open (location1.c_str ());

      if (!loc.is_open ())
      {
        if (TOP_LEFT1!=0 && TOP_LEFT2!=0)
        {
        	x1=TOP_LEFT2;
        	y1=TOP_LEFT1;
        }
        else
        {
      	  cout << "Error! There is no associated off-set location file." << endl;
          return 0;
        }
      }
      else
      {
        char comma;
        loc >> x1 >> comma >> y1;
      }
      loc.close();
    }

    // perform erosion to remove the noises around the boundary
    erosion (mask, 4);

    vector<KeyPoint> key2d;
    PointCloud pointcloud;
    PointCloud keypoints;

    // read data and perform surf feature matching
    read (rgbimg, depimg, mask, pointcloud, keypoints, key2d, x1, y1, HESSIAN);
    filter (pointcloud, 0.001, pointcloud);
    data.push_back (pointcloud);
    key3data.push_back (keypoints);
    key2data.push_back (key2d);

    // find the pairwise correspondecnes based on surf feature
    if (i >= 1)
    {
      pcl::CorrespondencesPtr correspondences (new pcl::Correspondences ());  //may have bugs??
      vector<DMatch> good_matches;
      match (rgbimg, previousrgb, key2d, previouskey2d, good_matches,
          correspondences);
      correspondencesdata.push_back (correspondences);
      correspondences.reset (new pcl::Correspondences);
    }
    previousrgb = rgbimg;
    previouskey2d = key2d;

    i++;
  }

  index.close();

  printf ("Loaded %d datasets.\n", (int) data.size ());

  PointCloudPtr final (new PointCloud);
  vector<Eigen::Matrix4f, Eigen::aligned_allocator<Eigen::Matrix4f> > matrix_buffer;

  // registration: initial+fine+global

  //initialize the PCL viewer
  p = new pcl::visualization::PCLVisualizer (argc, argv, "3D Moeling example");
  p->setSize (480, 640);
  p->setPosition (480, 200);

  // registration
  int first = 0;
  int last = 0;
  if (!LOOP_CLOSURE)
  {
	  last=(int) data.size ()-1;
  }

  // need to change back
  double threshold = MIN_DISTANCE;
  for (int z = 0; z < GITER; z++)
  {
    vector<PointCloud, Eigen::aligned_allocator<PointCloud> > out;

    PointCloudPtr output (new PointCloud);
    PointCloudPtr final_temp (new PointCloud);

    if (GLOBAL_FLAG)
    {

      vector<PointCloud, Eigen::aligned_allocator<PointCloud> > out_global;

      if (z == 0)
        pairwiseAlign (data, correspondencesdata, key3data, *output, out,
            matrix_buffer, 1, threshold);
      else
      {
        pairwiseAlign (data, correspondencesdata, key3data, *output, out,
            matrix_buffer, 0, threshold);
      }

      // global optimization
      cout << endl << "Global optimization begins ... at iteration " <<z+1<< endl;
      globalAlign (out, final_temp, matrix_buffer, out_global, first, last);
      data = out_global;
    }
    else
    {
      pairwiseAlign (data, correspondencesdata, key3data, *output, out,
          matrix_buffer, 1, threshold);
      final = output;
      break;
    }
    threshold = threshold * 0.9;

    final = final_temp;
    showScene (*p, matrix_buffer[(int) data.size ()-1], *final);
  }

  if (PRINT_TRANSFORM)
  {
    for (int i = 0; i < data.size (); i++)
    {
      cout << "# " << i + 1 << ": " << endl;
      for (int k = 0; k < 3; k++)
      {
        for (int j = 0; j < 4; j++)
        {
          cout << matrix_buffer[i] (k, j) << " ";
        }
        cout << endl;
      }
      cout << "-------------------------------------------" << endl;
    }
  }

  // de-noising
  PointCloudPtr denoised (new PointCloud);
  if (DENOISING)
  {
    cout<<"statistical outlier removal begins ..."<<endl;
    pcl::StatisticalOutlierRemoval<PointT> sor;
    sor.setInputCloud (final);
    sor.setMeanK (50);
    sor.setStddevMulThresh (1);
    sor.filter (*denoised);
    showScene (*p, matrix_buffer[(int) data.size ()-1], *denoised);
  }
  else
    denoised = final;

  // if voxel down-sampling is needed
  PointCloudPtr downsampled (new PointCloud);
  if (DEPTH_AVG)
  {
    cout<<"Down sampling the point cloud..."<<endl;
    pcl::VoxelGrid<PointT> vf;
    vf.setInputCloud (denoised);
    vf.setLeafSize (DOWN_SAMPLE, DOWN_SAMPLE, DOWN_SAMPLE);
    vf.filter (*downsampled);
    showScene (*p, matrix_buffer[(int) data.size ()-1], *downsampled);
  }
  else
    downsampled=denoised;

  // moving least square smoothing
  PointCloudPtr mls_points (new PointCloud);
  if (MLS && DEPTH_AVG)
  // if the pcd is not downsampled, the mls process will be much slower
  {
    cout<<"moving least square begins ..."<<endl;
    pcl::search::KdTree<PointT>::Ptr tree (new pcl::search::KdTree<PointT>);
    // Output has the PointNormal type in order to store the normals calculated by MLS
    pcl::MovingLeastSquares<PointT, PointT> mls;
    mls.setComputeNormals (true);
    mls.setInputCloud (downsampled);
    mls.setPolynomialFit (true);
    mls.setSearchMethod (tree);
    mls.setSearchRadius (0.01);
    mls.process (*mls_points);

    showScene (*p, matrix_buffer[(int) data.size ()-1], *mls_points);
  }
  else
    mls_points = downsampled;

  //=============EOD of The Process========================

  cout<<"Number of points: "<<mls_points->size()<<endl;

  // save the final pcd
  string ply_name = root_dir + "/" + "result.ply";
  cout << "saving data to " << ply_name << endl;
  save2ply(*mls_points, ply_name.c_str());

  cout << "done!" << endl;

  // show the result
  while (!p->wasStopped ())
  {
    p->spin ();
    boost::this_thread::sleep (boost::posix_time::microseconds (100000));
  }

  return (0);

}

void
showScene (pcl::visualization::PCLVisualizer& viewer, Eigen::Matrix4f m, PointCloud output)
{

    viewer.removeAllPointClouds ();
    viewer.addPointCloud (output.makeShared ());

     Eigen::Matrix3f r;
     r<< m(0, 0), m(0, 1), m(0,2), m(1,0), m(1,1), m(1,2), m(2,0), m(2,1), m(2,2);
     Eigen::Vector3f t(m(0,3),m(1,3),m(2,3));

     Eigen::Affine3f pose = viewer.getViewerPose();
     pose.linear()=r;
     pose.translation()=t;

     Eigen::Vector3f pos_vector = pose * Eigen::Vector3f (0, 0, 0);
     Eigen::Vector3f look_at_vector = pose.rotation () * Eigen::Vector3f (0, 0, 1) + pos_vector;
     Eigen::Vector3f up_vector = pose.rotation () * Eigen::Vector3f (0, 1, 0);
     viewer.setCameraPosition(pos_vector[0],pos_vector[1],pos_vector[2],look_at_vector[0],look_at_vector[1],look_at_vector[2],
                                                   up_vector[0],up_vector[1],up_vector[2]);
     viewer.spinOnce ();
}


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef NT2_TOOLBOX_TRIGONOMETRIC_FUNCTIONS_SCALAR_SINCOS_HPP_INCLUDED
#define NT2_TOOLBOX_TRIGONOMETRIC_FUNCTIONS_SCALAR_SINCOS_HPP_INCLUDED
#include <nt2/toolbox/trigonometric/functions/sincos.hpp>
#include <boost/fusion/tuple.hpp>
#include <nt2/toolbox/trigonometric/functions/scalar/impl/trigo.hpp>
/////////////////////////////////////////////////////////////////////////////
// reference based Implementation
/////////////////////////////////////////////////////////////////////////////
namespace nt2 { namespace ext
{
  NT2_FUNCTOR_IMPLEMENTATION(nt2::tag::sincos_, tag::cpu_,
			     (A0)(A1),
			     (scalar_ < arithmetic_<A0> > )
			     (scalar_ < floating_<A1> > )
			     (scalar_ < floating_<A1> > )
			     )
  {
    typedef int result_type;
    inline int operator()(A0 const& a0,A1 & a1,A1 & a2) const
    {
      a1 = impl::trig_base <A1,radian_tag,tag::not_simd_type>::sincosa(A1(a0),a2);
      return 0;
    }
  };

  NT2_FUNCTOR_IMPLEMENTATION(nt2::tag::sincos_, tag::cpu_,
			     (A0)(A1),
			     (scalar_ < arithmetic_<A0> > )
			     (scalar_ < floating_<A1> > )
			     )
  {
    typedef A1 result_type;
    inline A1 operator()(A0 const& a0,A1 & a2) const
    {
      return impl::trig_base <A1,radian_tag,tag::not_simd_type>::sincosa(A1(a0),a2);
    }
  };

  /////////////////////////////////////////////////////////////////////////////
  // Implementation when type  is floating_
  /////////////////////////////////////////////////////////////////////////////
  NT2_FUNCTOR_IMPLEMENTATION(nt2::tag::sincos_, tag::cpu_,
			     (A0),
			     (scalar_ < arithmetic_<A0> > )
			     )
  {
    typedef typename boost::dispatch::meta::as_floating<A0>::type etype;
    typedef boost::fusion::tuple<etype, etype>           result_type;

    NT2_FUNCTOR_CALL(1)
      {
	result_type res;
	boost::fusion::at_c<0>(res) = impl::trig_base <etype,radian_tag
	  ,  tag::not_simd_type
	  >::sincosa( tofloat(a0)
		      , boost::fusion::at_c<1>(res)
		      );
	return res;
      }
  };

  /////////////////////////////////////////////////////////////////////////////
  // Implementation when type  is floating_
  /////////////////////////////////////////////////////////////////////////////

  NT2_FUNCTOR_IMPLEMENTATION(nt2::tag::sincos_, tag::cpu_,
			     (A0),
			     (scalar_ < floating_<A0> > )
			     )
  {
    typedef typename meta::strip<A0>::type                 etype;
    typedef boost::fusion::tuple<etype, etype>       result_type;

    NT2_FUNCTOR_CALL(1)
      {
	result_type res;
	boost::fusion::at_c<0>(res) = impl::trig_base < A0,radian_tag
	  ,  tag::not_simd_type
	  >::sincosa( a0
		      , boost::fusion::at_c<1>(res)
		      );
        return res;
      }
    };
} }
#endif


#define BOOST_TEST_MAIN
#include <boost/test/unit_test.hpp>

#include <tvgutil/numbers/RandomNumberGenerator.h>
using namespace tvgutil;

BOOST_AUTO_TEST_SUITE(test_RandomNumberGenerator)

BOOST_AUTO_TEST_CASE(generate_int_from_uniform_test)
{
  RandomNumberGenerator rng(1234);
  BOOST_CHECK_EQUAL(rng.generate_int_from_uniform(23,23), 23);
}

BOOST_AUTO_TEST_SUITE_END()


/*
 * Pointcloudtogrid.hpp
 *
 *  Created on: Aug 1, 2019
 *      Author: Peter XXC
 *	 Institute: ZJU, Robotics 104
 */

#pragma once

// Elevation Mapping
#include "elevation_mapping/ElevationMap.hpp"
#include "elevation_mapping/RobotMotionMapUpdater.hpp"
#include "elevation_mapping/sensor_processors/SensorProcessorBase.hpp"
#include "elevation_mapping/WeightedEmpiricalCumulativeDistributionFunction.hpp"

// Grid Map
#include <grid_map_msgs/GetGridMap.h>
#include <grid_map_msgs/ProcessFile.h>
#include "grid_map_core/TypeDefs.hpp"
#include "grid_map_core/SubmapGeometry.hpp"
#include "grid_map_core/BufferRegion.hpp"

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

// PCL
#include <pcl/conversions.h>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/PCLPointCloud2.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/filters/voxel_grid.h>

// ROS
#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
#include <message_filters/cache.h>
#include <message_filters/subscriber.h>
#include <tf/transform_listener.h>
#include <geometry_msgs/PoseWithCovarianceStamped.h>
#include <std_srvs/Empty.h>

// Boost
#include <boost/thread.hpp>

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

#define NOT_OCCUPIED 0
#define OCCUPIED 1
#define OUT_OF_SCOPE 2

using namespace grid_map;
using namespace pcl;

struct Cell_feature
{
    //short cell_position[2] = {-1, -1};
    float highest = 0.0;
    float lowest = 0.0;
    float h_var = 0.0;
    float l_var = 0.0;
    int flag = NOT_OCCUPIED;
};

namespace elevation_mapping {

class PointcloudToGrid
{

public:

    /*!
   * Constructor.
   */
    PointcloudToGrid();

    /*!
    * Destructor.
    */
    virtual ~PointcloudToGrid();

    /*!
    * Callback function for processing pointcloud to grid map.
    * @param pointCloud the point cloud to be processed.
    */

    void pointCloudProcess(const pcl::PointCloud<pcl::PointXYZRGB>::ConstPtr &InputCloud, Eigen::VectorXf& variances, struct Cell_feature cell_with_pos[]);

    /*!
    * Set the geometry of the grid map.
    * @param length the side lengths in x, and y-direction of the grid map [m].
    * @param resolution the cell size in [m/cell].
    */
    void setScale(const Length& length, const double resolution);

private:

    //! Side length of the map in x- and y-direction [m].
    Length length_;

    //! Map resolution in xy plane [m/cell].
    double resolution_;

    //! Map grid cell number.
    Index grid_num;

    //! Size of the buffer (rows and cols of the data structure).
    Size size_;

    //! cell struct.
    //Cell_feature *cell_with_pos;

};

}


#define METRICS_PRINT_RAW_COUNTS

#include <boost/range/iterator_range.hpp>
#include <gtest/gtest.h>

#include "../LSTMCausalityTagger.h"
#include "../Metrics.h"

using namespace std;
using boost::make_iterator_range;

constexpr unsigned CAUSE_INDEX = CausalityRelation::CAUSE;
constexpr unsigned EFFECT_INDEX = CausalityRelation::EFFECT;

TEST(MetricsTest, JaccardIndexWorks) {
  BecauseRelation::IndexList gold = {1, 2, 3, 4, 5, 6, 7};
  BecauseRelation::IndexList predicted = {1, 3, 5, 7, 9, 11};
  double jaccard = CausalityMetrics::CalculateJaccard(gold, predicted, 0);
  EXPECT_DOUBLE_EQ(4/9., jaccard);
  jaccard = CausalityMetrics::CalculateJaccard(gold, predicted, 2);
  EXPECT_DOUBLE_EQ(4/11., jaccard);
}

TEST(MetricsTest, F1Works) {
  unsigned tp = 1770;
  unsigned fp = 150;
  unsigned fn = 330;
  double f1 = ClassificationMetrics::CalculateF1(
      ClassificationMetrics::CalculatePrecision(tp, fp),
      ClassificationMetrics::CalculateRecall(tp, fn));
  EXPECT_DOUBLE_EQ(0.880597014925373, f1);
}

TEST(MetricsTest, F1WorksInMetricsClasses) {
  unsigned tp = 1770;
  unsigned fp = 150;
  unsigned fn = 330;

  ArgumentMetrics am(tp, fp, fn, 0, tp + fp + fn);
  EXPECT_DOUBLE_EQ(0.880597014925373, am.spans->GetF1());

  ClassificationMetrics cm(tp, fp, fn);
  EXPECT_DOUBLE_EQ(0.880597014925373, cm.GetF1());
}

struct DataMetricsTest: public ::testing::Test {
protected:
  static void SetUpCorpus(const string& data_path,
                          unique_ptr<BecauseOracleTransitionCorpus>* corpus,
                          vector<vector<CausalityRelation>>* relations) {
    corpus->reset(
        new BecauseOracleTransitionCorpus(vocab.get(), data_path, true, false));
    for (unsigned i = 0; i < (*corpus)->sentences.size(); ++i) {
      const lstm_parser::Sentence& sentence = (*corpus)->sentences[i];
      const vector<unsigned>& actions = (*corpus)->correct_act_sent[i];
      relations->push_back(LSTMCausalityTagger::Decode(sentence, actions,
                                                       false, true));
      /*
      cerr << sentence << endl;
      for (const auto& rel : relations->back()) {
        cerr << rel << endl;
      }
      //*/
    }
  }

  inline CausalityMetrics CompareCorpora(
      const BecauseOracleTransitionCorpus& corpus1,
      const BecauseOracleTransitionCorpus& corpus2,
      const vector<vector<CausalityRelation>>& rels1,
      const vector<vector<CausalityRelation>>& rels2,
      bool record_differences = false, double match_threshold = 1.0) {
    CausalityMetrics total_metrics;

    for (unsigned i = 0; i < corpus1.sentences.size(); ++i) {
      const lstm_parser::Sentence& sentence = corpus1.sentences[i];
      const vector<CausalityRelation>& original_rels = rels1[i];
      const vector<CausalityRelation>& modified_rels = rels2[i];

      SpanTokenFilter filter = {false, sentence, corpus1.pos_is_punct};
      GraphEnhancedParseTree pseudo_parse(sentence);
      CausalityMetrics sentence_metrics(original_rels, modified_rels, corpus1,
                                        pseudo_parse, filter, 0, {},
                                        record_differences, false,
                                        match_threshold);
      total_metrics += sentence_metrics;
    }

    return total_metrics;
  }

  inline CausalityMetrics CompareOriginalAndModified(
      bool record_differences = false, double match_threshold = 1.0) {
    return CompareCorpora(*original_corpus, *modified_corpus,
                          original_relations, modified_relations,
                          record_differences, match_threshold);
  }

  static void SetUpTestCase() {
    vocab.reset(new lstm_parser::CorpusVocabulary);
    SetUpCorpus("lstm-causality/tests/data/original", &original_corpus,
                &original_relations);
    SetUpCorpus("lstm-causality/tests/data/modified", &modified_corpus,
                &modified_relations);
  }

  static void TearDownTestCase() {
    original_corpus.release();
  }

  static unique_ptr<lstm_parser::CorpusVocabulary> vocab;
  static unique_ptr<BecauseOracleTransitionCorpus> original_corpus;
  static vector<vector<CausalityRelation>> original_relations;
  static unique_ptr<BecauseOracleTransitionCorpus> modified_corpus;
  static vector<vector<CausalityRelation>> modified_relations;
};

unique_ptr<lstm_parser::CorpusVocabulary> DataMetricsTest::vocab;
unique_ptr<BecauseOracleTransitionCorpus> DataMetricsTest::original_corpus;
vector<vector<CausalityRelation>> DataMetricsTest::original_relations;
unique_ptr<BecauseOracleTransitionCorpus> DataMetricsTest::modified_corpus;
vector<vector<CausalityRelation>> DataMetricsTest::modified_relations;

#define TEST_METRICS(calculated_metrics, correct_connective_metrics, \
                     correct_cause_metrics, correct_effect_metrics) \
  EXPECT_EQ(correct_connective_metrics, \
            *calculated_metrics.connective_metrics); \
  EXPECT_EQ(correct_cause_metrics, \
            *calculated_metrics.argument_metrics[CAUSE_INDEX]); \
  EXPECT_EQ(correct_effect_metrics, \
            *calculated_metrics.argument_metrics[EFFECT_INDEX]);

TEST_F(DataMetricsTest, SameAnnotationGetsPerfectScores) {
  CausalityMetrics self_metrics = CompareCorpora(
      *original_corpus, *original_corpus, original_relations,
      original_relations);

  ClassificationMetrics correct_connective_metrics(7, 0, 0);
  ArgumentMetrics correct_cause_metrics(6, 0, 0, 1, 7);  // 1 causeless instance
  ArgumentMetrics correct_effect_metrics(7, 0, 0, 1, 7);
  ASSERT_EQ(correct_connective_metrics, correct_connective_metrics);
  ASSERT_EQ(correct_cause_metrics, correct_cause_metrics);
  TEST_METRICS(self_metrics, correct_connective_metrics, correct_cause_metrics,
               correct_effect_metrics);
}

TEST_F(DataMetricsTest, ModifiedAnnotationsGivesLessPerfectScores) {
  CausalityMetrics compared_metrics = CompareOriginalAndModified();
  ClassificationMetrics correct_connective_metrics(5, 2, 2);
  // One cause changed -> arg FP+FN; one cause deleted -> arg FN. Plus the 2 FPs
  // and FNs from connectives.
  ArgumentMetrics correct_cause_metrics(2, 3, 4, (1 + 1 + 0 + 0 + 1) / 5.);
  ArgumentMetrics correct_effect_metrics(4, 3, 3, 33/35.);
  TEST_METRICS(compared_metrics, correct_connective_metrics,
               correct_cause_metrics, correct_effect_metrics);
}

TEST_F(DataMetricsTest, ModifiedAnnotationsPartialGivesBetterScores) {
  CausalityMetrics compared_metrics = CompareOriginalAndModified(false, 0.5);
  ClassificationMetrics correct_connective_metrics(6, 1, 1);
  // One cause changed -> arg FP+FN; one cause deleted -> arg FN. Plus the 1 FP
  // and FN from connectives.
  ArgumentMetrics correct_cause_metrics(3, 2, 3, (1 + 1 + 1 + 0 + 0 + 1) / 6.);
  ArgumentMetrics correct_effect_metrics(6, 1, 1, (5/7. + 5)/6);
  TEST_METRICS(compared_metrics, correct_connective_metrics,
               correct_cause_metrics, correct_effect_metrics);
}

TEST_F(DataMetricsTest, AddingMetricsWorks) {
  CausalityMetrics compared_metrics = CompareOriginalAndModified();
  CausalityMetrics tweaked_metrics = compared_metrics;
  tweaked_metrics.argument_metrics[CAUSE_INDEX]->jaccard_index = 0.3;
  tweaked_metrics.argument_metrics[EFFECT_INDEX]->jaccard_index = 1.0;
  CausalityMetrics summed_metrics = compared_metrics + tweaked_metrics;

  ClassificationMetrics correct_connective_metrics(10, 4, 4);
  ArgumentMetrics correct_cause_metrics(4, 6, 8, 0.45);
  ArgumentMetrics correct_effect_metrics(8, 6, 6, 34/35.);  // 34 = mean(33, 35)
  TEST_METRICS(summed_metrics, correct_connective_metrics,
               correct_cause_metrics, correct_effect_metrics);
}

TEST_F(DataMetricsTest, AveragingMetricsWorks) {
  CausalityMetrics compared_metrics = CompareOriginalAndModified();

  vector<CausalityMetrics> repeated = {compared_metrics, compared_metrics,
                                       compared_metrics};
  AveragedCausalityMetrics avged_same(make_iterator_range(repeated));
  TEST_METRICS(avged_same, *compared_metrics.connective_metrics,
               *compared_metrics.argument_metrics[CAUSE_INDEX],
               *compared_metrics.argument_metrics[EFFECT_INDEX]);

  CausalityMetrics self_metrics = CompareCorpora(
      *original_corpus, *original_corpus, original_relations,
      original_relations);
  vector<CausalityMetrics> to_average = {compared_metrics, self_metrics};
  AveragedCausalityMetrics averaged_metrics(make_iterator_range(to_average));
  AveragedCausalityMetrics correct_metrics(  // start off averaging nothing
      make_iterator_range(to_average.begin(), to_average.begin()));

  // Manually set everything for the correct averaged metrics.
  AveragedClassificationMetrics* correct_conn_metrics =
      static_cast<AveragedClassificationMetrics*>(
          correct_metrics.connective_metrics.get());
  double precision = ClassificationMetrics::CalculatePrecision(6, 1);
  double recall = ClassificationMetrics::CalculateRecall(6, 1);
  double f1 = ClassificationMetrics::CalculateF1(precision, recall);
  tie(correct_conn_metrics->tp, correct_conn_metrics->fp,
      correct_conn_metrics->fn, correct_conn_metrics->avg_accuracy,
      correct_conn_metrics->avg_precision,
      correct_conn_metrics->avg_recall, correct_conn_metrics->avg_f1) =
          make_tuple(6, 1, 1, nan(""), precision, recall, f1);

  AveragedArgumentMetrics* correct_cause_metrics =
      static_cast<AveragedArgumentMetrics*>(
          correct_metrics.argument_metrics[CAUSE_INDEX].get());
  AveragedClassificationMetrics *cause_spans =
      static_cast<AveragedClassificationMetrics *>(
          correct_cause_metrics->spans.get());
  tie(cause_spans->tp, cause_spans->fp, cause_spans->fn,
      correct_cause_metrics->jaccard_index) =
          make_tuple(4, 1.5, 2, (1 + 0.6) / 2);
  precision = (1 + 2/5.) / 2;
  recall = (1 + 2/6.) / 2;
  f1 = (1 + ClassificationMetrics::CalculateF1(2/5., 2/6.)) / 2;
  tie(cause_spans->avg_precision, cause_spans->avg_recall, cause_spans->avg_f1)
      = make_tuple(precision, recall, f1);

  AveragedArgumentMetrics* correct_effect_metrics =
      static_cast<AveragedArgumentMetrics*>(
          correct_metrics.argument_metrics[EFFECT_INDEX].get());
  AveragedClassificationMetrics *effect_spans =
      static_cast<AveragedClassificationMetrics *>(
          correct_effect_metrics->spans.get());
  tie(effect_spans->tp, effect_spans->fp, effect_spans->fn,
      correct_effect_metrics->jaccard_index) =
          make_tuple(5.5, 1.5, 1.5, 34 / 35.);
  double p_r_f1 = (4/7. + 1) / 2;
  tie(effect_spans->avg_precision, effect_spans->avg_recall,
      effect_spans->avg_f1) = make_tuple(p_r_f1, p_r_f1, p_r_f1);

  TEST_METRICS(averaged_metrics, *correct_conn_metrics,
               *correct_cause_metrics, *correct_effect_metrics);
}

TEST_F(DataMetricsTest, CorrectNumberOfDifferencesRecorded) {
  CausalityMetrics self_metrics = CompareCorpora(
      *original_corpus, *original_corpus, original_relations,
      original_relations, true);
  EXPECT_EQ(7, self_metrics.GetArgumentMatches().size());
  EXPECT_EQ(0, self_metrics.GetArgumentMismatches().size());
  EXPECT_EQ(0, self_metrics.GetFPs().size());
  EXPECT_EQ(0, self_metrics.GetFNs().size());

  CausalityMetrics modified_metrics = CompareCorpora(
      *original_corpus, *modified_corpus, original_relations,
      modified_relations, true);
  // 2 totally correct; 2 FPs; 2 FNs; 2 correct connectives with incorrect
  // arguments, with 3 incorrect arguments between them.
  EXPECT_EQ(2, modified_metrics.GetArgumentMatches().size());
  EXPECT_EQ(3, modified_metrics.GetArgumentMismatches().size());
  EXPECT_EQ(2, modified_metrics.GetFPs().size());
  EXPECT_EQ(2, modified_metrics.GetFNs().size());
}