#define BOOST_TEST_MODULE test_fs_n_fold_adult_transformed_dataset
#include <boost/test/included/unit_test.hpp>

//#define TIMERS
//#define TRACE_EVALUATOR

#include "../lib/classifier_creator_fs_n_fold.h"
#include "../lib/classifier_creator_train_tmc.h"

BOOST_AUTO_TEST_CASE( test_fs_n_fold_adult_transformed_dataset_roc  ) {
    const std::string names_file("data/adult_transformed/adult_transformed.names");
    const std::string data_file("data/adult_transformed/adult_transformed.data");
    const std::string eval_file("data/adult_transformed/adult_transformed.test");
    BOOST_TEST_MESSAGE("Creating train sample from file: " << data_file);
    BOOST_TEST_MESSAGE("Creating evaluation sample from file: " << eval_file);

    std::shared_ptr<FeatureNames> pFN = std::make_shared<FeatureNames>(names_file);

    SampleCreator sample_creator;
    sample_creator.set_feature_names(pFN);

    std::shared_ptr<Sample> pSample = sample_creator.from_file(data_file);
    std::shared_ptr<Sample> pSampleEval = sample_creator.from_file(eval_file);

    BOOST_CHECK_EQUAL(pSample->dim(), 13);
    BOOST_CHECK_EQUAL(pSample->size(), 26314);
    BOOST_CHECK_EQUAL(pSampleEval->size(), 13906);
    std::shared_ptr<ClassifierCreatorTrain> pTC_aux = std::make_shared<ClassifierCreatorTrainTmc>();
    std::shared_ptr<ClassifierCreatorFsNfold> pCCFS = std::make_shared<ClassifierCreatorFsNfold>();
    pCCFS->verbose(true);
    (*pCCFS).set_classifier_creator_train(pTC_aux).init(pSample);
    (*pCCFS).set_n_folds(5).train();

    std::shared_ptr<Classifier> pFsClTmc = pCCFS->get_classifier();

    auto pEvaluator = std::make_shared<Evaluator>();
    (*pEvaluator).set_conf_type(Evaluator::number).set_classifier(pFsClTmc);
    std::pair<std::pair<double, double>, std::pair<double, double>> confusion_matrix_train, confusion_matrix_eval;
    double roc_err_train, roc_err_eval, err_rate_train, err_rate_eval;

        (*pEvaluator).set_sample(pSample);
        confusion_matrix_train = pEvaluator->get_confusion_matrix();
        roc_err_train = pEvaluator->get_roc_error();
        err_rate_train = pEvaluator->get_error_rate();

        (*pEvaluator).set_sample(pSampleEval);
        confusion_matrix_eval = pEvaluator->get_confusion_matrix();
        roc_err_eval = pEvaluator->get_roc_error();
        err_rate_eval = pEvaluator->get_error_rate();

        BOOST_TEST_MESSAGE("");
        BOOST_TEST_MESSAGE("########################");
        BOOST_TEST_MESSAGE("Train conf matr ### Eval conf matr");
        BOOST_TEST_MESSAGE(confusion_matrix_train.first.first << "  " << confusion_matrix_train.first.second << "    ###    "
                                                              << confusion_matrix_eval.first.first << "   " << confusion_matrix_eval.first.second);
        BOOST_TEST_MESSAGE(confusion_matrix_train.second.first << "  " << confusion_matrix_train.second.second << "    ###    "
                                                               << confusion_matrix_eval.second.first << "   " << confusion_matrix_eval.second.second);
        BOOST_TEST_MESSAGE("------------------------------------------------");
        BOOST_TEST_MESSAGE("Train rank err ### Eval rank err");
        BOOST_TEST_MESSAGE( roc_err_train << "    ###    " << roc_err_eval);
        BOOST_TEST_MESSAGE("------------------------------------------------");
        BOOST_TEST_MESSAGE("Train err rate ### Eval err rate");
        BOOST_TEST_MESSAGE( err_rate_train << "    ###    " << err_rate_eval);
        BOOST_TEST_MESSAGE("################################################");

    BOOST_CHECK_EQUAL(roc_err_train, 0.074392899680209762);
    BOOST_CHECK_EQUAL(roc_err_eval, 0.09294518600907764);
    BOOST_CHECK_EQUAL(err_rate_train, 0.13433037482690777);
    BOOST_CHECK_EQUAL(err_rate_eval, 0.14824828129875775);
    std::pair<std::pair<double, double>, std::pair<double, double>> confusion_matrix_train_({{908866308,265218900},{503763608,4046712803}});
    BOOST_CHECK(confusion_matrix_train == confusion_matrix_train_);
    std::pair<std::pair<double, double>, std::pair<double, double>> confusion_matrix_eval_({{427794236,152200867},{271140266,2004487167}});
    BOOST_CHECK(confusion_matrix_eval == confusion_matrix_eval_);

    BOOST_CHECK_EQUAL(confusion_matrix_train.first.first,confusion_matrix_train_.first.first);
    BOOST_CHECK_EQUAL(confusion_matrix_train.first.second,confusion_matrix_train_.first.second);
    BOOST_CHECK_EQUAL(confusion_matrix_train.second.first,confusion_matrix_train_.second.first);
    BOOST_CHECK_EQUAL(confusion_matrix_train.second.second,confusion_matrix_train_.second.second);

    BOOST_CHECK_EQUAL(confusion_matrix_eval.first.first,confusion_matrix_eval_.first.first);
    BOOST_CHECK_EQUAL(confusion_matrix_eval.first.second,confusion_matrix_eval_.first.second);
    BOOST_CHECK_EQUAL(confusion_matrix_eval.second.first,confusion_matrix_eval_.second.first);
    BOOST_CHECK_EQUAL(confusion_matrix_eval.second.second,confusion_matrix_eval_.second.second);
}

BOOST_AUTO_TEST_CASE( test_fs_n_fold_adult_transformed_dataset_classerr  ) {
    const std::string names_file("data/adult_transformed/adult_transformed.names");
    const std::string data_file("data/adult_transformed/adult_transformed.data");
    const std::string eval_file("data/adult_transformed/adult_transformed.test");
    BOOST_TEST_MESSAGE("Creating train sample from file: " << data_file);
    BOOST_TEST_MESSAGE("Creating evaluation sample from file: " << eval_file);

    std::shared_ptr<FeatureNames> pFN = std::make_shared<FeatureNames>(names_file);

    SampleCreator sample_creator;
    sample_creator.set_feature_names(pFN);

    std::shared_ptr<Sample> pSample = sample_creator.from_file(data_file);
    std::shared_ptr<Sample> pSampleEval = sample_creator.from_file(eval_file);

    BOOST_CHECK_EQUAL(pSample->dim(), 13);
    BOOST_CHECK_EQUAL(pSample->size(), 26314);
    BOOST_CHECK_EQUAL(pSampleEval->size(), 13906);
    std::shared_ptr<ClassifierCreatorTrain> pTC_aux = std::make_shared<ClassifierCreatorTrainTmc>();
    std::shared_ptr<ClassifierCreatorFsNfold> pCCFS = std::make_shared<ClassifierCreatorFsNfold>();
    pCCFS->verbose(true);
    (*pCCFS).set_classifier_creator_train(pTC_aux).init(pSample);
    (*pCCFS).set_n_folds(2).set_kpi_type(ClassifierCreatorFsNfold::class_err).train();

    std::shared_ptr<Classifier> pFsClTmc = pCCFS->get_classifier();

    auto pEvaluator = std::make_shared<Evaluator>();
    (*pEvaluator).set_conf_type(Evaluator::number).set_classifier(pFsClTmc);
    std::pair<std::pair<double, double>, std::pair<double, double>> confusion_matrix_train, confusion_matrix_eval;
    double roc_err_train, roc_err_eval, err_rate_train, err_rate_eval;

    (*pEvaluator).set_sample(pSample);
    confusion_matrix_train = pEvaluator->get_confusion_matrix();
    roc_err_train = pEvaluator->get_roc_error();
    err_rate_train = pEvaluator->get_error_rate();

    (*pEvaluator).set_sample(pSampleEval);
    confusion_matrix_eval = pEvaluator->get_confusion_matrix();
    roc_err_eval = pEvaluator->get_roc_error();
    err_rate_eval = pEvaluator->get_error_rate();

    BOOST_TEST_MESSAGE("");
    BOOST_TEST_MESSAGE("########################");
    BOOST_TEST_MESSAGE("Train conf matr ### Eval conf matr");
    BOOST_TEST_MESSAGE(confusion_matrix_train.first.first << "  " << confusion_matrix_train.first.second << "    ###    "
                                                          << confusion_matrix_eval.first.first << "   " << confusion_matrix_eval.first.second);
    BOOST_TEST_MESSAGE(confusion_matrix_train.second.first << "  " << confusion_matrix_train.second.second << "    ###    "
                                                           << confusion_matrix_eval.second.first << "   " << confusion_matrix_eval.second.second);
    BOOST_TEST_MESSAGE("------------------------------------------------");
    BOOST_TEST_MESSAGE("Train rank err ### Eval rank err");
    BOOST_TEST_MESSAGE( roc_err_train << "    ###    " << roc_err_eval);
    BOOST_TEST_MESSAGE("------------------------------------------------");
    BOOST_TEST_MESSAGE("Train err rate ### Eval err rate");
    BOOST_TEST_MESSAGE( err_rate_train << "    ###    " << err_rate_eval);
    BOOST_TEST_MESSAGE("################################################");

    BOOST_CHECK_EQUAL(roc_err_train, 0.089554665836618294);
    BOOST_CHECK_EQUAL(roc_err_eval, 0.095041575784745022);
    std::pair<std::pair<double, double>, std::pair<double, double>> confusion_matrix_train_({{822729363,233218553},{589900553,4078713150}});
    BOOST_CHECK(confusion_matrix_train == confusion_matrix_train_);
    std::pair<std::pair<double, double>, std::pair<double, double>> confusion_matrix_eval_({{401599020,118273444},{297335482,2038414590}});
    BOOST_CHECK(confusion_matrix_eval == confusion_matrix_eval_);

    BOOST_CHECK_EQUAL(confusion_matrix_train.first.first,confusion_matrix_train_.first.first);
    BOOST_CHECK_EQUAL(confusion_matrix_train.first.second,confusion_matrix_train_.first.second);
    BOOST_CHECK_EQUAL(confusion_matrix_train.second.first,confusion_matrix_train_.second.first);
    BOOST_CHECK_EQUAL(confusion_matrix_train.second.second,confusion_matrix_train_.second.second);

    BOOST_CHECK_EQUAL(confusion_matrix_eval.first.first,confusion_matrix_eval_.first.first);
    BOOST_CHECK_EQUAL(confusion_matrix_eval.first.second,confusion_matrix_eval_.first.second);
    BOOST_CHECK_EQUAL(confusion_matrix_eval.second.first,confusion_matrix_eval_.second.first);
    BOOST_CHECK_EQUAL(confusion_matrix_eval.second.second,confusion_matrix_eval_.second.second);
}


#include "stc_ml_c.h"

#include <xmmintrin.h>
#include <cmath>
#include <limits>
#include <algorithm>
#include <sstream>
#include <fstream>
#include <iomanip>
#include <string.h> // due to memcpy

#include <boost/random/uniform_int.hpp>       // this is required for Marsene-Twister random number generator
#include <boost/random/variate_generator.hpp>
#include <boost/random/mersenne_twister.hpp>

#include "stc_embed_c.h"
#include "stc_extract_c.h"
#include "sse_mathfun.h"    // library with optimized functions obtained from http://gruntthepeon.free.fr/ssemath/

// {{{ write_vector_to_file()
template< class T > void write_vector_to_file( uint n, T *ptr, const char* file_name ) {

    std::ofstream f( file_name );
    for ( uint i = 0; i < n; i++ )
        f << std::left << std::setw( 20 ) << i << std::left << std::setw( 20 ) << ptr[i] << std::endl;
    f.close();
}
// }}}

// {{{ write_matrix_to_file()
// write column-ordered matrix into file
template< class T > void write_matrix_to_file( uint rows, uint columns, T *ptr, const char* file_name ) {

    std::ofstream f( file_name );
    for ( uint i = 0; i < rows; i++ ) {
        f << std::left << std::setw( 20 ) << i;
        for ( uint j = 0; j < columns; j++ )
            f << std::left << std::setw( 20 ) << ptr[j * rows + i];
        f << std::endl;
    }
    f.close();
}
// }}}

// {{{ align_*()
// Templates to handle aligned version of new and delete operators.
// These functions are necessary for creating arrays aligned address of certain multiples, such as 16.
template< class T > T* align_new( unsigned int n, unsigned int align_size ) {
    char *ptr, *ptr2, *aligned_ptr;
    int align_mask = align_size - 1;

    ptr = new char[n * sizeof(T) + align_size + sizeof(int)];
    if ( ptr == 0 ) return 0;

    ptr2 = ptr + sizeof(int);
    aligned_ptr = ptr2 + (align_size - ((size_t) ptr2 & align_mask));

    ptr2 = aligned_ptr - sizeof(int);
    *((int*) ptr2) = (int) (aligned_ptr - ptr);

    return (T*) aligned_ptr;
}

template< class T > void align_delete( T *ptr ) {
    int *ptr2 = (int*) ptr - 1;
    char *p;

    p = (char*) ptr;
    p -= *ptr2;
    delete[] p;
}
// }}}

// {{{ randperm()
/* Generates random permutation of length n based on the MT random number generator with seed 'seed'. */
void randperm( uint n, uint seed, uint* perm ) {

    boost::mt19937 *generator = new boost::mt19937( seed );
    boost::variate_generator< boost::mt19937, boost::uniform_int< > > *randi = new boost::variate_generator< boost::mt19937,
        boost::uniform_int< > >( *generator, boost::uniform_int< >( 0, INT_MAX ) );

    // generate random permutation - this is used to shuffle cover pixels to randomize the effect of different neighboring pixels
    for ( uint i = 0; i < n; i++ )
        perm[i] = i;
    for ( uint i = 0; i < n; i++ ) {
        uint j = (*randi)() % (n - i);
        uint tmp = perm[i];
        perm[i] = perm[i + j];
        perm[i + j] = tmp;
    }

    delete generator;
    delete randi;
}
// }}}

// {{{ sum_inplace()
inline float sum_inplace( __m128 x ) {
    float y;
    // add all 4 terms from x together
    x = _mm_add_ps( x, _mm_shuffle_ps(x,x,_MM_SHUFFLE(1,0,3,2)) );
    x = _mm_add_ps( x, _mm_shuffle_ps(x,x,_MM_SHUFFLE(2,3,0,1)) );
    _mm_store_ss( &y, x );
    return y;
}
// }}}

// {{{ calc_entropy()
float calc_entropy( uint n, uint k, float* costs, float lambda ) {

    float const LOG2 = log( 2.0 );
    __m128 inf = _mm_set1_ps( F_INF );
    __m128 v_lambda = _mm_set1_ps( -lambda );
    __m128 z, d, rho, p, entr, mask;

    entr = _mm_setzero_ps();
    for ( uint i = 0; i < n / 4; i++ ) {
        z = _mm_setzero_ps();
        d = _mm_setzero_ps();
        for ( uint j = 0; j < k; j++ ) {
            rho = _mm_load_ps( costs + j * n + 4 * i ); // costs array must be aligned in memory
            p = exp_ps( _mm_mul_ps( v_lambda, rho ) );
            z = _mm_add_ps( z, p );

            mask = _mm_cmpeq_ps( rho, inf ); // if p<eps, then do not accumulate it to d since x*exp(-x) tends to zero
            p = _mm_mul_ps( rho, p );
            p = _mm_andnot_ps( mask, p ); // apply mask
            d = _mm_add_ps( d, p );
        }
        entr = _mm_sub_ps( entr, _mm_div_ps( _mm_mul_ps( v_lambda, d ), z ) );
        entr = _mm_add_ps( entr, log_ps( z ) );
    }
    return sum_inplace( entr ) / LOG2;
}
// }}}

// {{{ get_lambda_entropy()
float get_lambda_entropy( uint n, uint k, float *costs, float payload, float initial_lambda = 10 ) {

    float p1, p2, p3, lambda1, lambda2, lambda3;
    int j = 0;
    uint iterations = 0;

    lambda1 = 0;
    p1 = n * log( (float)k ) / log( 2.0f );
    lambda3 = initial_lambda;
    p3 = payload + 1; // this is just an initial value
    lambda2 = initial_lambda;
    while ( p3 > payload ) {
        lambda3 *= 2;
        p3 = calc_entropy( n, k, costs, lambda3 );
        j++;
        iterations++;
        // beta is probably unbounded => it seems that we cannot find beta such that
        // relative payload will be smaller than requested. Binary search does not make sence here.
        if ( j > 10 ) {
            return lambda3;
        }
    }
    while ( (p1 - p3) / n > payload / n * 1e-2 ) { // binary search for parameter lambda
        lambda2 = lambda1 + (lambda3 - lambda1) / 2;
        p2 = calc_entropy( n, k, costs, lambda2 );
        if ( p2 < payload ) {
            lambda3 = lambda2;
            p3 = p2;
        } else {
            lambda1 = lambda2;
            p1 = p2;
        }
        iterations++; // this is for monitoring the number of iterations
    }
    return lambda1 + (lambda3 - lambda1) / 2;
}
// }}}

// {{{ calc_distortion()
float calc_distortion( uint n, uint k, float* costs, float lambda ) {

    __m128 eps = _mm_set1_ps( std::numeric_limits< float >::epsilon() );
    __m128 v_lambda = _mm_set1_ps( -lambda );
    __m128 z, d, rho, p, dist, mask;

    dist = _mm_setzero_ps();
    for ( uint i = 0; i < n / 4; i++ ) { // n must be multiple of 4
        z = _mm_setzero_ps();
        d = _mm_setzero_ps();
        for ( uint j = 0; j < k; j++ ) {
            rho = _mm_load_ps( costs + j * n + 4 * i ); // costs array must be aligned in memory
            p = exp_ps( _mm_mul_ps( v_lambda, rho ) );
            z = _mm_add_ps( z, p );
            mask = _mm_cmplt_ps( p, eps ); // if p<eps, then do not accumulate it to d since x*exp(-x) tends to zero
            p = _mm_mul_ps( rho, p );
            p = _mm_andnot_ps( mask, p );
            d = _mm_add_ps( d, p );
        }
        dist = _mm_add_ps( dist, _mm_div_ps( d, z ) );
    }
    return sum_inplace( dist );
}
// }}}

// {{{ get_lambda_distortion()
float get_lambda_distortion( uint n, uint k, float *costs, float distortion, float initial_lambda = 10, float precision = 1e-3,
        uint iter_limit = 30 ) {

    float dist1, dist2, dist3, lambda1, lambda2, lambda3;
    int j = 0;
    uint iterations = 0;

    lambda1 = 0;
    dist1 = calc_distortion( n, k, costs, lambda1 );
    lambda3 = initial_lambda;
    dist2 = F_INF; // this is just an initial value
    lambda2 = initial_lambda;
    dist3 = distortion + 1;
    while ( dist3 > distortion ) {
        lambda3 *= 2;
        dist3 = calc_distortion( n, k, costs, lambda3 );
        j++;
        iterations++;
        // beta is probably unbounded => it seems that we cannot find beta such that
        // relative payload will be smaller than requested. Binary search cannot converge.
        if ( j > 10 ) {
            return lambda3;
        }
    }
    while ( (fabs( dist2 - distortion ) / n > precision) && (iterations < iter_limit) ) { // binary search for parameter lambda
        lambda2 = lambda1 + (lambda3 - lambda1) / 2;
        dist2 = calc_distortion( n, k, costs, lambda2 );
        if ( dist2 < distortion ) {
            lambda3 = lambda2;
            dist3 = dist2;
        } else {
            lambda1 = lambda2;
            dist1 = dist2;
        }
        iterations++; // this is for monitoring the number of iterations
    }
    return lambda1 + (lambda3 - lambda1) / 2;
}
// }}}

// {{{ binary_entropy_array()
float binary_entropy_array( uint n, float *prob ) {

    float h = 0;
    float const LOG2 = log( 2.0 );
    float const EPS = std::numeric_limits< float >::epsilon();

    for ( uint i = 0; i < n; i++ )
        if ( (prob[i] > EPS) && (1 - prob[i] > EPS) ) h -= prob[i] * log( prob[i] ) + (1 - prob[i]) * log( 1 - prob[i] );

    return h / LOG2;
}
// }}}

// {{{ entropy_array()
float entropy_array( uint n, float* prob ) {

    double h = 0;
    double const LOG2 = log( 2.0 );
    double const EPS = std::numeric_limits< double >::epsilon();

    for ( uint i = 0; i < n; i++ )
        if ( prob[i] > EPS ) h -= prob[i] * log( prob[i] );

    return h / LOG2;
}
// }}}

// {{{ mod()
inline uint mod( int x, int m ) {
    int tmp = x - (x / m) * m + m;
    return tmp % m;
}
// }}}

/* EMBEDDING ALGORITHMS */

// {{{ stc_embed_trial()
void stc_embed_trial( uint n, float* cover_bit_prob0, u8* message, uint stc_constraint_height, uint &num_msg_bits, uint* perm, u8* stego,
        uint &trial, uint max_trials, const char* debugging_file = "cost.txt" ) {

    bool success = false;
    u8* cover = new u8[n];
    double* cost = new double[n];
    while ( !success ) {
        randperm( n, num_msg_bits, perm );
        for ( uint i = 0; i < n; i++ ) {
            cover[perm[i]] = (cover_bit_prob0[i] < 0.5) ? 1 : 0;
            cost[perm[i]] = -log( (1 / std::max( cover_bit_prob0[i], 1 - cover_bit_prob0[i] )) - 1 );
            if ( cost[perm[i]] != cost[perm[i]] ) // if p20[i]>1 due to numerical error (this is possible due to float data type)
            cost[perm[i]] = D_INF; // then cost2[i] is NaN, it should be Inf
        }
        memcpy( stego, cover, n ); // initialize stego array by cover array
        // debugging
        // write_vector_to_file<double>(n, cost, debugging_file);
        try {
            if ( num_msg_bits != 0 ) stc_embed( cover, n, message, num_msg_bits, (void*) cost, true, stego, stc_constraint_height );
            success = true;
        } catch ( stc_exception& e ) {
            if ( e.error_id != 4 ) { // error_id=4 means No solution exists, thus we try to embed with different permutation.
                delete[] cost;
                delete[] cover;
                throw e;
            }
            num_msg_bits--; // by decreasing the number of  bits, we change the permutation used to shuffle the bits
            trial++;
            if ( trial > max_trials ) {
                delete[] cost;
                delete[] cover;
                throw stc_exception( "Maximum number of trials in layered construction exceeded (2).", 6 );
            }
        }
    }
    delete[] cost;
    delete[] cover;
}
// }}}

// {{{ check_costs()
// SANITY CHECKS for cost arrays
void check_costs( uint n, uint k, float *costs ) {

    bool test_nan, test_non_inf, test_minus_inf;
    for ( uint i = 0; i < n; i++ ) {
        test_nan = false; // Is any element NaN? Should be FALSE
        test_non_inf = false; // Is any element finite? Should be TRUE
        test_minus_inf = false; // Is any element minus Inf? should be FALSE
        for ( uint j = 0; j < k; j++ ) {
            test_nan |= (costs[k * i + j] != costs[k * i + j]);
            test_non_inf |= ((costs[k * i + j] != -F_INF) & (costs[k * i + j] != F_INF));
            test_minus_inf |= (costs[k * i + j] == -F_INF);
        }
        if ( test_nan ) {
            std::stringstream ss;
            ss << "Incorrect cost array." << i << "-th element contains NaN value. This is not a valid cost.";
            throw stc_exception( ss.str(), 6 );
        }
        if ( !test_non_inf ) {
            std::stringstream ss;
            ss << "Incorrect cost array." << i << "-th element does not contain any finite cost value. This is not a valid cost.";
            throw stc_exception( ss.str(), 6 );
        }
        if ( test_minus_inf ) {
            std::stringstream ss;
            ss << "Incorrect cost array." << i << "-th element contains -Inf value. This is not a valid cost.";
            throw stc_exception( ss.str(), 6 );
        }
    }
}
// }}}

// {{{ stc_pm1_pls_embed()
// MULTI-LAYERED EMBEDDING for plus/minus one changes
// payload limited case - returns distortion
float stc_pm1_pls_embed( uint cover_length, int* cover, float* costs, uint message_length, u8* message, // input variables
                         uint stc_constraint_height, float wet_cost,                                    // other input parameters
                         int* stego, uint* num_msg_bits, uint &max_trials, float* coding_loss ) {       // output variables

    return stc_pm1_dls_embed( cover_length, cover, costs, message_length, message, F_INF, stc_constraint_height, 0, wet_cost, stego,
            num_msg_bits, max_trials, coding_loss );
}
// }}}

// {{{ stc_pm1_dls_embed()
// distortion limited case - returns distortion
float stc_pm1_dls_embed( uint cover_length, int* cover, float* costs, uint message_length, u8* message, float target_distortion, // input variables
                         uint stc_constraint_height, float expected_coding_loss, float wet_cost,   // other input parameters
                         int* stego, uint* num_msg_bits, uint &max_trials, float* coding_loss ) {  // output variables

    check_costs( cover_length, 3, costs );
    float dist = 0;

    int *stego_values = new int[4 * cover_length];
    float *costs_ml2 = new float[4 * cover_length];
    for ( uint i = 0; i < cover_length; i++ ) {
        costs_ml2[4 * i + mod( (cover[i] - 1 + 4), 4 )] = costs[3 * i + 0]; // set cost of changing by -1
        stego_values[4 * i + mod( (cover[i] - 1 + 4), 4 )] = cover[i] - 1;
        costs_ml2[4 * i + mod( (cover[i] + 0 + 4), 4 )] = costs[3 * i + 1]; // set cost of changing by 0
        stego_values[4 * i + mod( (cover[i] + 0 + 4), 4 )] = cover[i];
        costs_ml2[4 * i + mod( (cover[i] + 1 + 4), 4 )] = costs[3 * i + 2]; // set cost of changing by +1
        stego_values[4 * i + mod( (cover[i] + 1 + 4), 4 )] = cover[i] + 1;
        costs_ml2[4 * i + mod( (cover[i] + 2 + 4), 4 )] = wet_cost; // set cost of changing by +2
        stego_values[4 * i + mod( (cover[i] + 2 + 4), 4 )] = cover[i] + 2;
    }

    // run general 2 layered embedding in distortion limited regime
    dist = stc_ml2_embed( cover_length, costs_ml2, stego_values, message_length, message, target_distortion, stc_constraint_height,
            expected_coding_loss, stego, num_msg_bits, max_trials, coding_loss );
    delete[] costs_ml2;
    delete[] stego_values;

    return dist;
}
// }}}

// {{{ stc_pm2_dls_embed()
// MULTI-LAYERED EMBEDDING for plus/minus one and two changes
// payload limited case - returns distortion
float stc_pm2_pls_embed( uint cover_length, int* cover, float* costs, uint message_length, u8* message, // input variables
        uint stc_constraint_height, float wet_cost, // other input parameters
        int* stego, uint* num_msg_bits, uint &max_trials, float* coding_loss ) { // output variables

    return stc_pm2_dls_embed( cover_length, cover, costs, message_length, message, F_INF, stc_constraint_height, 0, wet_cost, stego,
            num_msg_bits, max_trials, coding_loss );
}
// }}}

// {{{ stc_pm2_dls_embed()
// distortion limited case - returns distortion
float stc_pm2_dls_embed( uint cover_length, int* cover, float* costs, uint message_length, u8* message, float target_distortion, // input variables
        uint stc_constraint_height, float expected_coding_loss, float wet_cost, // other input parameters
        int* stego, uint* num_msg_bits, uint &max_trials, float* coding_loss ) { // output variables

    check_costs( cover_length, 5, costs );
    int *stego_values = new int[8 * cover_length];
    float* costs_ml3 = new float[8 * cover_length];
    std::fill_n( costs_ml3, 8 * cover_length, wet_cost ); // initialize new cost array

    for ( uint i = 0; i < cover_length; i++ ) {
        costs_ml3[8 * i + mod( (cover[i] - 2 + 8), 8 )] = costs[5 * i + 0]; // set cost of changing by -2
        stego_values[8 * i + mod( (cover[i] - 2 + 8), 8 )] = cover[i] - 2;
        costs_ml3[8 * i + mod( (cover[i] - 1 + 8), 8 )] = costs[5 * i + 1]; // set cost of changing by -1
        stego_values[8 * i + mod( (cover[i] - 1 + 8), 8 )] = cover[i] - 1;
        costs_ml3[8 * i + mod( (cover[i] + 0 + 8), 8 )] = costs[5 * i + 2]; // set cost of changing by 0
        stego_values[8 * i + mod( (cover[i] + 0 + 8), 8 )] = cover[i] + 0;
        costs_ml3[8 * i + mod( (cover[i] + 1 + 8), 8 )] = costs[5 * i + 3]; // set cost of changing by +1
        stego_values[8 * i + mod( (cover[i] + 1 + 8), 8 )] = cover[i] + 1;
        costs_ml3[8 * i + mod( (cover[i] + 2 + 8), 8 )] = costs[5 * i + 4]; // set cost of changing by +2
        stego_values[8 * i + mod( (cover[i] + 2 + 8), 8 )] = cover[i] + 2;
        stego_values[8 * i + mod( (cover[i] + 3 + 8), 8 )] = cover[i] + 3; // these values are not used and are defined
        stego_values[8 * i + mod( (cover[i] + 4 + 8), 8 )] = cover[i] + 4; // just to have the array complete
        stego_values[8 * i + mod( (cover[i] + 5 + 8), 8 )] = cover[i] + 5; //
    }

    // run general 3 layered embedding in distortion limited regime
    float dist = stc_ml3_embed( cover_length, costs_ml3, stego_values, message_length, message, target_distortion, stc_constraint_height,
            expected_coding_loss, stego, num_msg_bits, max_trials, coding_loss );
    delete[] costs_ml3;
    delete[] stego_values;

    return dist;
}
// }}}

// GENERAL MULTI-LAYERED EMBEDDING

// {{{ stc_ml1_embed()
// algorithm for embedding into 1 layer, both payload- and distortion-limited case
float stc_ml1_embed( uint cover_length, int* cover, short* direction, float* costs, uint message_length, u8* message,
        float target_distortion,// input variables
        uint stc_constraint_height, float expected_coding_loss, // other input parameters
        int* stego, uint* num_msg_bits, uint &max_trials, float* coding_loss ) { // output variables

    float distortion, lambda = 0, m_max = 0;
    bool success = false;
    uint m_actual = 0;
    uint n = cover_length + 4 - (cover_length % 4); // cover length rounded to multiple of 4
    uint *perm1 = new uint[n];

    float* c = align_new< float > ( 2 * n, 16 );
    std::fill_n( c, 2 * n, F_INF );
    std::fill_n( c, n, 0 );
    for ( uint i = 0; i < cover_length; i++ ) { // copy and transpose data for better reading via SSE instructions
        c[mod( cover[i], 2 ) * n + i] = 0; // cost of not changing the element
        c[mod( (cover[i] + 1), 2 ) * n + i] = costs[i]; // cost of changing the element
    }

    if ( target_distortion != F_INF ) { // distortion-limited sender
        lambda = get_lambda_distortion( n, 2, c, target_distortion, 2 ); //
        m_max = (1 - expected_coding_loss) * calc_entropy( n, 2, c, lambda ); //
        m_actual = std::min( message_length, (uint) floor( m_max ) ); //
    }
    if ( (target_distortion == F_INF) || (m_actual < floor( m_max )) ) { // payload-limited sender
        m_actual = std::min( cover_length, message_length ); // or distortion-limited sender with
    }

    /* SINGLE LAYER OF 1ST LSBs */
    num_msg_bits[0] = m_actual;
    uint trial = 0;
    u8* cover1 = new u8[cover_length];
    double* cost1 = new double[cover_length];
    u8* stego1 = new u8[cover_length];
    while ( !success ) {
        randperm( cover_length, num_msg_bits[0], perm1 );
        for ( uint i = 0; i < cover_length; i++ ) {
            cover1[perm1[i]] = mod( cover[i], 2 );
            cost1[perm1[i]] = costs[i];
            if ( cost1[perm1[i]] != cost1[perm1[i]] ) cost1[perm1[i]] = D_INF;
        }
        memcpy( stego1, cover1, cover_length ); // initialize stego array by cover array
        // debugging
        // write_vector_to_file<double>(n, cost, debugging_file);
        try {
            if ( num_msg_bits[0] != 0 ) stc_embed( cover1, cover_length, message, num_msg_bits[0], (void*) cost1, true, stego1,
                    stc_constraint_height );
            success = true;
        } catch ( stc_exception& e ) {
            if ( e.error_id != 4 ) { // error_id=4 means No solution exists, thus we try to embed with different permutation.
                delete[] cost1;
                delete[] cover1;
                delete[] stego1;
                delete[] perm1;
                delete[] c;
                throw e;
            }
            num_msg_bits[0]--; // by decreasing the number of  bits, we change the permutation used to shuffle the bits
            trial++;
            if ( trial > max_trials ) {
                delete[] cost1;
                delete[] cover1;
                delete[] stego1;
                delete[] perm1;
                delete[] c;
                throw stc_exception( "Maximum number of trials in layered construction exceeded (1).", 6 );
            }
        }
    }

    /* FINAL CALCULATIONS */
    distortion = 0;
    for ( uint i = 0; i < cover_length; i++ ) {
        stego[i] = (stego1[perm1[i]] == cover1[perm1[i]]) ? cover[i] : cover[i] + direction[i];
        distortion += (stego1[perm1[i]] == cover1[perm1[i]]) ? 0 : costs[i];
    }
    if ( coding_loss != 0 ) {
        float lambda_dist = get_lambda_distortion( n, 2, c, distortion, lambda, 0, 20 ); // use 20 iterations to make lambda_dist precise
        float max_payload = calc_entropy( n, 2, c, lambda_dist );
        (*coding_loss) = (max_payload - m_actual) / max_payload; // fraction of max_payload lost due to practical coding scheme
    }
    max_trials = trial;

    delete[] cost1;
    delete[] cover1;
    delete[] stego1;
    delete[] perm1;
    align_delete< float > ( c );

    return distortion;
}
// }}}

// {{{ stc_ml2_embed()
// algorithm for embedding into 2 layers with possibility to use only 1 layer, both payload- and distortion-limited cases
float stc_ml2_embed( uint cover_length, float* costs, int* stego_values, uint message_length, u8* message, float target_distortion, // input variables
        uint stc_constraint_height, float expected_coding_loss, // other input parameters
        int* stego, uint* num_msg_bits, uint &max_trials, float* coding_loss ) { // output and optional variables

    float distortion, dist_coding_loss, lambda = 0, m_max = 0;
    uint m_actual = 0;
    uint n = cover_length + 4 - (cover_length % 4); // cover length rounded to multiple of 4

    check_costs( cover_length, 4, costs );
    // if only binary embedding is sufficient, then use only 1st LSB layer
    bool lsb1_only = true;
    for ( uint i = 0; i < cover_length; i++ ) {
        uint n_finite_costs = 0; // number of finite cost values
        uint lsb_xor = 0;
        for ( uint k = 0; k < 4; k++ )
            if ( costs[4 * i + k] != F_INF ) {
                n_finite_costs++;
                lsb_xor ^= (k % 2);
            }
        lsb1_only &= ((n_finite_costs <= 2) & (lsb_xor == 1));
    }
    if ( lsb1_only ) { // use stc_ml1_embed method
        distortion = 0;
        int *cover = new int[cover_length];
        short *direction = new short[cover_length];
        float *costs_ml1 = new float[cover_length];
        for ( uint i = 0; i < cover_length; i++ ) { // normalize such that minimal element is 0 - this helps numerical stability
            uint min_id = 0;
            float f_min = F_INF;
            for ( uint j = 0; j < 4; j++ )
                if ( f_min > costs[4 * i + j] ) {
                    f_min = costs[4 * i + j]; // minimum value
                    min_id = j; // index of the minimal entry
                }
            costs_ml1[i] = F_INF;
            cover[i] = stego_values[4 * i + min_id];
            for ( uint j = 0; j < 4; j++ )
                if ( (costs[4 * i + j] != F_INF) && (min_id != j) ) {
                    distortion += f_min;
                    costs_ml1[i] = costs[4 * i + j] - f_min;
                    direction[i] = stego_values[4 * i + j] - cover[i];
                }
        }

        distortion += stc_ml1_embed( cover_length, cover, direction, costs_ml1, message_length, message, target_distortion,
                stc_constraint_height, expected_coding_loss, stego, num_msg_bits, max_trials, coding_loss );
        delete[] direction;
        delete[] costs_ml1;
        delete[] cover;
        return distortion;
    }

    // copy and transpose data for faster reading via SSE instructions
    float* c = align_new< float > ( 4 * n, 16 );
    std::fill_n( c, 4 * n, F_INF );
    std::fill_n( c, n, 0 );
    for ( uint i = 0; i < 4 * cover_length; i++ )
        c[n * (i % 4) + i / 4] = costs[i];
    // write_matrix_to_file<float>(n, 4, c, "cost_ml2.txt");
    for ( uint i = 0; i < n; i++ ) { // normalize such that minimal element is 0 - this helps numerical stability
        float f_min = F_INF;
        for ( uint j = 0; j < 4; j++ )
            f_min = std::min( f_min, c[j * n + i] );
        for ( uint j = 0; j < 4; j++ )
            c[j * n + i] -= f_min;
    }

    if ( target_distortion != F_INF ) {
        lambda = get_lambda_distortion( n, 4, c, target_distortion, 2 );
        m_max = (1 - expected_coding_loss) * calc_entropy( n, 4, c, lambda );
        m_actual = std::min( message_length, (uint) floor( m_max ) );
    }
    if ( (target_distortion == F_INF) || (m_actual < floor( m_max )) ) {
        m_actual = std::min( 2 * cover_length, message_length );
        lambda = get_lambda_entropy( n, 4, c, m_actual, 2 );
    }
    /*
     p = exp(-lambda*costs);
     p = p./(ones(4,1)*sum(p));
     */
    float* p = align_new< float > ( 4 * n, 16 );
    __m128 v_lambda = _mm_set1_ps( -lambda );
    for ( uint i = 0; i < n / 4; i++ ) {
        __m128 sum = _mm_setzero_ps();
        for ( uint j = 0; j < 4; j++ ) {
            __m128 x = _mm_load_ps( c + j * n + 4 * i );
            x = exp_ps( _mm_mul_ps( v_lambda, x ) );
            _mm_store_ps( p + j * n + 4 * i, x );
            sum = _mm_add_ps( sum, x );
        }
        for ( uint j = 0; j < 4; j++ ) {
            __m128 x = _mm_load_ps( p + j * n + 4 * i );
            x = _mm_div_ps( x, sum );
            _mm_store_ps( p + j * n + 4 * i, x );
        }
    }
    // this is for debugging purposes
    // float payload_dbg = entropy_array(4*n, p);

    uint trial = 0;
    float* p10 = new float[cover_length];
    float* p20 = new float[cover_length];
    u8* stego1 = new u8[cover_length];
    u8* stego2 = new u8[cover_length];
    uint *perm1 = new uint[cover_length];
    uint *perm2 = new uint[cover_length];

    /* LAYER OF 2ND LSBs */
    for ( uint i = 0; i < cover_length; i++ )
        p20[i] = p[i] + p[i + n]; // p20 = p(1,:)+p(2,:);         % probability of 2nd LSB of stego equal 0
    //num_msg_bits[1] = (uint) floor( binary_entropy_array( cover_length, p20 ) ); // msg_bits(2) = floor(sum(binary_entropy(p20)));    % number of msg bits embedded into 2nd LSBs
    num_msg_bits[1] = (uint) (message_length/2 /*+ message_length%2*/ ); // XXX

    try {
        stc_embed_trial( cover_length, p20, message, stc_constraint_height, num_msg_bits[1], perm2, stego2, trial, max_trials, "cost2.txt" );
    } catch ( stc_exception& e ) {
        delete[] p10;
        delete[] p20;
        delete[] perm1;
        delete[] perm2;
        delete[] stego1;
        delete[] stego2;
        align_delete< float > ( c );
        align_delete< float > ( p );
        throw e;
    }

    /* LAYER OF 1ST LSBs */
    for ( uint i = 0; i < cover_length; i++ ) //
        if ( stego2[perm2[i]] == 0 ) // % conditional probability of 1st LSB of stego equal 0 given LSB2=0
        p10[i] = p[i] / (p[i] + p[i + n]); // p10(i) = p(1,i)/(p(1,i)+p(2,i));
        else // % conditional probability of 1st LSB of stego equal 0 given LSB2=1
        p10[i] = p[i + 2 * n] / (p[i + 2 * n] + p[i + 3 * n]); // p10(i) = p(3,i)/(p(3,i)+p(4,i));
    num_msg_bits[0] = m_actual - num_msg_bits[1]; // msg_bits(1) = m_actual-msg_bits(2); % number of msg bits embedded into 1st LSBs
    try {
        stc_embed_trial( cover_length, p10, message + num_msg_bits[1], stc_constraint_height, num_msg_bits[0], perm1, stego1, trial,
                max_trials, "cost1.txt" );
    } catch ( stc_exception& e ) {
        delete[] p10;
        delete[] p20;
        delete[] perm1;
        delete[] perm2;
        delete[] stego1;
        delete[] stego2;
        align_delete< float > ( c );
        align_delete< float > ( p );
        throw e;
    }
    delete[] p10;
    delete[] p20;

    /* FINAL CALCULATIONS */
    distortion = 0;
    for ( uint i = 0; i < cover_length; i++ ) {
        stego[i] = stego_values[4 * i + 2 * stego2[perm2[i]] + stego1[perm1[i]]];
        distortion += costs[4 * i + 2 * stego2[perm2[i]] + stego1[perm1[i]]];
    }
    if ( coding_loss != 0 ) {
        dist_coding_loss = 0;
        for ( uint i = 0; i < cover_length; i++ )
            dist_coding_loss += c[i + n * (2 * stego2[perm2[i]] + stego1[perm1[i]])];
        float lambda_dist = get_lambda_distortion( n, 4, c, dist_coding_loss, lambda, 0, 20 ); // use 20 iterations to make lambda_dist precise
        float max_payload = calc_entropy( n, 4, c, lambda_dist );
        (*coding_loss) = (max_payload - m_actual) / max_payload; // fraction of max_payload lost due to practical coding scheme
    }
    max_trials = trial;

    delete[] stego1;
    delete[] stego2;
    delete[] perm1;
    delete[] perm2;
    align_delete< float > ( c );
    align_delete< float > ( p );

    return distortion;
}
// }}}

// {{{ stc_ml3_embed()
// algorithm for embedding into 3 layers, both payload- and distortion-limited case
float stc_ml3_embed( uint cover_length, float* costs, int* stego_values, uint message_length, u8* message, float target_distortion, // input variables
        uint stc_constraint_height, float expected_coding_loss, // other input parameters
        int* stego, uint* num_msg_bits, uint &max_trials, float* coding_loss ) { // output and optional variables

    float distortion, dist_coding_loss, lambda = 0, m_max = 0;
    uint m_actual = 0;
    uint n = cover_length + 4 - (cover_length % 4); // cover length rounded to multiple of 4

    check_costs( cover_length, 8, costs );
    float* c = align_new< float > ( 8 * n, 16 );
    std::fill_n( c, 8 * n, F_INF );
    std::fill_n( c, n, 0 );
    for ( uint i = 0; i < 8 * cover_length; i++ )
        c[n * (i % 8) + i / 8] = costs[i]; // copy and transpose data for better reading via SSE instructions
    // write_matrix_to_file<float>(n, 8, c, "cost_ml3.txt");
    for ( uint i = 0; i < n; i++ ) { // normalize such that minimal element is 0 - this helps numerical stability
        float f_min = F_INF;
        for ( uint j = 0; j < 8; j++ )
            f_min = std::min( f_min, c[j * n + i] );
        for ( uint j = 0; j < 8; j++ )
            c[j * n + i] -= f_min;
    }

    if ( target_distortion != F_INF ) {
        lambda = get_lambda_distortion( n, 8, c, target_distortion, 2.0 );
        m_max = (1 - expected_coding_loss) * calc_entropy( n, 8, c, lambda );
        m_actual = std::min( message_length, (uint) floor( m_max ) );
    }
    if ( (target_distortion == F_INF) || (m_actual < floor( m_max )) ) {
        m_actual = std::min( 3 * cover_length, message_length );
        lambda = get_lambda_entropy( n, 8, c, m_actual, 2.0 );
    }
    /*
     p = exp(-lambda*costs);
     p = p./(ones(8,1)*sum(p));
     */
    float* p = align_new< float > ( 8 * n, 16 );
    __m128 v_lambda = _mm_set1_ps( -lambda );
    for ( uint i = 0; i < n / 4; i++ ) {
        __m128 sum = _mm_setzero_ps();
        for ( uint j = 0; j < 8; j++ ) {
            __m128 x = _mm_load_ps( c + j * n + 4 * i );
            x = exp_ps( _mm_mul_ps( v_lambda, x ) );
            _mm_store_ps( p + j * n + 4 * i, x );
            sum = _mm_add_ps( sum, x );
        }
        for ( uint j = 0; j < 8; j++ ) {
            __m128 x = _mm_load_ps( p + j * n + 4 * i );
            x = _mm_div_ps( x, sum );
            _mm_store_ps( p + j * n + 4 * i, x );
        }
    }
    // this is for debugging
    // float payload_dbg = entropy_array(8*n, p);

    uint trial = 0;
    float* p10 = new float[cover_length];
    float* p20 = new float[cover_length];
    float* p30 = new float[cover_length];
    u8* stego1 = new u8[cover_length];
    u8* stego2 = new u8[cover_length];
    u8* stego3 = new u8[cover_length];
    uint *perm1 = new uint[cover_length];
    uint *perm2 = new uint[cover_length];
    uint *perm3 = new uint[cover_length];

    /* LAYER OF 3RD LSBs */
    for ( uint i = 0; i < cover_length; i++ )
        p30[i] = p[i] + p[i + n] + p[i + 2 * n] + p[i + 3 * n]; //
    num_msg_bits[2] = (uint) floor( binary_entropy_array( cover_length, p30 ) ); //
    try {
        stc_embed_trial( cover_length, p30, message, stc_constraint_height, num_msg_bits[2], perm3, stego3, trial, max_trials, "cost3.txt" );
    } catch ( stc_exception& e ) {
        delete[] p10;
        delete[] p20;
        delete[] p30;
        delete[] perm1;
        delete[] perm2;
        delete[] perm3;
        delete[] stego1;
        delete[] stego2;
        delete[] stego3;
        align_delete< float > ( c );
        align_delete< float > ( p );
        throw e;
    }

    /* LAYER OF 2ND LSBs */
    for ( uint i = 0; i < cover_length; i++ ) { //
        int s = 4 * stego3[perm3[i]]; // % conditional probability of 2nd LSB of stego equal 0 given LSB3
        p20[i] = (p[i + s * n] + p[i + (s + 1) * n]) / (p[i + s * n] + p[i + (s + 1) * n] + p[i + (s + 2) * n] + p[i + (s + 3) * n]);
    }
    num_msg_bits[1] = (uint) floor( binary_entropy_array( cover_length, p20 ) );// msg_bits(2) = floor(sum(binary_entropy(p20)));    % number of msg bits embedded into 2nd LSBs
    try {
        stc_embed_trial( cover_length, p20, message + num_msg_bits[2], stc_constraint_height, num_msg_bits[1], perm2, stego2, trial,
                max_trials, "cost2.txt" );
    } catch ( stc_exception& e ) {
        delete[] p10;
        delete[] p20;
        delete[] p30;
        delete[] perm1;
        delete[] perm2;
        delete[] perm3;
        delete[] stego1;
        delete[] stego2;
        delete[] stego3;
        align_delete< float > ( c );
        align_delete< float > ( p );
        throw e;
    }

    /* LAYER OF 1ST LSBs */
    for ( uint i = 0; i < cover_length; i++ ) { //
        int s = 4 * stego3[perm3[i]] + 2 * stego2[perm2[i]]; // % conditional probability of 1st LSB of stego equal 0 given LSB3 and LSB2
        p10[i] = p[i + s * n] / (p[i + s * n] + p[i + (s + 1) * n]);
    }
    num_msg_bits[0] = m_actual - num_msg_bits[1] - num_msg_bits[2]; // msg_bits(1) = m_actual-msg_bits(2)-msg_bits(3); % number of msg bits embedded into 1st LSBs
    try {
        stc_embed_trial( cover_length, p10, message + num_msg_bits[1] + num_msg_bits[2], stc_constraint_height, num_msg_bits[0], perm1,
                stego1, trial, max_trials, "cost1.txt" );
    } catch ( stc_exception& e ) {
        delete[] p10;
        delete[] p20;
        delete[] p30;
        delete[] perm1;
        delete[] perm2;
        delete[] perm3;
        delete[] stego1;
        delete[] stego2;
        delete[] stego3;
        align_delete< float > ( c );
        align_delete< float > ( p );
        throw e;
    }
    delete[] p10;
    delete[] p20;
    delete[] p30;
    max_trials = trial;

    /* FINAL CALCULATIONS */
    distortion = 0;
    for ( uint i = 0; i < cover_length; i++ ) {
        stego[i] = stego_values[8 * i + 4 * stego3[perm3[i]] + 2 * stego2[perm2[i]] + stego1[perm1[i]]];
        distortion += costs[8 * i + 4 * stego3[perm3[i]] + 2 * stego2[perm2[i]] + stego1[perm1[i]]];
    }
    if ( coding_loss != 0 ) {
        dist_coding_loss = 0;
        for ( uint i = 0; i < cover_length; i++ )
            dist_coding_loss += c[i + n * (4 * stego3[perm3[i]] + 2 * stego2[perm2[i]] + stego1[perm1[i]])];
        float lambda_dist = get_lambda_distortion( n, 8, c, dist_coding_loss, lambda, 0, 20 ); // use 20 iterations to make lambda_dist precise
        float max_payload = calc_entropy( n, 8, c, lambda_dist );
        (*coding_loss) = (max_payload - m_actual) / max_payload; // fraction of max_payload lost due to practical coding scheme
    }

    delete[] perm1;
    delete[] perm2;
    delete[] perm3;
    delete[] stego1;
    delete[] stego2;
    delete[] stego3;
    align_delete< float > ( c );
    align_delete< float > ( p );

    return distortion;
}
// }}}

/* EXTRACTION ALGORITHMS */

// {{{ stc_ml_extract()
/** Extraction algorithm for any l-layered construction.
 @param stego_length - ...
 @param stego - ...
 @param msg_bits - ...
 @param stc_constraint_height - ...
 @param message - ...
 */
void stc_ml_extract( uint stego_length, int* stego, uint num_of_layers, uint* num_msg_bits, // input variables
                     uint stc_constraint_height, // other input parameters
                     u8* message ) { // output variables

    u8* stego_bits = new u8[stego_length];
    u8* msg_ptr = message;
    uint *perm = new uint[stego_length];

    for ( uint l = num_of_layers; l > 0; l-- ) { // extract message from every layer starting from most significant ones
        // extract bits from l-th LSB plane
        if ( num_msg_bits[l - 1] > 0 ) {
            randperm( stego_length, num_msg_bits[l - 1], perm );
            for ( uint i = 0; i < stego_length; i++ )
                stego_bits[perm[i]] = mod( stego[i], (1 << l) ) >> (l - 1);
            stc_extract( stego_bits, stego_length, msg_ptr, num_msg_bits[l - 1], stc_constraint_height );
            msg_ptr += num_msg_bits[l - 1];
        }
    }

    delete[] stego_bits;
    delete[] perm;
}
// }}}


//////////////////////////////////////////////////////////////////////////////
//
// (C) Copyright Alejandro Cabrera 2011.
// Distributed under the Boost
// Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt or
// copy at http://www.boost.org/LICENSE_1_0.txt)
//
// See http://www.boost.org/libs/bloom_filter for documentation.
//
//////////////////////////////////////////////////////////////////////////////

#ifndef BOOST_BLOOM_FILTER_DYNAMIC_COUNTING_BLOOM_FILTER_HPP
#define BOOST_BLOOM_FILTER_DYNAMIC_COUNTING_BLOOM_FILTER_HPP 1

#include <cmath>
#include <vector>

#include <boost/config.hpp>

#include <boost/mpl/vector.hpp>
#include <boost/mpl/size.hpp>

#include <boost/static_assert.hpp>
#include <boost/type_traits/is_integral.hpp>
#include <boost/type_traits/is_unsigned.hpp>

#include <boost/bloom_filter/detail/counting_apply_hash.hpp>
#include <boost/bloom_filter/hash/default.hpp>

namespace boost {
  namespace bloom_filters {
    template <typename T,
	      size_t BitsPerBin = 4,
	      class HashFunctions = mpl::vector<boost_hash<T> >,
	      typename Block = size_t,
	      typename Allocator = std::allocator<Block> >
    class dynamic_counting_bloom_filter {

      // Block needs to be an integral type
      BOOST_STATIC_ASSERT( boost::is_integral<Block>::value == true);

      // Block needs to be an unsigned type
      BOOST_STATIC_ASSERT( boost::is_unsigned<Block>::value == true);

      // BitsPerBin needs to be greater than 0
      BOOST_STATIC_ASSERT( BitsPerBin > 0);

      // it doesn't make sense to ever support using a BitsPerBin value larger
      // than the number of bits per Block. In that case, the user shouldn't
      // be using a Bloom filter to represent their data.
      BOOST_STATIC_ASSERT( (BitsPerBin < (sizeof(Block) * 8) ) );

      // because of the nature of this implementation, the Bloom filter
      // can have internal fragmentation if the calculation for
      // bins_per_slot has a remainder. The severity of the  internal
      // fragmentation is equal to the remainder * the number of slots.
      // This check prevents internal fragmentation.
      // This also necessarily limits to bin sizes to one of:
      // [1,2,4,8,16,32(64-bit system only)] bits
      BOOST_STATIC_ASSERT( ((sizeof(Block) * 8) % BitsPerBin) == 0);

    public:
      typedef T value_type;
      typedef T key_type;
      typedef HashFunctions hash_function_type;
      typedef Block block_type;
      typedef Allocator allocator_type;
      typedef dynamic_counting_bloom_filter<T, BitsPerBin,
					    HashFunctions,
					    Block, Allocator> this_type;

      typedef std::vector<Block, Allocator> bucket_type;
      typedef typename bucket_type::iterator bucket_iterator;
      typedef typename bucket_type::const_iterator bucket_const_iterator;

      static const size_t slot_bits = sizeof(block_type) * 8;
      static const size_t default_num_bins = 32;

    private:
      size_t bucket_size(const size_t requested_bins) const {
	const size_t bin_bits = requested_bins * BitsPerBin;
	return bin_bits / slot_bits + 1;
      }

      typedef detail::counting_apply_hash<mpl::size<HashFunctions>::value - 1,
					  this_type> apply_hash_type;

    public:
      //* constructors
      dynamic_counting_bloom_filter()
	: bits(bucket_size(default_num_bins)),
	  _num_bins(default_num_bins)
      {
      }

      explicit dynamic_counting_bloom_filter(const size_t requested_bins)
	: bits(bucket_size(requested_bins)),
	  _num_bins(requested_bins)
      {
      }

      template <typename InputIterator>
      dynamic_counting_bloom_filter(const InputIterator start,
				    const InputIterator end)
	: bits(bucket_size(std::distance(start, end) * 4)),
	  _num_bins(std::distance(start, end) * 4)
      {
	for (InputIterator i = start; i != end; ++i)
	  this->insert(*i);
      }

      //* meta functions
      size_t num_bins() const
      {
	return this->_num_bins;
      }

      static BOOST_CONSTEXPR size_t bits_per_bin()
      {
	return BitsPerBin;
      }

      static BOOST_CONSTEXPR size_t bins_per_slot()
      {
	return sizeof(block_type) * 8 / BitsPerBin;
      }

      static BOOST_CONSTEXPR size_t mask()
      {
	return static_cast<Block>(0 - 1) >> (sizeof(Block) * 8 - BitsPerBin);
      }

      size_t bit_capacity() const
      {
        return this->num_bins() * BitsPerBin;
      }

      static BOOST_CONSTEXPR size_t num_hash_functions()
      {
        return mpl::size<HashFunctions>::value;
      }

      double false_positive_rate() const
      {
        const double n = static_cast<double>(this->count());
        static const double k = static_cast<double>(num_hash_functions());
        static const double m = static_cast<double>(this->num_bins());
        static const double e =
	  2.718281828459045235360287471352662497757247093699959574966;
        return std::pow(1 - std::pow(e, -k * n / m), k);
      }

      //? returns the number of bins that have at least 1 bit set
      size_t count() const
      {
	size_t ret = 0;

	for (bucket_const_iterator i = this->bits.begin(),
	       end = this->bits.end();
	     i != end; ++i) {
	  for (size_t bin = 0; bin < this->bins_per_slot(); ++bin) {
	    const size_t offset_bits = bin * BitsPerBin;
	    const size_t target_bits = (*i >> offset_bits) & this->mask();

	    if (target_bits > 0)
	      ++ret;
	  }
	}

        return ret;
      }

      bool empty() const
      {
	return this->count() == 0;
      }

      const bucket_type&
      data() const
      {
	return this->bits;
      }

      //* core ops
      void insert(const T& t)
      {
	apply_hash_type::insert(t,
				this->bits,
				this->num_bins());
      }

      template <typename InputIterator>
      void insert(const InputIterator start, const InputIterator end)
      {
	for (InputIterator i = start; i != end; ++i) {
	  this->insert(*i);
	}
      }

      void remove(const T& t)
      {
	apply_hash_type::remove(t,
				this->bits,
				this->num_bins());
      }

      template <typename InputIterator>
      void remove(const InputIterator start, const InputIterator end)
      {
	for (InputIterator i = start; i != end; ++i) {
	  this->remove(*i);
	}
      }

      bool probably_contains(const T& t) const
      {
	return apply_hash_type::contains(t,
					 this->bits,
					 this->num_bins());
      }

      //* auxiliary ops
      void clear()
      {
	for (bucket_iterator i = bits.begin(), end = bits.end();
	     i != end; ++i)
	  *i = 0;
      }

      void swap(dynamic_counting_bloom_filter& other)
      {
	dynamic_counting_bloom_filter tmp = other;
	other = *this;
	*this = tmp;
      }

      //* equality comparison operators
      template <typename _T, size_t _BitsPerBin,
		typename _HashFns, typename _Block, typename _Allocator>
      friend bool
      operator==(const dynamic_counting_bloom_filter<_T, _BitsPerBin,
						     _HashFns, _Block,
						     _Allocator>& lhs,
		 const dynamic_counting_bloom_filter<_T, _BitsPerBin,
						     _HashFns, _Block,
						     _Allocator>& rhs);

      template <typename _T, size_t _BitsPerBin,
		typename _HashFns, typename _Block, typename _Allocator>
      friend bool
      operator!=(const dynamic_counting_bloom_filter<_T, _BitsPerBin,
						     _HashFns, _Block,
						     _Allocator>& lhs,
		 const dynamic_counting_bloom_filter<_T, _BitsPerBin,
						     _HashFns, _Block,
						     _Allocator>& rhs);

    private:
      bucket_type bits;
      size_t _num_bins;
    };

    template<class T, size_t BitsPerBin, class HashFunctions,
	     typename Block, typename Allocator>
    void
    swap(dynamic_counting_bloom_filter<T, BitsPerBin,
				       HashFunctions, Block,
				       Allocator>& lhs,
	 dynamic_counting_bloom_filter<T, BitsPerBin,
				       HashFunctions, Block,
				       Allocator>& rhs)

    {
      lhs.swap(rhs);
    }

    template<class T, size_t BitsPerBin, class HashFunctions,
	     typename Block, typename Allocator>
    bool
    operator==(const dynamic_counting_bloom_filter<T, BitsPerBin,
						   HashFunctions,
						   Block,
						   Allocator>& lhs,
	       const dynamic_counting_bloom_filter<T, BitsPerBin,
						   HashFunctions,
						   Block,
						   Allocator>& rhs)
    {
      if (lhs.bit_capacity() != rhs.bit_capacity())
	throw detail::incompatible_size_exception();

      return (lhs.bits == rhs.bits);
    }

    template<class T, size_t BitsPerBin, class HashFunctions,
	     typename Block, typename Allocator>
    bool
    operator!=(const dynamic_counting_bloom_filter<T, BitsPerBin,
						   HashFunctions,
						   Block,
						   Allocator>& lhs,
	       const dynamic_counting_bloom_filter<T, BitsPerBin,
						   HashFunctions,
						   Block,
						   Allocator>& rhs)
    {
      if (lhs.bit_capacity() != rhs.bit_capacity())
	throw detail::incompatible_size_exception();

      return !(lhs == rhs);
    }

  } // namespace bloom_filter
} // namespace boost
#endif


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

#include "data.hpp"
#include "../helper.hpp"

#include "g2o/core/sparse_optimizer.h"
#include "g2o/core/block_solver.h"
#include "g2o/core/solver.h"
#include "g2o/core/optimization_algorithm_levenberg.h"
#include "g2o/core/base_vertex.h"
#include "g2o/core/base_unary_edge.h"
#include "g2o/solvers/csparse/linear_solver_csparse.h"

using namespace std;

typedef g2o::BlockSolver< g2o::BlockSolverTraits<2, 1> >  KLTBlockSolver;
typedef g2o::LinearSolverCSparse<KLTBlockSolver::PoseMatrixType> KLTLinearSolver;

class VertexKLT : public g2o::BaseVertex<2, Eigen::Vector2d>
{
  public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW;
    VertexKLT()
    {
    }

    virtual bool read(std::istream& /*is*/)
    {
      cerr << __PRETTY_FUNCTION__ << " not implemented yet" << endl;
      return false;
    }

    virtual bool write(std::ostream& /*os*/) const
    {
      cerr << __PRETTY_FUNCTION__ << " not implemented yet" << endl;
      return false;
    }

    virtual void setToOriginImpl()
    {
      cerr << __PRETTY_FUNCTION__ << " not implemented yet" << endl;
    }

    virtual void oplusImpl(const double* update)
    {
      Eigen::Vector2d::ConstMapType v(update);
      _estimate += v;
    }
};

class EdgeKLT : public g2o::BaseUnaryEdge<1, Eigen::VectorXd, VertexKLT>
{
  public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW
    EdgeKLT(const cv::Mat& _img1, const cv::Mat& _Ix, const cv::Mat& _Iy)
    :img1(_img1), Ix(_Ix), Iy(_Iy){}

    virtual bool read(std::istream& /*is*/)
    {
      cerr << __PRETTY_FUNCTION__ << " not implemented yet" << endl;
      return false;
    }
    virtual bool write(std::ostream& /*os*/) const
    {
      cerr << __PRETTY_FUNCTION__ << " not implemented yet" << endl;
      return false;
    }

    void computeError();
    // void linearizeOplus();

  private:
    const cv::Mat& img1;
    const cv::Mat& Ix;
    const cv::Mat& Iy;
};


/* Siconos is a program dedicated to modeling, simulation and control
 * of non smooth dynamical systems.
 *
 * Copyright 2021 INRIA.
 *
 * 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.
*/

#include "SiconosConfig.h"

#pragma GCC diagnostic ignored "-Wunused-local-typedefs"

#include <boost/numeric/ublas/lu.hpp>
#include <boost/numeric/ublas/operation.hpp>
#include <boost/numeric/ublas/operation_sparse.hpp>

//#define BIND_FORTRAN_LOWERCASE_UNDERSCORE
#include <boost/numeric/bindings/ublas/vector_proxy.hpp>
#include <boost/numeric/bindings/ublas/matrix_proxy.hpp>
#include <boost/numeric/bindings/trans.hpp>
#include <boost/numeric/bindings/blas.hpp>
#include <boost/numeric/bindings/lapack.hpp>
#include <boost/numeric/bindings/ublas/vector.hpp>
#include <boost/numeric/bindings/ublas/matrix.hpp>
#include <boost/numeric/bindings/std/vector.hpp>

namespace lapack = boost::numeric::bindings::lapack;

#include "SiconosVector.hpp"
#include "cholesky.hpp"
#include "SimpleMatrix.hpp"
#include "BlockMatrixIterators.hpp"
#include "BlockMatrix.hpp"

#include "SiconosAlgebra.hpp"

#include "NumericsMatrix.h"
#include "NumericsSparseMatrix.h"
#include "CSparseMatrix.h"

//#define DEBUG_MESSAGES
#include "siconos_debug.h"

#ifdef DEBUG_MESSAGES
#include "NumericsVector.h"
#include <cs.h>
#endif

using namespace Siconos;

void SimpleMatrix::PLUFactorizationInPlace()
{
  if (_isPLUFactorized)
  {
    std::cout << "SimpleMatrix::PLUFactorizationInPlace warning: this matrix is already PLUFactorized. " << std::endl;
    return;
  }
  if (_num == Siconos::DENSE)
  {
    if (!_ipiv)
      _ipiv.reset(new VInt(size(0)));
    else
      _ipiv->resize(size(0));
    int info = lapack::getrf(*mat.Dense, *_ipiv);
    if (info != 0)
    {
      _isPLUFactorized = false;
      _isPLUFactorizedInPlace = true;
      THROW_EXCEPTION("SimpleMatrix::PLUFactorizationInPlace failed: the matrix is singular.");
    }
    else
    {
      _isPLUFactorized = true;
      _isPLUFactorizedInPlace = true;
    }
  }
  else
  {
    int info = cholesky_decompose(*sparse());
    // \warning: VA 24/11/2010: work only for symmetric matrices. Should be replaced by efficient implementatation (e.g. mumps )
    if (info != 0)
    {
      display();
      _isPLUFactorized = false;
      _isPLUFactorizedInPlace = false;
      std::cout << "Problem in Cholesky Decomposition for the row number" << info   << std::endl;
      THROW_EXCEPTION("SimpleMatrix::PLUFactorizationInPlace failed. ");
    }
    else
    {
      _isPLUFactorized = true;
      _isPLUFactorizedInPlace = false;
    }
  }

}

void SimpleMatrix::PLUInverseInPlace()
{
  if(!_isPLUFactorized)
    PLUFactorizationInPlace();
  if(_num != Siconos::DENSE)
    THROW_EXCEPTION(" SimpleMatrix::PLUInverseInPlace: only implemented for dense matrices.");

#if defined(HAS_LAPACK_dgetri)
  int info = lapack::getri(*mat.Dense, *_ipiv);   // solve from factorization

  if(info != 0)
    THROW_EXCEPTION("SimpleMatrix::PLUInverseInPlace failed, the matrix is singular.");

  _isPLUInversed = true;
#else
  THROW_EXCEPTION("SimpleMatrix::PLUInverseInPlace not implemented with lapack.");
#endif
}

void SimpleMatrix::PLUForwardBackwardInPlace(SiconosMatrix &B)
{
  if(B.isBlock())
    THROW_EXCEPTION("SimpleMatrix PLUForwardBackwardInPlace(B) failed at solving Ax = B. Not yet implemented for a BlockMatrix B.");
  int info = 0;

  if(_num == Siconos::DENSE)
  {
    if(!_isPLUFactorized)  // call gesv => LU-factorize+solve
    {
      // solve system:
      if(!_ipiv)
        _ipiv.reset(new VInt(size(0)));
      else
        _ipiv->resize(size(0));
      info = lapack::gesv(*mat.Dense, *_ipiv, *(B.dense()));
      _isPLUFactorized = true;

      /*
        ublas::vector<double> S(std::max(size(0),size(1)));
        ublas::matrix<double, ublas::column_major> U(size(0),size(1));
        ublas::matrix<double, ublas::column_major> VT(size(0),size(1));

        int ierr = lapack::gesdd(*mat.Dense, S, U, VT);
        printf("info = %d, ierr = %d, emax = %f, emin = %f , cond = %f\n",info,ierr,S(0),S(2),S(0)/S(2));
      */
      // B now contains solution:
    }
    else // call getrs: only solve using previous lu-factorization
      if(B.num() == DENSE)
        info = lapack::getrs(*mat.Dense, *_ipiv, *(B.dense()));
      else
        THROW_EXCEPTION(" SimpleMatrix::PLUInverseInPlace: only implemented for dense matrices in RHS.");

  }
  else
  {
    if(!_isPLUFactorized)  // call first PLUFactorizationInPlace
    {
      PLUFactorizationInPlace();
    }
    // and then solve
    if(B.num() == Siconos::DENSE)
    {
      inplace_solve(*sparse(), *(B.dense()), ublas::lower_tag());
      inplace_solve(ublas::trans(*sparse()), *(B.dense()), ublas::upper_tag());
    }
    else if(B.num() == Siconos::SPARSE)
    {
      inplace_solve(*sparse(), *(B.sparse()), ublas::lower_tag());
      inplace_solve(ublas::trans(*sparse()), *(B.sparse()), ublas::upper_tag());
    }
    else
      THROW_EXCEPTION(" SimpleMatrix::PLUInverseInPlace: only implemented for dense ans sparse matrices in RHS.");
    info = 0 ;
  }
  //  THROW_EXCEPTION(" SimpleMatrix::PLUInverseInPlace: only implemented for dense matrices.");

  if(info != 0)
    THROW_EXCEPTION("SimpleMatrix::PLUForwardBackwardInPlace failed.");
}

void SimpleMatrix::PLUForwardBackwardInPlace(SiconosVector &B)
{
  DenseMat tmpB(B.size(), 1);
  ublas::column(tmpB, 0) = *(B.dense()); // Conversion of vector to matrix. Temporary solution.
  int info;

  if(_num == Siconos::DENSE)
  {
    if(!_isPLUFactorized)  // call gesv => LU-factorize+solve
    {
      // solve system:
      if(!_ipiv)
        _ipiv.reset(new VInt(size(0)));
      else
        _ipiv->resize(size(0));

      info = lapack::gesv(*mat.Dense, *_ipiv, tmpB);
      _isPLUFactorized = true;

      /*
        ublas::matrix<double> COPY(*mat.Dense);
        ublas::vector<double> S(std::max(size(0),size(1)));
        ublas::matrix<double, ublas::column_major> U(size(0),size(1));
`        ublas::matrix<double, ublas::column_major> VT(size(0),size(1));

        int ierr = lapack::gesdd(COPY, S, U, VT);
        printf("info = %d, ierr = %d, emax = %f, emin = %f , cond = %f\n",info,ierr,S(0),S(2),S(0)/S(2));
      */
      // B now contains solution:
    }
    else // call getrs: only solve using previous lu-factorization
      info = lapack::getrs(*mat.Dense, *_ipiv, tmpB);
  }
  else
  {
    if(!_isPLUFactorized)  // call first PLUFactorizationInPlace
    {
      PLUFactorizationInPlace();
    }
    // and then solve
    inplace_solve(*sparse(), tmpB, ublas::lower_tag());
    inplace_solve(ublas::trans(*sparse()), tmpB, ublas::upper_tag());
    info = 0;
  }
  if(info != 0)
    THROW_EXCEPTION("SimpleMatrix::PLUForwardBackwardInPlace failed.");

  noalias(*(B.dense())) = ublas::column(tmpB, 0);

}

void SimpleMatrix::resetLU()
{
  if(_ipiv) _ipiv->clear();
  _isPLUFactorized = false;
  _isPLUFactorizedInPlace = false;
  _isPLUInversed = false;
}

void SimpleMatrix::resetCholesky()
{
  _isCholeskyFactorized = false;
  _isCholeskyFactorizedInPlace = false;
}

void SimpleMatrix::resetQR()
{
  _isQRFactorized = false;

}

void SimpleMatrix::resetFactorizationFlags()
{
  resetLU();
  resetCholesky();
  resetQR();
}

// const SimpleMatrix operator * (const SimpleMatrix & A, const SimpleMatrix& B )
// {
//   return (DenseMat)prod(*A.dense() , *B.dense());
//   //  return A;
// }

void SimpleMatrix::SolveByLeastSquares(SiconosMatrix &B)
{
  if(B.isBlock())
    THROW_EXCEPTION("SimpleMatrix::SolveByLeastSquares(Siconos Matrix &B) failed. Not yet implemented for M being a BlockMatrix.");
  int info = 0;
#ifdef USE_OPTIMAL_WORKSPACE
  info += lapack::gels(*mat.Dense, *(B.dense()), lapack::optimal_workspace());
#endif
#ifdef USE_MINIMAL_WORKSPACE
  info += lapack::gels(*mat.Dense, *(B.dense()), lapack::minimal_workspace());
#endif
  if(info != 0)
    THROW_EXCEPTION("SimpleMatrix::SolveByLeastSquares failed.");
}

void SimpleMatrix::SolveByLeastSquares(SiconosVector &B)
{
  DenseMat tmpB(B.size(), 1);
  ublas::column(tmpB, 0) = *(B.dense()); // Conversion of vector to matrix. Temporary solution.
  int info = 0;

#ifdef USE_OPTIMAL_WORKSPACE
  info += lapack::gels(*mat.Dense, tmpB, lapack::optimal_workspace());
#endif
#ifdef USE_MINIMAL_WORKSPACE
  info += lapack::gels(*mat.Dense, tmpB, lapack::minimal_workspace());
#endif
  if(info != 0)
  {
    std::cout << "info = " << info << std::endl;
    THROW_EXCEPTION("SimpleMatrix::SolveByLeastSquares failed.");
  }
  else
  {
    noalias(*(B.dense())) = ublas::column(tmpB, 0);
  }

}

/*
void polePlacement(const SiconosMatrix& A, const SiconosVector& B, SiconosVector& P, bool transpose)
{
  unsigned int n = A.size(0);
  DenseMat AA(n, n);
  DenseMat Q(n, n);
  DenseVect tau(n);
  DenseVect BB(n);
  noalias(AA) = (*A.dense());
  lapack::gehrd(1, n, AA, tau);
  lapack::orghr(n, 1, n, Q, tau);
  noalias(BB) = prod(Q, *B.dense());
}
*/

void SimpleMatrix::Factorize()
{
  DEBUG_BEGIN("void SimpleMatrix::Factorize()\n");
  if(isFactorized())
  {
    std::cout << "SimpleMatrix::PLUFactorize warning: this matrix is already Factorized. " << std::endl;
    return;
  }

  /* set the numericsMatrix */

  updateNumericsMatrix();
  NumericsMatrix * NM = numericsMatrix();

  /* Factorization calling the right method in Numerics */
  int info =1;
  if (isSymmetric())
  {
    if (isPositiveDefinite()) // Cholesky Factorization
    {
      //std::cout << "Cholesky Factorize"<< std::endl;
      info  = NM_Cholesky_factorize(NM);

      if(info != 0)
      {
        _isCholeskyFactorized = false;
        THROW_EXCEPTION("SimpleMatrix::Factorize failed (Cholesky)");
      }
      else
      {
        _isCholeskyFactorized = true;
      }
    }
    else  //  LDLT Factorization
    {
      THROW_EXCEPTION("SimpleMatrix::Factorize failed: LDL^T not yet implemented.");
    }
  }
  else //  LU Factorization  by default
  {
    info  = NM_LU_factorize(NM);

    if(info != 0)
    {
      _isPLUFactorized = false;
      if(_num == DENSE)
        _isPLUFactorizedInPlace = true;
      THROW_EXCEPTION("SimpleMatrix::PLUFactorize failed: the matrix is singular.");
    }
    else
    {
      _isPLUFactorized = true;
       if(_num == DENSE)
         _isPLUFactorizedInPlace = true;
    }
  }
  DEBUG_END("void SimpleMatrix::Factorize()\n");
}

//#define SPARSE_RHS_COPY_TO_DENSE 1

void SimpleMatrix::Solve(SiconosMatrix &B)
{
  if(B.isBlock())
    THROW_EXCEPTION("SimpleMatrix Solve(B) failed at solving Ax = B. Not yet implemented for a BlockMatrix B.");

  int info = 1;
  if(!isFactorized())
  {
    Factorize();
  }
  // and then solve
  NumericsMatrix * NM = _numericsMatrix.get();

#ifdef SPARSE_RHS_COPY_TO_DENSE

  double * b;
  SP::SimpleMatrix Bdense;

  if(B.num() == DENSE)
  {
    b = B.getArray();
  }
  else if(B.num() == SPARSE)
  {
    // First way.
    // We copy to dense since our sparse solver is not able to take
    // into account for sparse r.h.s, yet.
    Bdense.reset (new SimpleMatrix(size(0),size(1)));
    * Bdense= B ;                                                // copy to dense
    b = &(*Bdense->getArray());

    // Second way
    // use inplace_solve of ublas (see above with SolveInPlace)
    // For that, we need to fill our factorization given by NM_LU_factorize
    // into a ublas sparse matrix
    // inplace_solve(*sparse(), *(B.sparse()), ublas::lower_tag());
    // inplace_solve(ublas::trans(*sparse()), *(B.sparse()), ublas::upper_tag());

  }
  else
    THROW_EXCEPTION(" SimpleMatrix::Solve: only implemented for dense and sparse matrices in RHS.");

#else
  B.updateNumericsMatrix();
  NumericsMatrix * NM_B = B.numericsMatrix();
  //NM_display(NM_B);
#endif

  if (isSymmetric())
  {
    if (isPositiveDefinite()) // Cholesky Solving
    {
      //std::cout << "Cholesky Solve"<< std::endl;
#ifdef SPARSE_RHS_COPY_TO_DENSE
      info  = NM_Cholesky_solve(NM, b, B.size(1));
#else
      info  = NM_Cholesky_solve_matrix_rhs(NM, NM_B);
#endif
      if(info != 0)
      {
        THROW_EXCEPTION("SimpleMatrix::Solve failed (Cholesky)");
      }
    }
    else  //  LDLT Factorization
    {
      THROW_EXCEPTION("SimpleMatrix::Solve failed: LDL^T not yet implemented.");
    }
  }
  else //  LU Factorization  by default
  {
#ifdef SPARSE_RHS_COPY_TO_DENSE
    info  = NM_LU_solve(NM, b, B.size(1));
#else
    info  = NM_LU_solve_matrix_rhs(NM, NM_B);
#endif
    if(info != 0)
    {
      THROW_EXCEPTION("SimpleMatrix::PLUFactorize failed: the matrix is singular.");
    }
  }

  if(B.num() == SPARSE)
  {
#ifdef SPARSE_RHS_COPY_TO_DENSE
    B = *Bdense ; // we copy back to sparse.
    //B.displayExpert();
#else
    /* we need to fill back again */
    //B.fromCSC(NM_csc(NM_B));
#endif
  }

  if(info != 0)
    THROW_EXCEPTION("SimpleMatrix::Solve failed.");
}

void SimpleMatrix::Solve(SiconosVector &B)
{
  DEBUG_BEGIN("SimpleMatrix::PLUSolve(SiconosVector &B)\n");

  if(!isFactorized())
  {
    Factorize();
  }

  // and then solve
  int info =1;
  NumericsMatrix * NM;
  double * b;
  SP::SiconosVector Bdense;

  if(B.num() == DENSE)
  {
    NM = _numericsMatrix.get();
    b = B.getArray();
  }
  else if(B.num() == SPARSE)
  {
    // First way. We copy to dense since our sparse solver is not able to take into account sparse r.h.s
    Bdense.reset (new SiconosVector(size(0)));
    * Bdense= B ;
    b = &(*Bdense->getArray());
    NM = _numericsMatrix.get();

    // Second way use inplace_solve of ublas
    // For that, we need to fill our factorization given by NM_LU_factorize into a ublas sparse matrix
    //inplace_solve(*sparse(), *(B.sparse()), ublas::lower_tag());
    //inplace_solve(ublas::trans(*sparse()), *(B.sparse()), ublas::upper_tag());
  }
  else
    THROW_EXCEPTION(" SimpleMatrix::Solve: only implemented for dense and sparse matrices in RHS.");

  if (isSymmetric())
  {
    if (isPositiveDefinite()) // Cholesky Factorization
    {
      //std::cout << "Cholesky Solve"<< std::endl;
      info  = NM_Cholesky_solve(NM, b, 1);

      if(info != 0)
      {
        THROW_EXCEPTION("SimpleMatrix::Solve failed (Cholesky)");
      }
    }
    else  //  LDLT Factorization
    {
      THROW_EXCEPTION("SimpleMatrix::Solve failed: LDL^T not yet implemented.");
    }
  }
  else //  LU Factorization  by default
  {
    info  = NM_LU_solve(NM, b, 1);

    if(info != 0)
    {
      THROW_EXCEPTION("SimpleMatrix::Solve failed (LU)");
    }
  }

  if(B.num() == SPARSE)
  {
    B = *Bdense ;                                                // we copy back to sparse.
  }

  if(info != 0)
    THROW_EXCEPTION("SimpleMatrix::Solve failed.");
  // else
  // {
  //   noalias(*(B.dense())) = ublas::column(tmpB, 0);
  // }
  DEBUG_END("SimpleMatrix::Solve(SiconosVector &B)\n");
}


/*
 * L2, L2-Zupt
 * @file test_gnss_l2.cpp
 * @authors Ryan Watson, Shounak Das, Cagri Kilic
 */

// GTSAM related includes.
#include <gtsam/slam/dataset.h>
#include <gtsam/nonlinear/ISAM2.h>
#include <gtsam/inference/Symbol.h>
#include <gtsam/slam/PriorFactor.h>
#include <gtsam/slam/BetweenFactor.h>
#include <gtsam/gnssNavigation/GnssData.h>
#include <gtsam/gnssNavigation/GnssTools.h>
#include <gtsam/gnssNavigation/nonBiasStates.h>
#include <gtsam/nonlinear/NonlinearFactorGraph.h>
#include <gtsam/gnssNavigation/GNSSMultiModalFactor.h>
#include <gtsam/nonlinear/LevenbergMarquardtOptimizer.h>

// LibCluster
#include <libcluster/merge.h>
#include <libcluster/probutils.h>
#include <libcluster/libcluster.h>
#include <libcluster/distributions.h>

// BOOST
#include <boost/filesystem.hpp>
#include <boost/program_options.hpp>
#include <boost/serialization/export.hpp>
#include <boost/archive/binary_iarchive.hpp>
#include <boost/archive/binary_oarchive.hpp>

// GPSTK
#include <gpstk/ConfDataReader.hpp>

// STD
#include <chrono>
#include <fstream>
#include <iostream>
#include <algorithm>
#include <vector>

using namespace std;
using namespace gtsam;
using namespace gpstk;
using namespace boost;
using namespace merge;
using namespace std::chrono;
using namespace libcluster;
using namespace distributions;
namespace NM = gtsam::noiseModel;
namespace po = boost::program_options;
typedef noiseModel::Diagonal diagNoise;

// Intel Threading Building Block
#include <tbb/tbb.h>

using symbol_shorthand::X; // nonBiasStates ( dx, dy, dz, trop, cb )
using symbol_shorthand::G;   // bias states ( Phase Biases )

int main(int argc, char* argv[])
{
        // define std out print color
        bool skipped_update(false);
        vector<int> prn_vec;
        vector<int> factor_count_vec;
        vector<rnxData> data;
        // vector<faultyRnxData> data;
        const string red("\033[0;31m");
        const string green("\033[0;32m");
        string confFile, gnssFile, station;
        double xn, yn, zn, range, phase, rho, gnssTime, prev_time;
        int startKey(0), currKey, startEpoch(0), svn, numBatch(0), state_count(0), update_count(0);
        int nThreads(-1), phase_break, break_count(0), nextKey, factor_count(-1), res_count(-1);
        bool printECEF, printENU, printAmb, first_ob(true);
        Eigen::MatrixXd residuals;
        vector<mixtureComponents> globalMixtureModel;
        int num_zupts = 0;

        // string out_file = "/home/navlab-shounak/Desktop/Fusion/t10_clean_results_latest/l2_t10_w500_Fmod.xyz";
        string out_file = "results/l2_t11.xyz";
        gnssFile = "data/out11sat4F.gtsam";
        ifstream inputFile("data/zupt_Tags_t11.txt");
        //t9 nominal ECEF values
        // xn = 859154.0695;
        // yn = -4836304.2164;
        // zn = 4055377.5475;

        //t10 nominal ECEF values
        // xn = 859153.0167;
        // yn = -4836303.7245;
        // zn = 4055378.4991;

        // t11 nominal ECEF values
        xn = 859156.4189;
        yn = -4836305.5491;
        zn = 4055375.2899;

        // 859154.0695, -4836304.2164, 4055377.5475 - t9
        // 859153.0167, -4836303.7245, 4055378.4991 - t10
        // 859156.4189, -4836305.5491, 4055375.2899 - t11

        ofstream out_os(out_file);
        cout.precision(12);

        out_os.precision(12);

        po::options_description desc("Available options");
        po::variables_map vm;
        po::store(po::command_line_parser(argc, argv).options(desc).run(), vm);
        po::notify(vm);

        // gnssFile = "data/out9sat4F.gtsam";
        // gnssFile = "/home/navlab-shounak/Desktop/Fusion/gtsam_data_t10/out10sat4F.gtsam";

        //----------------------------------------------------------------------

        //read the zupt times from a file (a set?, they can use the method count())

        // open file
        // ifstream inputFile("data/zupt_Tags_t9.txt");
        vector<double> zupt_tags;

        // test file open
        if (inputFile) {
            double val;

            // read the elements in the file into a vector
            while ( inputFile >> val ) {
                zupt_tags.push_back(val);
            }
        }

        printENU = false;
        printAmb = false;
        printECEF = true;

        Point3 nomXYZ(xn, yn, zn);
        Point3 prop_xyz = nomXYZ;

        try {data = readGNSS_SingleFreq(gnssFile); }
        catch(std::exception& e)
        {
                cout << red << "\n\n Cannot read GNSS data file " << endl;
                exit(1);
        }

        std::auto_ptr<tbb::task_scheduler_init> init;
        if(nThreads > 0) {
                init.reset(new tbb::task_scheduler_init(nThreads));
        }
        else
                cout << green << " \n\n Using threads for all processors" << endl;

        ISAM2DoglegParams doglegParams;
        ISAM2Params parameters;
        parameters.relinearizeThreshold = 0.01;
        parameters.relinearizeSkip = 1000;
        ISAM2 isam(parameters);

        double output_time = 0.0;
        double rangeWeight = 0.5;  //2.5;
        double phaseWeight = 0.05; //0.25

        ifstream file(gnssFile.c_str());
        string value;

        nonBiasStates prior_nonBias = (gtsam::Vector(5) << 0.0, 0.0, 0.0, 0.0, 0.0).finished();

        // zupt noise model

        noiseModel::Diagonal::shared_ptr zuptNoise = noiseModel::Diagonal::Variances((gtsam::Vector(5) << 1e-3, 1e-3, 1e-3, 1e3, 1e-3).finished());

        // non-zupt noise model

        noiseModel::Diagonal::shared_ptr non_zuptNoise = noiseModel::Diagonal::Variances((gtsam::Vector(5) << 0.5, 0.5, 0.5, 1e3, 1e-3).finished());

        phaseBias bias_state(Z_1x1);
        gnssStateVector phase_arc(Z_34x1);
        gnssStateVector bias_counter(Z_34x1);
        for (int i=1; i<34; i++) {bias_counter(i) = bias_counter(i-1) + 10000; }

        nonBiasStates initEst(Z_5x1);
        nonBiasStates between_nonBias_State(Z_5x1);

        Values initial_values, result;

        noiseModel::Diagonal::shared_ptr nonBias_InitNoise = noiseModel::Diagonal::Variances((gtsam::Vector(5) << 0.2, 0.2, 0.2, 3e6, 1e-1).finished());

        noiseModel::Diagonal::shared_ptr nonBias_Reset = noiseModel::Diagonal::Variances((gtsam::Vector(5) << 3e4, 3e4, 3e4, 3e6, 1e-1).finished());

        noiseModel::Diagonal::shared_ptr nonBias_ProcessNoise = noiseModel::Diagonal::Variances((gtsam::Vector(5) << 1.0, 1.0, 1.0, 1e3, 1e-3).finished());

        noiseModel::Diagonal::shared_ptr initNoise = noiseModel::Diagonal::Variances((gtsam::Vector(1) << 3e6).finished());

        NonlinearFactorGraph *graph = new NonlinearFactorGraph();

        residuals.setZero(1000,2);

        // For the l2 model, we will use a mixture model with 1 component
        Eigen::RowVectorXd m(2);
        m << 0.0, 0.0;

        // Add component 1. (i.e., the assumed error covariance)
        Eigen::MatrixXd c(2,2);
        c<< std::pow(rangeWeight,2), 0.0, 0.0, std::pow(phaseWeight,2);
        globalMixtureModel.push_back(boost::make_tuple(0, 0, 0.0, m, c));

        int lastStep = get<0>(data.back());

        std::vector<int> num_obs (1000, 0);

        bool is_zupt;

        cout << " The number of epochs in the Gtsam data file -- " << data.size() << endl;

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

                is_zupt = false;

                auto start = high_resolution_clock::now();

                double gnssTime = get<0>(data[i]);
                int currKey = get<1>(data[i]);
                if (first_ob) {
                        first_ob=false;
                        startKey = currKey;
                        graph->add(PriorFactor<nonBiasStates>(X(currKey), initEst,  nonBias_InitNoise));
                        ++factor_count;
                        initial_values.insert(X(currKey), initEst);
                }
                int nextKey = get<1>(data[i+1]);
                int svn = get<2>(data[i]);
                Point3 satXYZ = get<3>(data[i]);
                double rho = get<4>(data[i]);
                double range = get<5>(data[i]);
                double phase = get<6>(data[i]);
                double phase_break = get<7>(data[i]);

                gtsam::Vector2 obs;
                obs << range-rho, phase-rho;

                if ( currKey == startKey || phase_arc[svn]!=phase_break)
                {
                        bias_state[0] = phase-range;
                        if (currKey > startKey) { bias_counter[svn] = bias_counter[svn] +1; }
                        initial_values.insert(G(bias_counter[svn]), bias_state);

                        graph->add(boost::make_shared<PriorFactor<phaseBias> >(G(bias_counter[svn]), bias_state,  initNoise));

                        phase_arc[svn] = phase_break;
                        ++factor_count;
                }

                graph->add(boost::make_shared<GNSSMultiModalFactor>(X(currKey), G(bias_counter[svn]), obs, satXYZ, nomXYZ, globalMixtureModel));

                prn_vec.push_back(svn);
                factor_count_vec.push_back(++factor_count);

                //---------------------------------------------------------------------
                // add between factor here
                //
                // if (i != startEpoch){
                //     int prevKey = get<1>(data[i-1]);
                //     double prevgnssTime = get<0>(data[i-1]);
                //
                //     for (int j = 0; j < zupt_tags.size()-1; j++){
                //         if ((std::abs(zupt_tags[j] - prevgnssTime) < 0.01) && (std::abs(zupt_tags[j+1] - gnssTime) < 0.01)){
                //             is_zupt = true;
                //             graph->add(BetweenFactor<nonBiasStates>(X(currKey),X(prevKey), between_nonBias_State, zuptNoise));
                //             ++factor_count;
                //             num_zupts = num_zupts+1;
                //             cout << " Zupt applied -- " << num_zupts <<  " between times " << prevgnssTime << " <--> " << gnssTime <<  endl;
                //             break;
                //         }
                //     }
                //
                //     if(is_zupt == false){
                //         graph->add(BetweenFactor<nonBiasStates>(X(currKey),X(prevKey),between_nonBias_State, non_zuptNoise));
                //         ++factor_count;
                //     }
                // }

                //---------------------------------------------------------------------

                if (currKey != nextKey && nextKey != 0) {

                        if (currKey > startKey ) {
                                if ( lastStep == nextKey ) { break; }

                                if (state_count < 600 || state_count > (lastStep-600))
                                {
                                        graph->add(PriorFactor<nonBiasStates>(X(currKey), initEst,  nonBias_InitNoise));

                                        ++factor_count;
                                }

                                graph->add(PriorFactor<nonBiasStates>(X(currKey), prior_nonBias, nonBias_Reset));
                                ++factor_count;

                        }

                        isam.update(*graph, initial_values);
                        isam.update();
                        isam.update();
                        result = isam.calculateEstimate();

                        prior_nonBias = result.at<nonBiasStates>(X(currKey));
                        Point3 delta_xyz = (gtsam::Vector(3) << prior_nonBias.x(), prior_nonBias.y(), prior_nonBias.z()).finished();
                        prop_xyz = nomXYZ - delta_xyz;

                        if (printECEF) {
                                out_os << gnssTime << " " << prop_xyz.x() << " " << prop_xyz.y() << " " << prop_xyz.z() << endl;
                        }

                        if (printENU) {
                                Point3 enu = xyz2enu(prop_xyz, nomXYZ);
                                out_os << "enu " << gnssTime << " " << enu.x() << " " << enu.y() << " " << enu.z() << endl;
                        }

                        if (printAmb) {
                                out_os << "gps " << " " << gnssTime << " ";
                                for (int k=0; k<prn_vec.size(); k++) {
                                        out_os << result.at<phaseBias>(G(bias_counter[prn_vec[k]])) << " ";
                                }
                                out_os << endl;
                        }

                        ++state_count;
                        output_time = output_time +1;

                        factor_count_vec.clear();
                        factor_count = -1;

                        graph->resize(0);
                        initial_values.clear();
                        prn_vec.clear();

                        auto stop = high_resolution_clock::now();
                        auto duration = duration_cast<microseconds>(stop - start);

                        // cout << "Delta time: "
                        //      << duration.count() << " microseconds" << endl;

                        initial_values.insert(X(nextKey), prior_nonBias);
                }

        }

        return 0;
}


#pragma once
#include "../argcheck.hpp"
#include <Eigen/Core>
#include <Eigen/Sparse>
#include <cstddef>
#include <future>
#include <iostream>
#include <random>
#include <stdexcept>
#include <thread>
#include <vector>

namespace KNN {
/**
 * Given R \times C sparse matrix X, compute
 * r x R matrix T * X^T
 * against r x C matrix T.
 * where each *ROW* is supposed to contain only top_k non-zero elements.
 */
template <typename Real, class SimilarityType> struct KNNComputer {

  using Triplet = typename Eigen::Triplet<Real>;
  using CSRMatrix = Eigen::SparseMatrix<Real, Eigen::RowMajor>;
  using CSCMatrix = Eigen::SparseMatrix<Real, Eigen::ColMajor>;

  using CSRIterType = typename CSRMatrix::InnerIterator;
  using CSCIterType = typename CSCMatrix::InnerIterator;
  using DenseVector = Eigen::Matrix<Real, Eigen::Dynamic, 1>;
  KNNComputer(const CSRMatrix &X_arg, Real shrinkage, size_t n_threads,
              size_t max_chunk_size)
      : X_t(X_arg.transpose()), n_threads(n_threads), N(X_arg.rows()),
        n_features(X_arg.cols()), shrinkage(shrinkage), norms(X_arg.rows()),
        max_chunk_size(max_chunk_size) {
    irspack::check_arg_lower_bounded<Real>(shrinkage, 0, "shrinkage");
    irspack::check_arg_lower_bounded<size_t>(n_threads, 1, "n_threads");
    irspack::check_arg_lower_bounded<size_t>(max_chunk_size, 1,
                                             "max_chunk_size");

    X_t.makeCompressed();
  }

  CSRMatrix compute_similarity(const CSRMatrix &target, size_t top_k) const {
    if (target.cols() != this->n_features)
      throw std::invalid_argument("illegal # of feature.");
    CSRMatrix result(target.rows(), this->N);
    if (n_threads == 1) {

      auto triplets =
          this->compute_similarity_triple(target, 0, target.rows(), top_k);
      result.setFromTriplets(triplets.begin(), triplets.end());
      return result;
    } else {
      const size_t per_block_size = target.rows() / n_threads;
      const size_t remnant = target.rows() % n_threads;
      size_t job_start = 0;
      std::vector<std::future<std::vector<Triplet>>> thread_results;
      std::vector<Triplet> accumulated_result;
      for (size_t i = 0; i < n_threads; i++) {
        size_t block_size = per_block_size;
        if (i < remnant) {
          block_size += 1;
        }
        thread_results.push_back(
            std::async(std::launch::async,
                       [&target, job_start, block_size, this, top_k]() {
                         return this->compute_similarity_triple(
                             target, job_start, job_start + block_size, top_k);
                       }));
        job_start += block_size;
      }
      assert(job_start == target.rows());
      for (auto &thread_result : thread_results) {
        std::vector<Triplet> thread_result_ret = thread_result.get();
        accumulated_result.insert(accumulated_result.end(),
                                  thread_result_ret.begin(),
                                  thread_result_ret.end());
      }
      result.setFromTriplets(accumulated_result.begin(),
                             accumulated_result.end());
      return result;
    }
  }

  std::vector<Triplet> compute_similarity_triple(const CSRMatrix &target,
                                                 size_t start, size_t end,
                                                 size_t top_k) const {
    /*
     We want to compute
     target * X^T
     with each *ROW* restricted by top-k strategy.
     We devide the problem row-wise.
     target[start:end, :] * X^T
    */

    using IndexType = typename CSCMatrix::StorageIndex;

    std::vector<Triplet> triples;
    std::vector<IndexType> buffer(this->N);

    for (size_t cursor = start; cursor < end; cursor += this->max_chunk_size) {
      size_t end_local = std::min(cursor + this->max_chunk_size, end);
      CSRMatrix block_result_row =
          this->compute_sim_block(target, cursor, end_local);
      block_result_row.makeCompressed();

      const Real *data_ptr = block_result_row.valuePtr();
      auto index_ptr = block_result_row.innerIndexPtr();
      auto index_start_ptr = block_result_row.outerIndexPtr();

      for (IndexType row = 0; row < block_result_row.rows(); row++) {
        size_t nz_size = index_start_ptr[row + 1] - index_start_ptr[row];
        size_t col_size = std::min(nz_size, top_k);
        auto data_start = data_ptr + index_start_ptr[row];
        auto index_start = index_ptr + index_start_ptr[row];
        for (size_t i = 0; i < nz_size; i++) {
          buffer[i] = i;
        }
        std::sort(buffer.begin(), buffer.begin() + nz_size,
                  [&data_start](IndexType &col1, IndexType &col2) {
                    return data_start[col1] > data_start[col2];
                  });
        std::sort(buffer.begin(), buffer.begin() + col_size);
        for (size_t j = 0; j < col_size; j++) {
          triples.emplace_back(row + cursor,
                               static_cast<IndexType>(index_start[buffer[j]]),
                               data_start[buffer[j]]);
        }
      }
    }
    return triples;
  }

public:
  CSCMatrix X_t;
  const size_t n_threads;
  const int N, n_features;

protected:
  Real shrinkage;
  DenseVector norms;

private:
  inline CSRMatrix compute_sim_block(const CSRMatrix &target, size_t start,
                                     size_t end) const {
    return static_cast<const SimilarityType &>(*this).compute_similarity_imple(
        target, start, end);
  }

protected:
  size_t max_chunk_size = 128;
};

} // namespace KNN


/**
 ** Isaac Genome Alignment Software
 ** Copyright (c) 2010-2017 Illumina, Inc.
 ** All rights reserved.
 **
 ** This software is provided under the terms and conditions of the
 ** GNU GENERAL PUBLIC LICENSE Version 3
 **
 ** You should have received a copy of the GNU GENERAL PUBLIC LICENSE Version 3
 ** along with this program. If not, see
 ** <https://github.com/illumina/licenses/>.
 **
 ** \file RestOfGenomeCorrection.hh
 **
 ** \brief Component to encapsulate probability of read randomly aligning to the genome.
 **
 ** \author Roman Petrovski
 **/

#ifndef iSAAC_ALIGNMENT_REST_OF_GENOME_CORRECTION_HH
#define iSAAC_ALIGNMENT_REST_OF_GENOME_CORRECTION_HH

#include <boost/foreach.hpp>

#include "common/Debug.hh"
#include "reference/Contig.hh"

namespace isaac
{
namespace alignment
{

class RestOfGenomeCorrection
{
public:
    RestOfGenomeCorrection() : rogCorrection_(0.0)
    {
    }

    RestOfGenomeCorrection(
        const reference::ContigList &contigList,
        const flowcell::ReadMetadataList &readMetadataList)
    {
        setGenome(contigList, readMetadataList);
    }

    void setGenome(
        const reference::ContigList &contigList,
        const flowcell::ReadMetadataList &readMetadataList)
    {
        getRestOfGenomeCorrection(contigList, readMetadataList, rogCorrectionList_);
        rogCorrection_ = Quality::restOfGenomeCorrection(reference::genomeLength(contigList),
                                                         flowcell::getTotalReadLength(readMetadataList));
        // can't have 0.0 in it as it turns the alignment score into 0
        rogCorrection_ = std::max(rogCorrection_, std::numeric_limits<double>::min());
    }

    double getReadRogCorrection(const unsigned readIndex) const
    {
        ISAAC_ASSERT_MSG(sizeof(rogCorrectionList_) > readIndex, "Only up to 2 reads supported");
        return rogCorrectionList_[readIndex];
    }
    double getRogCorrection() const {return rogCorrection_;}
private:
    static const unsigned READS_MAX = 2;
    /// Rest-of-genome correction for individual fragments
    boost::array<double, READS_MAX> rogCorrectionList_;
    /// Rest-of-genome correction for the template when all fragments match
    double rogCorrection_;

    static void getRestOfGenomeCorrection(
        const reference::ContigList &contigList,
        const flowcell::ReadMetadataList &readMetadataList,
        boost::array<double, READS_MAX> &rogCorrectionList)
    {
        const size_t genomeLength = reference::genomeLength(contigList);
        BOOST_FOREACH(const flowcell::ReadMetadata &readMetadata, readMetadataList)
        {
            rogCorrectionList[readMetadata.getIndex()] = Quality::restOfGenomeCorrection(genomeLength, readMetadata.getLength());
            // can't have 0.0 in it as it turns the alignment score into 0
            rogCorrectionList[readMetadata.getIndex()] = std::max(rogCorrectionList[readMetadata.getIndex()], std::numeric_limits<double>::min());
        }
    }
};

} // namespace alignment
} // namespace isaac

#endif // #ifndef iSAAC_ALIGNMENT_REST_OF_GENOME_CORRECTION_HH