#include <cmath>
#include <cstdlib>
#include <iostream>
#include <list>
#include <bitset>
#include <map>
#include <vector>
#include <string>
#include <sstream>
#include <iomanip>
#include <assert.h>
#include <cmath>
#include "mem/ruby/common/Address.hh"
#include "mem/ruby/system/System.hh"
#include "mem/ruby/slicc_interface/AbstractCacheEntry.hh"
#include "debug/TagTable.hh"
#include "debug/TagTable1.hh"
//#include "CacheMemoryPF.hh"
#include "tagtable_stats.hpp"

#define BLOCKSIZE 64
#define PAGESIZE 4096
#define ADDR_LEN 48
#define DEPTH 4
#define MAXTRACES 8

extern uint64_t CACHESIZE;// (64*MEGA)
extern uint64_t BLOCKS;// (CACHESIZE/BLOCKSIZE)	//number of blocks in cache
extern uint64_t PAGES;// (CACHESIZE/PAGESIZE)	//number of pages in cache

//#include <boost/circular_buffer.hpp>
//extern uint64_t Address;

 // Maximum Number of trace files that can be specified

//block to be used for insertion and sorting ("block" may be a misnomer since it actually represents a contiguous set of cache blocks)
struct block { //contiguous block of data
  int offset;
  int len;
  short page_offset;
  AbstractCacheEntry* PFentry_ptr;
  friend std::ostream& operator<< (std::ostream & out, const block &blk) {
    out << "Block offset = 0x" << std::hex << blk.offset << std::dec << ", length = " << blk.len
	<< ", page_offset = 0x" << std::hex << blk.page_offset << std::dec << std::endl;
    return out;
  }
};

bool block_ordering(block first, block second);
void print_assigned(const std::vector<uint64_t> assigned);
typedef struct l3map_entry_t {
  uint seq_no;  //indicates offset - from page table's base - of this particular entry-block
  uint size;    //indicates number of entries/tags it currently maps
  std::vector<uint64_t> meta_tags;//[(BLOCKSIZE/ENTRYSIZE)];      //tags for finding metadata in L3 (matches a LEVEL_PTR entry's tag)
  std::vector<uint64_t> data_tags;//[(BLOCKSIZE/ENTRYSIZE)];      //tags for path on pagetable associated with metadata block (equal to meta_tags in SST_ENTRYs)
  std::vector<bool> is_lvl;//[(BLOCKSIZE/ENTRYSIZE)];   //indicate which entries are associated with lvl_ptrs (makes eviction faster since I can directly use the tag for everything else, i.e., don't have search every entry to find it)

  l3map_entry_t(uint32_t _s) :
    meta_tags(_s, 0),
    data_tags(_s, 0),
    is_lvl(_s, false)
  { }
} l3map_entry;

struct sst_field {
  bool presence;		//NOT NECESSARY, len == 0 IMPLIES NOT PRESENT... 1 bit:  are the blocks in this range present? (or is this an exception, huh?  can't be an exception, this is covered by the sst_entry itself, not the field, right?)
  short page_offset;		//log(PAGESIZE/BLOCKSIZE) bits:  subblock w/in page associated with data corresponding to "offset" address (TODO:  implies all other subblocks covered by this entry are contiguous in page)
  short len;			//log(PAGESIZE/BLOCKSIZE) bits:  length beyond this range that presence is valid
  short offset;			//log(PAGESIZE/BLOCKSIZE) bits
  AbstractCacheEntry* PFentry_ptr;
};

//Sorted Segment Table (SST) entry (size = 1 + FIELDS * (1 + 3*log(PAGESIZE/BLOCKSIZE)))
struct sst_entry {
  // bool is_match;			//1 bit:  is presence a match or an exception? (i.e., do the chunks encode the present blocks or the non-present blocks?  Tracking gaps isn't implemented => should always be true)
  std::vector<sst_field> fields;	//FIELDS * (1 + 3*log(PAGESIZE/BLOCKSIZE)) bits
  std::bitset<(PAGESIZE/BLOCKSIZE)> prefetch_vector;	//0 bits (not actually present in physical implementation - just a stat):  flag blocks that were brought in speculatively (due to negotiation to prevent an entry outgrowing # of FIELDS)
//  boost::circular_buffer<uint32_t> field_lru;		//LRU of fields for eviction

  sst_entry(uint _num_fields) :
    fields(_num_fields)
     {
      //initialize LRU
      for(uint32_t f = 0; f < _num_fields; ++f) {
//	field_lru.push_front(f);
      }
    }
};

typedef enum {
  SST_ENTRY,            /* Entry of FIELDS fields defining chunks of space */
  LEVEL_PTR,            /* Pointer to next lowest level in the table */
  STATUS_VECTOR,        /* Vector defining block-level status */
  ENTRY_TYPE_NUM
} ENTRY_TYPE;

struct entry;

class entry_level {
public:
  entry_level( int w, int d, uint32_t _num_chunks, int proot_depth, entry_level *p_root);
  ~entry_level();
  entry* lookup(int index);
  int get_depth(){return depth;}
  int get_mask(){return mask;}
  entry_level* get_pageroot(){return page_root;}
  uint get_evictions();	//Accumulates and returns evictions of all children levels
  void update_evictions(int removed) { evictions += removed; }
  uint get_l3_evictions();
  short replace(int preferred, int &expansions, int &merges);
  void evict_block(short victim, int &expansions, int &merges);
  bool evict_entry(uint64_t tag);				//evict SST entry associated with a tag
  void update_tag(uint64_t path_tag, uint64_t old_tag, uint64_t new_tag);	//reassign tag from old to new (needed to handle LEVEL_PTRs found in 'evict_entry')
  friend std::ostream& operator<< (std::ostream & out, const entry_level & e_lvl);
  std::ostream & print_level(std::ostream & out) const;
  void print_level() const;
  int get_active_entries();	//walk structure and determine number of valid entries (either status vector or sst entries)
  //Break down how blocks are tracked for space efficiency determination
  int get_sst_tracked(int &entries, int &lvlptrs);
  int get_vector_tracked(int &entries);
  void verify_page_offsets(int offset_array[]);	//function to verify occupancy & presence vector periodically
  entry *copy_entry(entry *e);
  void get_size(int size[], int lvl_count[]);
  int getindex(uint64_t addr, int lvl, int mask);
  void defragment();	//Defragment entries (i.e., re-build all entries to maximize contiguous blocks
  uint64_t get_tag(uint64_t addr);
  //Translation Cache Functions
  int get_tc_occupancy();		//get number of paths present in the translation cache
  void print_tc_status(/*int check*/);	//print the current Translation Cache contents
  void evict_tc_entry();		//find and evict a translation cache entry to keep within capacity
  bool verify_tc_status(int check);	//Checks that a given TC subtree has the correct number of paths
  int fix_tc_paths();			//fix tc_paths for every entry in the level based on its children
  void dec_tc_paths(uint64_t tag);	//Targeted fixing of Translation Cache paths, only the entries that can be on the path of 'tag' are affected

  //Kill the status vector
  bool negotiate_new_assignment(std::pair<short, short> changes, int &expansions, int &merges);
  void clear_presence(std::pair<short, short>);		//clear presence bits for evicted blocks
  void set_presence(std::pair<short, short>);		//set presence bits for prefetched blocks
  void mark_evicted(int row_offset);			//Mark specified row offset as being evicted (triggered by pagetable)
  int prefetch(uint64_t addr, short offset, bool _evict);	//Rearrange blocks in row to make entry associated with 'addr' have <= FIELDS fields
							// return number of blocks prefetched (evict instead of prefetch if _evict)
  void find_and_evict(std::pair<short, short> range);	//find blocks with row offsets given in 'range' and evict them
  void print_presence_vector();

  //Prediction functions
  uint32_t getDistance(uint64_t addr);		//Determine minimum distance this block is from an existing chunk (*NOT* its own chunk however)
  void getChunkSizeStats(tagtable_stats &chunk_size_stats);	//Create histogram of chunk sizes for all existing valid entries

  //Unit Test Interfaces
  uint32_t getSSTEntries();
  uint32_t getLvlPtrs();
  uint32_t getChunks();

private:

  uint m_num_fields;
  int width;			//number of bits used to index array
  int pageroot_depth;	//depth at which page roots reside (i.e., roots of subtrees beneath which all blocks are on the same page (row) in the cache) - passed in by page table on construction
  int mask;			//mask used to determine index for this level
  entry **array;		//array of entries
  int depth;			//depth in page table where this level resides
  int occupancy;		//number of blocks captured by last level table (important for ensuring <= 64 entries mapped per page)
  int offsets_present[(PAGESIZE/BLOCKSIZE)];	//array to verify occupancy & presence vector periodically
  //  short next_replace;		//victim for next replacement (if page is full)
  std::bitset<PAGESIZE/BLOCKSIZE> presence_vector;	//offsets already taken
  int moves;			//moves necessitated to keep SST entries contiguous (will need to replace with function later)
  entry_level *page_root;	//entry level that forms the root of the blocks that fit on a single page
  std::list<block> fix_entry(int i);
  void populate_entry(int index, std::list<block> blocks, int &expansions, int &merges);
  //  int update_next_replace();
  short get_victim();
  void print_entry(int index, int indent) const;
  std::ostream & print_entry(int index, std::ostream & out, int indent) const;
  void find_offender(int page);
  std::list<block> build_entry(std::list<block> present_blocks);
  void bulk_update_presence();
  void accumulate_page_offsets(std::vector<int> &offset_array);

  //Statistics
  uint evictions;
  uint l3_evictions;	//evictions caused by evicting data from the L3

  //Defragment functions
  void add_to_vector(std::vector<uint64_t> &assigned);
  void populate_sst_entry(uint32_t index, uint32_t &curr_field, uint32_t &page_off, std::vector<uint64_t> assigned);
  void populate_level(uint32_t &curr_field, uint32_t &page_off, std::vector<uint64_t> assigned);
};

struct entry {
  bool valid;		//1 bit:  is this entry valid?
  int tc_path;		//log(#tcEntries) bits:  number of translation paths this entry is on
  ENTRY_TYPE type;	//1 bit (SST or PTR):  what type of entry is this?
  //below, either or (i.e., length is max(length(addr,sst_e))
  entry_level *addr;	//# bits to access next level of table in L3
  sst_entry sst_e;	//FIELDS * (1 + 3*log(PAGESIZE/BLOCKSIZE)))
  short status_vector[(PAGESIZE/BLOCKSIZE)];  //not used:  status vector stores either -1 (not present) or the block's offset in the page
  uint64_t tag;		//(ADDR_LEN - block offset - page offset - bits to identify row) bits:  address bits needed to disambiguate this entry (when at leaf this isn't necessary, only when entry exists at a higher level)

  entry(uint _num_fields) :
  sst_e(_num_fields)
  { }

  friend std::ostream& operator<< (std::ostream & out, const entry &entry) {
    int entry_width = 97;
    assert(entry.type == SST_ENTRY);
    out << std::setfill(' ') << std::setw(5) << "" << std::setfill('-') << std::setw(entry_width) << "" << std::endl << std::hex;
      for(const auto &fld : entry.sst_e.fields) {
	out << " | o:" << fld.offset << " (po: " << fld.page_offset << ") p:" << fld.presence << " l:" << fld.len;
      }
      out << "|(" << entry.tc_path << ") tag 0x" << std::hex << entry.tag << std::dec << std::endl;
      out << std::setfill(' ') << std::setw(5) << "" << std::setfill(' ') << std::setw(0) << "" << std::setfill('-') << std::setw(entry_width) << "" << std::endl;
    return out;
  }
};

//int m_entry_size=uint32_t(pow(2, (CeilLog2((20 + _num_chunks*19)/8))));


//
// Created by eliane on 04/03/19.
//

#include "VinaLike.h"

#include <exception>

#include <Structures/Atom.h>

#include <Engines/Internals/InternalsUtilityFunctions.h>

#include <boost/log/trivial.hpp>

#include "VinaLikeCommon.h"

namespace SmolDock::Score {

    double force_Instantiate_VinaLikeIntermolecularScoringFunction(const iConformer &conformer, iTransform &transform,
                                                                   const iProtein &protein,
                                                                   std::array<double, VinaLike_numCoefficients> nonDefaultCoeffs) {
        return VinaLikeIntermolecularScoringFunction<true, false>(conformer, transform, protein, nonDefaultCoeffs);
    }

    double force_Instantiate_VinaLikeIntermolecularScoringFunction_vectorized(
            const iConformer_Vectorized &conformer, iTransform &transform,
            const iProtein_vectorized &protein,
            std::array<double, VinaLike_numCoefficients> nonDefaultCoeffs) {
        return VinaLikeIntermolecularScoringFunction_vectorized<true, false>
                (conformer, transform, protein, nonDefaultCoeffs
                );
    }

    const std::array<std::string, VinaLike::numCoefficients>
            VinaLike::coefficientsNames = {"Gauss1", "Gauss2", "RepulsionExceptCovalent", "Hydrophobic", "Hydrogen"};

    template<bool OnlyIntermolecular, bool useNonDefaultCoefficients>
    double VinaLikeIntermolecularScoringFunction(const iConformer &ligand_, iTransform &transform,
                                                 const iProtein &protein,
                                                 std::array<double, VinaLike_numCoefficients> nonDefaultCoeffs) {

        BOOST_ASSERT(!ligand_.x.empty());
        BOOST_ASSERT(!protein.x.empty());

        transform.doHousekeeping();
        if (std::abs(transform.rota.norm() - 1) > 0.1) {
            transform.rota.normalize();
        }

        BOOST_ASSERT(transform.bondRotationsAngles.size() == ligand_.num_rotatable_bond);

        double score_raw = 0;

        double distance_total = 0.0;
        double gauss1_total = 0.0;
        double gauss2_total = 0.0;
        double repuls_total = 0.0;
        double hydrophobic_total = 0.0;
        double hydrogen_total = 0.0;

        iConformer ligand = ligand_;
        //applyBondRotationInPlace(ligand, transform);

        Eigen::Vector3d ProtCenterPosition = {protein.center_x, protein.center_y, protein.center_z};

        if constexpr(!OnlyIntermolecular) // C++17
        {
            for (unsigned int idxLig = 0; idxLig < ligand.x.size(); idxLig++) {
                for (unsigned int idxLig2 = idxLig; idxLig2 < ligand.x.size(); idxLig2++) {
                    Eigen::Vector3d LigDistance = {ligand.x[idxLig] - ligand.x[idxLig2],
                                                   ligand.y[idxLig] - ligand.y[idxLig2],
                                                   ligand.z[idxLig] - ligand.z[idxLig2]};

                    double distance_raw = LigDistance.norm();
                    const double distance = distanceFromRawDistance(distance_raw, ligand.atomicRadius[idxLig],
                                                                    ligand.atomicRadius[idxLig2]);

//                score_raw += VinaClassic::coeff_gauss1      * vinaGaussComponent(distance, 0.0, 0.5);
//                score_raw += VinaClassic::coeff_gauss2      * vinaGaussComponent(distance, 3.0, 2.0);
                    score_raw += VinaClassic::coeff_repulsion * vinaRepulsionComponent(distance, 0.0);
                }
            }
        }

        for (unsigned int idxLig = 0; idxLig < ligand.x.size(); idxLig++) {
            for (unsigned int idxProt = 0; idxProt < protein.x.size(); idxProt++) {

                Eigen::Vector3d LigPosition = {ligand.x[idxLig], ligand.y[idxLig], ligand.z[idxLig]};
                applyRigidTransformInPlace(LigPosition, transform);

                Eigen::Vector3d ProtPosition = {protein.x[idxProt], protein.y[idxProt], protein.z[idxProt]};
                Eigen::Vector3d distToCenterVector = LigPosition - ProtCenterPosition;

                double distanceToProteinCenter = distToCenterVector.norm();

                if (distanceToProteinCenter > (protein.radius - 1)) {
                    score_raw += std::pow((distanceToProteinCenter - protein.radius), 4) + 10;
                    continue;
                }

                Eigen::Vector3d distVect = ProtPosition - LigPosition;

                const double rawDist = distVect.norm();

                if (rawDist >= VinaClassic::interaction_cutoff)
                    continue;

                const double atomicRadiusLig = ligand.atomicRadius[idxLig];
                const double atomicRadiusProt = protein.atomicRadius[idxProt];

                const double radToRemove = (atomicRadiusLig + atomicRadiusProt);

                const double distance = rawDist - radToRemove;

                distance_total += distance;

                const unsigned int atom1AtomicNumber = ligand.type[idxLig];
                const unsigned int atom1AtomVariant = ligand.variant[idxLig];
                const unsigned int atom2AtomicNumber = protein.type[idxProt];
                const unsigned int atom2AtomVariant = protein.variant[idxProt];

                if constexpr(useNonDefaultCoefficients) {
                    score_raw += nonDefaultCoeffs[0] * vinaGaussComponent(distance, 0.0, 0.5);
                    score_raw += nonDefaultCoeffs[1] * vinaGaussComponent(distance, 3.0, 2.0);
                    score_raw += nonDefaultCoeffs[2] * vinaRepulsionComponent(distance, 0.0);
                    score_raw += nonDefaultCoeffs[3] * vinaHydrophobicComponent(distance,
                                                                                atom1AtomicNumber, atom1AtomVariant,
                                                                                atom2AtomicNumber, atom2AtomVariant);

                    score_raw += nonDefaultCoeffs[4] * vinaHydrogenComponent(distance,
                                                                             atom1AtomicNumber, atom1AtomVariant,
                                                                             atom2AtomicNumber, atom2AtomVariant);
                } else {

                    gauss1_total += vinaGaussComponent(distance, 0.0, 0.5);
                    gauss2_total += vinaGaussComponent(distance, 3.0, 2.0);
                    repuls_total += vinaRepulsionComponent(distance, 0.0);
                    hydrophobic_total += vinaHydrophobicComponent(distance,
                                                                  atom1AtomicNumber, atom1AtomVariant,
                                                                  atom2AtomicNumber, atom2AtomVariant);
                    hydrogen_total += vinaHydrogenComponent(distance,
                                                            atom1AtomicNumber, atom1AtomVariant,
                                                            atom2AtomicNumber, atom2AtomVariant);

                    score_raw += VinaClassic::coeff_gauss1 * vinaGaussComponent(distance, 0.0, 0.5);
                    score_raw += VinaClassic::coeff_gauss2 * vinaGaussComponent(distance, 3.0, 2.0);
                    score_raw += VinaClassic::coeff_repulsion * vinaRepulsionComponent(distance, 0.0);
                    score_raw += VinaClassic::coeff_hydrophobic * vinaHydrophobicComponent(distance,
                                                                                atom1AtomicNumber,
                                                                                atom1AtomVariant,
                                                                                atom2AtomicNumber,
                                                                                atom2AtomVariant);

                    score_raw += VinaClassic::coeff_hydrogen * vinaHydrogenComponent(distance,
                                                                                atom1AtomicNumber,
                                                                                atom1AtomVariant,
                                                                                atom2AtomicNumber,
                                                                                atom2AtomVariant);
                }

            } // for
        } // for

        double final_score = score_raw / (1 + (VinaClassic::coeff_entropic * ligand.num_rotatable_bond));

/*
        std::cout << "\n Non-Vectorized : \n";
        std::cout << "dist    : " << distance_total << std::endl;
        std::cout << "gauss1  : " << gauss1_total << std::endl;
        std::cout << "gauss2  : " << gauss2_total << std::endl;
        std::cout << "repuls  : " << repuls_total << std::endl;
        std::cout << "hydroph : " << hydrophobic_total << std::endl;
        std::cout << "hydrog  : " << hydrogen_total << std::endl;
        std::cout << "\n\n";
*/
        return final_score;
    }

    template<bool OnlyIntermolecular, bool useNonDefaultCoefficients>
    double VinaLikeIntermolecularScoringFunction_vectorized(const iConformer_Vectorized &ligand_, iTransform &transform,
                                                            const iProtein_vectorized &protein,
                                                            std::array<double, VinaLike_numCoefficients> nonDefaultCoeffs) {

        BOOST_ASSERT(ligand_.x.entriesCount() != 0);
        BOOST_ASSERT(protein.x.entriesCount() != 0);

        transform.doHousekeeping();
        if (std::abs(transform.rota.norm() - 1) > 0.1) {
            transform.rota.normalize();
        }

        BOOST_ASSERT(transform.bondRotationsAngles.size() == ligand_.num_rotatable_bond);

        double score_raw = 0;

        double distance_total = 0.0;
        double gauss1_total = 0.0;
        double gauss2_total = 0.0;
        double repuls_total = 0.0;
        double hydrophobic_total = 0.0;
        double hydrogen_total = 0.0;

        iConformer_Vectorized ligand = ligand_;
        //applyBondRotationInPlace(ligand, transform);

        const double xProtCenter = protein.center_x;
        const double yProtCenter = protein.center_y;
        const double zProtCenter = protein.center_z;

        if constexpr(!OnlyIntermolecular) // C++17
        {

            for (unsigned int idxLig1 = 0; idxLig1 < ligand.x.vectorsCount(); ++idxLig1) {
                for (unsigned int idxLig2 = idxLig1; idxLig2 < ligand.x.vectorsCount(); ++idxLig2) {

                    const Vc::Vector<double> x_diff = ligand.x.vector(idxLig1) - ligand.x.vector(idxLig2);
                    const Vc::Vector<double> y_diff = ligand.y.vector(idxLig1) - ligand.y.vector(idxLig2);
                    const Vc::Vector<double> z_diff = ligand.z.vector(idxLig1) - ligand.z.vector(idxLig2);
                    const Vc::Vector<double> squared_sum = (x_diff * x_diff) + (y_diff * y_diff) + (z_diff * z_diff);
                    const Vc::Vector<double> distances_raw = Vc::sqrt(squared_sum);

                    const Vc::Vector<double> distances = distanceFromRawDistance(distances_raw,
                                                                                ligand.atomicRadius.vector(idxLig1),
                                                                                ligand.atomicRadius.vector(idxLig2));

                    score_raw += VinaClassic::coeff_repulsion * vinaRepulsionComponent(distances, 0.0).sum();

                }

            }
        }

        int atomicNumber_VectorIdx_ligand = -1;
        int atomicNumber_VectorIdx_protein = -1;
        int atomVariant_VectorIdx_ligand = -1;
        int atomVariant_VectorIdx_protein = -1;

        Vc::Memory<Vc::Vector<unsigned char>, Vc::Vector<double>::Size> ligandAtomicNumInformation;
        Vc::Memory<Vc::Vector<unsigned int>, Vc::Vector<double>::Size> ligandAtomVariantInformation;

        Vc::Memory<Vc::Vector<unsigned char>, Vc::Vector<double>::Size> proteinAtomicNumInformation;
        Vc::Memory<Vc::Vector<unsigned int>, Vc::Vector<double>::Size> proteinAtomVariantInformation;

        for (unsigned int idxLig = 0; idxLig < ligand.x.vectorsCount(); ++idxLig) {
            const unsigned int globalIdx_Ligand = (idxLig * Vc::Vector<double>::Size);
            const unsigned char offsetUChar_Ligand = globalIdx_Ligand % Vc::Vector<unsigned char>::Size;
            const unsigned char offsetUInt_Ligand = globalIdx_Ligand % Vc::Vector<unsigned int>::Size;

            const unsigned int offset_UChar_ligand = offsetUChar_Ligand % Vc::Vector<double>::Size;
            if (offset_UChar_ligand == 0) {
                // we are at the start of a
                atomicNumber_VectorIdx_ligand++;
                ligandAtomicNumInformation = ligand.type.vector(atomicNumber_VectorIdx_ligand);
            }

            const unsigned int offset_UInt_ligand = offsetUInt_Ligand % Vc::Vector<double>::Size;
            if (offset_UInt_ligand == 0) {
                // we are at the start of a
                atomVariant_VectorIdx_ligand++;
                ligandAtomVariantInformation = ligand.variant.vector(atomVariant_VectorIdx_ligand);
            }

            Vc::Vector<double> ligandHydrophobicMaskMaker;
            for (unsigned int k = 0; k < Vc::Vector<double>::Size; ++k) {
                ligandHydrophobicMaskMaker[k] = isHydrophobic_prepareMask(
                        ligandAtomicNumInformation[offset_UChar_ligand + k],
                        ligandAtomVariantInformation[offset_UChar_ligand + k]
                );
            }

            Vc::Mask<double> ligandHydrophobicMask = (ligandHydrophobicMaskMaker == Vc::Vector<double>::One());

            for (unsigned int idxProt = 0; idxProt < protein.x.vectorsCount(); idxProt++) {
                Vc::Vector<double> xLig, yLig, zLig;
                xLig = ligand.x.vector(idxLig);
                yLig = ligand.y.vector(idxLig);
                zLig = ligand.z.vector(idxLig);

                applyRigidTransformInPlace(xLig, yLig, zLig, transform);

                const Vc::Vector<double> xsquared_disttoProtCenter = (xLig - xProtCenter) * (xLig - xProtCenter);
                const Vc::Vector<double> ysquared_disttoProtCenter = (yLig - yProtCenter) * (yLig - yProtCenter);
                const Vc::Vector<double> zsquared_disttoProtCenter = (zLig - zProtCenter) * (zLig - zProtCenter);

                const Vc::Vector<double> distancesToProtCenter = Vc::sqrt(xsquared_disttoProtCenter
                                                                          + ysquared_disttoProtCenter
                                                                          + zsquared_disttoProtCenter);

                const Vc::Vector<double> tooFarFromCenterPenalty = Vc::iif(
                        distancesToProtCenter > (protein.radius - 1),
                        Vc::exp(4 * Vc::log(distancesToProtCenter - protein.radius)) + 10,
                        Vc::Vector<double>(Vc::Zero));

                score_raw += tooFarFromCenterPenalty.sum();

                const Vc::Vector<double> xProt = protein.x.vector(idxProt);
                const Vc::Vector<double> yProt = protein.y.vector(idxProt);
                const Vc::Vector<double> zProt = protein.z.vector(idxProt);

                const Vc::Vector<double> xsquared = (xLig - xProt) * (xLig - xProt);
                const Vc::Vector<double> ysquared = (yLig - yProt) * (yLig - yProt);
                const Vc::Vector<double> zsquared = (zLig - zProt) * (zLig - zProt);

                const Vc::Vector<double> rawDistances = Vc::sqrt(xsquared + ysquared + zsquared);

                if (Vc::all_of(rawDistances >= VinaClassic::interaction_cutoff)) {
                    /*
                     * Theoretically, this helps because atoms which are close in idx are ~ close in position
                     * So all_of will be true a non negigible fraction of the time
                     * TODO : benchmark this versus just the mask
                     */
                    continue;
                }

                // If we have some higher and some lower than the cutoff, we we mask them at the end in .sum()
                auto cutOffMask = rawDistances < VinaClassic::interaction_cutoff;

                const Vc::Vector<double> distances = distanceFromRawDistance(rawDistances,
                                                                             ligand.atomicRadius.vector(idxLig),
                                                                             protein.atomicRadius.vector(idxProt));

                //std::cout << cutOffMask << " -> " << distances << std::endl;

                distance_total += distances.sum(cutOffMask);

                const unsigned int globalIdx_Prot = (idxProt * Vc::Vector<double>::Size);
                const unsigned char offsetUChar_Prot = globalIdx_Prot % Vc::Vector<unsigned char>::Size;
                const unsigned char offsetUInt_Prot = globalIdx_Prot % Vc::Vector<unsigned int>::Size;

                const unsigned int offset_UChar_protein = offsetUChar_Prot % Vc::Vector<double>::Size;
                if (offset_UChar_protein == 0) {
                    // we are at the start of a
                    atomicNumber_VectorIdx_protein++;
                    proteinAtomicNumInformation = protein.type.vector(atomicNumber_VectorIdx_protein);
                }

                const unsigned int offset_UInt_protein = offsetUInt_Prot % Vc::Vector<double>::Size;
                if (offset_UInt_protein == 0) {
                    // we are at the start of a
                    atomVariant_VectorIdx_protein++;
                    proteinAtomVariantInformation = protein.variant.vector(atomVariant_VectorIdx_ligand);
                }

                Vc::Vector<double> proteinHydrophobicMaskMaker;
                for (unsigned int k = 0; k < Vc::Vector<double>::Size; ++k) {
                    proteinHydrophobicMaskMaker[k] = isHydrophobic_prepareMask(
                            proteinAtomicNumInformation[offset_UChar_protein + k],
                            proteinAtomVariantInformation[offset_UInt_protein + k]
                    );
                }

                Vc::Vector<double> HydrogenMaskMaker;
                for (unsigned int k = 0; k < Vc::Vector<double>::Size; ++k) {
                    ligandHydrophobicMaskMaker[k] = hydrogenBondingPossible_prepareMask(
                            ligandAtomicNumInformation[offset_UChar_ligand + k],
                            ligandAtomVariantInformation[offset_UChar_ligand + k],
                            proteinAtomicNumInformation[offset_UChar_ligand + k],
                            proteinAtomVariantInformation[offset_UChar_ligand + k]
                    );
                }

                Vc::Mask<double> hydrophobicMask =
                        (proteinHydrophobicMaskMaker == Vc::Vector<double>::One()) && ligandHydrophobicMask;
                Vc::Mask<double> hydrogenMask = (HydrogenMaskMaker == Vc::Vector<double>::One());

                if constexpr(useNonDefaultCoefficients) {

                    score_raw += nonDefaultCoeffs[0] * vinaGaussComponent(distances, 0.0, 0.5).sum(cutOffMask);
                    score_raw += nonDefaultCoeffs[1] * vinaGaussComponent(distances, 3.0, 2.0).sum(cutOffMask);
                    score_raw += nonDefaultCoeffs[2] * vinaRepulsionComponent(distances, 0.0).sum(cutOffMask);
                    score_raw +=
                            nonDefaultCoeffs[3] * vinaHydrophobicComponent(distances, hydrophobicMask).sum(cutOffMask);
                    score_raw += nonDefaultCoeffs[4] * vinaHydrogenComponent(distances, hydrogenMask).sum(cutOffMask);
                } else {
                    gauss1_total += vinaGaussComponent(distances, 0.0, 0.5).sum(cutOffMask);
                    gauss2_total += vinaGaussComponent(distances, 3.0, 2.0).sum(cutOffMask);
                    repuls_total += vinaRepulsionComponent(distances, 0.0).sum(cutOffMask);
                    hydrophobic_total += vinaHydrophobicComponent(distances, hydrophobicMask).sum(cutOffMask);
                    hydrogen_total += vinaHydrogenComponent(distances, hydrogenMask).sum(cutOffMask);

                    score_raw += VinaClassic::coeff_gauss1 * vinaGaussComponent(distances, 0.0, 0.5).sum(cutOffMask);
                    score_raw += VinaClassic::coeff_gauss2 * vinaGaussComponent(distances, 3.0, 2.0).sum(cutOffMask);
                    score_raw += VinaClassic::coeff_repulsion * vinaRepulsionComponent(distances, 0.0).sum(cutOffMask);
                    score_raw += VinaClassic::coeff_hydrophobic *
                                 vinaHydrophobicComponent(distances, hydrophobicMask).sum(cutOffMask);
                    score_raw += VinaClassic::coeff_hydrogen *
                                 vinaHydrogenComponent(distances, hydrogenMask).sum(cutOffMask);
                }

            }
        }
        double final_score = score_raw / (1 + (VinaClassic::coeff_entropic * ligand.num_rotatable_bond));

        std::cout << "\n Vectorized : \n";
        std::cout << "dist    : " << distance_total << std::endl;
        std::cout << "gauss1  : " << gauss1_total << std::endl;
        std::cout << "gauss2  : " << gauss2_total << std::endl;
        std::cout << "repuls  : " << repuls_total << std::endl;
        std::cout << "hydroph : " << hydrophobic_total << std::endl;
        std::cout << "hydrog  : " << hydrogen_total << std::endl;
        std::cout << "\n\n";

        return final_score;
    }

    VinaLike::VinaLike(const iConformer &startingConformation_,
                       const iProtein &p,
                       const iTransform &initialTransform_,
                       double differential_epsilon_,
                       bool useNonDefaultCoefficient) :
            useNonDefaultCoefficient(useNonDefaultCoefficient),
            startingConformation(startingConformation_),
            prot(p),
            initialTransform(initialTransform_),
            differential_epsilon(differential_epsilon_) {
        this->numberOfRotatableBonds = this->startingConformation.num_rotatable_bond;
        this->numberOfParamInState = 7 + (this->numberOfRotatableBonds);

        if (this->initialTransform.bondRotationsAngles.size() != this->numberOfRotatableBonds) {
            BOOST_LOG_TRIVIAL(error)
                << "Discrepency between the number of rotatable bonds in the iConformer and iTransform ("
                << this->numberOfRotatableBonds << " != " << this->initialTransform.bondRotationsAngles.size() << ")";
            std::terminate();
        }

        if (this->useNonDefaultCoefficient) {
                this->nonDefaultCoefficients[0] = VinaClassic::coeff_gauss1;
                this->nonDefaultCoefficients[1] = VinaClassic::coeff_gauss2;
                this->nonDefaultCoefficients[2] = VinaClassic::coeff_repulsion;
                this->nonDefaultCoefficients[3] = VinaClassic::coeff_hydrophobic;
                this->nonDefaultCoefficients[4] = VinaClassic::coeff_hydrogen;
        }

    }

    double VinaLike::Evaluate(const arma::mat &x) {
        BOOST_ASSERT(x.n_rows == this->numberOfParamInState);

        iTransform tr = this->internalToExternalRepr(x);

        normalizeQuaternionInPlace(tr.rota);

        double score_ = this->useNonDefaultCoefficient ?
                        VinaLikeIntermolecularScoringFunction<false, true>(this->startingConformation, tr, this->prot,
                                                                           this->nonDefaultCoefficients)
                                                       : VinaLikeIntermolecularScoringFunction<false, false>(
                        this->startingConformation, tr, this->prot);

        return score_;
    }

    double VinaLike::EvaluateWithGradient(const arma::mat &x, arma::mat &grad) {

        BOOST_ASSERT(!x.has_nan());
        BOOST_ASSERT(!grad.has_nan());
        BOOST_ASSERT(x.n_rows == this->numberOfParamInState);
        BOOST_ASSERT(grad.n_rows == this->numberOfParamInState);

        iTransform tr = this->internalToExternalRepr(x);
        tr.doHousekeeping();

        double score_ = this->useNonDefaultCoefficient ?
                        VinaLikeIntermolecularScoringFunction<false, true>(this->startingConformation, tr, this->prot,
                                                                           this->nonDefaultCoefficients)
                                                       : VinaLikeIntermolecularScoringFunction<false, false>(
                        this->startingConformation, tr, this->prot);

        // Translation
        {
            iTransform transform_dx = tr;
            transform_dx.transl.x() += this->differential_epsilon;
            const double gradScore = this->useNonDefaultCoefficient ?
                                     VinaLikeIntermolecularScoringFunction<false, true>(this->startingConformation,
                                                                                transform_dx, this->prot,
                                                                                this->nonDefaultCoefficients)
                                                                    : VinaLikeIntermolecularScoringFunction<false, false>(
                            this->startingConformation, transform_dx, this->prot);
            grad[0] = gradScore - score_;
        }

        {
            iTransform transform_dy = tr;
            transform_dy.transl.y() += this->differential_epsilon;
            const double gradScore = this->useNonDefaultCoefficient ?
                                     VinaLikeIntermolecularScoringFunction<false, true>(this->startingConformation,
                                                                                transform_dy, this->prot,
                                                                                this->nonDefaultCoefficients)
                                                                    : VinaLikeIntermolecularScoringFunction<false, false>(
                            this->startingConformation, transform_dy, this->prot);
            grad[1] = gradScore - score_;
        }

        {
            iTransform transform_dz = tr;
            transform_dz.transl.z() += this->differential_epsilon;
            const double gradScore = this->useNonDefaultCoefficient ?
                                     VinaLikeIntermolecularScoringFunction<false, true>(this->startingConformation,
                                                                                transform_dz, this->prot,
                                                                                this->nonDefaultCoefficients)
                                                                    : VinaLikeIntermolecularScoringFunction<false, false>(
                            this->startingConformation, transform_dz, this->prot);
            grad[2] = gradScore - score_;
        }

        // Rotation

        {
            iTransform transform_dqs = tr;
            transform_dqs.rota.w() += this->differential_epsilon;
            transform_dqs.doHousekeeping();
            const double gradScore = this->useNonDefaultCoefficient ?
                                     VinaLikeIntermolecularScoringFunction<false, true>(this->startingConformation,
                                                                                transform_dqs, this->prot,
                                                                                this->nonDefaultCoefficients)
                                                                    : VinaLikeIntermolecularScoringFunction<false, false>(
                            this->startingConformation, transform_dqs, this->prot);
            grad[3] = gradScore - score_;
        }

        {
            iTransform transform_dqx = tr;
            transform_dqx.rota.x() += this->differential_epsilon;
            transform_dqx.doHousekeeping();
            const double gradScore = this->useNonDefaultCoefficient ?
                                     VinaLikeIntermolecularScoringFunction<false, true>(this->startingConformation,
                                                                                transform_dqx, this->prot,
                                                                                this->nonDefaultCoefficients)
                                                                    : VinaLikeIntermolecularScoringFunction<false, false>(
                            this->startingConformation, transform_dqx, this->prot);
            grad[4] = gradScore - score_;

        }

        {
            iTransform transform_dqy = tr;
            transform_dqy.rota.x() += this->differential_epsilon;
            transform_dqy.doHousekeeping();
            const double gradScore = this->useNonDefaultCoefficient ?
                                     VinaLikeIntermolecularScoringFunction<false, true>(this->startingConformation,
                                                                                transform_dqy, this->prot,
                                                                                this->nonDefaultCoefficients)
                                                                    : VinaLikeIntermolecularScoringFunction<false, false>(
                            this->startingConformation, transform_dqy, this->prot);
            grad[5] = gradScore - score_;
        }

        {
            iTransform transform_dqz = tr;
            transform_dqz.rota.x() += this->differential_epsilon;
            transform_dqz.doHousekeeping();
            const double gradScore = this->useNonDefaultCoefficient ?
                                     VinaLikeIntermolecularScoringFunction<false, true>(this->startingConformation,
                                                                                transform_dqz, this->prot,
                                                                                this->nonDefaultCoefficients)
                                                                    : VinaLikeIntermolecularScoringFunction<false, false>(
                            this->startingConformation, transform_dqz, this->prot);
            grad[6] = gradScore - score_;
        }

        for (unsigned int i = 0; i < this->numberOfRotatableBonds; i++) {
            iTransform transform_dbondrot = tr;
            transform_dbondrot.bondRotationsAngles[i] += this->differential_epsilon;
            const double gradScore = this->useNonDefaultCoefficient ?
                                     VinaLikeIntermolecularScoringFunction<false, true>(this->startingConformation,
                                                                                transform_dbondrot, this->prot,
                                                                                this->nonDefaultCoefficients)
                                                                    : VinaLikeIntermolecularScoringFunction<false, false>(
                            this->startingConformation, transform_dbondrot, this->prot);
            grad[7 + i] = gradScore - score_;

        }

        /*
        BOOST_LOG_TRIVIAL(debug) << "Transform: " << x.t();
        BOOST_LOG_TRIVIAL(debug) << "Score: " << score_;
        BOOST_LOG_TRIVIAL(debug) << "Gradient";
        BOOST_LOG_TRIVIAL(debug) << "     ds: " << grad[3];
        BOOST_LOG_TRIVIAL(debug) << "     du: " << grad[4] << "   dx: " << grad[0];
        BOOST_LOG_TRIVIAL(debug) << "     dv: " << grad[5] << "   dy: " << grad[1];
        BOOST_LOG_TRIVIAL(debug) << "     dt: " << grad[6] << "   dx: " << grad[2];
        //*/

        BOOST_ASSERT(score_ == score_); // catches NaN
        return score_;
    }

    double VinaLike::getDifferentialEpsilon() const {
        return this->differential_epsilon;
    }

    arma::mat VinaLike::getStartingConditions() const {
        return this->externalToInternalRepr(this->initialTransform);
    }

    iConformer VinaLike::getConformerForParamMatrix(const arma::mat &x) {
        BOOST_ASSERT(x.n_rows == this->numberOfParamInState);

        iTransform tr = this->internalToExternalRepr(x);
        tr.doHousekeeping();

        iConformer ret = this->startingConformation;
        applyBondRotationInPlace(ret, tr);
        applyRigidTransformInPlace(ret, tr);

        return ret;
    }

    unsigned int VinaLike::getParamVectorDimension() const {
        return this->numberOfParamInState;
    }

    std::vector<std::tuple<std::string, double>> VinaLike::EvaluateSubcomponents(const arma::mat &x) {
        std::vector<std::tuple<std::string, double>> ret;

        BOOST_ASSERT(x.n_rows == this->numberOfParamInState);

        iTransform tr = this->internalToExternalRepr(x);
        tr.doHousekeeping();

        BOOST_ASSERT(!this->startingConformation.x.empty());
        BOOST_ASSERT(!this->prot.x.empty());
        BOOST_ASSERT(tr.bondRotationsAngles.size() == this->startingConformation.num_rotatable_bond);

        if (std::abs(tr.rota.norm() - 1) > 0.1) {
            tr.rota.normalize();
        }

        double gauss1_total = 0.0;
        double gauss2_total = 0.0;
        double repulsion_total = 0.0;
        double hydrogen_total = 0.0;
        double hydrophobic_total = 0.0;
        double score_raw = 0.0;

        double intramolecular_repuls_total = 0.0;
        double intramolecular_score = 0.0;

        iConformer ligand = this->startingConformation;
        applyBondRotationInPlace(ligand, tr);

        Eigen::Vector3d ProtCenterPosition = {this->prot.center_x, this->prot.center_y, this->prot.center_z};

        for (unsigned int idxLig = 0; idxLig < ligand.x.size(); idxLig++) {
            for (unsigned int idxLig2 = idxLig; idxLig2 < ligand.x.size(); idxLig2++) {

                if (idxLig == idxLig2)
                    continue;

                Eigen::Vector3d LigDistance = {ligand.x[idxLig] - ligand.x[idxLig2],
                                               ligand.y[idxLig] - ligand.y[idxLig2],
                                               ligand.z[idxLig] - ligand.z[idxLig2]};

                double distance_raw = LigDistance.norm();
                const double distance = distanceFromRawDistance(distance_raw, ligand.atomicRadius[idxLig],
                                                                ligand.atomicRadius[idxLig2]);

//                score_raw += VinaClassic::coeff_gauss1      * vinaGaussComponent(distance, 0.0, 0.5);
//                score_raw += VinaClassic::coeff_gauss2      * vinaGaussComponent(distance, 3.0, 2.0);
                intramolecular_repuls_total += vinaRepulsionComponent(distance, 0.0);

            }
        }

        intramolecular_score = VinaClassic::coeff_repulsion * intramolecular_repuls_total;

        for (unsigned int idxLig = 0; idxLig < ligand.x.size(); idxLig++) {
            for (unsigned int idxProt = 0; idxProt < this->prot.x.size(); idxProt++) {

                Eigen::Vector3d LigPosition = {ligand.x[idxLig], ligand.y[idxLig], ligand.z[idxLig]};
                applyRigidTransformInPlace(LigPosition, tr);

                Eigen::Vector3d ProtPosition = {this->prot.x[idxProt], this->prot.y[idxProt], this->prot.z[idxProt]};
                Eigen::Vector3d distToCenterVector = LigPosition - ProtCenterPosition;

                double distanceToProteinCenter = distToCenterVector.norm();

                if (distanceToProteinCenter > (this->prot.radius - 1)) {
                    score_raw += std::pow((distanceToProteinCenter - this->prot.radius), 4) + 10;
                    continue;
                }

                Eigen::Vector3d distVect = ProtPosition - LigPosition;

                double rawDist = distVect.norm();

                if (rawDist >= VinaClassic::interaction_cutoff)
                    continue;

                double distance = distanceFromRawDistance(rawDist, ligand.atomicRadius[idxLig],
                                                          this->prot.atomicRadius[idxProt]);

                const unsigned int atom1AtomicNumber = ligand.type[idxLig];
                const unsigned int atom1AtomVariant = ligand.variant[idxLig];
                const unsigned int atom2AtomicNumber = this->prot.type[idxProt];
                const unsigned int atom2AtomVariant = this->prot.variant[idxProt];

                gauss1_total += vinaGaussComponent(distance, 0.0, 0.5);
                gauss2_total += vinaGaussComponent(distance, 3.0, 2.0);
                repulsion_total += vinaRepulsionComponent(distance, 0.0);
                hydrophobic_total += vinaHydrophobicComponent(distance,
                                                              atom1AtomicNumber, atom1AtomVariant,
                                                              atom2AtomicNumber, atom2AtomVariant);
                hydrogen_total += vinaHydrogenComponent(distance,
                                                        atom1AtomicNumber, atom1AtomVariant,
                                                        atom2AtomicNumber, atom2AtomVariant);

            } // for
        } // for

        double score_sum = 0.0;

        if (this->useNonDefaultCoefficient) {
            ret.emplace_back(std::make_tuple("NonDefaultCoeffs", 1.0));
            score_sum = this->nonDefaultCoefficients[0] * gauss1_total
                        + this->nonDefaultCoefficients[1] * gauss2_total
                        + this->nonDefaultCoefficients[2] * repulsion_total
                        + this->nonDefaultCoefficients[3] * hydrophobic_total
                        + this->nonDefaultCoefficients[4] * hydrogen_total;
        } else {
            score_sum = VinaClassic::coeff_gauss1 * gauss1_total
                        + VinaClassic::coeff_gauss2 * gauss2_total
                        + VinaClassic::coeff_repulsion * repulsion_total
                        + VinaClassic::coeff_hydrophobic * hydrophobic_total
                        + VinaClassic::coeff_hydrogen * hydrogen_total;
        }

        double final_score_fromSum = score_sum / (1 + (VinaClassic::coeff_entropic * ligand.num_rotatable_bond));

        ret.emplace_back(std::make_tuple("Gauss1", gauss1_total));
        ret.emplace_back(std::make_tuple("Gauss2", gauss2_total));
        ret.emplace_back(std::make_tuple("Repulsion", repulsion_total));
        ret.emplace_back(std::make_tuple("Hydrophobic", hydrophobic_total));
        ret.emplace_back(std::make_tuple("Hydrogen", hydrogen_total));
        ret.emplace_back(std::make_tuple("Intra_Repuls", intramolecular_repuls_total));
        ret.emplace_back(std::make_tuple("Intra_Score", intramolecular_score));
        ret.emplace_back(std::make_tuple("ScoreRawSum", score_sum));
        ret.emplace_back(std::make_tuple("Score", final_score_fromSum));
        return ret;
    }

    double VinaLike::EvaluateOnlyIntermolecular(const arma::mat &x) {
        BOOST_ASSERT(x.n_rows == this->numberOfParamInState);

        iTransform tr = this->internalToExternalRepr(x);

        tr.doHousekeeping();

        // Template parameter controls whether onlyIntermolecular interaction are taken into account. Here we want true
        double score_ = this->useNonDefaultCoefficient ?
                        VinaLikeIntermolecularScoringFunction<true, true>(this->startingConformation, tr, this->prot,
                                                                           this->nonDefaultCoefficients)
                                                       : VinaLikeIntermolecularScoringFunction<true, false>(
                        this->startingConformation, tr, this->prot);

        return score_;
    }

    unsigned int VinaLike::getCoefficientsVectorWidth() {
        return this->numCoefficients;
    }

    std::vector<std::string> VinaLike::getCoefficientsNames() {
        return std::vector<std::string>(this->coefficientsNames.begin(), this->coefficientsNames.end());
    }

    std::vector<double> VinaLike::getCurrentCoefficients() {
        if (this->useNonDefaultCoefficient) {
            return std::vector<double>(this->nonDefaultCoefficients.begin(), this->nonDefaultCoefficients.end());
        }
        return {VinaClassic::coeff_gauss1, VinaClassic::coeff_gauss2,
                VinaClassic::coeff_repulsion, VinaClassic::coeff_hydrophobic,
                VinaClassic::coeff_hydrogen};
    }

    bool VinaLike::setNonDefaultCoefficients(std::vector<double> coeffs) {
        if (coeffs.size() != this->numCoefficients) {
            BOOST_LOG_TRIVIAL(error) << "Trying to set " << this->numCoefficients << " coefficients with vector of "
                                     << coeffs.size() << " values.";
            return false;
        }
        if (this->useNonDefaultCoefficient == false) {
            BOOST_LOG_TRIVIAL(error)
                << "Trying to set non default coefficient, but this scoring function was constructed with default coefficients only.";
            BOOST_LOG_TRIVIAL(error) << "Check the parameters passed to the scoring function constructor.";
            return false;
        }
        for (unsigned int j = 0; j < this->nonDefaultCoefficients.size(); ++j) {
            this->nonDefaultCoefficients[j] = coeffs[j];
        }
        return true;
    }

}


//
// Copyright (c) 2015-2019 CNRS INRIA
// Copyright (c) 2015 Wandercraft, 86 rue de Paris 91400 Orsay, France.
//

#ifndef __pinocchio_data_hxx__
#define __pinocchio_data_hxx__

#include "pinocchio/spatial/fwd.hpp"
#include "pinocchio/multibody/model.hpp"
#include "pinocchio/utils/string-generator.hpp"
#include "pinocchio/multibody/liegroup/liegroup-algo.hpp"

#include <boost/bind.hpp>
#include <boost/utility.hpp>

/// @cond DEV

namespace pinocchio
{

  template<typename Scalar, int Options, template<typename,int> class JointCollectionTpl>
  inline DataTpl<Scalar,Options,JointCollectionTpl>::
  DataTpl(const Model & model)
  : joints(0)
  , a((std::size_t)model.njoints)
  , oa((std::size_t)model.njoints)
  , a_gf((std::size_t)model.njoints)
  , oa_gf((std::size_t)model.njoints)
  , v((std::size_t)model.njoints)
  , ov((std::size_t)model.njoints)
  , f((std::size_t)model.njoints)
  , of((std::size_t)model.njoints)
  , h((std::size_t)model.njoints)
  , oh((std::size_t)model.njoints)
  , oMi((std::size_t)model.njoints)
  , liMi((std::size_t)model.njoints)
  , tau(model.nv)
  , nle(model.nv)
  , g(model.nv)
  , oMf((std::size_t)model.nframes)
  , Ycrb((std::size_t)model.njoints)
  , dYcrb((std::size_t)model.njoints)
  , M(model.nv,model.nv)
  , Minv(model.nv,model.nv)
  , C(model.nv,model.nv)
  , dHdq(6,model.nv)
  , dFdq(6,model.nv)
  , dFdv(6,model.nv)
  , dFda(6,model.nv)
  , SDinv(6,model.nv)
  , UDinv(6,model.nv)
  , IS(6,model.nv)
  , vxI((std::size_t)model.njoints)
  , Ivx((std::size_t)model.njoints)
  , oYcrb((std::size_t)model.njoints)
  , doYcrb((std::size_t)model.njoints)
  , ddq(model.nv)
  , Yaba((std::size_t)model.njoints)
  , u(model.nv)
  , Ag(6,model.nv)
  , dAg(6,model.nv)
  , Fcrb((std::size_t)model.njoints)
  , lastChild((std::size_t)model.njoints)
  , nvSubtree((std::size_t)model.njoints)
  , U(model.nv,model.nv)
  , D(model.nv)
  , Dinv(model.nv)
  , tmp(model.nv)
  , parents_fromRow((std::size_t)model.nv)
  , nvSubtree_fromRow((std::size_t)model.nv)
  , J(6,model.nv)
  , dJ(6,model.nv)
  , dVdq(6,model.nv)
  , dAdq(6,model.nv)
  , dAdv(6,model.nv)
  , dtau_dq(MatrixXs::Zero(model.nv,model.nv))
  , dtau_dv(MatrixXs::Zero(model.nv,model.nv))
  , ddq_dq(MatrixXs::Zero(model.nv,model.nv))
  , ddq_dv(MatrixXs::Zero(model.nv,model.nv))
  , iMf((std::size_t)model.njoints)
  , com((std::size_t)model.njoints)
  , vcom((std::size_t)model.njoints)
  , acom((std::size_t)model.njoints)
  , mass((std::size_t)model.njoints)
  , Jcom(3,model.nv)
  , JMinvJt()
  , llt_JMinvJt()
  , lambda_c()
  , sDUiJt(model.nv,model.nv)
  , torque_residual(model.nv)
  , dq_after(model.nv)
  , impulse_c()
  , staticRegressor(3,4*(model.njoints-1))
  {
    typedef typename Model::JointIndex JointIndex;

    /* Create data strcture associated to the joints */
    for(JointIndex i=0;i<(JointIndex)(model.njoints);++i)
      joints.push_back(CreateJointData<Scalar,Options,JointCollectionTpl>::run(model.joints[i]));

    /* Init for CRBA */
    M.setZero(); Minv.setZero();
    for(JointIndex i=0;i<(JointIndex)(model.njoints);++i)
    { Fcrb[i].resize(6,model.nv); }

    computeLastChild(model);

    /* Init for Coriolis */
    C.setZero();

    /* Init for Cholesky */
    U.setIdentity();
    computeParents_fromRow(model);

    /* Init Jacobian */
    J.setZero();
    Ag.setZero();

    /* Init universe states relatively to itself */

    a[0].setZero();
    oa[0].setZero();
    v[0].setZero();
    ov[0].setZero();
    a_gf[0] = -model.gravity;
    f[0].setZero();
    h[0].setZero();
    oMi[0].setIdentity();
    liMi[0].setIdentity();
    oMf[0].setIdentity();

    Yaba[0].setZero();
    Ycrb[0].setZero();
  }

  template<typename Scalar, int Options, template<typename,int> class JointCollectionTpl>
  inline void DataTpl<Scalar,Options,JointCollectionTpl>::
  computeLastChild(const Model & model)
  {
    typedef typename Model::Index Index;

    std::fill(lastChild.begin(),lastChild.end(),-1);
    for( int i=model.njoints-1;i>=0;--i )
    {
      if(lastChild[(Index)i] == -1) lastChild[(Index)i] = i;
      const Index & parent = model.parents[(Index)i];
      lastChild[parent] = std::max(lastChild[(Index)i],lastChild[parent]);

      nvSubtree[(Index)i]
      = idx_v(model.joints[(Index)lastChild[(Index)i]]) + nv(model.joints[(Index)lastChild[(Index)i]])
      - idx_v(model.joints[(Index)i]);
    }
  }

  template<typename Scalar, int Options, template<typename,int> class JointCollectionTpl>
  inline void DataTpl<Scalar,Options,JointCollectionTpl>
  ::computeParents_fromRow(const Model & model)
  {
    typedef typename Model::Index Index;

    for(Index joint=1;joint<(Index)(model.njoints);joint++)
    {
      const Index & parent = model.parents[joint];
      const int nvj    = nv   (model.joints[joint]);
      const int idx_vj = idx_v(model.joints[joint]);

      if(parent>0) parents_fromRow[(Index)idx_vj] = idx_v(model.joints[parent])+nv(model.joints[parent])-1;
      else         parents_fromRow[(Index)idx_vj] = -1;
      nvSubtree_fromRow[(Index)idx_vj] = nvSubtree[joint];

      for(int row=1;row<nvj;++row)
      {
        parents_fromRow[(Index)(idx_vj+row)] = idx_vj+row-1;
        nvSubtree_fromRow[(Index)(idx_vj+row)] = nvSubtree[joint]-row;
      }
    }
  }

} // namespace pinocchio

/// @endcond

#endif // ifndef __pinocchio_data_hxx__


#pragma once
//===----------------------------------------------------------------------===//
#include "bah_patterns.hpp"
#include "bah_types.hpp"

#include <dtl/bitmap/util/plain_bitmap.hpp>
#include <dtl/dtl.hpp>

#include <boost/dynamic_bitset.hpp>

#include <cstddef>
#include <iomanip>
#include <type_traits>
#include <vector>
//===----------------------------------------------------------------------===//
namespace dtl {
//===----------------------------------------------------------------------===//
/// Implementation of the Byte Aligned Hybrid bitmap compression technique as
/// described in the paper 'BAH: A Bitmap Index Compression Algorithm for Fast
/// Data Retrieval' by Li et al.
class bah {
  static constexpr std::size_t word_bitlength = sizeof(bah_word_t) * 8;

  /// The encoded bitmap.
  std::vector<bah_byte_t> main_;
  std::vector<bah_word_t> data_;
  std::vector<bah_byte_t> index_;
  std::vector<bah_word_t> counter_;
  /// The length of the encoded bitmap.
  std::size_t encoded_bitmap_length_ = 0;

  /// The different word kinds.
  enum class word_kind {
    ZERO,
    ENCODABLE,
    LITERAL
  };

  /// Determines the word kind for a single word of the plain bitmap.
  static word_kind
  get_word_kind(bah_word_t word) {
    if (word == 0) {
      return word_kind::ZERO;
    }
    if (get_oep_code(word) != dtl::bah_code_not_found) {
      return word_kind::ENCODABLE;
    }
    if (get_tep_code(word) != dtl::bah_code_not_found) {
      return word_kind::ENCODABLE;
    }
    return word_kind::LITERAL;
  }

public:
  bah() = default;

  explicit bah(const boost::dynamic_bitset<$u32>& in) {
    auto* bitmap_begin = in.m_bits.data();
    auto* bitmap_end = in.m_bits.data() + in.m_bits.size();
    const auto word_cnt = bitmap_end - bitmap_begin; // TODO handle the case when the bitmap size is not a multiple of the word size

    // The plain bitmap is segmented into words of the same kind.
    // Keep track of the first and last word of the current segment.
    std::size_t seg_begin = 0;
    std::size_t seg_end = 0;

    // Encode the bitmap segment by segment.
    while (seg_begin < word_cnt) {
      auto& w = bitmap_begin[seg_begin];
      auto seg_kind = get_word_kind(w);
      seg_end = seg_begin + 1;
      // Find the end of the current 'segment'.
      while (seg_end < word_cnt
          && get_word_kind(bitmap_begin[seg_end]) == seg_kind) {
        auto seg_len = seg_end - seg_begin;
        if ((seg_kind == word_kind::LITERAL && seg_len == 63)
            || (seg_kind == word_kind::ENCODABLE)) {
          break;
        }
        ++seg_end;
      }
      // Encode the segment.
      if (seg_kind == word_kind::ZERO) {
        auto len = seg_end - seg_begin;
        if (len > 252) {
          main_.push_back(bah_byte_t(0));
          counter_.push_back(len);
        }
        else {
          while (len > 0) {
            auto l = std::min(len, 63ul);
            main_.push_back(static_cast<bah_byte_t>(l));
            len -= l;
          }
        }
      }
      else if (seg_kind == word_kind::LITERAL) {
        auto len = seg_end - seg_begin;
        main_.push_back(static_cast<bah_byte_t>(len) | (bah_byte_t(0b01) << 6));
        for (std::size_t i = seg_begin; i < seg_end; ++i) {
          data_.push_back(bitmap_begin[i]);
        }
      }
      else { /* ENCODABLE */
        auto code = get_oep_code(w);
        if (code != dtl::bah_code_not_found) {
          main_.push_back(
              static_cast<bah_byte_t>(code) | (bah_byte_t(0b10) << 6));
        }
        else {
          code = get_tep_code(w);
          bah_byte_t code_lo = code & 0b111111;
          bah_byte_t code_hi = code >> 6;
          main_.push_back(code_lo | (bah_byte_t(0b11) << 6));
          index_.push_back(code_hi);
        }
      }

      // Next.
      seg_begin = seg_end;
    }
    encoded_bitmap_length_ = in.size();

    // Try to reduce the memory consumption.
    shrink();
  }

  ~bah() = default;
  bah(const bah& other) = default;
  bah(bah&& other) noexcept = default;
  bah& operator=(const bah& other) = default;
  bah& operator=(bah&& other) noexcept = default;

  /// Return the size in bytes.
  std::size_t __forceinline__
  size_in_bytes() const {
    return main_.size() * sizeof(bah_byte_t)
        + data_.size() * sizeof(bah_word_t)
        + index_.size() * sizeof(bah_byte_t)
        + counter_.size() * sizeof(bah_word_t)
        + 3 * sizeof(void*) /* pointers/offsets to the three other arrays */
        + sizeof(encoded_bitmap_length_); /* bit-length of the original bitmap */
  }

  /// Returns the size of the bitmap.
  std::size_t __forceinline__
  size() const {
    return encoded_bitmap_length_;
  }

  /// Conversion to an plain bitmap. // TODO remove
  dtl::plain_bitmap<$u64>
  to_plain_bitmap() {
    dtl::plain_bitmap<$u64> ret(encoded_bitmap_length_, false);
    // The write position.
    std::size_t i = 0;

    // The read position(s).
    std::size_t data_idx = 0;
    std::size_t index_idx = 0;
    std::size_t counter_idx = 0;
    for (std::size_t main_idx = 0; main_idx < main_.size(); ++main_idx) {
      const auto& w = main_[main_idx];
      const auto type = w >> 6;
      const auto n = w & 0b111111;
      switch (type) {
        case 0b00: { /* ZERO */
          std::size_t count = 0;
          if (n == 0) {
            count = counter_[counter_idx++];
          }
          else {
            count = n;
          }
          auto l = count * word_bitlength;
          ret.clear(i, std::min(i + l, encoded_bitmap_length_));
          i += l;
          break;
        }
        case 0b01: { /* LITERAL */
          for (std::size_t j = 0; j < n; ++j) {
            auto lit = data_[data_idx++];
            ret.store_bits(i,
                std::min(i + word_bitlength, encoded_bitmap_length_), lit);
            i += word_bitlength;
          }
          break;
        }
        case 0b10: { /* ENCODABLE (one byte) */
          auto pattern = get_oep(n);
          ret.store_bits(i,
              std::min(i + word_bitlength, encoded_bitmap_length_), pattern);
          i += word_bitlength;
          break;
        }
        case 0b11: { /* ENCODABLE (two bytes) */
          auto code = bah_word_t(n) | (bah_word_t(index_[index_idx++]) << 6);
          auto pattern = get_tep(code);
          ret.store_bits(i,
              std::min(i + word_bitlength, encoded_bitmap_length_), pattern);
          i += word_bitlength;
          break;
        }
      }
    }
    if (i < encoded_bitmap_length_) {
      ret.clear(i, encoded_bitmap_length_);
    }
    return std::move(ret);
  }

  static std::string
  name() {
    return "bah";
  }

  /// Returns the value of the bit at the position pos.
  u1 __forceinline__
  test(const std::size_t pos) const {
    // The current range spanned by the current encoded word.
    std::size_t b = 0;
    std::size_t e = 0;

    // The read position(s).
    std::size_t data_idx = 0;
    std::size_t index_idx = 0;
    std::size_t counter_idx = 0;
    for (std::size_t main_idx = 0; main_idx < main_.size(); ++main_idx) {
      const auto& w = main_[main_idx];
      const auto type = w >> 6;
      const auto n = w & 0b111111;
      switch (type) {
        case 0b00: { /* ZERO */
          std::size_t count = 0;
          if (n == 0) {
            count = counter_[counter_idx++];
          }
          else {
            count = n;
          }
          auto l = count * word_bitlength;
          e = std::min(b + l, encoded_bitmap_length_);
          if (b <= pos && pos < e) {
            return false;
          }
          break;
        }
        case 0b01: { /* LITERAL */
          e = std::min(b + n * word_bitlength, encoded_bitmap_length_);
          if (b <= pos && pos < e) {
            auto offset = pos - b;
            auto lit_word = data_[data_idx + (offset / word_bitlength)];
            return dtl::bits::bit_test(lit_word, offset % word_bitlength);
          }
          data_idx += n;
          break;
        }
        case 0b10: { /* ENCODABLE (one byte) */
          e = std::min(b + word_bitlength, encoded_bitmap_length_);
          if (b <= pos && pos < e) {
            auto pattern = get_oep(n);
            return dtl::bits::bit_test(pattern, pos % word_bitlength);
          }
          break;
        }
        case 0b11: { /* ENCODABLE (two bytes) */
          e = std::min(b + word_bitlength, encoded_bitmap_length_);
          auto code = bah_word_t(n) | (bah_word_t(index_[index_idx++]) << 6);
          if (b <= pos && pos < e) {
            auto pattern = get_tep(code);
            return dtl::bits::bit_test(pattern, pos % word_bitlength);
          }
          break;
        }
      }
      // Next.
      b = e;
    }
    return false;
  }

  /// Try to reduce the memory consumption. This function is supposed to be
  /// called after the bitmap has been modified.
  __forceinline__ void
  shrink() {
    main_.shrink_to_fit();
    data_.shrink_to_fit();
    index_.shrink_to_fit();
    counter_.shrink_to_fit();
  }

  //===--------------------------------------------------------------------===//
  /// 1-run iterator
  class iter {
    const bah& outer_;

    /// The read positions.
    std::size_t main_idx_;
    std::size_t data_idx_;
    std::size_t index_idx_;
    std::size_t counter_idx_;

    //    word_kind segment_kind_;
    bah_word_t literal_;
    std::size_t literal_words_remaining_;
    std::size_t literal_in_word_idx_;

    /// Points to the beginning of a 1-run.
    $u64 pos_;
    /// The length of the current 1-run.
    $u64 length_;

  public:
    explicit __forceinline__
    iter(const bah& outer)
        : outer_(outer),
          main_idx_(0),
          data_idx_(0),
          index_idx_(0),
          counter_idx_(0),
          literal_(0),
          literal_words_remaining_(0),
          literal_in_word_idx_(0),
          pos_(0), length_(0) {
      // Find the first 1-run.
      const auto main_cnt = outer_.main_.size();
      while (main_idx_ < main_cnt) {
        pos_ += length_;
        length_ = 0;

        const auto& w = outer_.main_[main_idx_++];
        const auto type = w >> 6;
        const auto n = w & 0b111111;

        switch (type) {
          case 0b00: { /* ZERO */
            std::size_t count = 0;
            if (n == 0) {
              count = outer_.counter_[counter_idx_++];
            }
            else {
              count = n;
            }
            auto l = count * word_bitlength;
            length_ = l;
            break; // Skip over zeros.
          }
          case 0b01: { /* LITERAL */
            literal_ = outer_.data_[data_idx_++];
            assert(literal_ != 0);
            literal_words_remaining_ = n - 1;
            assert(literal_words_remaining_ < n);
            const std::size_t b = dtl::bits::tz_count(literal_);
            std::size_t e = b + 1;
            for (; e < word_bitlength; ++e) {
              u1 is_set = dtl::bits::bit_test(literal_, e);
              if (!is_set) break;
            }
            literal_in_word_idx_ = e;
            pos_ += b;
            length_ = e - b;
            return;
          }
          case 0b10: { /* ENCODABLE (one byte) */
            literal_ = get_oep(n);
            assert(literal_ != 0);
            literal_words_remaining_ = 0;
            const std::size_t b = dtl::bits::tz_count(literal_);
            std::size_t e = b + 1;
            for (; e < word_bitlength; ++e) {
              u1 is_set = dtl::bits::bit_test(literal_, e);
              if (!is_set) break;
            }
            literal_in_word_idx_ = e;
            pos_ += b;
            length_ = e - b;
            return;
          }
          case 0b11: { /* ENCODABLE (two bytes) */
            auto code =
                bah_word_t(n) | (bah_word_t(outer_.index_[index_idx_++]) << 6);
            literal_ = get_tep(code);
            assert(literal_ != 0);
            literal_words_remaining_ = 0;
            const std::size_t b = dtl::bits::tz_count(literal_);
            std::size_t e = b + 1;
            for (; e < word_bitlength; ++e) {
              u1 is_set = dtl::bits::bit_test(literal_, e);
              if (!is_set) break;
            }
            literal_in_word_idx_ = e;
            pos_ += b;
            length_ = e - b;
            return;
          }
        }
      }
      pos_ = outer_.encoded_bitmap_length_;
      length_ = 0;
    }

    /// Forward the iterator to the next 1-run.
    void __forceinline__
    next() {
      const auto main_cnt = outer_.main_.size();
      while (main_idx_ < main_cnt || literal_ != 0) {
        pos_ += length_;
        length_ = 0;

        // Check, whether we are currently in a literal segment.
        if (literal_ != 0) {
          if (literal_in_word_idx_ < word_bitlength) {
            // Find the next set bit.
            std::size_t b = literal_in_word_idx_;
            for (; b < word_bitlength; b++) { // TODO optimize
              u1 is_set = dtl::bits::bit_test(literal_, b);
              if (is_set) break;
            }
            if (b < word_bitlength) {
              // Determine the length of the current 1-run and return.
              std::size_t e = b + 1;
              for (; e < word_bitlength; ++e) { // TODO optimize
                u1 is_set = dtl::bits::bit_test(literal_, e);
                if (!is_set) break;
              }
              pos_ += b - literal_in_word_idx_;
              length_ = e - b;
              literal_in_word_idx_ = e;
              return;
            }
            pos_ += b - literal_in_word_idx_;
          }
          // Fetch the next literal word, if any.
          if (literal_words_remaining_ > 0) {
            literal_ = outer_.data_[data_idx_++];
            --literal_words_remaining_;
            literal_in_word_idx_ = 0;
            continue;
          }
          else {
            // Reached the end of the literal segment.
            literal_ = 0;
            literal_in_word_idx_ = 0;
          }
        }

        assert(literal_ == 0);

        if (main_idx_ >= main_cnt) break;

        const auto& w = outer_.main_[main_idx_++];
        const auto type = w >> 6;
        const auto n = w & 0b111111;
        switch (type) {
          case 0b00: { /* ZERO */
            std::size_t count = 0;
            if (n == 0) {
              count = outer_.counter_[counter_idx_++];
            }
            else {
              count = n;
            }
            auto l = count * word_bitlength;
            length_ = l;
            break; // Skip over zeros.
          }
          case 0b01: { /* LITERAL */
            literal_ = outer_.data_[data_idx_++];
            assert(literal_ != 0);
            literal_words_remaining_ = n - 1;
            assert(literal_words_remaining_ < n);
            const std::size_t b = dtl::bits::tz_count(literal_);
            std::size_t e = b + 1;
            for (; e < word_bitlength; ++e) {
              u1 is_set = dtl::bits::bit_test(literal_, e);
              if (!is_set) break;
            }
            literal_in_word_idx_ = e;
            pos_ += b;
            length_ = e - b;
            return;
          }
          case 0b10: { /* ENCODABLE (one byte) */
            literal_ = get_oep(n);
            assert(literal_ != 0);
            literal_words_remaining_ = 0;
            const std::size_t b = dtl::bits::tz_count(literal_);
            std::size_t e = b + 1;
            for (; e < word_bitlength; ++e) { // TODO optimize
              u1 is_set = dtl::bits::bit_test(literal_, e);
              if (!is_set) break;
            }
            literal_in_word_idx_ = e;
            pos_ += b;
            length_ = e - b;
            return;
          }
          case 0b11: { /* ENCODABLE (two bytes) */
            auto code =
                bah_word_t(n) | (bah_word_t(outer_.index_[index_idx_++]) << 6);
            literal_ = get_tep(code);
            assert(literal_ != 0);
            literal_words_remaining_ = 0;
            const std::size_t b = dtl::bits::tz_count(literal_);
            std::size_t e = b + 1;
            for (; e < word_bitlength; ++e) { // TODO optimize
              u1 is_set = dtl::bits::bit_test(literal_, e);
              if (!is_set) break;
            }
            literal_in_word_idx_ = e;
            pos_ += b;
            length_ = e - b;
            return;
          }
        }
      }
      pos_ = outer_.encoded_bitmap_length_;
      length_ = 0;
    }

    /// Forward the iterator to the desired position.
    void __forceinline__
    skip_to(const std::size_t to_pos) {
      assert(pos_ <= to_pos);
      if (to_pos >= outer_.encoded_bitmap_length_) {
        pos_ = outer_.encoded_bitmap_length_;
        length_ = 0;
        return;
      }
      // Call next until the desired position has been reached.
      while (!end() && pos() + length() <= to_pos) {
        next();
      }
      // Adjust the current position and run length.
      if (!end() && pos() < to_pos) {
        length_ -= to_pos - pos_;
        pos_ = to_pos;
      }
    }

    u1 __forceinline__
    end() const noexcept {
      return length_ == 0;
    }

    u64 __forceinline__
    pos() const noexcept {
      return pos_;
    }

    u64 __forceinline__
    length() const noexcept {
      return length_;
    }
  };
  //===--------------------------------------------------------------------===//

  using skip_iter_type = iter;
  using scan_iter_type = iter;

  /// Returns a 1-run iterator.
  skip_iter_type __forceinline__
  it() const {
    return skip_iter_type(*this);
  }

  /// Returns a 1-run iterator.
  scan_iter_type __forceinline__
  scan_it() const {
    return scan_iter_type(*this);
  }

  /// Returns the name of the instance including the most important parameters
  /// in JSON.
  std::string
  info() const {
    return "{\"name\":\"" + name() + "\""
        + ",\"n\":" + std::to_string(encoded_bitmap_length_)
        + ",\"size\":" + std::to_string(size_in_bytes())
        + ",\"word_size\":" + std::to_string(sizeof(bah_word_t))
        + "}";
  }

  // For debugging purposes.
  void
  print(std::ostream& os) const {
    os << "main idx | word type | content" << std::endl;
    os << "---------|-----------|---------------------------------" << std::endl;
    // The read position(s).
    std::size_t data_idx = 0;
    std::size_t index_idx = 0;
    std::size_t counter_idx = 0;
    for (std::size_t main_idx = 0; main_idx < main_.size(); ++main_idx) {
      os << std::setw(8) << main_idx << " | ";
      auto& w = main_[main_idx];
      const auto type = w >> 6;
      const auto n = w & 0b111111;
      switch (type) {
        case 0b00: { /* ZERO */
          std::size_t count = 0;
          if (n == 0) {
            count = counter_[counter_idx++];
          }
          else {
            count = n;
          }
          os << count << " x ZERO";
          break;
        }
        case 0b01: { /* LITERAL */
          os << n << " x LITERAL";
          for (std::size_t i = 0; i < n; ++i) {
            os << std::endl;
            os << "         |           | "
               << std::bitset<word_bitlength>(data_[data_idx + i]);
          }
          data_idx += n;
          break;
        }
        case 0b10: { /* ENCODABLE (one byte) */
          os << "1"
             << " x ENCODABLE (one byte)";
          os << std::endl;
          os << "         |           | "
              << std::bitset<word_bitlength>(get_oep(n));
          break;
        }
        case 0b11: { /* ENCODABLE (two bytes) */
          os << "1"
             << " x ENCODABLE (two bytes)";
          auto code = bah_word_t(n) | (bah_word_t(index_[index_idx++]) << 6);
          os << std::endl;
          os << "         |           | "
              << std::bitset<word_bitlength>(get_tep(code));
          break;
        }
      }
      os << std::endl;
    }
  }
};
//===----------------------------------------------------------------------===//
} // namespace dtl


/*********************************************************************
 * Software License Agreement (BSD License)
 *
 *  Copyright (c) 2020,
 *  Max Planck Institute for Intelligent Systems (MPI-IS).
 *  All rights reserved.
 *
 *  Redistribution and use in source and binary forms, with or without
 *  modification, are permitted provided that the following conditions
 *  are met:
 *
 *   * Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *   * Redistributions in binary form must reproduce the above
 *     copyright notice, this list of conditions and the following
 *     disclaimer in the documentation and/or other materials provided
 *     with the distribution.
 *   * Neither the name of the MPI-IS nor the names
 *     of its contributors may be used to endorse or promote products
 *     derived from this software without specific prior written
 *     permission.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 *  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 *  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 *  FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 *  COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 *  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 *  BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 *  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 *  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 *  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 *  ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 *  POSSIBILITY OF SUCH DAMAGE.
 *********************************************************************/

/* Author: Andreas Orthey, Sohaib Akbar */

#include <ompl/multilevel/planners/sparse/SMLRImpl.h>
#include <ompl/tools/config/SelfConfig.h>
#include <boost/foreach.hpp>
#include <ompl/datastructures/NearestNeighbors.h>
#include "ompl/datastructures/PDF.h"
#include <boost/math/constants/constants.hpp>
#include <ompl/multilevel/datastructures/pathrestriction/PathRestriction.h>

#define foreach BOOST_FOREACH

using namespace ompl::multilevel;

SMLRImpl::SMLRImpl(const ompl::base::SpaceInformationPtr &si, BundleSpace *parent_) : BaseT(si, parent_)
{
    setName("SMLRImpl" + std::to_string(id_));
    randomWorkStates_.resize(5);
    getBundle()->allocStates(randomWorkStates_);

    setMetric("geodesic");
    setGraphSampler("visibilityregion");

    firstRun_ = true;
    isInfeasible_ = false;
}

void SMLRImpl::clear()
{
    BaseT::clear();
    firstRun_ = true;
    isInfeasible_ = false;
}

SMLRImpl::~SMLRImpl()
{
    getBundle()->freeStates(randomWorkStates_);
}

double SMLRImpl::getImportance() const
{
    return 1.0 / (consecutiveFailures_ + 1);
}

void SMLRImpl::grow()
{
    if (firstRun_)
    {
        init();
        firstRun_ = false;

        getGoalPtr()->sampleGoal(qGoal_->state);
        addConfiguration(qGoal_);
        goalConfigurations_.push_back(qGoal_);

        findSection();
    }

    // if(pis_.getSampledGoalsCount() < getGoalPtr()->maxSampleCount())
    // {
    //     const base::State *state = pis_.nextGoal();
    //     Configuration qgoal = new Configuration(getBundle(), state);
    //     qgoal->isGoal = true;
    //     goalConfigurations_.push_back(qgoal);
    // }

    if (!sampleBundleValid(xRandom_->state))
    {
        return;
    }

    Configuration *xNew = new Configuration(getBundle(), xRandom_->state);

    addConfigurationConditional(xNew);

    if (!hasSolution_)
    {
        bool same_component = sameComponent(getStartIndex(), getGoalIndex());
        if (same_component)
        {
            hasSolution_ = true;
        }
    }
}

bool SMLRImpl::hasConverged()
{
    bool progressFailure = (consecutiveFailures_ >= maxFailures_);
    if (progressFailure)
    {
        OMPL_INFORM("Converged with probability %f (no valid samples for %d rounds).",
        (1.0 - 1.0 / (double)consecutiveFailures_),
        consecutiveFailures_);
    }
    return progressFailure;
}

bool SMLRImpl::isInfeasible()
{
    bool progressFailure = ((consecutiveFailures_ >= maxFailures_) && !hasSolution_);
    if (progressFailure)
    {
        OMPL_INFORM("Infeasibility detected with probability %f (no valid samples for %d rounds).",
                    (1.0 - 1.0 / (double)consecutiveFailures_), consecutiveFailures_);
        isInfeasible_ = true;
    }
    return progressFailure;
}


/*
 * KMajority.cpp
 *
 *  Created on: Aug 28, 2013
 *      Author: andresf
 */

#include <KMajority.h>
#include <CentersChooser.h>

#include <boost/iostreams/filter/gzip.hpp>
#include <boost/iostreams/filtering_stream.hpp>

#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/flann/random.h>
#include <opencv2/flann/dist.h>

#include <iostream>
#include <bitset>
#include <fstream>
#include <functional>

namespace vlr {

KMajority::KMajority(vlr::Mat& data, const cvflann::IndexParams& params,
		const cvflann::IndexParams& nnIndexParams) :
		m_dataset(data), m_dim(data.cols), m_nnIndex(NULL), m_nnIndexParams(
				nnIndexParams) {

	// Attributes initialization
	m_numClusters = cvflann::get_param<int>(params, "num.clusters");
	m_maxIterations = cvflann::get_param<int>(params, "max.iterations");
	m_centersInitMethod = cvflann::get_param<cvflann::flann_centers_init_t>(
			params, "centers.init.method");
	m_nnType = cvflann::get_param<vlr::indexType>(params, "nn.type");
	m_numDatapoints = m_dataset.rows;

	// Initially all transactions belong to any cluster
	m_belongsTo.clear();
	m_belongsTo.resize(m_dataset.rows, m_numClusters);

	// Initially all transactions are at the farthest possible distance
	// i.e. m_dim*8 the max Hamming distance
	m_distanceTo.clear();
	m_distanceTo.resize(m_dataset.rows, data.cols * 8);

	// Initially no transaction is assigned to any cluster
	m_clusterCounts.clear();
	m_clusterCounts.resize(m_numClusters, 0);

}

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

KMajority::~KMajority() {
	delete m_nnIndex;
}

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

void KMajority::build() {

	if (m_dataset.type() != CV_8U) {
		throw std::runtime_error(
				"[KMajority::build] Descriptors matrix is not binary");
	}

	if (m_dataset.empty()) {
		throw std::runtime_error("[KMajority::build] Descriptors is empty");
	}

	// Trivial case: less data than clusters, assign one data point per cluster
	if (m_numDatapoints <= m_numClusters) {
		m_centroids.create(m_numClusters, m_dim, m_dataset.type());
		for (int i = 0; i < m_numDatapoints; ++i) {
			m_dataset.row(i).copyTo(
					m_centroids(cv::Range(i, i + 1), cv::Range(0, m_dim)));
			m_belongsTo[i] = i;
		}
		return;
	}

#if KMAJVERBOSE
	printf("-- Bootstrapping clustering process\n");
#endif

	// Randomly generate clusters
#if KMAJVERBOSE
	printf("   Initializing clusters centers.\n");
#endif
	initCentroids();

	// Update nearest neighbors index upon new centers
#if KMAJVERBOSE
	printf("   Updating nearest neighbors index.\n");
#endif
	updateIndex();

	// Assign data to clusters
#if KMAJVERBOSE
	printf("   Quantizing data into clusters.\n");
#endif
	quantize();

	bool converged = false;
	int iteration = 0;

	while (converged == false && iteration < m_maxIterations) {

		++iteration;

#if KMAJVERBOSE
		printf("-- Iteration=[%d]\n", iteration);
		fflush(stdout);
#endif

		// Compute the new clusters centers
#if KMAJVERBOSE
		printf("   Computing new clusters centers.\n");
#endif
		computeCentroids();

		// Update nearest neighbors index upon new centers
#if KMAJVERBOSE
		printf("   Updating nearest neighbors index.\n");
#endif
		updateIndex();

		// Reassign data to clusters
#if KMAJVERBOSE
		printf("   Quantizing data into clusters.\n");
#endif
		converged = quantize();

		// Handle empty clusters case
#if KMAJVERBOSE
		printf("   Handling empty clusters case.\n");
#endif
		handleEmptyClusters();
	}

}

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

void KMajority::initCentroids() {

	// Initializing variables useful for obtaining indexes of random chosen center
	std::vector<int> centers_idx(m_numClusters);
	int centers_length;

	// Array of indices indicating data points involved in the clustering process
	int* indices = new int[m_dataset.rows];
	for (int i = 0; i < m_dataset.rows; ++i) {
		indices[i] = i;
	}

	// Randomly chose centers
	CentersChooser<Distance::ElementType, cv::Hamming>::create(
			m_centersInitMethod)->chooseCenters(m_numClusters, indices,
			m_numDatapoints, centers_idx, centers_length, m_dataset);
	CV_Assert(centers_length == m_numClusters);

	std::sort(centers_idx.begin(), centers_idx.end());

	delete[] indices;

	// Assign centers based on the chosen indexes
	m_centroids.create(centers_length, m_dim, m_dataset.type());
	for (int i = 0; i < centers_length; ++i) {
		m_dataset.row(centers_idx[i]).copyTo(
				m_centroids(cv::Range(i, i + 1), cv::Range(0, m_dim)));
	}

}

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

bool KMajority::quantize() {

	bool converged = true;

	// Number of nearest neighbors
	int knn = 1;

	// The indices of the nearest neighbors found (numQueries X numNeighbors)
	cvflann::Matrix<int> indices(new int[1 * knn], 1, knn);

	// Distances to the nearest neighbors found (numQueries X numNeighbors)
	cvflann::Matrix<DistanceType> distances(new DistanceType[1 * knn], 1, knn);

	for (int i = 0; i < m_numDatapoints; ++i) {
		std::fill(indices.data, indices.data + indices.rows * indices.cols, 0);
		std::fill(distances.data,
				distances.data + distances.rows * distances.cols, 0.0f);

		cvflann::Matrix<Distance::ElementType> descriptor(
				(Distance::ElementType*) m_dataset.row(i).data, 1,
				m_dataset.cols);

		/* Get new cluster it belongs to */
		m_nnIndex->knnSearch(descriptor, indices, distances, knn,
				cvflann::SearchParams());

		/* Check if cluster assignment changed */
		// If it did then algorithm hasn't converged yet
		if (m_belongsTo[i] != indices[0][0]) {
			converged = false;
		}

		/* Update cluster assignment and cluster counts */
		// Decrease cluster count in case it was assigned to some valid cluster before.
		// Recall that initially all transaction are assigned to kth cluster which
		// is not valid since valid clusters run from 0 to k-1 both inclusive.
		if (m_belongsTo[i] != m_numClusters) {
			--m_clusterCounts[m_belongsTo[i]];
		}
		m_belongsTo[i] = indices[0][0];
		++m_clusterCounts[indices[0][0]];
		m_distanceTo[i] = distances[0][0];
	}

	delete[] indices.data;
	delete[] distances.data;

	return converged;
}

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

void KMajority::computeCentroids() {

	// Warning: using matrix of integers, there might be an overflow when summing too much descriptors
	cv::Mat bitwiseCount(m_numClusters, m_dim * 8, cv::DataType<int>::type);
	// Zeroing matrix of cumulative bits
	bitwiseCount = cv::Scalar::all(0);
	// Zeroing all cluster centers dimensions
	m_centroids = cv::Scalar::all(0);

	// Bitwise summing the data into each center
	for (int i = 0; i < m_numDatapoints; ++i) {
		cv::Mat b = bitwiseCount.row(m_belongsTo[i]);
		KMajority::cumBitSum(m_dataset.row(i), b);
	}

	// Bitwise majority voting
	for (int j = 0; j < m_numClusters; j++) {
		cv::Mat centroid = m_centroids.row(j);
		KMajority::majorityVoting(bitwiseCount.row(j), centroid,
				m_clusterCounts[j]);
	}
}

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

void KMajority::save(const std::string& filename) const {

	if (m_centroids.empty()) {
		throw std::runtime_error("[KMajority::save] Tree is empty");
	}

	cv::FileStorage fs(filename.c_str(), cv::FileStorage::WRITE);

	if (fs.isOpened() == false) {
		throw std::runtime_error("[KMajority::save] "
				"Unable to open file [" + filename + "] for writing");
	}

	fs << "type" << "AKMAJ";
	fs << "Centers" << m_centroids;

	fs.release();

}

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

void KMajority::load(const std::string& filename) {

	enum nodeFields {
		header, start, rows, cols, dt, data
	};
	std::string nodeFieldsNames[] = { "YAML", "Centers", "rows:", "cols:",
			"dt:", "data:" };

	std::ifstream inputZippedFileStream;
	boost::iostreams::filtering_istream inputFileStream;

	std::string line, field;
	std::stringstream ss;

	// Open file
	inputZippedFileStream.open(filename.c_str(),
			std::fstream::in | std::fstream::binary);

	// Check file
	if (inputZippedFileStream.good() == false) {
		throw std::runtime_error("[KMajority::load] "
				"Unable to open file [" + filename + "] for reading");
	}

	int _rows = -1;
	int _cols = -1;
	std::string _type;
	int elemIdx = -1;
	unsigned int elem;

	try {
		inputFileStream.push(boost::iostreams::gzip_decompressor());
		inputFileStream.push(inputZippedFileStream);

		while (getline(inputFileStream, line)) {
			ss.clear();
			ss.str(line);
			ss >> field;
			if (field.compare(nodeFieldsNames[header]) == 0) {
				continue;
			} else if (field.compare(nodeFieldsNames[start]) == 0) {
				continue;
			} else if (field.compare(nodeFieldsNames[rows]) == 0) {
				ss >> _rows;
			} else if (field.compare(nodeFieldsNames[cols]) == 0) {
				ss >> _cols;
			} else if (field.compare(nodeFieldsNames[dt]) == 0) {
				ss >> _type;
			} else {
				if (field.compare(nodeFieldsNames[data]) == 0) {
					m_centroids = cv::Mat::zeros(_rows, _cols,
							_type.compare("f") == 0 ? CV_32F : CV_8U);
					line.replace(line.find(nodeFieldsNames[data]), 5, " ");
				}

				std::replace(line.begin(), line.end(), '[', ' ');
				std::replace(line.begin(), line.end(), ',', ' ');
				std::replace(line.begin(), line.end(), ']', ' ');

				ss.clear();
				ss.str(line);

				while ((ss >> elem).fail() == false) {
					++elemIdx;
					int row = floor(elemIdx / _cols);
					int col = elemIdx % _cols;
					m_centroids.at<uchar>(row, col) = elem;
				}

			}
		}

	} catch (const boost::iostreams::gzip_error& e) {
		throw std::runtime_error("[KMajority::load] "
				"Got error while parsing file [" + std::string(e.what()) + "]");
	}

	// Close file
	inputZippedFileStream.close();

}

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

void KMajority::cumBitSum(const cv::Mat& data, cv::Mat& accVector) {

	// cumResult and data must be row vectors
	if (data.rows != 1 || accVector.rows != 1) {
		throw std::runtime_error(
				"[KMajority::cumBitSum] data and cumResult parameters must be row vectors\n");
	}
	// cumResult and data must be same length
	if (data.cols * 8 != accVector.cols) {
		throw std::runtime_error(
				"[KMajority::cumBitSum] number of columns in cumResult must be that of data times 8\n");
	}

	uchar byte = 0;
	for (int l = 0; l < accVector.cols; l++) {
		// bit: 7-(l%8) col: (int)l/8 descriptor: i
		// Load byte every 8 bits
		if ((l % 8) == 0) {
			byte = *(data.col((int) l / 8).data);
		}
		// Note: ignore maybe-uninitialized warning because loop starts with l=0 that means byte gets a value as soon as the loop start
		// bit at ith position is mod(bitleftshift(byte,i),2) where ith position is 7-mod(l,8) i.e 7, 6, 5, 4, 3, 2, 1, 0
		accVector.at<int>(0, l) += ((int) ((byte >> (7 - (l % 8))) % 2));
	}

}

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

void KMajority::majorityVoting(const cv::Mat& accVector, cv::Mat& result,
		const int& threshold) {

	// cumResult and data must be a row vectors
	if (accVector.rows != 1 || result.rows != 1) {
		throw std::runtime_error(
				"[KMajority::majorityVoting] 'accVector' and 'result' parameters must be row vectors\n");
	}

	// cumResult and data must be same length
	if (result.cols * 8 != accVector.cols) {
		throw std::runtime_error(
				"[KMajority::majorityVoting] number of columns in 'accVector' must be that of 'result' times 8\n");
	}

	// In this point I already have stored in bitwiseCount the bitwise sum of all data assigned to jth cluster
	for (int l = 0; l < accVector.cols; ++l) {
		// If the bitcount for jth cluster at dimension l is greater than half of the data assigned to it
		// then set lth centroid bit to 1 otherwise set it to 0 (break ties randomly)
		bool bit;
		// There is a tie if the number of data assigned to jth cluster is even
		// AND the number of bits set to 1 in lth dimension is the half of the data assigned to jth cluster
		if (threshold % 2 == 1
				&& 2 * accVector.at<int>(0, l) == (int) threshold) {
			bit = rand() % 2;
		} else {
			bit = 2 * accVector.at<int>(0, l) > (int) (threshold);
		}
		// Stores the majority voting result from the LSB to the MSB
		result.at<unsigned char>(0, (int) (accVector.cols - 1 - l) / 8) += (bit)
				<< ((accVector.cols - 1 - l) % 8);
	}
}

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

void KMajority::handleEmptyClusters() {

	// If some cluster appeared to be empty then:
	// 1. Find the biggest cluster.
	// 2. Find farthest point in the biggest cluster
	// 3. Exclude the farthest point from the biggest cluster and form a new 1-point cluster.

	for (int k = 0; k < m_numClusters; ++k) {
		if (m_clusterCounts[k] != 0) {
			continue;
		}

		// 1. Find the biggest cluster
		int max_k = 0;
		for (int k1 = 1; k1 < m_numClusters; ++k1) {
			if (m_clusterCounts[max_k] < m_clusterCounts[k1])
				max_k = k1;
		}

		// 2. Find farthest point in the biggest cluster
		DistanceType maxDist(-1);
		int idxFarthestPt = -1;
		for (int i = 0; i < m_numDatapoints; ++i) {
			if (m_belongsTo[i] == max_k) {
				if (maxDist < m_distanceTo[i]) {
					maxDist = m_distanceTo[i];
					idxFarthestPt = i;
				}
			}
		}

		// 3. Exclude the farthest point from the biggest cluster and form a new 1-point cluster
		--m_clusterCounts[max_k];
		++m_clusterCounts[k];
		m_belongsTo[idxFarthestPt] = k;
	}
}

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

void KMajority::updateIndex() {

	m_nnIndex = vlr::createIndexByType(
			cvflann::Matrix<Distance::ElementType>(
					(Distance::ElementType*) m_centroids.data, m_centroids.rows,
					m_centroids.cols), m_nnType, m_nnIndexParams);

	double mytime = cv::getTickCount();
        m_nnIndex->buildIndex();
	mytime = ((double) cv::getTickCount() - mytime) / cv::getTickFrequency() * 1000;

	printf("   Index built in [%lf] ms\n", mytime);

}

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

const cv::Mat& KMajority::getCentroids() const {
	return m_centroids;
}

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

const std::vector<int>& KMajority::getClusterCounts() const {
	return m_clusterCounts;
}

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

const std::vector<int>& KMajority::getClusterAssignments() const {
	return m_belongsTo;
}

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

cvflann::NNIndex<Distance>* createIndexByType(
		const cvflann::Matrix<typename Distance::ElementType>& dataset,
		vlr::indexType type, const cvflann::IndexParams& userDefParams) {

	cvflann::IndexParams::const_iterator it;

	cvflann::IndexParams params;
	cvflann::NNIndex<Distance>* nnIndex;

	double mytime = cv::getTickCount();

	switch (type) {
	case vlr::LINEAR:
		printf("-- Creating [Linear] index\n");
		params = cvflann::LinearIndexParams();
		// Do not copy any parameters, linear index doesn't need any
		nnIndex = new cvflann::LinearIndex<Distance>(dataset, params,
				Distance());
		break;
	case vlr::HIERARCHICAL:
		printf("-- Creating [HierarchicalClustering] index\n");
		params = cvflann::HierarchicalClusteringIndexParams();
		for (it = userDefParams.begin(); it != userDefParams.end(); ++it) {
			int value = it->second.cast<int>();
			params[it->first] = value;
		}
		nnIndex = new cvflann::HierarchicalClusteringIndex<Distance>(dataset,
				params, Distance());
		break;
	default:
		throw std::runtime_error("Unknown index type");
	}
	mytime = ((double) cv::getTickCount() - mytime) / cv::getTickFrequency()
			* 1000;

	printf("   Index created in [%lf] ms\n", mytime);

	return nnIndex;
}

} /* namespace vlr */