#include "../sigma/params.h"
#include "../sigma/coinspend.h"
#include "../sigma/coin.h"
#include "../main.h"
#include "../secp256k1/include/Scalar.h"
#include "../sigma.h"
#include "./test_bitcoin.h"
#include "../wallet/wallet.h"

#include "test/fixtures.h"
#include "test/testutil.h"

#include <boost/test/unit_test.hpp>

#include <stdlib.h>

BOOST_FIXTURE_TEST_SUITE(sigma_state_tests, ZerocoinTestingSetup200)

static const uint256 txHash = uint256S("a64bf7b459d3bb09653e444d75a942e9848ed8e1f30e2890f999426ed6dd4a2c");

CBlockIndex CreateBlockIndex(int nHeight)
{
    CBlockIndex index;
    index.nHeight = nHeight;
    index.pprev = chainActive.Tip();
    index.phashBlock = new uint256();
    return index;
}

CBlock CreateBlockWithMints(const std::vector<sigma::PublicCoin> mints)
{
    CBlock block;
    block.sigmaTxInfo = std::make_shared<sigma::CSigmaTxInfo>();
    block.sigmaTxInfo->mints = mints;
    return block;
}

std::vector<sigma::PrivateCoin> generateCoins(
    const sigma::Params* params, int n, sigma::CoinDenomination denom)
{
    std::vector<sigma::PrivateCoin> privCoins;

    for(int i =0; i<n; i++)
        privCoins.push_back(sigma::PrivateCoin(params, denom));

    return privCoins;
}

std::vector<sigma::PublicCoin> getPubcoins(const std::vector<sigma::PrivateCoin> coins)
{
    std::vector<sigma::PublicCoin> pubCoins;

    BOOST_FOREACH(auto& coin, coins)
        pubCoins.push_back(coin.getPublicCoin());

    return pubCoins;
}

// Checking AddSpend
BOOST_AUTO_TEST_CASE(sigma_addspend)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    auto params = sigma::Params::get_default();

    const sigma::PrivateCoin privcoin(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin;
    pubcoin = privcoin.getPublicCoin();

    std::vector<sigma::PublicCoin> anonymity_set;
    anonymity_set.push_back(pubcoin);

    // Doesn't really matter what metadata we give here, it must pass.
    sigma::SpendMetaData metaData(0, uint256S("120"), uint256S("120"));

    sigma::CoinSpend coin(params, privcoin, anonymity_set, metaData, true);

    auto coinSerial = coin.getCoinSerialNumber();
    auto initSize = sigmaState->GetSpends().count(coinSerial);
    sigmaState->AddSpend(coinSerial, pubcoin.getDenomination(), 0);
    auto actSize = sigmaState->GetSpends().count(coinSerial);

    BOOST_CHECK_MESSAGE(initSize + 1 == actSize, "Serial was not added to usedCoinSerials.");
    sigmaState->Reset();
}

// Checking HasCoin when coin does not exist
BOOST_AUTO_TEST_CASE(sigma_hascoin_false)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    auto params = sigma::Params::get_default();

    const sigma::PrivateCoin privcoin(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin;
    pubcoin = privcoin.getPublicCoin();
    auto hasCoin = sigmaState->HasCoin(pubcoin);

    BOOST_CHECK_MESSAGE(!hasCoin, "The coin should not be in mintedPubCoins.");
    sigmaState->Reset();
}

// Checking HasCoin when coin exists
BOOST_AUTO_TEST_CASE(sigma_hascoin_true)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    CScript scriptPubKey2;
    auto params = sigma::Params::get_default();

    const sigma::PrivateCoin privcoin(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin;
    pubcoin = privcoin.getPublicCoin();
    CBlockIndex index = CreateBlockIndex(1);
    auto mintsBlock = CreateBlockWithMints({pubcoin});

    sigmaState->AddMintsToStateAndBlockIndex(&index, &mintsBlock);
    auto hasCoin = sigmaState->HasCoin(pubcoin);

    BOOST_CHECK_MESSAGE(hasCoin, "The coin should not be in mintedPubCoins.");
    sigmaState->Reset();
}

// Checking GetMintedCoinHeightAndId when coin exists
BOOST_AUTO_TEST_CASE(sigma_getmintcoinheightandid_true)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    CScript scriptPubKey2;
    auto params = sigma::Params::get_default();

    const sigma::PrivateCoin privcoin(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin;
    pubcoin = privcoin.getPublicCoin();
    CBlockIndex index = CreateBlockIndex(1);
    auto mintsBlock = CreateBlockWithMints({pubcoin});

    sigmaState->AddMintsToStateAndBlockIndex(&index, &mintsBlock);
    auto cnData = sigmaState->GetMintedCoinHeightAndId(pubcoin);

    BOOST_CHECK_MESSAGE(cnData.first == 1, "Unexpected minted coin height.");
    BOOST_CHECK_MESSAGE(cnData.second == 1, "Unexpected minted coin id.");
    sigmaState->Reset();
}

// Checking GetMintedCoinHeightAndId when coin does not exist
BOOST_AUTO_TEST_CASE(sigma_get_mintcoin_height_and_id_false)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    CScript scriptPubKey2;
    auto params = sigma::Params::get_default();

    const sigma::PrivateCoin privcoin(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin;
    pubcoin = privcoin.getPublicCoin();

    auto cnData = sigmaState->GetMintedCoinHeightAndId(pubcoin);
    BOOST_CHECK_MESSAGE(cnData.first == -1, "Unexpected minted coin height.");
    BOOST_CHECK_MESSAGE(cnData.second == -1, "Unexpected minted coin id.");

    sigmaState->Reset();
}

// Checking AddMint two times with same coin
BOOST_AUTO_TEST_CASE(sigma_addmint_double)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    CScript scriptPubKey2;
    auto params = sigma::Params::get_default();

    const sigma::PrivateCoin privcoin(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin;
    pubcoin = privcoin.getPublicCoin();
    CBlockIndex index = CreateBlockIndex(1);
    auto mintsBlock = CreateBlockWithMints({pubcoin});

    sigmaState->AddMintsToStateAndBlockIndex(&index, &mintsBlock);
    auto mintedPubCoin = sigmaState->GetMints();

    BOOST_CHECK_MESSAGE(mintedPubCoin.size() == 1,
        "Unexpected mintedPubCoin size after first call.");

    sigmaState->AddMintsToStateAndBlockIndex(&index, &mintsBlock);
    mintedPubCoin = sigmaState->GetMints();

    BOOST_CHECK_MESSAGE(mintedPubCoin.size() == 1,
         "Unexpected mintedPubCoin size after second call.");

    sigmaState->Reset();
}

// Checking AddMint two different coins on one block
BOOST_AUTO_TEST_CASE(sigma_addmint_two)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    CScript scriptPubKey2;
    auto params1 = sigma::Params::get_default();
    const sigma::PrivateCoin privcoin1(params1, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin1;
    pubcoin1 = privcoin1.getPublicCoin();

    auto params2 = sigma::Params::get_default();
    const sigma::PrivateCoin privcoin2(params2, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin2;
    pubcoin2 = privcoin2.getPublicCoin();

    auto mintsBlock = CreateBlockWithMints({pubcoin1, pubcoin2});

    CBlockIndex index = CreateBlockIndex(1);
    sigmaState->AddMintsToStateAndBlockIndex(&index, &mintsBlock);

    auto mintedPubCoin = sigmaState->GetMints();

    BOOST_CHECK_MESSAGE(mintedPubCoin.size() == 2, "Unexpected mintedPubCoin size.");

    sigmaState->Reset();
}

// Checking AddMint ZC_SPEND_V3_COINSPERID_LIMIT and check group id increase.
BOOST_AUTO_TEST_CASE(sigma_addmints_coinperid_limit)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    const sigma::CoinDenomination testDenomination = sigma::CoinDenomination::SIGMA_DENOM_0_05;
    const auto testDenomStr = sigma::DenominationToString(testDenomination);

    // To make sure have coin more than ZC_SPEND_V3_COINSPERID in first group.
    auto mintsPerBlock = 100;

    std::string strError;
    // Make sure that transactions get to mempool
    pwalletMain->SetBroadcastTransactions(true);

    // Create 400-200+1 = 201 new empty blocks. // consensus.nMintV3SigmaStartBlock = 400
    CreateAndProcessEmptyBlocks(201, scriptPubKey);

    // Generate mints ZC_SPEND_V3_COINSPERID - 1, last group ID should be 1.
    int allMints = 0;
    while (allMints < ZC_SPEND_V3_COINSPERID_LIMIT - 1) {
        auto mintThisBlock = std::min(ZC_SPEND_V3_COINSPERID_LIMIT - 1 - allMints, mintsPerBlock);
        allMints += mintThisBlock;

        BOOST_CHECK_MESSAGE(pwalletMain->CreateZerocoinMintModel(
            strError, {{testDenomStr, mintThisBlock}}, SIGMA), strError + " - Create Mint failed");

        BOOST_CHECK_MESSAGE(mempool.size() == 1, "Mint was not added to mempool");
        CreateAndProcessBlock({}, scriptPubKey);
        BOOST_CHECK_MESSAGE(mempool.size() == 0, "Mempool did not get empty.");
    }

    sigma::CSigmaState::SigmaCoinGroupInfo group1;
    sigmaState->GetCoinGroupInfo(testDenomination, 1, group1);
    BOOST_CHECK_EQUAL(group1.nCoins, ZC_SPEND_V3_COINSPERID_LIMIT - 1);
    BOOST_CHECK_EQUAL(sigmaState->GetLatestCoinIds().find(testDenomination)->second, 1);

    // Try to generate more coins to make exceed hardcap, new coins should be push to new group instead.
    auto exceedHardCapAmount = ZC_SPEND_V3_COINSPERID_LIMIT + 1;
    auto moreMintsToMakeExceedLimit =  exceedHardCapAmount - group1.nCoins;

    BOOST_CHECK_MESSAGE(pwalletMain->CreateZerocoinMintModel(
        strError, {{testDenomStr, moreMintsToMakeExceedLimit}}, SIGMA), strError + " - Create Mint failed");
    BOOST_CHECK_MESSAGE(mempool.size() == 1, "Mint was not added to mempool");
    CreateAndProcessBlock({}, scriptPubKey);
    BOOST_CHECK_MESSAGE(mempool.size() == 0, "Mempool did not get empty.");

    // New Mints should not be added to first group.
    sigmaState->GetCoinGroupInfo(testDenomination, 1, group1);
    BOOST_CHECK_EQUAL(group1.nCoins, ZC_SPEND_V3_COINSPERID_LIMIT - 1);

    // New Mints should be added to news group.
    BOOST_CHECK_EQUAL(sigmaState->GetLatestCoinIds().find(testDenomination)->second, 2);
    sigma::CSigmaState::SigmaCoinGroupInfo group2;
    sigmaState->GetCoinGroupInfo(testDenomination, 2, group2);
    BOOST_CHECK_EQUAL(group2.nCoins, moreMintsToMakeExceedLimit);

    // Remove last block, coin ID should be decrease back.
    DisconnectBlocks(1);

    BOOST_CHECK_EQUAL(sigmaState->GetLatestCoinIds().find(testDenomination)->second, 1);

    sigmaState->Reset();
}

// Checking RemoveSpendFromMempool, when coin is in mempool
BOOST_AUTO_TEST_CASE(sigma_remove_spend_from_mempool_coin_in)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    auto params = sigma::Params::get_default();

    const sigma::PrivateCoin privcoin(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin;
    pubcoin = privcoin.getPublicCoin();

    std::vector<sigma::PublicCoin> anonymity_set;
    anonymity_set.push_back(pubcoin);

    // Doesn't really matter what metadata we give here, it must pass.
    sigma::SpendMetaData metaData(0, uint256S("120"), uint256S("120"));

    sigma::CoinSpend coin(params, privcoin, anonymity_set, metaData, true);

    auto coinSerial = coin.getCoinSerialNumber();

    sigmaState->AddSpendToMempool(coinSerial, txHash);
    BOOST_CHECK_MESSAGE(sigmaState->GetMempoolCoinSerials().size() == 1,
      "Unexpected mempoolCoinSerials size after call AddSpendToMempool.");

    sigmaState->RemoveSpendFromMempool(coinSerial);
    BOOST_CHECK_MESSAGE(sigmaState->GetMempoolCoinSerials().size() == 0,
      "Unexpected mempoolCoinSerials size after call AddSpendToMempool.");
    sigmaState->Reset();
}

// Checking RemoveSpendFromMempool, when coin is not in mempool
BOOST_AUTO_TEST_CASE(sigma_remove_spend_from_mempool_coin_not_in)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    auto params = sigma::Params::get_default();

    const sigma::PrivateCoin privcoin(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin;
    pubcoin = privcoin.getPublicCoin();

    std::vector<sigma::PublicCoin> anonymity_set;
    anonymity_set.push_back(pubcoin);

    // Doesn't really matter what metadata we give here, it must pass.
    sigma::SpendMetaData metaData(0, uint256S("120"), uint256S("120"));

    sigma::CoinSpend coin(params, privcoin, anonymity_set, metaData, true);

    auto coinSerial = coin.getCoinSerialNumber();

    sigmaState->RemoveSpendFromMempool(coinSerial);
    BOOST_CHECK_MESSAGE(sigmaState->GetMempoolCoinSerials().size() == 0,
      "Unexpected mempoolCoinSerials size after call AddSpendToMempool.");
    sigmaState->Reset();
}

// Checking AddSpendToMempool, when coin was used (in usedCoinSerials)
BOOST_AUTO_TEST_CASE(sigma_addspend_to_mempool_coin_used)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    auto params = sigma::Params::get_default();

    const sigma::PrivateCoin privcoin(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin;
    pubcoin = privcoin.getPublicCoin();

    // Doesn't really matter what metadata we give here, it must pass.
    sigma::SpendMetaData metaData(0, uint256S("120"), uint256S("120"));

    std::vector<sigma::PublicCoin> anonymity_set;
    anonymity_set.push_back(pubcoin);
    sigma::CoinSpend coin(params, privcoin, anonymity_set, metaData, true);

    auto coinSerial = coin.getCoinSerialNumber();

    sigmaState->AddSpend(coinSerial, pubcoin.getDenomination(), 0);
    BOOST_CHECK_MESSAGE(sigmaState->GetMempoolCoinSerials().size() == 0,
      "Unexpected mempoolCoinSerials size before call AddSpendToMempool.");

    sigmaState->AddSpendToMempool(coinSerial, txHash);
    BOOST_CHECK_MESSAGE(sigmaState->GetMempoolCoinSerials().size() == 0,
      "Unexpected mempoolCoinSerials size after call AddSpendToMempool.");

    sigmaState->Reset();
}

// Checking AddSpendToMempool, when coin was not used
BOOST_AUTO_TEST_CASE(sigma_addspendtomempool)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    auto params = sigma::Params::get_default();

    const sigma::PrivateCoin privcoin(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin;
    pubcoin = privcoin.getPublicCoin();

    std::vector<sigma::PublicCoin> anonymity_set;
    anonymity_set.push_back(pubcoin);

    // Doesn't really matter what metadata we give here, it must pass.
    sigma::SpendMetaData metaData(0, uint256S("120"), uint256S("120"));

    sigma::CoinSpend coin(params, privcoin, anonymity_set, metaData, true);

    auto coinSerial = coin.getCoinSerialNumber();

    sigmaState->AddSpendToMempool(coinSerial, txHash);
    BOOST_CHECK_MESSAGE(sigmaState->GetMempoolCoinSerials().size() == 1,
      "Unexpected mempoolCoinSerials size after call AddSpendToMempool.");

    sigmaState->Reset();
}

// Checking AddSpendToMempool, when coin is already in mempool
BOOST_AUTO_TEST_CASE(sigma_addspendtomempool_coinin)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    auto params = sigma::Params::get_default();

    const sigma::PrivateCoin privcoin(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin;
    pubcoin = privcoin.getPublicCoin();

    std::vector<sigma::PublicCoin> anonymity_set;
    anonymity_set.push_back(pubcoin);

    // Doesn't really matter what metadata we give here, it must pass.
    sigma::SpendMetaData metaData(0, uint256S("120"), uint256S("120"));

    sigma::CoinSpend coin(params, privcoin, anonymity_set, metaData, true);

    auto coinSerial = coin.getCoinSerialNumber();

    sigmaState->AddSpendToMempool(coinSerial, txHash);

    BOOST_CHECK_MESSAGE(sigmaState->GetMempoolCoinSerials().size() == 1,
      "Unexpected mempoolCoinSerials size after first call AddSpendToMempool.");
    sigmaState->AddSpendToMempool(coinSerial, txHash);
    BOOST_CHECK_MESSAGE(sigmaState->GetMempoolCoinSerials().size() == 1,
      "Unexpected mempoolCoinSerials size after second call AddSpendToMempool.");

    sigmaState->Reset();
}

// Checking CanAddSpendToMempool, when coin is already in mempool
BOOST_AUTO_TEST_CASE(sigma_canaddspendtomempool_inmempool)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    auto params = sigma::Params::get_default();

    const sigma::PrivateCoin privcoin(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin;
    pubcoin = privcoin.getPublicCoin();

    std::vector<sigma::PublicCoin> anonymity_set;
    anonymity_set.push_back(pubcoin);

    // Doesn't really matter what metadata we give here, it must pass.
    sigma::SpendMetaData metaData(0, uint256S("120"), uint256S("120"));

    sigma::CoinSpend coin(params, privcoin, anonymity_set, metaData, true);

    auto coinSerial = coin.getCoinSerialNumber();

    BOOST_CHECK_MESSAGE(sigmaState->CanAddSpendToMempool(coinSerial),
      "CanAddSpendToMempool return false, which means coin already in use, but should not.");

    sigmaState->AddSpendToMempool(coinSerial, txHash);

    BOOST_CHECK_MESSAGE(!sigmaState->CanAddSpendToMempool(coinSerial),
      "CanAddSpendToMempool return true, which means coin not in use, but should be.");

    sigmaState->Reset();
}

// Checking CanAddSpendToMempool, when coin is already used
BOOST_AUTO_TEST_CASE(sigma_canaddspendtomempool_used)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    auto params = sigma::Params::get_default();

    const sigma::PrivateCoin privcoin(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin;
    pubcoin = privcoin.getPublicCoin();

    std::vector<sigma::PublicCoin> anonymity_set;
    anonymity_set.push_back(pubcoin);

    // Doesn't really matter what metadata we give here, it must pass.
    sigma::SpendMetaData metaData(0, uint256S("120"), uint256S("120"));

    sigma::CoinSpend coin(params, privcoin, anonymity_set, metaData, true);

    auto coinSerial = coin.getCoinSerialNumber();

    sigmaState->AddSpend(coinSerial, pubcoin.getDenomination(), 0);

    BOOST_CHECK_MESSAGE(!sigmaState->CanAddSpendToMempool(coinSerial),
      "CanAddSpendToMempool return true, which means coin not in use, but should be.");

    sigmaState->Reset();
}

// Checking Reset
BOOST_AUTO_TEST_CASE(sigma_reset)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    CScript scriptPubKey2;
    auto params = sigma::Params::get_default();

    const sigma::PrivateCoin privcoin(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin;
    pubcoin = privcoin.getPublicCoin();
    CBlockIndex index = CreateBlockIndex(1);

    // Let's add data to zerocoinstate before reset

    auto mintsBlock = CreateBlockWithMints({pubcoin});
    sigmaState->AddMintsToStateAndBlockIndex(&index, &mintsBlock);

    BOOST_CHECK_MESSAGE(sigmaState->GetMints().size() == 1,
      "Unexpected mintedPubCoin size before reset.");
    BOOST_CHECK_MESSAGE(sigmaState->GetCoinGroups().size() == 1,
      "Unexpected coinGroups size before reset.");

    std::vector<sigma::PublicCoin> anonymity_set;
    anonymity_set.push_back(pubcoin);

    // Doesn't really matter what metadata we give here, it must pass.
    sigma::SpendMetaData metaData(0, uint256S("120"), uint256S("120"));

    sigma::CoinSpend coin(params, privcoin, anonymity_set, metaData, true);

    auto coinSerial = coin.getCoinSerialNumber();

    sigmaState->AddSpendToMempool(coinSerial, txHash);

    BOOST_CHECK_MESSAGE(sigmaState->GetMempoolCoinSerials().size() == 1,
      "Unexpected mempoolCoinSerials size before reset.");

    sigmaState->AddSpend(coinSerial, pubcoin.getDenomination(), 0);

    BOOST_CHECK_MESSAGE(sigmaState->GetSpends().size() == 1,
      "Unexpected usedCoinSerials size before reset.");
    BOOST_CHECK_MESSAGE(sigmaState->GetLatestCoinIds().size() == 1,
      "Unexpected mintedPubCoin size before reset.");

    sigmaState->Reset();

    BOOST_CHECK_MESSAGE(sigmaState->GetMints().size() == 0,
      "Unexpected mintedPubCoin size after reset.");
    BOOST_CHECK_MESSAGE(sigmaState->GetCoinGroups().size() == 0,
      "Unexpected coinGroups size after reset.");
    BOOST_CHECK_MESSAGE(sigmaState->GetSpends().size() == 0,
      "Unexpected usedCoinSerials size after reset.");
    BOOST_CHECK_MESSAGE(sigmaState->GetLatestCoinIds().size() == 0,
      "Unexpected mintedPubCoin size after reset.");
    BOOST_CHECK_MESSAGE(sigmaState->GetMempoolCoinSerials().size() == 0,
      "Unexpected mintedPubCoin size after reset.");
}

// Checking GetCoinGroupInfo, when coingroup is exist
BOOST_AUTO_TEST_CASE(sigma_getcoingroupinfo_existing)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    CScript scriptPubKey2;
    auto params = sigma::Params::get_default();

    const sigma::PrivateCoin privcoin(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin;
    pubcoin = privcoin.getPublicCoin();
    CBlockIndex index = CreateBlockIndex(1);
    auto mintsBlock = CreateBlockWithMints({pubcoin});

    sigmaState->AddMintsToStateAndBlockIndex(&index, &mintsBlock);
    auto mintedPubCoin = sigmaState->GetMints();

    BOOST_CHECK_MESSAGE(mintedPubCoin.size() == 1,
        "Unexpected mintedPubCoin size after first call.");

    sigma::CSigmaState::SigmaCoinGroupInfo result;
    sigmaState->GetCoinGroupInfo(pubcoin.getDenomination(), 1, result);
    BOOST_CHECK_MESSAGE(result.nCoins == 1,
        "Unexpected number of coins in group.");
    BOOST_CHECK_MESSAGE(result.firstBlock->mintedPubCoins.size() == index.mintedPubCoins.size(),
        "Unexpected first block index for Group info.");
    BOOST_CHECK_MESSAGE(result.lastBlock->mintedPubCoins.size() == index.mintedPubCoins.size(),
        "Unexpected last block index for Group info.");

    sigmaState->Reset();
}

// Checking GetCoinGroupInfo, when coingroup is not minted
BOOST_AUTO_TEST_CASE(sigma_getcoingroupinfo_not_minted)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    CScript scriptPubKey2;
    auto params = sigma::Params::get_default();

    const sigma::PrivateCoin privcoin(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin;
    pubcoin = privcoin.getPublicCoin();
    CBlockIndex index = CreateBlockIndex(1);

    sigma::CSigmaState::SigmaCoinGroupInfo result;
    sigmaState->GetCoinGroupInfo(pubcoin.getDenomination(), 1, result);
    BOOST_CHECK_MESSAGE(result.nCoins == 0,
        "Unexpected number of coins in group.");

    sigmaState->Reset();
}

BOOST_AUTO_TEST_CASE(zerocoin_sigma_addblock_nonexist_index)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    CScript scriptPubKey2;
    auto params = sigma::Params::get_default();

    const sigma::PrivateCoin privcoin(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin;
    pubcoin = privcoin.getPublicCoin();
    CBlockIndex index = CreateBlockIndex(1);

	sigmaState->AddBlock(&index);
	BOOST_CHECK_MESSAGE(sigmaState->GetMints().size() == 0,
	  "Unexpected mintedPubCoins size, add new block without minted txs.");

	BOOST_CHECK_MESSAGE(sigmaState->GetSpends().size() == 0,
	  "Unexpected usedCoinSerials size, add new block without spend txs.");

    sigmaState->Reset();
}

BOOST_AUTO_TEST_CASE(zerocoin_sigma_addblock_minted_spend)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    CScript scriptPubKey2;
    auto params = sigma::Params::get_default();

    const sigma::PrivateCoin privcoin1(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin1;
    pubcoin1 = privcoin1.getPublicCoin();

    const sigma::PrivateCoin privcoin2(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin2;
    pubcoin2 = privcoin2.getPublicCoin();

    CBlockIndex index = CreateBlockIndex(1);
    std::pair<sigma::CoinDenomination, int> denomination1Group1(
        sigma::CoinDenomination::SIGMA_DENOM_1,1);

	index.sigmaMintedPubCoins[denomination1Group1].push_back(pubcoin1);
	index.sigmaMintedPubCoins[denomination1Group1].push_back(pubcoin2);

	sigmaState->AddBlock(&index);
	BOOST_CHECK_MESSAGE(sigmaState->GetMints().size() == 2,
	  "Unexpected mintedPubCoins size, add new block with 2 minted txs.");

	BOOST_CHECK_MESSAGE(sigmaState->GetSpends().size() == 0,
	  "Unexpected usedCoinSerials size, add new block without spend txs.");

	// spend
    std::vector<sigma::PublicCoin> anonymity_set;
    anonymity_set.push_back(pubcoin1);
    anonymity_set.push_back(pubcoin2);

    // Doesn't really matter what metadata we give here, it must pass.
    sigma::SpendMetaData metaData(0, uint256S("120"), uint256S("120"));

    sigma::CoinSpend coinSpend(params,privcoin1,anonymity_set, metaData, true);

	auto spendSerial = coinSpend.getCoinSerialNumber();

    CBlockIndex index2 = CreateBlockIndex(2);
	index2.sigmaSpentSerials.clear();
	index2.sigmaSpentSerials.insert(std::make_pair(spendSerial, sigma::CSpendCoinInfo::make(coinSpend.getDenomination(), 0)));
	sigmaState->AddBlock(&index2);
	BOOST_CHECK_MESSAGE(sigmaState->GetMints().size() == 2,
	  "Unexpected mintedPubCoins size, add new block without additional minted.");

	BOOST_CHECK_MESSAGE(sigmaState->GetSpends().size() == 1,
	  "Unexpected usedCoinSerials size, add new block with 1 spend txs.");

    // minted more coin
    const sigma::PrivateCoin privcoin3(params, sigma::CoinDenomination::SIGMA_DENOM_1);
    sigma::PublicCoin pubcoin3;
    pubcoin3 = privcoin3.getPublicCoin();
    CBlockIndex index3 = CreateBlockIndex(3);

    index3.sigmaMintedPubCoins[denomination1Group1].push_back(pubcoin3);
    sigmaState->AddBlock(&index3);
    BOOST_CHECK_MESSAGE(sigmaState->GetMints().size() == 3,
	  "Unexpected mintedPubCoins size, add new block with one more minted.");

	BOOST_CHECK_MESSAGE(sigmaState->GetSpends().size() == 1,
	  "Unexpected usedCoinSerials size, add new block without new spend");

    sigmaState->Reset();
}

BOOST_AUTO_TEST_CASE(zerocoin_sigma_removeblock_remove)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    CScript scriptPubKey2;
    auto params = sigma::Params::get_default();

    // add index 1 with 10 minted
    auto coins = generateCoins(params, 10, sigma::CoinDenomination::SIGMA_DENOM_1);
    auto pubCoins = getPubcoins(coins);

    auto index1 = CreateBlockIndex(1);
    std::pair<sigma::CoinDenomination, int> denomination1Group1(sigma::CoinDenomination::SIGMA_DENOM_1, 1);
    index1.sigmaMintedPubCoins[denomination1Group1] = pubCoins;

    // add index 2 with 10 minted and 1 spend
    auto coins2 = generateCoins(params,10, sigma::CoinDenomination::SIGMA_DENOM_1);
    auto pubCoins2 = getPubcoins(coins2);

    auto index2 = CreateBlockIndex(2);
    std::pair<sigma::CoinDenomination, int> denomination1Group2(sigma::CoinDenomination::SIGMA_DENOM_1, 2);
    index2.sigmaMintedPubCoins[denomination1Group2] = pubCoins2;

    // Doesn't really matter what metadata we give here, it must pass.
    sigma::SpendMetaData metaData(0, uint256S("120"), uint256S("120"));

    sigma::CoinSpend coinSpend(params, coins[0], pubCoins, metaData, true);

    index2.sigmaSpentSerials.clear();
    index2.sigmaSpentSerials.insert(std::make_pair(coinSpend.getCoinSerialNumber(), sigma::CSpendCoinInfo::make(coinSpend.getDenomination(), 0)));

    sigmaState->AddBlock(&index1);
    sigmaState->AddBlock(&index2);

    BOOST_CHECK_MESSAGE(sigmaState->GetLatestCoinIds().find(sigma::CoinDenomination::SIGMA_DENOM_1)->second == 2,
      "Unexpected lastestcoinId");
    BOOST_CHECK_MESSAGE(sigmaState->HasCoin(pubCoins2[0]),
      "Coin isn't in state before remove index 2");

    // remove one
    sigmaState->RemoveBlock(&index2);
    BOOST_CHECK_MESSAGE(sigmaState->GetMints().size() == 10,
	  "Unexpected mintedPubCoins size, remove index contain 10 minteds.");

    BOOST_CHECK_MESSAGE(sigmaState->GetSpends().size() == 0,
      "Unexpected usedCoinSerials size, remove index contain 1 spend.");

    BOOST_CHECK_MESSAGE(sigmaState->GetLatestCoinIds().find(sigma::CoinDenomination::SIGMA_DENOM_1)->second == 1,
      "Unexpected lastestcoinId");

    BOOST_CHECK_MESSAGE(sigmaState->HasCoin(pubCoins[0]),
      "Coin isn't in state before remove index 1");
    BOOST_CHECK_MESSAGE(!sigmaState->HasCoin(pubCoins2[0]),
      "Coin is in state after remove index 2");

    // remove all
    sigmaState->RemoveBlock(&index1);
    BOOST_CHECK_MESSAGE(sigmaState->GetMints().size() == 0,
	  "Unexpected mintedPubCoins size, remove index contain 10 minteds.");

	BOOST_CHECK_MESSAGE(sigmaState->GetSpends().size() == 0,
	  "Unexpected usedCoinSerials size, remove index contain no spend.");

    BOOST_CHECK_MESSAGE(sigmaState->GetLatestCoinIds().find(sigma::CoinDenomination::SIGMA_DENOM_1) == sigmaState->GetLatestCoinIds().end(),
      "Unexpected lastestcoinId remove all");

    BOOST_CHECK_MESSAGE(!sigmaState->HasCoin(pubCoins[0]),
      "Coin is in state after remove index 1");

    sigmaState->Reset();
}

BOOST_AUTO_TEST_CASE(getmempoolconflictingtxhash_added_no)
{
    sigma::CSigmaState state;

    secp_primitives::Scalar serial;
    serial.randomize();

    BOOST_CHECK(state.GetMempoolConflictingTxHash(serial) == uint256());
}

BOOST_AUTO_TEST_CASE(getmempoolconflictingtxhash_added_yes)
{
    sigma::CSigmaState state;

    secp_primitives::Scalar serial;
    serial.randomize();

    auto txid = uint256S("c8cdacf6b51275a3de9496073c75708abde26cb2f6cb164e0a1a0ed942c3c6e7");

    BOOST_CHECK(state.AddSpendToMempool(serial, txid));
    BOOST_CHECK(state.GetMempoolConflictingTxHash(serial) == txid);
}

BOOST_AUTO_TEST_CASE(zerocoingetspendserialnumberv3_valid_tx_valid_vin)
{
    // setup
    CScript scriptPubKey2;
    auto params = sigma::Params::get_default();

    // spend
    auto coins = generateCoins(params, 10, sigma::CoinDenomination::SIGMA_DENOM_0_1);
    auto pubCoins = getPubcoins(coins);

    // Doesn't really matter what metadata we give here, it must pass.
    sigma::SpendMetaData metaData(0, uint256S("120"), uint256S("120"));

    sigma::CoinSpend coinSpend(params, coins[0], pubCoins, metaData, true);

    CDataStream serializedCoinSpend(SER_NETWORK, PROTOCOL_VERSION);
    serializedCoinSpend << coinSpend;

    // create tx and vin

    CTxIn newTxIn;
    newTxIn.scriptSig = CScript();
    newTxIn.prevout.SetNull();

    CScript tmp = CScript() << OP_SIGMASPEND;
    tmp.insert(tmp.end(),serializedCoinSpend.begin(),serializedCoinSpend.end());

    newTxIn.scriptSig.assign(tmp.begin(),tmp.end());

    CMutableTransaction newtx;
    newtx.vin.clear();
    newtx.vout.clear();

    newtx.vin.push_back(newTxIn);

    // mock vout
    CTxOut newTxOut(0, CScript());
    newtx.vout.push_back(newTxOut);

    CTransaction ctx(newtx);

    // check
    BOOST_CHECK_MESSAGE(sigma::GetSigmaSpendSerialNumber(ctx,newTxIn) != Scalar(uint64_t(0)),
      "Expect serial number, got 0");

    // add more spend vin
    sigma::CoinSpend coinSpend2(params, coins[1], pubCoins, metaData, true);

    CDataStream serializedCoinSpend2(SER_NETWORK, PROTOCOL_VERSION);
    serializedCoinSpend2 << coinSpend2;

    CTxIn newTxIn2;
    newTxIn2.scriptSig = CScript();
    newTxIn2.prevout.SetNull();

    CScript tmp2 = CScript() << OP_SIGMASPEND;
    tmp2.insert(tmp2.end(), serializedCoinSpend2.begin(), serializedCoinSpend2.end());

    newTxIn2.scriptSig.assign(tmp2.begin(), tmp2.end());
    newtx.vin.push_back(newTxIn);

    CTransaction ctx2(newtx);

    // check
    BOOST_CHECK_MESSAGE(sigma::GetSigmaSpendSerialNumber(ctx2,newTxIn2) != Scalar(uint64_t(0)),
      "2 vin, Expect serial number, got 0");

    // not allow unspend vin
    // add unspend vin
    CTxIn newTxVin3;
    newTxVin3.scriptSig = CScript();
    newTxVin3.prevout.SetNull();

    CScript tmp3 = CScript() << OP_RETURN;
    newTxVin3.scriptSig.assign(tmp3.begin(), tmp3.end());

    newtx.vin.push_back(newTxVin3);

    CTransaction ctx3(newtx);

    BOOST_CHECK_MESSAGE(sigma::GetSigmaSpendSerialNumber(ctx3, newTxVin3) == Scalar(uint64_t(0)),
      "Expect 0 got serial");
}

BOOST_AUTO_TEST_CASE(zerocoingetspendserialnumberv3_invalid_script)
{
    // setup
    CScript scriptPubKey2;
    auto params = sigma::Params::get_default();

    // spend
    auto coins = generateCoins(params, 10, sigma::CoinDenomination::SIGMA_DENOM_0_1);
    auto pubCoins = getPubcoins(coins);

    // Doesn't really matter what metadata we give here, it must pass.
    sigma::SpendMetaData metaData(0, uint256S("120"), uint256S("120"));

    sigma::CoinSpend coinSpend(params, coins[0], pubCoins, metaData, true);

    CDataStream serializedCoinSpend(SER_NETWORK, PROTOCOL_VERSION);
    serializedCoinSpend << coinSpend;

    // create tx and vin

    CTxIn newTxIn;
    newTxIn.scriptSig = CScript();
    newTxIn.prevout.SetNull();

    CScript tmp = CScript() << OP_SIGMASPEND;
    auto itr = serializedCoinSpend.begin();
    // ignore first byte to make it invalid
    tmp.insert(tmp.end(),++itr,serializedCoinSpend.end());

    CMutableTransaction newtx;

    newtx.vin.push_back(newTxIn);

    // mock vout
    CTxOut newTxOut(0, CScript());
    newtx.vout.push_back(newTxOut);

    CTransaction ctx(newtx);

    // check
    BOOST_CHECK_MESSAGE(sigma::GetSigmaSpendSerialNumber(ctx,newTxIn) == Scalar(uint64_t(0)),
      "Expect 0 got serial, Wrong script");
}

BOOST_AUTO_TEST_CASE(getzerocoinstate_not_null)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    BOOST_CHECK_MESSAGE(sigmaState != NULL, "sigma::CSigmaState::GetState() returned null");
}

BOOST_AUTO_TEST_CASE(sigma_build_state)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    auto params = sigma::Params::get_default();

    CBlockIndex index0 = CreateBlockIndex(0);
    chainActive.SetTip(&index0);

    CBlockIndex index1 = CreateBlockIndex(1);

    // add index 1 with 10 SIGMA_DENOM_1
    auto coins = generateCoins(params, 10, sigma::CoinDenomination::SIGMA_DENOM_1);
    auto pubCoins = getPubcoins(coins);
    std::pair<sigma::CoinDenomination, int> denomination1Group1(sigma::CoinDenomination::SIGMA_DENOM_1, 1);
    std::pair<sigma::CoinDenomination, int> denomination10Group1(sigma::CoinDenomination::SIGMA_DENOM_10, 1);

    index1.sigmaMintedPubCoins[denomination1Group1] = pubCoins;

    chainActive.SetTip(&index1);

    CBlockIndex index2 = CreateBlockIndex(2);
    // add index 2 with 1 DENOMINATION_1  mints and 1 spend
    auto coins2 = generateCoins(params, 1, sigma::CoinDenomination::SIGMA_DENOM_1);
    auto pubCoins2 = getPubcoins(coins2);
    auto coins3 = generateCoins(params, 1, sigma::CoinDenomination::SIGMA_DENOM_10);
    auto pubCoins3 = getPubcoins(coins3);

    // mock spend
    secp_primitives::Scalar serial;
    serial.randomize();

    index2.sigmaSpentSerials.insert(std::make_pair(serial, sigma::CSpendCoinInfo::make(sigma::CoinDenomination::SIGMA_DENOM_1, 0)));

    index2.sigmaMintedPubCoins[denomination1Group1] = pubCoins2;
    index2.sigmaMintedPubCoins[denomination10Group1] = pubCoins3;

    chainActive.SetTip(&index2);

    for(int i =3 ;i<=100;i++){
		CBlockIndex index = CreateBlockIndex(i);
        chainActive.SetTip(&index);
	}

    sigma::BuildSigmaStateFromIndex(&chainActive);

    // check group
    sigma::CSigmaState::SigmaCoinGroupInfo group;
    sigmaState->GetCoinGroupInfo(sigma::CoinDenomination::SIGMA_DENOM_1, 1, group);
    BOOST_CHECK_MESSAGE(group.firstBlock->nHeight == index1.nHeight, "Expect firstBlock == index1");
    BOOST_CHECK_MESSAGE(group.lastBlock->nHeight == index2.nHeight, "Expect lastBlock == index2");
    BOOST_CHECK_MESSAGE(group.nCoins == 11, "Expect nCoins == 11");

    sigma::CSigmaState::SigmaCoinGroupInfo group2;
    sigmaState->GetCoinGroupInfo(sigma::CoinDenomination::SIGMA_DENOM_10, 1, group2);
    BOOST_CHECK_MESSAGE(group2.firstBlock->nHeight == index2.nHeight, "Expect firstBlock == index2");
    BOOST_CHECK_MESSAGE(group2.lastBlock->nHeight == index2.nHeight, "Expect lastBlock == index2");
    BOOST_CHECK_MESSAGE(group2.nCoins == 1, "Expect nCoins == 1");

    // check serial
    secp_primitives::Scalar notFoundSerial;
    notFoundSerial.randomize();
    BOOST_CHECK_MESSAGE(sigmaState->IsUsedCoinSerial(serial), "Expect found serial");
    BOOST_CHECK_MESSAGE(!sigmaState->IsUsedCoinSerial(notFoundSerial), "Expect not found serial");

    // has coin
    auto notFoundCoins = generateCoins(params, 10, sigma::CoinDenomination::SIGMA_DENOM_0_1);
    auto notFoundPubCoins = getPubcoins(notFoundCoins);
    BOOST_CHECK_MESSAGE(sigmaState->HasCoin(pubCoins[0]), "Expect found pubcoin");
    BOOST_CHECK_MESSAGE(sigmaState->HasCoin(pubCoins3[0]), "Expect found pubcoin");
    BOOST_CHECK_MESSAGE(!sigmaState->HasCoin(notFoundPubCoins[0]), "Expect not found pubcoin");

    // get mint coin heigh and id
    std::pair<int,int> groupHeightAndID = sigmaState->GetMintedCoinHeightAndId(pubCoins2[0]);
    BOOST_CHECK_MESSAGE(groupHeightAndID.first == index2.nHeight, "Expect pubcoin on index2");
    BOOST_CHECK_MESSAGE(groupHeightAndID.second == 1, "Expect pubcoin id == 1");

    std::pair<int,int> notFoundGroupHeightAndID = sigmaState->GetMintedCoinHeightAndId(notFoundPubCoins[0]);
    BOOST_CHECK_MESSAGE(notFoundGroupHeightAndID == std::make_pair(-1,-1),"Expect not found return -1,-1");

    sigmaState->Reset();
}

BOOST_AUTO_TEST_CASE(sigma_build_state_no_sigma)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    auto params = sigma::Params::get_default();

    std::vector<CBlockIndex> indexs;
    indexs.resize(101);
    for(int i =0 ;i<=100;i++){
		CBlockIndex index = CreateBlockIndex(i);
        chainActive.SetTip(&index);
        indexs[i] = index;
	}

    sigma::BuildSigmaStateFromIndex(&chainActive);

    // check group
    sigma::CSigmaState::SigmaCoinGroupInfo group;
    bool found = sigmaState->GetCoinGroupInfo(sigma::CoinDenomination::SIGMA_DENOM_1, 1, group);
    BOOST_CHECK_MESSAGE(!found, "Expect group not found");

    // check serial
    secp_primitives::Scalar notFoundSerial;
    notFoundSerial.randomize();
    BOOST_CHECK_MESSAGE(!sigmaState->IsUsedCoinSerial(notFoundSerial), "Expect not found serial");

    // has coin
    auto notFoundCoins = generateCoins(params,10, sigma::CoinDenomination::SIGMA_DENOM_0_1);
    auto notFoundPubCoins = getPubcoins(notFoundCoins);
    BOOST_CHECK_MESSAGE(!sigmaState->HasCoin(notFoundPubCoins[0]), "Expect not found pubcoin");

    // get mint coin heigh and id
    std::pair<int,int> notFoundGroupHeightAndID = sigmaState->GetMintedCoinHeightAndId(notFoundPubCoins[0]);
    BOOST_CHECK_MESSAGE(notFoundGroupHeightAndID == std::make_pair(-1,-1),"Expect not found return -1,-1");

    sigmaState->Reset();
}

BOOST_AUTO_TEST_CASE(sigma_getcoinsetforspend)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();
    sigma::Params* params = sigma::Params::get_default();
    int nextIndex = 0;
    std::vector<CBlockIndex> indexes;
    indexes.resize(101);

    // nextIndex = 0
    indexes[nextIndex] = CreateBlockIndex(nextIndex);
    chainActive.SetTip(&indexes[nextIndex]);

    // nextIndex = 1
    nextIndex++;
    indexes[nextIndex] = CreateBlockIndex(nextIndex);

    std::pair<sigma::CoinDenomination, int> denomination1Group1(sigma::CoinDenomination::SIGMA_DENOM_1, 1);
    std::pair<sigma::CoinDenomination, int> denomination10Group1(sigma::CoinDenomination::SIGMA_DENOM_10, 1);

    // add index 1 with 10 SIGMA_DENOM_1
    auto coins = generateCoins(params, 10, sigma::CoinDenomination::SIGMA_DENOM_1);
    auto pubCoins = getPubcoins(coins);

    // add index 2 with 1 DENOM_1  mints and 1 spend
    auto coins2 = generateCoins(params, 1, sigma::CoinDenomination::SIGMA_DENOM_1);
    auto pubCoins2 = getPubcoins(coins2);
    auto coins3 = generateCoins(params, 5, sigma::CoinDenomination::SIGMA_DENOM_10);
    auto pubCoins3 = getPubcoins(coins3);

    indexes[nextIndex].sigmaMintedPubCoins[denomination1Group1] = pubCoins;
    chainActive.SetTip(&indexes[nextIndex]);

    nextIndex++;
    indexes[nextIndex] = CreateBlockIndex(nextIndex);

    // mock spend
    secp_primitives::Scalar serial;
    serial.randomize();

    indexes[nextIndex].sigmaSpentSerials.insert(std::make_pair(serial, sigma::CSpendCoinInfo::make(sigma::CoinDenomination::SIGMA_DENOM_1, 0)));
    indexes[nextIndex].sigmaMintedPubCoins[denomination1Group1] = pubCoins2;
    indexes[nextIndex].sigmaMintedPubCoins[denomination10Group1] = pubCoins3;

    chainActive.SetTip(&indexes[nextIndex]);

    // nextIndex = 3
    nextIndex++;
    for( ; nextIndex<=100; nextIndex++){
		indexes[nextIndex] = CreateBlockIndex(nextIndex);
        chainActive.SetTip(&indexes[nextIndex]);
	}

    sigma::BuildSigmaStateFromIndex(&chainActive);

    uint256 blockHash_out;
    uint256 blockHash_empty;

    std::vector<sigma::PublicCoin> coins_out3;
    // maxheight < blockheight
    auto coins_amount = sigmaState->GetCoinSetForSpend(&chainActive, 1,
    sigma::CoinDenomination::SIGMA_DENOM_10, 1, blockHash_out, coins_out3);
    BOOST_CHECK_MESSAGE(coins_amount == 0, "Unexpected Coins amount for spend, should be 0.");
    BOOST_CHECK_MESSAGE(coins_out3.size() == 0, "Unexpected coins out, should be 0.");
    BOOST_CHECK_MESSAGE(blockHash_out == blockHash_empty , "Unexpected blockhash for small height.");

    std::vector<sigma::PublicCoin> coins_out1;
    // maxheight > blockheight
    coins_amount = sigmaState->GetCoinSetForSpend(&chainActive, nextIndex,
    sigma::CoinDenomination::SIGMA_DENOM_10, 1, blockHash_out, coins_out1);
    BOOST_CHECK_MESSAGE(coins_amount == 5, "Unexpected Coins amount for spend, should be 5.");
    BOOST_CHECK_MESSAGE(coins_out1 == pubCoins3, "Unexpected coins out for denom 10.");
    BOOST_CHECK_MESSAGE(blockHash_out == indexes[2].GetBlockHash(), "Unexpected blockhash for denom 10.");

    std::vector<sigma::PublicCoin> coins_out2;
    // maxheight > blockheight another denom
    coins_amount = sigmaState->GetCoinSetForSpend(&chainActive, nextIndex,
    sigma::CoinDenomination::SIGMA_DENOM_1, 1, blockHash_out, coins_out2);
    BOOST_CHECK_MESSAGE(coins_amount == 11, "Unexpected Coins amount for spend, should be 11.");
    BOOST_CHECK_MESSAGE(coins_out2.size() == pubCoins2.size() + pubCoins.size(), "Unexpected coins out for denom 1.");
    BOOST_CHECK_MESSAGE(blockHash_out == indexes[2].GetBlockHash(), "Unexpected blockhash for denom 1.");

    sigmaState->Reset();
}

namespace {
    Scalar generateSpend(sigma::CoinDenomination denom) {
        auto params = sigma::Params::get_default();

        const sigma::PrivateCoin privcoin(params, denom);
        sigma::PublicCoin pubcoin;
        pubcoin = privcoin.getPublicCoin();

        std::vector<sigma::PublicCoin> anonymity_set;
        anonymity_set.push_back(pubcoin);

        // Doesn't really matter what metadata we give here, it must pass.
        sigma::SpendMetaData metaData(0, uint256S("120"), uint256S("120"));

        sigma::CoinSpend coin(params, privcoin, anonymity_set, metaData, true);

        return coin.getCoinSerialNumber();
    }

    CBlock generateMint(sigma::CoinDenomination denom) {
        auto params = sigma::Params::get_default();

        const sigma::PrivateCoin privcoin(params, denom);

        return CreateBlockWithMints({privcoin.getPublicCoin()});
    }
}

BOOST_AUTO_TEST_CASE(sigma_surge_detection_positive)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();

    CBlockIndex index = CreateBlockIndex(1);
    CBlock mintsBlock = generateMint(sigma::CoinDenomination::SIGMA_DENOM_1);
    sigmaState->AddMintsToStateAndBlockIndex(&index, &mintsBlock);
    BOOST_CHECK(sigmaState->IsSurgeConditionDetected() == false);

    sigmaState->AddSpend(generateSpend(sigma::CoinDenomination::SIGMA_DENOM_1), sigma::CoinDenomination::SIGMA_DENOM_1, 1);
    BOOST_CHECK(sigmaState->IsSurgeConditionDetected() == false);

    sigmaState->Reset();
}

BOOST_AUTO_TEST_CASE(sigma_surge_detection_reset)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();

    sigmaState->Reset();
    BOOST_CHECK(sigmaState->IsSurgeConditionDetected() == false);

    sigmaState->AddSpend(generateSpend(sigma::CoinDenomination::SIGMA_DENOM_1), sigma::CoinDenomination::SIGMA_DENOM_1, 1);
    BOOST_CHECK(sigmaState->IsSurgeConditionDetected() == true);

    CBlockIndex index = CreateBlockIndex(1);
    CBlock mintsBlock = generateMint(sigma::CoinDenomination::SIGMA_DENOM_1);
    sigmaState->AddMintsToStateAndBlockIndex(&index, &mintsBlock);
    BOOST_CHECK(sigmaState->IsSurgeConditionDetected() == false);

    sigmaState->Reset();
}

// Check that failure in any [denom, group] makes it failed for all other
// combinations of [denom,group]
BOOST_AUTO_TEST_CASE(sigma_surge_detection_failure_anywhere)
{
    sigma::CSigmaState *sigmaState = sigma::CSigmaState::GetState();

    sigmaState->Reset();
    BOOST_CHECK(sigmaState->IsSurgeConditionDetected() == false);

    sigmaState->AddSpend(generateSpend(sigma::CoinDenomination::SIGMA_DENOM_1), sigma::CoinDenomination::SIGMA_DENOM_100, 1);
    BOOST_CHECK(sigmaState->IsSurgeConditionDetected() == true);

    CBlockIndex index = CreateBlockIndex(1);
    CBlock mintsBlock1 = generateMint(sigma::CoinDenomination::SIGMA_DENOM_1);
    sigmaState->AddMintsToStateAndBlockIndex(&index, &mintsBlock1);
    BOOST_CHECK(sigmaState->IsSurgeConditionDetected() == true);

    sigmaState->AddSpend(generateSpend(sigma::CoinDenomination::SIGMA_DENOM_1), sigma::CoinDenomination::SIGMA_DENOM_1, 1);
    BOOST_CHECK(sigmaState->IsSurgeConditionDetected() == true);

    CBlock mintsBlock100 = generateMint(sigma::CoinDenomination::SIGMA_DENOM_100);
    sigmaState->AddMintsToStateAndBlockIndex(&index, &mintsBlock100);
    BOOST_CHECK(sigmaState->IsSurgeConditionDetected() == false);

    sigmaState->Reset();
}

BOOST_AUTO_TEST_SUITE_END()


/*
 * Copyright (c) 2009 Carnegie Mellon University.
 *     All rights reserved.
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing,
 *  software distributed under the License is distributed on an "AS
 *  IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 *  express or implied.  See the License for the specific language
 *  governing permissions and limitations under the License.
 *
 * For more about this software visit:
 *
 *      http://www.graphlab.ml.cmu.edu
 *
 */

/**
 *
 * \brief This application performs MAP inference on Markov Nets
 * provided in standard UAI file format via Dual-Decomposition.
 *
 *
 *  \authors Dhruv Batra, André Martins, Aroma Mahendru
 */

#ifndef _AD3_QP_HPP_
#define _AD3_QP_HPP_

#include <Eigen/Eigenvalues>
#include <math.h>
#include <limits>
#include "dd_grlab.hpp"

#define NEARLY_ZERO_TOL(a,tol) (((a)<=(tol)) && ((a)>=(-(tol))))
#define NEARLY_EQ_TOL(a,b,tol) (((a)-(b))*((a)-(b))<=(tol))
#define num_max_iterations_QP_ 10

/**
 * \brief ADMM vertex program  general implements Alternating direction dual
 * decomposition for general dense factors.
 */

struct admm_vertex_program_general:public admm_vertex_program {

/**
 * \brief Maximize returns the maximum value and configuration with reference to
 * input additional and variable log potentials. addtional log potential corresponds
 * factor potentials and variable potential corresponds to sum of lagrange
 * multipliers and unary potentials divided by degree of the unary vertex.
 */

void Maximize(vertex_type& vertex, vec additional_log_potentials, vec variable_log_potentials,
                Configuration &configuration,
                double *value) {

    vector <Configuration> states(vertex.data().nvars,-1);
    int best = -1;
    *value = -1e12;
    for (int index = 0;
         index < additional_log_potentials.size();
         ++index) {
      double score = additional_log_potentials[index];
      get_configuration_states(vertex,index, &states);
      int offset = 0;
      for (int i = 0; i < vertex.data().nvars; ++i) {
        score += variable_log_potentials[offset+states[i]];
        offset = vertex.data().cards[i];

      }

      if (configuration < 0 || score > *value) {
        configuration = index;
        *value = score;
      }
    }
    assert(configuration >= 0);

  }

 void DeleteConfiguration(Configuration &configuration) {
    configuration = -1;
  }

/**
 * \brief InvertAfterInsertion function is used to invert the matrix. Used in solveQP
 */

 bool InvertAfterInsertion(vertex_type& vertex, vector <double> & inverse_A_,
        const vector<Configuration> &active_set, const Configuration &inserted_element) {

  vector<double> inverse_A = inverse_A_;
  int size_A = active_set.size() + 1;
  vector<double> r(size_A);

  r[0] = 1.0;
  for (int i = 0; i < active_set.size(); ++i) {
    // Count how many variable values the new assignment
    // have in common with the i-th assignment.
    int num_common_values = CountCommonValues(vertex, active_set[i], inserted_element);
    r[i+1] = static_cast<double>(num_common_values);
  }

  double r0 = static_cast<double>(CountCommonValues(vertex,
      inserted_element, inserted_element));
  double s = r0;
  for (int i = 0; i < size_A; ++i) {
    if (r[i] == 0.0) continue;
    s -= r[i] * r[i] * inverse_A[i * size_A + i];
    for (int j = i+1; j < size_A; ++j) {
      if (r[j] == 0.0) continue;
      s -= 2 * r[i] * r[j] * inverse_A[i * size_A + j];
    }
  }

    if (NEARLY_ZERO_TOL(s, 1e-9)) {
         if (opts.verbose> 2) {
      cout << "Warning: updated matrix will become singular after insertion."
           << endl;
    }
    return false;
  }

  double invs = 1.0 / s;
  vector<double> d(size_A, 0.0);
  for (int i = 0; i < size_A; ++i) {
    if (r[i] == 0.0) continue;
    for (int j = 0; j < size_A; ++j) {
      d[j] += inverse_A[i * size_A + j] * r[i];
    }
  }

  int size_A_after = size_A + 1;
  inverse_A_.resize(size_A_after * size_A_after);
  for (int i = 0; i < size_A; ++i) {
    for (int j = 0; j < size_A; ++j) {
      inverse_A_[i * size_A_after + j] = inverse_A[i * size_A + j] +
          invs * d[i] * d[j];
    }
    inverse_A_[i * size_A_after + size_A] = -invs * d[i];
    inverse_A_[size_A * size_A_after + i] = -invs * d[i];
  }
  inverse_A_[size_A * size_A_after + size_A] = invs;

  return true;
}

/**
 * \brief InvertAfterRemoval function is used to invert the matrix. Used in solveQP
 */
void InvertAfterRemoval(vector <double> &inverse_A_,const vector<Configuration> &active_set,
                                       int removed_index) {
  vector<double> inverse_A = inverse_A_;
  int size_A = active_set.size() + 1;
  vector<double> r(size_A);

  ++removed_index; // Index in A has an offset of 1.
  double invs = inverse_A[removed_index * size_A + removed_index];
  assert(!NEARLY_ZERO_TOL(invs, 1e-12));
  double s = 1.0 / invs;
  vector<double> d(size_A - 1, 0.0);
  int k = 0;
  for (int i = 0; i < size_A; ++i) {
    if (i == removed_index) continue;
    d[k] = -s * inverse_A[removed_index * size_A + i];
    ++k;
  }

  int size_A_after = size_A - 1;
  inverse_A_.resize(size_A_after * size_A_after);
  k = 0;
  for (int i = 0; i < size_A; ++i) {
    if (i == removed_index) continue;
    int l = 0;
    for (int j = 0; j < size_A; ++j) {
      if (j == removed_index) continue;
      inverse_A_[k * size_A_after + l] = inverse_A[i * size_A + j] -
          invs * d[k] * d[l];
      ++l;
    }
    ++k;
  }
}

/**
 * \brief ComputeActiveSetSimilarities computes Mnz'*Mnz. Used in solveQP
 */
void ComputeActiveSetSimilarities(vertex_type& vertex,
    const vector<Configuration> &active_set,
    vector<double> *similarities) {
  int size = active_set.size();

  // Compute similarity matrix.
  similarities->resize(size * size);
  (*similarities)[0] = 0.0;
  for (int i = 0; i < active_set.size(); ++i) {
    (*similarities)[i*size + i] = static_cast<double>(
        CountCommonValues(vertex,active_set[i], active_set[i]) );
    for (int j = i+1; j < active_set.size(); ++j) {
      // Count how many variable values the i-th and j-th
      // assignments have in common.
      int num_common_values = CountCommonValues(vertex,active_set[i], active_set[j]);
      (*similarities)[i*size + j] = num_common_values;
      (*similarities)[j*size + i] = num_common_values;
    }
  }
}

/**
 * \brief  ComputeMarginalsFromSparseDistribution computes marginalvalues for unary
 * factor from given factor distribution.
 */

void ComputeMarginalsFromSparseDistribution( vertex_type& vertex,
    const vector<Configuration> &active_set,
    const vector<double> &distribution,
    vec  &variable_posteriors,
    vec &additional_posteriors) {
    variable_posteriors.setZero();
    additional_posteriors.setZero();
    for (int i = 0; i < active_set.size(); ++i) {
    UpdateMarginalsFromConfiguration(vertex,active_set[i],
                                       distribution[i],
                                       variable_posteriors,
                                       additional_posteriors);
    }
  }

   // Given a configuration with a probability (weight),
  // increment the vectors of variable and additional posteriors.
  void UpdateMarginalsFromConfiguration(vertex_type& vertex,
    const Configuration &configuration,
    double weight,
    vec &variable_posteriors,
    vec &additional_posteriors) {

     vector <Configuration> states(vertex.data().nvars, -1);
     get_configuration_states(vertex, configuration, &states);

            int offset = 0;

            for (int k = 0; k < vertex.data().nvars; ++k)
            {   variable_posteriors[offset + states[k]] += weight;
                offset += vertex.data().cards[k];
            }
    additional_posteriors[configuration] += weight;

  }
  // Count how many common values two configurations have.
  int CountCommonValues(vertex_type& vertex,Configuration configuration1,
                        Configuration configuration2) {

    //assert(states1->size() == states2->size());
    int count = 0;
    vector <Configuration> states1(vertex.data().nvars, -1);
    vector <Configuration> states2(vertex.data().nvars, -1);
    get_configuration_states(vertex, configuration1, &states1);
    get_configuration_states(vertex, configuration2, &states2);
    for(int i = 0; i< vertex.data().nvars; i++)
    {  if (states1[i] == states2[i])
      { count++;} }
    return count;
  }

/**
 * \brief Evaluate returns the maximum value  with reference to
 * input additional and variable log potentials and configuration. addtional
 * log potential corresponds factor potentials and variable potential corresponds
 * to sum of lagrange  * multipliers and unary potentials divided by degree of
 * the unary vertex.
 */

void Evaluate(vertex_type& vertex, vec additional_log_potentials, vec variable_log_potentials,
                const Configuration configuration,
                double *value) {

    vector<Configuration> states(vertex.data().nvars, -1);
    get_configuration_states(vertex, configuration, &states);
    *value = 0.0;
    int offset = 0;
    for (int i = 0;i<vertex.data().nvars; ++i) {
      *value += variable_log_potentials[offset + states[i]];
      offset = vertex.data().cards[i];
    }
    *value += additional_log_potentials[configuration];
  }

  void EigenDecompose(vector<double> *similarities,
                            vector<double> *eigenvalues) {

  int size = sqrt(similarities->size());

  Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> es;
  Eigen::MatrixXd sim(size, size);
  int t = 0;
  for (int i = 0; i < size; ++i) {
    for (int j = 0; j < size; ++j) {
      sim(i, j) = (*similarities)[t];
      ++t;
    }
  }
  es.compute(sim);
  const Eigen::VectorXd &eigvals = es.eigenvalues();
  eigenvalues->resize(size);
  for (int i = 0; i < size; ++i) {
    (*eigenvalues)[i] = eigvals[i];
  }
  const Eigen::MatrixXd &eigvectors = es.eigenvectors().transpose();
  t = 0;
  for (int i = 0; i < size; ++i) {
    for (int j = 0; j < size; ++j) {
      (*similarities)[t] = eigvectors(i, j);
      ++t;
    }
  }

}

// Function to solve each quadratic sub problem for dense factors.
// It uses active set method. Caching is deactivated
// TODO: Activate caching feature

void SolveQP_dense(vertex_type& vertex,const gather_type& total,
            vec& variable_posteriors, vec& additional_posteriors) {
   vertex_data& vdata = vertex.data();
   //cout<<"loop 1 in solveQP .."<<endl;
   vec additional_log_potentials = vdata.potentials;
   vec variable_log_potentials = total.neighbor_distribution + total.messages;
   vector <Configuration> active_set_;
   vector<double> distribution_;
   vector<double> inverse_A_;
  // Initialize the active set.

   if (active_set_.size() == 0) {
    variable_posteriors.resize(variable_log_potentials.size());
    additional_posteriors.resize(additional_log_potentials.size());
    distribution_.clear();
    // Initialize by solving the LP, discarding the quadratic
    // term.
    Configuration configuration = -1;
    double value;
    Maximize(vertex, additional_log_potentials, variable_log_potentials,
             configuration,
             &value);
    active_set_.push_back(configuration);
    distribution_.push_back(1.0);

    // Initialize inv(A) as [-M,1;1,0].
    inverse_A_.resize(4);
    inverse_A_[0] = static_cast<double>(
        -CountCommonValues(vertex,configuration, configuration));
    inverse_A_[1] = 1;
    inverse_A_[2] = 1;
    inverse_A_[3] = 0;
  }

  bool changed_active_set = true;
  vector<double> z;
  int num_max_iterations = num_max_iterations_QP_;
  double tau = 0;
  for (int iter = 0; iter < num_max_iterations; ++iter) {
    bool same_as_before = true;
    bool unbounded = false;
    if (changed_active_set) {
      // Recompute vector b.
      vector<double> b(active_set_.size() + 1, 0.0);
      b[0] = 1.0;
      for (int i = 0; i < active_set_.size(); ++i) {
        const Configuration &configuration = active_set_[i];
        double score;
        Evaluate(vertex, additional_log_potentials, variable_log_potentials,
                 configuration,
                 &score);
        b[i+1] = score;
      }
      // Solve the system Az = b.
      z.resize(active_set_.size());
      int size_A = active_set_.size() + 1;
      for (int i = 0; i < active_set_.size(); ++i) {
        z[i] = 0.0;
        for (int j = 0; j < size_A; ++j) {
          z[i] += inverse_A_[(i+1) * size_A + j] * b[j];
        }
      }
      tau = 0.0;
      for (int j = 0; j < size_A; ++j) {
        tau += inverse_A_[j] * b[j];
      }

      same_as_before = false;
    }

    if (same_as_before) {
      // Compute the variable marginals from the full distribution
      // stored in z.
      ComputeMarginalsFromSparseDistribution(vertex, active_set_,
                                             z,
                                             variable_posteriors,
                                             additional_posteriors);
      // Get the most violated constraint
      // (by calling the black box that computes the MAP).
      vec scores = variable_log_potentials;
      for (int i = 0; i < scores.size(); ++i) {
        scores[i] -= variable_posteriors[i];
      }
      Configuration configuration = -1;
      double value = 0.0;

      Maximize(vertex,
                additional_log_potentials, scores,
               configuration,
               &value);
      double very_small_threshold = 1e-9;
      if (value <= tau + very_small_threshold) { // value <= tau.
        // We have found the solution;
        // the distribution, active set, and inv(A) are cached for the next round.
        DeleteConfiguration(configuration);
        return;
      } else {
        for (int k = 0; k < active_set_.size(); ++k) {
          // This is expensive and should just be a sanity check.
          // However, in practice, numerical issues force an already existing
          // configuration to try to be added. Therefore, we always check
          // if a configuration already exists before inserting it.
          // If it does, that means the active set method converged to a
          // solution (but numerical issues had prevented us to see it.)
          if (active_set_[k] == configuration) {
            if (opts.verbose > 2) {
              cout << "Warning: value - tau = "
                   << value - tau << " " << value << " " << tau
                   << endl;
            }
            // We have found the solution;
            // the distribution, active set, and inv(A)
            // are cached for the next round.
            DeleteConfiguration(configuration);

            // Just in case, clean the cache.
            // This may prevent eventual numerical problems in the future.
            for (int j = 0; j < active_set_.size(); ++j) {
              if (j == k) continue; // This configuration was deleted already.
              DeleteConfiguration(active_set_[j]);
            }
            active_set_.clear();
            inverse_A_.clear();
            distribution_.clear();

            // Return.
            return;
          }
        }
        z.push_back(0.0);
        distribution_ = z;

        // Update inv(A).
        bool singular = !InvertAfterInsertion(vertex, inverse_A_, active_set_, configuration);
        if (singular) {
          // If adding a new configuration causes the matrix to be singular,
          // don't just add it. Instead, look for a configuration in the null
          // space and remove it before inserting the new one.
          // Right now, if more than one such configuration exists, we just
          // remove the first one we find. There's a chance this could cause
          // some cyclic behaviour. If that is the case, we should randomize
          // this choice.
          // Note: This step is expensive and requires an eigendecomposition.
          // TODO: I think there is a graph interpretation for this problem.
          // Maybe some specialized graph algorithm is cheaper than doing
          // the eigendecomposition.
          vector<double> similarities(active_set_.size() * active_set_.size());
          ComputeActiveSetSimilarities(vertex, active_set_, &similarities);

          //cout<<"compute active similarities in solveQP .."<<endl;
          vector<double> padded_similarities((active_set_.size()+2) *
                                             (active_set_.size()+2), 1.0);
          for (int i = 0; i < active_set_.size(); ++i) {
            for (int j = 0; j < active_set_.size(); ++j) {
              padded_similarities[(i+1)*(active_set_.size()+2) + (j+1)] =
                  similarities[i*active_set_.size() + j];
            }
          }
          padded_similarities[0] = 0.0;
          for (int i = 0; i < active_set_.size(); ++i) {
            double value = static_cast<double>(
                CountCommonValues(vertex, configuration, active_set_[i]));
            padded_similarities[(i+1)*(active_set_.size()+2) +
                                (active_set_.size()+1)] = value;
            padded_similarities[(active_set_.size()+1)*(active_set_.size()+2) +
                                (i+1)] = value;
          }
          double value = static_cast<double>(
              CountCommonValues(vertex, configuration, configuration));
          padded_similarities[(active_set_.size()+1)*(active_set_.size()+2) +
                              (active_set_.size()+1)] = value;

          vector<double> eigenvalues(active_set_.size()+2);
          EigenDecompose(&padded_similarities, &eigenvalues);
          int zero_eigenvalue = -1;
          for (int i = 0; i < active_set_.size()+2; ++i) {
            if (NEARLY_EQ_TOL(eigenvalues[i], 0.0, 1e-9)) {
              if (zero_eigenvalue >= 0) {
                // If this happens, something failed. Maybe a numerical problem
                // may cause this. In that case, just give up, clean the cache
                // and return. Hopefully the next iteration will fix it.
                cout << "Multiple zero eigenvalues: "
                     << eigenvalues[zero_eigenvalue] << " and "
                     << eigenvalues[i] << endl;
                cout << "Warning: Giving up." << endl;
                // Clean the cache.
                for (int j = 0; j < active_set_.size(); ++j) {
                  DeleteConfiguration(active_set_[j]);
                }
                active_set_.clear();
                inverse_A_.clear();
                distribution_.clear();
                return;
              }
              zero_eigenvalue = i;
            }
          }
          assert(zero_eigenvalue >= 0);
          vector<int> configurations_to_remove;
          for (int j = 1; j < active_set_.size()+1; ++j) {
            double value = padded_similarities[zero_eigenvalue*(active_set_.size()+2) + j];
            if (!NEARLY_EQ_TOL(value, 0.0, 1e-9)) {
              configurations_to_remove.push_back(j-1);
            }
          }
          if (opts.verbose > 2) {
            cout << "Pick a configuration to remove (" << configurations_to_remove.size()
                 << " out of " << active_set_.size() << ")." << endl;
          }

          assert(configurations_to_remove.size() >= 1);
          int j = configurations_to_remove[0];

          // Update inv(A).
          InvertAfterRemoval(inverse_A_, active_set_, j);

          // Remove blocking constraint from the active set.
          DeleteConfiguration(active_set_[j]); // Delete configutation.
          active_set_.erase(active_set_.begin() + j);

          singular = !InvertAfterInsertion(vertex, inverse_A_, active_set_, configuration);
          assert(!singular);
        }

        // Insert configuration to active set.
        if (opts.verbose > 2) {
          cout << "Inserted one element to the active set (iteration "
               << iter << ")." << endl;
        }
        active_set_.push_back(configuration);
        changed_active_set = true;
      }
    } else {
      // Solution has changed from the previous iteration.
      // Look for blocking constraints.
      int blocking = -1;
      bool exist_blocking = false;
      double alpha = 1.0;
      for (int i = 0; i < active_set_.size(); ++i) {
        assert(distribution_[i] >= -1e-12);
        if (z[i] >= distribution_[i]) continue;
        if (z[i] < 0) exist_blocking = true;
        double tmp = distribution_[i] / (distribution_[i] - z[i]);
        if (blocking < 0 || tmp < alpha) {
          alpha = tmp;
          blocking = i;
        }
      }

      if (!exist_blocking) {
        // No blocking constraints.
        assert(!unbounded);
        distribution_ = z;
        alpha = 1.0;
        changed_active_set = false;
      } else {
        if (alpha > 1.0 && !unbounded) alpha = 1.0;
        // Interpolate between factor_posteriors_[i] and z.
        if (alpha == 1.0) {
          distribution_ = z;
        } else {
          for (int i = 0; i < active_set_.size(); ++i) {
            z[i] = (1 - alpha) * distribution_[i] + alpha * z[i];
            distribution_[i] = z[i];
          }
        }

        // Update inv(A).
        InvertAfterRemoval(inverse_A_, active_set_, blocking);

        // Remove blocking constraint from the active set.
        if (opts.verbose > 2) {
          cout << "Removed one element to the active set (iteration "
               << iter << ")." << endl;
        }

        DeleteConfiguration(active_set_[blocking]); // Delete configutation.
        active_set_.erase(active_set_.begin() + blocking);

        z.erase(z.begin() + blocking);
        distribution_.erase(distribution_.begin() + blocking);
        changed_active_set = true;
        for (int i = 0; i < distribution_.size(); ++i) {
          assert(distribution_[i] > -1e-16);
        }
      }
    }
  }

  // Maximum number of iterations reached.
  // Return the best existing solution by computing the variable marginals
  // from the full distribution stored in z.
  //assert(false);
  ComputeMarginalsFromSparseDistribution(vertex, active_set_,
                                         z,
                                         variable_posteriors,
                                         additional_posteriors);
  }

  void InsertionSort(pair<double, int> arr[], int length) {
  int i, j;
  pair<double, int> tmp;

  for (i = 1; i < length; i++) {
    j = i;
    while (j > 0 && arr[j - 1].first > arr[j].first) {
      tmp = arr[j];
      arr[j] = arr[j - 1];
      arr[j - 1] = tmp;
      j--;
    }
  }
}

  int project_onto_budget_constraint_cached(vec& x,
                                          int d,
                                          double budget,
                                          vector<pair<double,int> >& y) {
  int j, k, l, level;
  double s = 0.0;
  double tau = 0.0, tightsum;
  double left, right = -std::numeric_limits<double>::infinity();

  // Load x into a reordered y (the reordering is cached).
  if (y.size() != d) {
    y.resize(d);
    for (j = 0; j < d; j++) {
      s -= x[j];
      y[j].first = -x[j];
      y[j].second = j;
    }
    sort(y.begin(), y.end());
  } else {
    for (j = 0; j < d; j++) {
      s -= x[j];
      y[j].first = -x[y[j].second];
    }
    // If reordering is cached, use a sorting algorithm
    // which is fast when the vector is almost sorted.
    InsertionSort(&y[0], d);
  }

  tightsum = s;
  s += budget;

  k = l = level = 0;
  bool found = false;
  double val_a, val_b;
  while (k < d && l < d) {
    if (level != 0) {
      tau = (s - tightsum) / static_cast<double>(level);
    }
    if (k < d) val_a = y[k].first;
    val_b = 1.0 + y[l].first;
    left = right;
    if (k == d || val_b <= val_a) {
      right = val_b;
    } else {
      right = val_a;
    }
    if ((level == 0 && s == tightsum) || (level != 0 && tau <= right)) {
      // Found the right split-point!
      found = true;
      break;
    }
    if (k == d || val_b <= val_a) {
      tightsum += val_b;
      --level;
      ++l;
    } else {
      tightsum -= val_a;
      ++level;
      ++k;
    }
  }

  if (!found) {
    left = right;
    right = std::numeric_limits<double>::infinity();
  }

  for (j = 0; j < d; j++) {
    if (-x[j] >= right) {
      x[j] = 0.0;
    } else if (1.0 - x[j] <= left) {
      x[j] = 1.0;
    } else {
      x[j] += tau;
    }
  }

  return 0;
}

  // Solve the QP subproblem for budget factor.
  // TODO Enable caching
void SolveQP_budget(vertex_type& vertex,const gather_type& total,
            vec& variable_posteriors, vec& additional_posteriors){

  vertex_data& vdata =  vertex.data();
  vec variable_log_potentials = total.neighbor_distribution + total.messages;
  vector<pair<double,int> > last_sort_;
  for (int f = 0; f < variable_log_potentials.size(); ++f) {
    variable_posteriors[f] = variable_log_potentials[f];
    if (variable_posteriors[f] < 0.0) {
      variable_posteriors[f] = 0.0;
    } else if (variable_posteriors[f] > 1.0) {
      variable_posteriors[f] = 1.0;
    }
  }

  double s = 0.0;
  for (int f = 0; f < vdata.nvars; ++f) {
    s += variable_posteriors[f];
  }

  if (s > static_cast<double>(vdata.budget)) {
    for (int f = 0; f < variable_log_potentials.size(); ++f) {
      variable_posteriors[f] = variable_log_potentials[f];
    }
    project_onto_budget_constraint_cached(variable_posteriors,
                                          variable_log_potentials.size(),
                                          static_cast<double>(vdata.budget),
                                          last_sort_);
  }

}

// Finds best configuration ofr budget factors
void SolveMAP_budget(vertex_type& vertex,const gather_type& total,
            vec& variable_posteriors, vec& additional_posteriors, double& value) {

   vertex_data& vdata = vertex.data();
  // Create a local copy of the log potentials.
   vec log_potentials(total.messages);
  double valaux;

  value = 0.0;

  int num_active = 0;
  double sum = 0.0;

  for (int f = 0; f < vdata.nvars; ++f) {
    valaux = log_potentials[f];
    if (valaux < 0.0) {
      variable_posteriors[f] =  0.0;
    } else {
      sum += valaux;
      variable_posteriors[f] = 1.0;
    }
    ++num_active;
  }
  if (num_active > vdata.budget) {
    vector<pair<double,int> > scores(vdata.nvars);
    for (int f = 0; f < vdata.nvars; ++f) {
      scores[f].first = -log_potentials[f];
      scores[f].second = f;
    }

    sort(scores.begin(), scores.end());
    num_active = 0;
    sum = 0.0;
    for (int k = 0; k < vdata.budget; ++k) {
      valaux = -scores[k].first;
      if (valaux < 0.0) break;
      int f = scores[k].second;
      variable_posteriors[f] = 1.0;
      sum += valaux;
      ++num_active;
    }

    for (int k = num_active; k < vdata.nvars; ++k) {
      int f = scores[k].second;
      variable_posteriors[f] = 0.0;
    }
   // cout<<"third inner loop ..."<<endl;
  }

  value += sum;

}

// Finds best configuration for dense factors
void SolveMAP_dense(vertex_type& vertex,const gather_type& total,
            vec& variable_posteriors, vec& additional_posteriors, double& value){
     vertex_data& vdata = vertex.data();
     vec beliefs = vdata.potentials;
     int num_configurations = vdata.potentials.size();
     for (int index_configuration = 0;
           index_configuration < num_configurations;
            ++index_configuration) {
        vector<int> states(vdata.nvars, -1);
        get_configuration_states(vertex, index_configuration, &states);
        int offset = 0;
        for (int k = 0; k < vdata.nvars; ++k) {
             beliefs[index_configuration] += total.messages[offset + states[k]];
             offset += vdata.cards[k];}
    }

        value = beliefs.maxCoeff();

 }

// Finds beliefs using dense and budget factors
void compute_beliefs(vertex_type& vertex,const gather_type& total,
            vec& variable_posteriors, vec& additional_posteriors){
switch(vertex.data().factor_type){

case 0: SolveQP_dense(vertex,total, variable_posteriors, additional_posteriors);
        break;
case 1: SolveQP_budget(vertex,total, variable_posteriors, additional_posteriors);

}
};

// General solvMAP function
void SolveMAP(vertex_type& vertex,const gather_type& total,
            vec& variable_posteriors, vec& additional_posteriors, double& value){
switch(vertex.data().factor_type){

case 0: SolveMAP_dense(vertex,total, variable_posteriors, additional_posteriors, value);
        break;
case 1: SolveMAP_budget(vertex,total, variable_posteriors, additional_posteriors, value);
  }
 };

};

#endif


//
// Original authors: Andrew Dowsey <andrew.dowsey@manchester.ac.uk>
//
// Copyright 2014 CADET Bioinformatics Laboratory
//                University of Manchester, UK
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#pragma once

#include <string>
#include <vector>
#include <fstream>
#include <boost/filesystem.hpp>
#include <boost/timer/timer.hpp>

#include "aux_Eigen.hpp"

using namespace std;
using namespace boost;

struct Coef
{
	int lx, ly, x, y;
	fp v;

	bool operator()(Coef& a, Coef& b);
};

class Reconstructer
{
protected:
	filesystem::path out_path;

	// accumulated image
	Matrix<fp,Dynamic,Dynamic> img;

	double mz_min;
	double mz_max;
	double rt_min;
	double rt_max;

	filesystem::path stream_path;
	ofstream         stream_ofs;
	int              stream_index;
	int              chunk_index;

    timer::cpu_timer read_time;
    timer::cpu_timer viz_time;
	size_t chunk_count;
	PNGWriter<fp> writer;

	int nx,ny,nxy;

public:
	Reconstructer(const string& out_dir, int width, int height);

	void next_stream(double mz_min, double mz_max, double rt_min, double rt_max);
	void next_chunk(const vector<Coef>& cs);
};


#include "Object.h"
#include "Camera.h"
#include "Light.h"
#include "read_json.h"
#include "write_ppm.h"
#include "viewing_ray.h"
#include "raycolor.h"
#include <Eigen/Core>
#include <vector>
#include <iostream>
#include <memory>
#include <limits>
#include <functional>
#include "read_rgba_from_png.h"
#include "rgba_to_rgb.h"

int main(int argc, char * argv[])
{
  Camera camera;
  std::vector< std::shared_ptr<Object> > objects;
  std::vector< std::shared_ptr<Light> > lights;

  std::vector<std::vector<unsigned char>> textures;
  std::vector<int> texture_height_offsets;
  bool use_textures = false;

  std::string outputname;

  if (argc == 2 && strcmp(argv[1], "texture") == 0)
  {
    outputname = "textured_seattle.ppm";
    use_textures = true;
    // Read a camera and scene description from given .json file
    read_json(
      "../shared/data/seattle-postcard.json",
      camera,
      objects,
      lights);

    // read texture files
    const char * default_input_file_names[] {
      "../images/seattle_skyline.png",
      "../images/floor_texture.png"
    };
    char ** input_file_names = const_cast<char **>(default_input_file_names);

    // background texture
    int background_width, background_height;
    std::vector<unsigned char> background_rgba;
    read_rgba_from_png(input_file_names[0], background_rgba, background_width, background_height);

    std::vector<unsigned char> background_rgb;
    rgba_to_rgb(background_rgba, background_width, background_height, background_rgb);
    textures.push_back(background_rgb);
    texture_height_offsets.push_back(background_height - 1);

    std::shared_ptr<Plane> background(new Plane());
    background->texture_id = 0;
    background->point = Eigen::Vector3d(0,0,-5);
    background->normal = Eigen::Vector3d(0,0,1).normalized();
    objects.push_back(background);

    // floor texture
    int floor_width, floor_height;
    std::vector<unsigned char> floor_rgba;
    read_rgba_from_png(input_file_names[1], floor_rgba, floor_width, floor_height);
    std::vector<unsigned char> floor_rgb;
    rgba_to_rgb(floor_rgba, floor_width, floor_height, floor_rgb);
    textures.push_back(floor_rgb);
    texture_height_offsets.push_back(floor_height - 1);

    std::shared_ptr<Plane> floor_plane(new Plane());
    floor_plane->texture_id = 1;
    floor_plane->point = Eigen::Vector3d(0,-1,0);
    floor_plane->normal = Eigen::Vector3d(0,1,0).normalized();
    objects.push_back(floor_plane);
  }
  else
  {
    // Read a camera and scene description from given .json file
    read_json(
      "../shared/data/untextured-seattle-postcard.json",
      camera,
      objects,
      lights);
    outputname = "untextured_seattle.ppm";
  }

  int width =  640;
  int height = 360;
  std::vector<unsigned char> rgb_image(3*width*height);
  // For each pixel (i,j)
  for(unsigned i=0; i<height; ++i)
  {
    for(unsigned j=0; j<width; ++j)
    {
      // Set background color
      Eigen::Vector3d rgb(0,0,0);

      // Compute viewing ray
      Ray ray;
      viewing_ray(camera,i,j,width,height,ray);

      // Shoot ray and collect color
      if (raycolor(ray,1.0,objects,lights,0, i, j, textures, texture_height_offsets, rgb))
      {
        // Write double precision color into image
        auto clamp = [](double s){ return std::max(std::min(s,1.0),0.0);};
        rgb_image[0+3*(j+width*i)] = 255.0*clamp(rgb(0));
        rgb_image[1+3*(j+width*i)] = 255.0*clamp(rgb(1));
        rgb_image[2+3*(j+width*i)] = 255.0*clamp(rgb(2));
      }
      else
      {
        rgb_image[0+3*(j+width*i)] = rgb(0);
        rgb_image[1+3*(j+width*i)] = rgb(1);
        rgb_image[2+3*(j+width*i)] = rgb(2);
      }
    }
  }

  write_ppm(outputname, rgb_image, width, height, 3);
}


#include <boost/config.hpp>
#include <ctime>
#include <fstream>
#include <iostream>
#include <iomanip>
#include <boost/graph/use_mpi.hpp>
#include <boost/graph/distributed/mpi_process_group.hpp>
#include <boost/graph/distributed/adjacency_list.hpp>
#include <boost/graph/distributed/connected_components_parallel_search.hpp>
#include <boost/graph/distributed/connected_components.hpp>
#include <boost/iostreams/filtering_stream.hpp>
#include <boost/iostreams/filtering_streambuf.hpp>
#include <boost/iostreams/filter/gzip.hpp>
#include <algorithm>
#include <vector>
#include <queue>
#include <limits>
#include <map>
int main (int argc, char ** argv){
	int gsiz = atoi(argv[1]);
	using namespace boost;
	using boost::graph::distributed::mpi_process_group;
	{
		mpi::environment env(argc, argv);
		mpi::communicator world;

		mpi_process_group pg;

	        int nproc = world.size();

		parallel::variant_distribution<mpi_process_group> distrib
			      = parallel::block(pg, gsiz);

		typedef adjacency_list < vecS,
				boost::distributedS<boost::graph::distributed::mpi_process_group, boost::vecS>,
				undirectedS > Graph;

		Graph G (gsiz);

                int pid = process_id (G .process_group());
                if (pid == 0) std::cerr << "start;" << time(0) << std::endl;
		if (1) {
			unsigned int a, b;
                        std::stringstream fnames;
                        fnames << argv[2] << "." << std::setw(5) << std::setfill('0') << pid << ".gz";
                        std::string fname = fnames.str ();
                        std::ifstream file(fname.data(), std::ios_base::in | std::ios_base::binary);
                        boost::iostreams::filtering_istream in;
                        in.push(boost::iostreams::gzip_decompressor());
                        in.push(file);
			in>>a; in>>b;
			//std::cout<<a<<b;
			while (!in.eof()){
				add_edge (vertex(a, G), vertex(b, G), G);
				in>>a;in>>b;
				//if (! ((j++)%1000000000)) std::cerr << j << std::endl;
			}
			std::cerr << "read;" << pid << ';' << time(0) << std::endl;
		}

		synchronize(G);
                if (pid == 0) std::cerr << "read;" << time(0) << std::endl;

		int nv = num_vertices(G);
		//std::cerr << "Created local graph " << process_id (G .process_group()) <<
		//	" with " << nv << " vertices" << std::endl;

		std::vector<long int> lcomponent (num_vertices(G));
		typedef iterator_property_map<std::vector<long int>::iterator,
				property_map<Graph, vertex_index_t>::type> ComponentMap;
		ComponentMap component (lcomponent.begin(), get (vertex_index, G));

                std::stringstream fnames;
                fnames << argv[3] << "." << std::setw(5) << std::setfill('0') << pid << ".gz";
                std::ofstream file(fnames.str().data(), std::ios_base::out | std::ios_base::binary);
                boost::iostreams::filtering_ostream out;
                out .push(boost::iostreams::gzip_compressor());
                out .push(file);

		//out << "Created components for local graph " << pid <<
		//	" with " << nv << " vertices" << std::endl;
		int num = connected_components_ps (G, component);
	        //synchronize(component);
                //out << "computed;" << pid << ';' << time(0) << std::endl;
		if (pid == 0){
                   component.set_max_ghost_cells(0);
                   synchronize(component);
                }else{
                   synchronize(component);
                }
		if (pid == 0) std::cerr << "connected;" << time(0) << ";components" << std::endl;
                //for (int i = 0; i < nv; i++){
                //for (int i = from; i < nv; i++){

                synchronize(G .process_group());
                synchronize(component);
                typedef graph_traits<Graph>::vertex_iterator vertex_iter;
                std::pair<vertex_iter, vertex_iter> lG;
                typedef property_map<Graph, vertex_index_t>::type IndexMap;
                IndexMap index = get(vertex_index, G);
                for (lG = vertices(G); lG.first != lG.second; ++lG.first){
                   out << index [*lG.first] << ';' << get (component, *lG.first) << std::endl;
                }
		//if (0){
                //int from = gsiz*pid/nproc;
                //int to = gsiz*(pid+1)/nproc;
                //if (pid == nproc-1) to = gsiz;
		//for (int i = from; i < to; i++){
		//	int c = get (component, vertex(i, G));
		//	out << i << ";" << c << std::endl;
		//}//}
                synchronize(G .process_group());
                if (pid == 0) std::cerr << "done;" << time(0)  << std::endl;
		//	synchronize(component);
		//}else{
		//	synchronize(component);
		//}
		//write_graphviz("cc.dot", G, paint_by_number(component));
	}
	return 0;
}


/*
 * polygon_ext.cc
 *
 *  Copyright (C) 2013 Diamond Light Source
 *
 *  Author: James Parkhurst
 *
 *  This code is distributed under the BSD license, a copy of which is
 *  included in the root directory of this package.
 */
#include <boost/python.hpp>
#include <boost/python/def.hpp>
#include <scitbx/vec2.h>
#include <dials/algorithms/polygon/area.h>
#include <dials/array_family/scitbx_shared_and_versa.h>

namespace dials { namespace algorithms { namespace polygon { namespace boost_python {

  using namespace boost::python;

  using scitbx::vec2;

  void export_area() {
    def("simple_area", &simple_area<af::shared<vec2<double> > >, (arg("poly")));
  }

  BOOST_PYTHON_MODULE(dials_algorithms_polygon_ext) {
    export_area();
  }

}}}}  // namespace dials::algorithms::polygon::boost_python


#include <iostream>
#include <fstream>
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/io.hpp>
#include <boost/program_options.hpp>
#include <boost/format.hpp>

#include "simulation.hpp"

#include "string"

#include "rngmit.hpp"
// this is polymer

using namespace std;
namespace po = boost::program_options;

extern unsigned long int rng_ia[55];
extern int rng_p,rng_pp;
extern int rng_mod[55];

string inFile;

bool CONTINUE;

int main(int argc , char** argv) {
	// Get Simulation Parameters and options
	double optDt,optRelGamma,optGamma0,optOmega;
	unsigned int optSteps;

	if(argc!=13 && argc!=3){cout << "ERROR: Wrong number of arguments" << "\n"; return 0;}
	else if(argc==3)
	{
		if (string(argv[1]) == "-con")
		{
			CONTINUE=1;
		}
		else{cout << "ERROR: only continue possible --> -con"<<"\n"; return 0;}

		inFile=argv[2];
	}
	else if(argc==13)
	{
		if (string(argv[1]) == "-con")
		{
			CONTINUE=1;
		}
		else if (string(argv[1]) == "-start")
		{
			CONTINUE=0;
		}
		else{cout << "ERROR: -start or -con!" << "\n"; return 0;}

		inFile=argv[2];

		for (int i = 3; i < argc; i++)
		{
			if (string(argv[i]) == "-dt") {
				optDt=stod( argv[i + 1]); i++;
			}
			else if (string(argv[i]) == "-s"){
				optSteps=(int)(stod( argv[i + 1])); i++;
			}
			else if (string(argv[i]) == "-rg"){
				optRelGamma=stod( argv[i + 1]); i++;
			}
			else if (string(argv[i]) == "-g0"){
				optGamma0=stod( argv[i + 1]); i++;
			}
			else if (string(argv[i]) == "-w"){
				optOmega=stod( argv[i + 1]); i++;
			}
			else {cout << "ERROR: wrong order of arguments!" << "\n"; return 0;}
		}
}

	//double optDt = 1e-11; //0.001 gives artifacts, default: 0.00001
	//unsigned int optSteps = 1e8;//4100000

	unsigned int optN = 1;
	int seed = 1;

  	Simulation simulation(optN, optDt, optRelGamma, optGamma0, optOmega, seed, inFile,CONTINUE);

  	simulation.integrate(optSteps);

	cout << "simulation finished" << endl;

	return 0;
}


#include <boost/range/algorithm/for_each.hpp>
#include "async_loop.h"
#include "ann.h"
#include "shuffled_range.h"
#include "utils.h"
#include <matrix/confusion_matrix.h>

namespace ann { namespace async {

ThreadCommandQueue::message::buffer serializeConfusionMatrix(const math::ConfusionMatrix::internalMatrix &matrix);
ThreadCommandQueue::message::buffer serializeTestStatus(const size_t numImages);
ThreadCommandQueue::message::buffer serializeTrainStatus(const IANN::trainData &trainData);

struct TrainState
{
    std::unique_ptr<ShuffledRange> mRangeGenerator;
    ShuffledRange::IndexedRange mCurrentRange;

    math::MatrixF mImages;
    math::MatrixF mLabels;

    float mLastLoss = 0.0f;

    size_t mCurrentIteration = 0;

    bool mCreateSnapshot = false;
};

struct TestState
{
    std::unique_ptr<ShuffledRange> mRangeGenerator;
    ShuffledRange::IndexedRange mCurrentRange;
    math::ConfusionMatrix mConfusionMatrix = math::ConfusionMatrix(10);

    size_t mCurrentIteration = 0;
};

AsyncLoop::AsyncLoop(ThreadCommandQueuePtr threadCommand, MainCommandQueuePtr mainCommand, db::DBInterfacePtr trainDb, db::DBInterfacePtr testDb)
    : mTrainDb(trainDb)
    , mTestDb(testDb)
    , mIncomingQueue(threadCommand)
    , mResultQueue(mainCommand)
{
}

void AsyncLoop::Run()
{
    std::unique_ptr<TrainState> trainState;
    std::unique_ptr<TestState> testState;

    std::cout << "starting message loop on the background thread" << std::endl;

    while (mInternalState != InternalState::QUIT)
    {
        // get command
        auto result = mIncomingQueue->getCommand(mInternalState == InternalState::IDLE);
        if (result)
        {
            switch (result->mCommand)
            {
                case commandToThread::CONFIGURE_NETWORK:
                    mNetwork = std::make_unique<ANN>();
                    mNetwork->configureNetwork(mTrainDb->getPixelCount());

                    trainState = std::make_unique<TrainState>();
                    trainState->mRangeGenerator.reset( new ShuffledRange(mTrainDb->getImageCount()) );

                    mInternalState = InternalState::IDLE;
                    break;
                case commandToThread::GO_IDLE:
                    mInternalState = InternalState::IDLE;
                    break;
                case commandToThread::STOP:
                    mInternalState = InternalState::QUIT;
                    break;
                case commandToThread::START_TRAINING:
                    mInternalState = InternalState::BUSY;
                    mNetworkState = NetworkState::TRAIN;
                    break;
                case commandToThread::GET_TRAINING_SNAPSHOT:
                    if (trainState)
                    {
                        trainState->mCreateSnapshot = true;
                    }
                    break;
                case commandToThread::START_TESTING:
                    mInternalState = InternalState::BUSY;
                    mNetworkState = NetworkState::RECOGNIZE;

                    testState = std::make_unique<TestState>();
                    testState->mRangeGenerator.reset( new ShuffledRange( mTestDb->getImageCount()) );
                    break;
            }
        }

        if (mInternalState == InternalState::BUSY)
        {
            switch(mNetworkState)
            {
                case NetworkState::TRAIN:
                    nextTrainStep(*mNetwork, *trainState);
                    if (trainState->mCreateSnapshot)
                    {
                        auto trainResult = mNetwork->collectTrainData();
                        boost::for_each(trainState->mCurrentRange, [&trainResult](size_t range) {
                            trainResult.push_back(static_cast<uint32_t>(range));
                        });
                        MainCommandQueue::message snapshotMessge{
                            commandToMain::TRAINING_SNAPSHOT,
                            serializeTrainStatus(trainResult)
                        };
                        mResultQueue->pushCommand(snapshotMessge);

                        trainState->mCreateSnapshot = false;
                    }
                    break;
                case NetworkState::RECOGNIZE:
                    {
                        auto hasMoreElements = nextTestStep(*mNetwork, *testState);

                        MainCommandQueue::message trainStatusMsg{
                            commandToMain::TESTING_PROGRESS_STATUS,
                            serializeTestStatus((testState->mCurrentIteration) * 50)
                        };
                        mResultQueue->pushCommand(trainStatusMsg);

                        if (!hasMoreElements)
                        {
                            mNetworkState = NetworkState::NONE;
                            mInternalState = InternalState::IDLE;
                            std::cout << testState->mConfusionMatrix.toString() << std::endl;

                            MainCommandQueue::message message {
                                commandToMain::TESTING_FINISHED,
                                serializeConfusionMatrix(testState->mConfusionMatrix.getMatrix())
                            };

                            mResultQueue->pushCommand(message);

                        }
                    }
                    break;

            }
        }
    }
    std::cout << "worker thread finished" << std::endl;
}

ThreadCommandQueue::message::buffer serializeConfusionMatrix(const math::ConfusionMatrix::internalMatrix &matrix)
{
    ThreadCommandQueue::message::buffer result;
    const uint32_t rowCount = static_cast<uint32_t>(matrix.getRowCount());
    result.resize(sizeof(rowCount) * (1 + rowCount * rowCount));

    auto offset = result.data();
    std::memcpy(offset, &rowCount, sizeof(uint32_t));
    offset += sizeof(uint32_t);
    for (auto value : matrix)
    {
        auto v = static_cast<uint32_t>(value);
        std::memcpy(offset, &v, sizeof(uint32_t));
        offset += sizeof(uint32_t);
    }
    return result;
}

ThreadCommandQueue::message::buffer serializeTestStatus(const size_t numImages)
{
    ThreadCommandQueue::message::buffer result;
    result.resize(sizeof(uint32_t));
    auto v = static_cast<uint32_t>(numImages);
    std::memcpy(result.data(), &v, sizeof(uint32_t));
    return result;
}

ThreadCommandQueue::message::buffer serializeTrainStatus(const IANN::trainData &trainData)
{
    ThreadCommandQueue::message::buffer result;

    result.resize(trainData.size() * sizeof(uint32_t));
    auto offset = result.data();
    boost::for_each(trainData, [&offset](uint32_t value){
        std::memcpy(offset, &value, sizeof(value));
        offset += sizeof(value);
    });

    return result;
}

ShuffledRange::IndexedRange generateNewRange(ShuffledRange &generator, size_t count)
{
    if (!generator.hasMoreElements())
    {
        generator.shuffle();
    }
    auto result = generator.getNextN(count);
    return result;
}

void AsyncLoop::nextTrainStep(ann::IANN &network, TrainState& state)
{
    const size_t batchSize = 50;

    std::cout << state.mCurrentIteration<< ": " << state.mLastLoss << std::endl;

    state.mCurrentRange = generateNewRange(*state.mRangeGenerator, batchSize);
    state.mImages = getImageDataBasedOnRange(*mTrainDb, state.mCurrentRange);
    state.mLabels = getImageClassBasedOnRange(*mTrainDb, state.mCurrentRange);

    state.mLastLoss = mNetwork->train(state.mImages, state.mLabels);
    state.mCurrentIteration += 1;
}

bool AsyncLoop::nextTestStep(ann::IANN &network, TestState& state)
{
    const size_t batchSize = 50;
    const size_t targetCount = 1000;

    state.mCurrentRange = generateNewRange(*state.mRangeGenerator, batchSize);
    auto images = getImageDataBasedOnRange(*mTestDb, state.mCurrentRange);
    auto labels = getImageClassBasedOnRange(*mTestDb, state.mCurrentRange);

    const auto& result = mNetwork->test(images);
    state.mConfusionMatrix.update(result, labels);
    state.mCurrentIteration += 1;

    return state.mCurrentIteration * batchSize < targetCount;
}

} } // namespace


#include "solver/CG.h"

#include "utility/DataPtr.h"
#include "volume/Volume.h"

#include "cuda/Volume.cuh"
#include "cuda/LinearSys.cuh"
#include "cuda/CudaUtility.h"
#include "cuda/CG.cuh"

#include "solver/LinearSystemDevice.h"
#include "solver/LinearSystemHost.h"

#include <Eigen/IterativeLinearSolvers>

#include "utility/Timer.h"

#include <iostream>
#include <fstream>

namespace fast {

	template <typename T>
	FAST_EXPORT CG<T>::CG(bool verbose, bool precondition) : Solver<T>(verbose), _precondition(precondition)
	{

	}

	template <typename T>
	FAST_EXPORT	size_t fast::CG<T>::requiredMemory(size_t nnz) const
	{
		return CGState<T>::requiredBytesize(nnz, _precondition);
	}

	template <typename T>
	FAST_EXPORT typename Solver<T>::Output fast::CG<T>::solve(LinearSystem * linearSys, const typename Solver<T>::SolveParams & solveParams)
	{

		assert(linearSys->type() == primitiveTypeof<T>());

		typename Solver<T>::Output out;
		LinearSystemDevice<T> * devicePtr = dynamic_cast<LinearSystemDevice<T>*>(linearSys);
		if (devicePtr) {

			//Try to allocate solver state
			CGState<T> state;
			try {
				state.resize(devicePtr->NNZ, _precondition);
			}
			catch (thrust::system_error & er) {
				std::cerr << er.what() << std::endl;
				return out;
			}

			bool res = CG_Solve<T>(
				*devicePtr,
				state,
				solveParams.tolerance,
				solveParams.maxIter,
				&out.error,
				&out.iterations,
				Solver<T>::isVerbose()
				);

			out.status = res ? Solver<T>::Status::SOLVER_STATUS_SUCCESS : Solver<T>::Status::SOLVER_STATUS_NO_CONVERGENCE;
			return out;
		}

		auto hostPtr = dynamic_cast<LinearSystemHost<T>*>(linearSys);
		if (hostPtr) {

			int threadNum = omp_get_num_procs();
			if (solveParams.cpuThreads > 0) {
				threadNum = solveParams.cpuThreads;
			}
			omp_set_num_threads(threadNum);
			Eigen::setNbThreads(threadNum);

			Eigen::ConjugateGradient<Eigen::SparseMatrix<T, Eigen::RowMajor>, Eigen::Lower | Eigen::Upper>
				_solver;

			_solver.compute(hostPtr->A);
			_solver.setTolerance(solveParams.tolerance);
			_solver.setMaxIterations(solveParams.maxIter);

			hostPtr->x = _solver.solveWithGuess(hostPtr->b, hostPtr->x);

			out.error = _solver.error();
			if (_solver.info() == Eigen::ComputationInfo::Success)
				out.status = Solver<T>::Status::SOLVER_STATUS_SUCCESS;
			else
				out.status = Solver<T>::Status::SOLVER_STATUS_NO_CONVERGENCE;

			out.iterations = _solver.iterations();
			//copyHostToDevice(_x.data(), sys.x);

			return out;
		}

		//Linear system derivation not implemented
		assert(false);
		return out;
	}

	template class CG<double>;
	template class CG<float>;
}


<% render(instance) %>
<%= includes((0..x).size) %>

#include <boost/hana/integral.hpp>
#include <boost/hana/iterable/iterable.hpp>

template <int> struct x { };

int main() {
    auto go = boost::hana::drop(
        boost::hana::int_<  <%= x / 2 %>    >,
        <%= list(
            (0..x).map { |i| "x<#{i}>" },
            (0..x).map { |i| "x<#{i}>{}" }
        ) %>
    );
}