#include "distance_measures.h"
#include "ProteinMatrix.h"
#include <cmath>
#include <string>

// #include <xmmintrin.h>
// #include <emmintrin.h>
// #include <pmmintrin.h>

#include <boost/assign/std/vector.hpp>
#include <boost/assign/list_of.hpp> // for 'list_of()'

using namespace boost::assign; // bring 'list_of()' into scope
//using namespace boost::assign; // bring 'operator+()' into scope

using namespace std;
using namespace uqam_doc;

//static function pointer calculators
//using uqam_doc::DistSeqFuncStatic;

//abstract class - interface
//dist_2seq(string seq_a, string seq_b)
using uqam_doc::DistSeqInterface;
// virtual functions implementations
using uqam_doc::Hamming;
using uqam_doc::Hamming2;

using uqam_doc::JukesKantorNucl;
using uqam_doc::JukesKantorProt;
using uqam_doc::KimProt;
using uqam_doc::ScoreDist;
using uqam_doc::Blosum80;
using uqam_doc::Blosum62;

/*
float DistSeqFuncStatic::hamming_2seq(string seq_a, string seq_b) {

    float d = 0.0;

    for (int i = 0; i < seq_a.length(); i++) {
        d += (seq_a.at(i) == seq_b.at(i)) ? 0.0 : 1.0;

    }
    d /= seq_a.length();

    return d;

}
 */

/*************************************************************************************
/*	Nom			:	p_Kimura2Parameter
/*	Parametres	:	pmatSites       (in) : Matrice de sites
/*					pmatDistances   (in) : Matrice de distances
/*					pintNbreEspeces (in) : Nombre d'esp�ces
/*					pintNbreSites   (in) : Nombre de sites
/*	Retour		:	-
/*	Objectif	:	Comparer les s�quences deux par deux, en utilisant la m�thode
/*					Kimura 2-Parameter, et stoker les distances entre elles dans la
/*					matrice des distances.
 **************************************************************************************/

/*
float DistSeqFuncStatic::kimprot_2seq(string seq_a, string seq_b) {

    //cout << "------------------k2p_2seq" << endl;
    int nb_sites = seq_a.length();
    float dist;

    int k; // Compteur
    float dblCorrespandance; // Nombre de correspondances entre les sites
    int intPosComparable; // Nombre de positions comparables, l� o� il n'y a pas de gaps
    int intGaps; // Les lacunes contenues dans les s�quences
    float dblD; // Uncorrected distance

    dblCorrespandance = 0;
    intPosComparable = 0;
    intGaps = 0;

    for (k = 1; k <= nb_sites; k++) {
        if ((seq_a[k] == seq_b[k]) && (seq_a[k] != '-') && (seq_a[k] != '?'))
            dblCorrespandance++;
        else if ((seq_a[k] == '-') || (seq_a[k] == '?') ||
                (seq_b[k] == '-') || (seq_b[k] == '?'))
            intGaps++;
    }
    intPosComparable = nb_sites - intGaps;

    if (intPosComparable != 0) {
        dblD = 1 - dblCorrespandance / intPosComparable;
        dblD = 1 - dblD - 0.2 * dblD * dblD;
        if (dblD <= 0) {
            //non comparable
            dist = INFINITY;
        } else {
            dist = -log(dblD);
        }
    } else {
        //non comparable
        dist = INFINITY;
    }

    return dist;

}
 */

DistAtom Hamming::dist_2seq(string seq_a, string seq_b) {

    float intDist = 0.0;
    float realDist = 0.0;

    for (int i = 0; i < seq_a.length(); i++) {

        intDist += (seq_a[i] == seq_b[i] ? 0 : 1);

    }

    realDist = intDist / seq_a.length();

    return DistAtom(intDist, seq_a.length(), 0, realDist);
}

void Hamming::testIntrin() {
    /**
    int i;

    __m128d a, b, c;
    float x0[2] __attribute__((aligned(16))) = {1.2, 3.5};
    float x1[2] __attribute__((aligned(16))) = {0.7, 2.6};
    a = _mm_load_pd(x0);
    b = _mm_load_pd(x1);
    c = _mm_add_pd(a, b);
    _mm_store_pd(x0, c);
    for (i = 0; i < 2; i++) {
       cout << x0[i] << endl;
    }
     **/

}

DistAtom Hamming2::dist_2seq(string seq_a, string seq_b) {

    float pos_comparable = 0;
    float prop_gaps = 0;

    float two_gap = 0;
    float one_gap = 0;
    float ident = 0;
    float subst = 0;

    float intDist = 0.0;
    float realDist = 0.0;

    for (int i = 0; i < seq_a.length(); i++) {

        //if no gaps
        if (seq_a[i] != '-' && seq_b[i] != '-') {

            if (seq_a[i] == seq_b[i]) { // identities
                ident++;

            } else { //substitutions
                subst++;

            }

        } else if (seq_a[i] == '-' && seq_b[i] == '-') {
            two_gap++;

        } else {
            one_gap++;
        }

        if (seq_a[i] != seq_b[i]) {
            intDist++;
        }

        /*
        //identities
        if ((seq_a[i] == seq_b[i]) && (seq_a[i] != '-') && (seq_a[i] != '?')) {
            intCorrespandance++;
        //gaps
        } else if ((seq_a[i] == '-') || (seq_a[i] == '?') ||
                (seq_b[i] == '-') || (seq_b[i] == '?')) {
            intGaps++;
        //substitutions
        } else {
            intDist++;
        }
         */

    }

    intDist = subst + one_gap;
    pos_comparable = ident + subst + one_gap;

    prop_gaps = (one_gap + two_gap) / float(seq_a.length());

    if (prop_gaps <= 0.5) {
        realDist = intDist /= pos_comparable;
        //realDist = intDist /= seq_a.length();
    } else {
        intDist = numeric_limits<float>::max();
        realDist = numeric_limits<float>::max();
    }

    /*

   intPosComparable = seq_a.length() - intGaps;
   //si trop de gaps, incomparables
   if (intGaps >= seq_a.length() * 0.9) {
       intDist = numeric_limits<float>::max();
       realDist = numeric_limits<float>::max();
   } else {
       realDist = intDist /= seq_a.length();
   }
     */
    return DistAtom(intDist, seq_a.length(), 0, realDist);
    //debug
    //return DistAtom(prop_gaps, seq_a.length(), 0, prop_gaps);
}

void JukesKantorNucl::set_penality(float penality) {
    penalty_ = penality;
}

JukesKantorNucl::JukesKantorNucl(float penality) {
    cout << "================JUKES CANTOR NUCL" << endl;
    penalty_ = penality;
}

JukesKantorNucl::~JukesKantorNucl() {

}

/*************************************************************************************
/*	Nom			:	p_JukesCantorNucleo
/*	Param�tres	:	pmatSites       (in) : Matrice de sites
/*					pmatDistances   (in) : Matrice de distances
/*					pintNbreEspeces (in) : Nombre d'esp�ces
/*					pintNbreSites   (in) : Nombre de sites
/*					pdblPenalite    (in) : P�nalit� donn�e par l'utilisateur
/*	Retour		:	-
/*	Objectif	:	Comparer les s�quences deux par deux, en utilisant la m�thode
/*					Jukes-Cantor pour nucl�otides, et stoker les distances entre elles
/*					dans la matrice des distances.
 **************************************************************************************/
DistAtom JukesKantorNucl::dist_2seq(string seq_a, string seq_b) {

    int nb_sites = seq_a.length();
    float dist;
    int k; // Compteurs
    int intCorrespandance; // Nombre de correspondances entre les sites
    int intPosComparable; // Nombre de positions comparables, la o� il n'y a pas de gaps
    int intGaps; // Les lacunes contenues dans les sequences
    float dblD; // Uncorrected distance
    float dblb; // Param�tre b

    dblb = 3. / 4;

    intCorrespandance = 0;
    intGaps = 0;

    for (k = 1; k <= nb_sites; k++) {
        if ((seq_a[k] == seq_b[k]) && (seq_a[k] != '-') && (seq_a[k] != '?'))
            intCorrespandance++;
        else if ((seq_a[k] == '-') || (seq_a[k] == '?') ||
                (seq_b[k] == '-') || (seq_b[k] == '?'))
            intGaps++;
    }
    intPosComparable = nb_sites - intGaps;

    if (intPosComparable != 0) {
        dblD = 1 - intCorrespandance / (intPosComparable + penalty_ * intGaps);
        dblD = dblD / dblb;
        if (1 - dblD <= 0) dist = numeric_limits<float>::max();
        else dist = -dblb * log(1 - dblD);
    } else {
        dist = numeric_limits<float>::max();
    }
    return DistAtom(numeric_limits<float>::max(), intPosComparable, intGaps, dist);

}

void JukesKantorProt::set_penality(float penality) {
    penality_ = penality;
}

JukesKantorProt::JukesKantorProt(float penality) {
    penality_ = penality;
}

JukesKantorProt::~JukesKantorProt() {

}

/*************************************************************************************
/*	Nom			:	p_JukesCantorProt
/*	Param�tres	:	pmatSites       (in) : Matrice de sites
/*					pmatDistances   (in) : Matrice de distances
/*					pintNbreEspeces (in) : Nombre d'esp�ces
/*					pintNbreSites   (in) : Nombre de sites
/*					pdblPenalite    (in) : P�nalit� donn�e par l'utilisateur
/*	Retour		:	-
/*	Objectif	:	Comparer les s�quences deux par deux, en utilisant la m�thode
/*					Jukes-Cantor pour prot�ines, et stoker les distances entre elles
/*					dans la matrice des distances.
 **************************************************************************************/
DistAtom JukesKantorProt::dist_2seq(string seq_a, string seq_b) {

    int nb_sites = seq_a.length();
    float dist;

    int k; // Compteurs
    int intCorrespandance; // Nombre de correspondances entre les sites
    int intPosComparable; // Nombre de positions comparables, l� o� il n'y a pas de gaps
    int intGaps; // Les lacunes contenues dans les s�quences
    float dblD; // Uncorrected distance
    float dblb; // Param�tre b

    dblb = 19. / 20;

    intCorrespandance = 0;
    intGaps = 0;

    for (k = 1; k <= nb_sites; k++) {
        if ((seq_a[k] == seq_b[k]) && (seq_a[k] != '-') && (seq_a[k] != '?'))
            intCorrespandance++;
        else if ((seq_a[k] == '-') || (seq_a[k] == '?') ||
                (seq_b[k] == '-') || (seq_b[k] == '?'))
            intGaps++;
    }
    intPosComparable = nb_sites - intGaps;

    if (intPosComparable != 0) {
        dblD = 1.0 - intCorrespandance / (intPosComparable + penality_ * intGaps);
        dblD = dblD / dblb;
        if (1 - dblD <= 0)
            dist = numeric_limits<float>::max();
        else
            dist = -dblb * log(1 - dblD);
    } else
        dist = numeric_limits<float>::max();

    return DistAtom(numeric_limits<float>::max(), intPosComparable, intGaps, dist);
    //if(pmatDistances[j][i] < 0 || pmatDistances[j][i] > 10)
    //			printf("\n%d %d %lf",i,j,dist);

}

/*************************************************************************************
/*	Nom			:	p_KimuraProtein
/*	Param�tres	:	pmatSites       (in) : Matrice de sites
/*					pmatDistances   (in) : Matrice de distances
/*					pintNbreEspeces (in) : Nombre d'esp�ces
/*					pintNbreSites   (in) : Nombre de sites
/*	Retour		:	-
/*	Objectif	:	Comparer les s�quences deux par deux, en utilisant la m�thode
/*					Kimura Protein, et stoker les distances entre elles dans la matrice
/*					des distances.
 **************************************************************************************/
DistAtom KimProt::dist_2seq(string seq_a, string seq_b) {
    //cout << "------------------k2p_2seq" << endl;
    int nb_sites = seq_a.length();
    float dist;

    int k; // Compteur
    float dblCorrespandance; // Nombre de correspondances entre les sites
    int intPosComparable; // Nombre de positions comparables, l� o� il n'y a pas de gaps
    int intGaps; // Les lacunes contenues dans les s�quences
    float dblD; // Uncorrected distance

    dblCorrespandance = 0;
    intPosComparable = 0;
    intGaps = 0;

    for (k = 1; k <= nb_sites; k++) {
        if ((seq_a[k] == seq_b[k]) && (seq_a[k] != '-') && (seq_a[k] != '?'))
            dblCorrespandance++;
        else if ((seq_a[k] == '-') || (seq_a[k] == '?') ||
                (seq_b[k] == '-') || (seq_b[k] == '?'))
            intGaps++;
    }
    intPosComparable = nb_sites - intGaps;

    if (intPosComparable != 0) {
        dblD = 1 - dblCorrespandance / intPosComparable;
        dblD = 1 - dblD - 0.2 * dblD * dblD;
        if (dblD <= 0) {
            //non comparable
            dist = numeric_limits<float>::max();
        } else {
            dist = -log(dblD);
        }
    } else {
        //non comparable
        dist = numeric_limits<float>::max();
    }

    return DistAtom(numeric_limits<float>::max(), intPosComparable, intGaps, dist);

}

//load inherited fields

Blosum80::Blosum80() {
    //ftp://ftp.ncbi.nih.gov/blast/matrices/BLOSUM80
    //matrices
    //mat_prot_ident_ += "A", "R", "N", "D", "C", "Q", "E", "G", "H", "I", "L", "K", "M", "F", "P", "S", "T", "W", "Y", "V", "B", "Z", "X", "_";
    mat_prot_ident_ = list_of("A") ("R") ("N") ("D") ("C") ("Q") ("E") ("G") ("H") ("I") ("L") ("K") ("M") ("F") ("P") ("S") ("T") ("W") ("Y") ("V") ("B") ("Z") ("X") ("*");

    mat_prot_ += list_of(7)(-3)(-3)(-3)(-1)(-2)(-2)(0)(-3)(-3)(-3)(-1)(-2)(-4)(-1)(2)(0)(-5)(-4)(-1)(-3)(-2)(-1)(-8);
    mat_prot_ += list_of(-3)(9)(-1)(-3)(-6)(1)(-1)(-4)(0)(-5)(-4)(3)(-3)(-5)(-3)(-2)(-2)(-5)(-4)(-4)(-2)(0)(-2)(-8);
    mat_prot_ += list_of(-3)(-1)(9)(2)(-5)(0)(-1)(-1)(1)(-6)(-6)(0)(-4)(-6)(-4)(1)(0)(-7)(-4)(-5)(5)(-1)(-2)(-8);
    mat_prot_ += list_of(-3)(-3)(2)(10)(-7)(-1)(2)(-3)(-2)(-7)(-7)(-2)(-6)(-6)(-3)(-1)(-2)(-8)(-6)(-6)(6)(1)(-3)(-8);
    mat_prot_ += list_of(-1)(-6)(-5)(-7)(13)(-5)(-7)(-6)(-7)(-2)(-3)(-6)(-3)(-4)(-6)(-2)(-2)(-5)(-5)(-2)(-6)(-7)(-4)(-8);
    mat_prot_ += list_of(-2)(1)(0)(-1)(-5)(9)(3)(-4)(1)(-5)(-4)(2)(-1)(-5)(-3)(-1)(-1)(-4)(-3)(-4)(-1)(5)(-2)(-8);
    mat_prot_ += list_of(-2)(-1)(-1)(2)(-7)(3)(8)(-4)(0)(-6)(-6)(1)(-4)(-6)(-2)(-1)(-2)(-6)(-5)(-4)(1)(6)(-2)(-8);
    mat_prot_ += list_of(0)(-4)(-1)(-3)(-6)(-4)(-4)(9)(-4)(-7)(-7)(-3)(-5)(-6)(-5)(-1)(-3)(-6)(-6)(-6)(-2)(-4)(-3)(-8);
    mat_prot_ += list_of(-3)(0)(1)(-2)(-7)(1)(0)(-4)(12)(-6)(-5)(-1)(-4)(-2)(-4)(-2)(-3)(-4)(3)(-5)(-1)(0)(-2)(-8);
    mat_prot_ += list_of(-3)(-5)(-6)(-7)(-2)(-5)(-6)(-7)(-6)(7)(2)(-5)(2)(-1)(-5)(-4)(-2)(-5)(-3)(4)(-6)(-6)(-2)(-8);
    mat_prot_ += list_of(-3)(-4)(-6)(-7)(-3)(-4)(-6)(-7)(-5)(2)(6)(-4)(3)(0)(-5)(-4)(-3)(-4)(-2)(1)(-7)(-5)(-2)(-8);
    mat_prot_ += list_of(-1)(3)(0)(-2)(-6)(2)(1)(-3)(-1)(-5)(-4)(8)(-3)(-5)(-2)(-1)(-1)(-6)(-4)(-4)(-1)(1)(-2)(-8);
    mat_prot_ += list_of(-2)(-3)(-4)(-6)(-3)(-1)(-4)(-5)(-4)(2)(3)(-3)(9)(0)(-4)(-3)(-1)(-3)(-3)(1)(-5)(-3)(-2)(-8);
    mat_prot_ += list_of(-4)(-5)(-6)(-6)(-4)(-5)(-6)(-6)(-2)(-1)(0)(-5)(0)(10)(-6)(-4)(-4)(0)(4)(-2)(-6)(-6)(-3)(-8);
    mat_prot_ += list_of(-1)(-3)(-4)(-3)(-6)(-3)(-2)(-5)(-4)(-5)(-5)(-2)(-4)(-6)(12)(-2)(-3)(-7)(-6)(-4)(-4)(-2)(-3)(-8);
    mat_prot_ += list_of(2)(-2)(1)(-1)(-2)(-1)(-1)(-1)(-2)(-4)(-4)(-1)(-3)(-4)(-2)(7)(2)(-6)(-3)(-3)(0)(-1)(-1)(-8);
    mat_prot_ += list_of(0)(-2)(0)(-2)(-2)(-1)(-2)(-3)(-3)(-2)(-3)(-1)(-1)(-4)(-3)(2)(8)(-5)(-3)(0)(-1)(-2)(-1)(-8);
    mat_prot_ += list_of(-5)(-5)(-7)(-8)(-5)(-4)(-6)(-6)(-4)(-5)(-4)(-6)(-3)(0)(-7)(-6)(-5)(16)(3)(-5)(-8)(-5)(-5)(-8);
    mat_prot_ += list_of(-4)(-4)(-4)(-6)(-5)(-3)(-5)(-6)(3)(-3)(-2)(-4)(-3)(4)(-6)(-3)(-3)(3)(11)(-3)(-5)(-4)(-3)(-8);
    mat_prot_ += list_of(-1)(-4)(-5)(-6)(-2)(-4)(-4)(-6)(-5)(4)(1)(-4)(1)(-2)(-4)(-3)(0)(-5)(-3)(7)(-6)(-4)(-2)(-8);
    mat_prot_ += list_of(-3)(-2)(5)(6)(-6)(-1)(1)(-2)(-1)(-6)(-7)(-1)(-5)(-6)(-4)(0)(-1)(-8)(-5)(-6)(6)(0)(-3)(-8);
    mat_prot_ += list_of(-2)(0)(-1)(1)(-7)(5)(6)(-4)(0)(-6)(-5)(1)(-3)(-6)(-2)(-1)(-2)(-5)(-4)(-4)(0)(6)(-1)(-8);
    mat_prot_ += list_of(-1)(-2)(-2)(-3)(-4)(-2)(-2)(-3)(-2)(-2)(-2)(-2)(-2)(-3)(-3)(-1)(-1)(-5)(-3)(-2)(-3)(-1)(-2)(-8);
    mat_prot_ += list_of(-8)(-8)(-8)(-8)(-8)(-8)(-8)(-8)(-8)(-8)(-8)(-8)(-8)(-8)(-8)(-8)(-8)(-8)(-8)(-8)(-8)(-8)(-8)(1);

    //cout << mat_prot_ident_.at(22) << endl;
    //cout << mat_prot_ident_.at(17) << endl;

    //ftp://ftp.ncbi.nih.gov/blast/matrices/BLOSUM80

    //cout << mat_prot_.at(22).at(17) << endl;

    mat_prot_expect_ = -0.7442;

}

//load inherited fields

Blosum62::Blosum62() {
    //ftp://ftp.ncbi.nih.gov/blast/matrices/BLOSUM80
    //matrices
    //mat_prot_ident_ += "A", "R", "N", "D", "C", "Q", "E", "G", "H", "I", "L", "K", "M", "F", "P", "S", "T", "W", "Y", "V", "B", "Z", "X", "_";
    mat_prot_ident_ = list_of("A") ("R") ("N") ("D") ("C") ("Q") ("E") ("G") ("H") ("I") ("L") ("K") ("M") ("F") ("P") ("S") ("T") ("W") ("Y") ("V") ("B") ("Z") ("X") ("*");

    mat_prot_ += list_of (4)(-1)(-2)(-2)(0)(-1)(-1)(0)(-2)(-1)(-1)(-1)(-1)(-2)(-1)(1)(0)(-3)(-2)(0)(-2)(-1)(0)(-4);
    mat_prot_ += list_of (-1)(5)(0)(-2)(-3)(1)(0)(-2)(0)(-3)(-2)(2)(-1)(-3)(-2)(-1)(-1)(-3)(-2)(-3)(-1)(0)(-1)(-4);
    mat_prot_ += list_of (-2)(0)(6)(1)(-3)(0)(0)(0)(1)(-3)(-3)(0)(-2)(-3)(-2)(1)(0)(-4)(-2)(-3)(3)(0)(-1)(-4);
    mat_prot_ += list_of (-2)(-2)(1)(6)(-3)(0)(2)(-1)(-1)(-3)(-4)(-1)(-3)(-3)(-1)(0)(-1)(-4)(-3)(-3)(4)(1)(-1)(-4);
    mat_prot_ += list_of (0)(-3)(-3)(-3)(9)(-3)(-4)(-3)(-3)(-1)(-1)(-3)(-1)(-2)(-3)(-1)(-1)(-2)(-2)(-1)(-3)(-3)(-2)(-4);
    mat_prot_ += list_of (-1)(1)(0)(0)(-3)(5)(2)(-2)(0)(-3)(-2)(1)(0)(-3)(-1)(0)(-1)(-2)(-1)(-2)(0)(3)(-1)(-4);
    mat_prot_ += list_of (-1)(0)(0)(2)(-4)(2)(5)(-2)(0)(-3)(-3)(1)(-2)(-3)(-1)(0)(-1)(-3)(-2)(-2)(1)(4)(-1)(-4);
    mat_prot_ += list_of (0)(-2)(0)(-1)(-3)(-2)(-2)(6)(-2)(-4)(-4)(-2)(-3)(-3)(-2)(0)(-2)(-2)(-3)(-3)(-1)(-2)(-1)(-4);
    mat_prot_  += list_of (-2)(0)(1)(-1)(-3)(0)(0)(-2)(8)(-3)(-3)(-1)(-2)(-1)(-2)(-1)(-2)(-2)(2)(-3)(0)(0)(-1)(-4);
    mat_prot_ += list_of (-1)(-3)(-3)(-3)(-1)(-3)(-3)(-4)(-3)(4)(2)(-3)(1)(0)(-3)(-2)(-1)(-3)(-1)(3)(-3)(-3)(-1)(-4);
    mat_prot_ += list_of (-1)(-2)(-3)(-4)(-1)(-2)(-3)(-4)(-3)(2)(4)(-2)(2)(0)(-3)(-2)(-1)(-2)(-1)(1)(-4)(-3)(-1)(-4);
    mat_prot_ += list_of (-1)(2)(0)(-1)(-3)(1)(1)(-2)(-1)(-3)(-2)(5)(-1)(-3)(-1)(0)(-1)(-3)(-2)(-2)(0)(1)(-1)(-4);
    mat_prot_ += list_of (-1)(-1)(-2)(-3)(-1)(0)(-2)(-3)(-2)(1)(2)(-1)(5)(0)(-2)(-1)(-1)(-1)(-1)(1)(-3)(-1)(-1)(-4);
    mat_prot_ += list_of (-2)(-3)(-3)(-3)(-2)(-3)(-3)(-3)(-1)(0)(0)(-3)(0)(6)(-4)(-2)(-2)(1)(3)(-1)(-3)(-3)(-1)(-4);
    mat_prot_ += list_of (-1)(-2)(-2)(-1)(-3)(-1)(-1)(-2)(-2)(-3)(-3)(-1)(-2)(-4)(7)(-1)(-1)(-4)(-3)(-2)(-2)(-1)(-2)(-4);
    mat_prot_ += list_of (1)(-1)(1)(0)(-1)(0)(0)(0)(-1)(-2)(-2)(0)(-1)(-2)(-1)(4)(1)(-3)(-2)(-2)(0)(0)(0)(-4);
    mat_prot_ += list_of (0)(-1)(0)(-1)(-1)(-1)(-1)(-2)(-2)(-1)(-1)(-1)(-1)(-2)(-1)(1)(5)(-2)(-2)(0)(-1)(-1)(0)(-4);
    mat_prot_ += list_of (-3)(-3)(-4)(-4)(-2)(-2)(-3)(-2)(-2)(-3)(-2)(-3)(-1)(1)(-4)(-3)(-2)(11)(2)(-3)(-4)(-3)(-2)(-4);
    mat_prot_ += list_of (-2)(-2)(-2)(-3)(-2)(-1)(-2)(-3)(2)(-1)(-1)(-2)(-1)(3)(-3)(-2)(-2)(2)(7)(-1)(-3)(-2)(-1)(-4);
    mat_prot_ += list_of (0)(-3)(-3)(-3)(-1)(-2)(-2)(-3)(-3)(3)(1)(-2)(1)(-1)(-2)(-2)(0)(-3)(-1)(4)(-3)(-2)(-1)(-4);
    mat_prot_ += list_of (-2)(-1)(3)(4)(-3)(0)(1)(-1)(0)(-3)(-4)(0)(-3)(-3)(-2)(0)(-1)(-4)(-3)(-3)(4)(1)(-1)(-4);
    mat_prot_ += list_of (-1)(0)(0)(1)(-3)(3)(4)(-2)(0)(-3)(-3)(1)(-1)(-3)(-1)(0)(-1)(-3)(-2)(-2)(1)(4)(-1)(-4);
    mat_prot_ += list_of (0)(-1)(-1)(-1)(-2)(-1)(-1)(-1)(-1)(-1)(-1)(-1)(-1)(-1)(-2)(0)(0)(-2)(-1)(-1)(-1)(-1)(-1)(-4);
    mat_prot_ += list_of (-4)(-4)(-4)(-4)(-4)(-4)(-4)(-4)(-4)(-4)(-4)(-4)(-4)(-4)(-4)(-4)(-4)(-4)(-4)(-4)(-4)(-4)(-4)(1);

    //cout << mat_prot_ident_.at(22) << endl;
    //cout << mat_prot_ident_.at(17) << endl;

    //ftp://ftp.ncbi.nih.gov/blast/matrices/BLOSUM62

    //cout << mat_prot_.at(22).at(17) << endl;

    mat_prot_expect_ = -0.5209;

}

//constructor

ScoreDist::ScoreDist(ApplicationData* appl_data) {

    // mat_prot_ = BLOSUM80_E;

    if (appl_data->protmatrix_ == eBLOSUM80) {
        matrix_ = new Blosum80();
    } else if (appl_data->protmatrix_ == eBLOSUM62) {
        matrix_ = new Blosum62();

    //cout << " mat_prot_ident_[1]: " << matrix_->getMatProtIdent().at(1) << endl;
    //cout << " mat_prot_[1][5]: " << matrix_->getMatProt().at(22).at(17) << endl;
    }

}

DistAtom ScoreDist::dist_2seq(string seq_a, string seq_b) {
    //cout << "------------------k2p_2seq" << endl;
    //cout << "seq_a: " << seq_a << endl;
    //cout << "seq_b: " << seq_b << endl;

    int len = seq_a.length();

    float sigma_s1_s2 = simple_dist_2seq(seq_a, seq_b);
    float sigma_s1_s1 = simple_dist_2seq(seq_a, seq_a);
    float sigma_s2_s2 = simple_dist_2seq(seq_b, seq_b);
    float sigma_r_l = matrix_->getMatProtExpect() * len;

    //correct lower_bound
    if (sigma_s1_s2 < sigma_r_l) {
        sigma_s1_s2 = sigma_r_l;
    }

    float sigma_n = sigma_s1_s2 - sigma_r_l;
    float sigma_u_s1_s2 = (sigma_s1_s1 + sigma_s2_s2) / 2;
    float sigma_un = sigma_u_s1_s2 - sigma_r_l;
    float dist_01 = sigma_n /sigma_un;

    //poisson process correction
    float d_r = - log (dist_01) * 100 ;
    //calibration constant
    float c = 1.3370;

    float d_s = c * d_r;

    /*Evolutionary distances of 250–300 PAM units are commonly considered as the maximum
     * for reasonable distance estimation and, therefore,
     * the Scoredist estimate ds is restricted to the interval [0, 300] PAM.
     */
    float dist = 0;

    float MAX_PAM = 1000;

    //dist = c * dist_01;

    if (d_s > MAX_PAM) {
       dist  = MAX_PAM;
    }

    // cout << "dist: " << dist << endl;
    //wrap d_s
    return DistAtom(floor(dist*seq_a.length()), seq_a.length(), 0, dist);

}

float ScoreDist::simple_dist_2seq(string seq_a, string seq_b) {
    //cout << "------------------k2p_2seq" << endl;
    int nb_sites = seq_a.length();
    float s = 0.0;

    //sum
    for (int k = 0; k < nb_sites; k++) {
        s += dist_2car(seq_a.substr(k, 1), seq_b.substr(k, 1));

    }

    return s;

}

float ScoreDist::dist_2car(string car_a, string car_b) {

    float one_dist = 0.0;

    //cout << "car_a: " << car_a << endl;
    //cout << "car_b: " << car_b << endl;

    int idx_a = std::find(matrix_->getMatProtIdent().begin(), matrix_->getMatProtIdent().end(), car_a) - matrix_->getMatProtIdent().begin();
    int idx_b = std::find(matrix_->getMatProtIdent().begin(), matrix_->getMatProtIdent().end(), car_b) - matrix_->getMatProtIdent().begin();

    //underscore or other defaults to *
    if (idx_a == matrix_->getMatProtIdent().size()) {
        car_a = "*";
        //and refind
        idx_a = std::find(matrix_->getMatProtIdent().begin(), matrix_->getMatProtIdent().end(), car_a) - matrix_->getMatProtIdent().begin();
    }

    if (idx_b == matrix_->getMatProtIdent().size()) {
        car_b = "*";
        //and refind
        idx_b = std::find(matrix_->getMatProtIdent().begin(), matrix_->getMatProtIdent().end(), car_b) - matrix_->getMatProtIdent().begin();
    }

    one_dist = matrix_->getMatProt().at(idx_a).at(idx_b);

    //cout << "one_dist: " << one_dist << endl;

    return one_dist;

}


/**
 * $Id$
 *
 * Software License Agreement (GNU General Public License)
 *
 *  Copyright (C) 2015:
 *
 *    Johann Prankl, prankl@acin.tuwien.ac.at
 *    Aitor Aldoma, aldoma@acin.tuwien.ac.at
 *
 *      Automation and Control Institute
 *      Vienna University of Technology
 *      Gusshausstraße 25-29
 *      1170 Vienn, Austria
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 * @author Johann Prankl, Aitor Aldoma
 *
 */

#ifndef KP_BA_COST_FUNCTIONS_HPP
#define KP_BA_COST_FUNCTIONS_HPP

#include <Eigen/Dense>
#include <ceres/rotation.h>
#include <v4r/common/impl/Vector.hpp>
#include <v4r/keypoints/impl/invPose.hpp>

namespace v4r
{

template<class T>
inline void print_values(const std::string &txt, const T &val1, const T &val2, const T &val3, const T &val4)
{
    //Do nothing
}

template<>
inline void print_values<double>(const std::string &txt, const double &val1, const double &val2, const double &val3, const double &val4)
{
  std::cout<<txt<<" "<<val1<<" "<<val2<<" "<<val3<<" "<<val4<<std::endl;
}

/**
 * Apply camera intrinsics to the normalized point to get image coordinates.
 * This applies the radial lens distortion to a point which is in normalized
 * camera coordinates (i.e. the principal point is at (0, 0)) to get image
 * coordinates in pixels. Templated for use with autodifferentiation.
 */
template <typename T>
inline void applyRadialDistortionCameraIntrinsics(const T &focal_length_x,
                                                  const T &focal_length_y,
                                                  const T &principal_point_x,
                                                  const T &principal_point_y,
                                                  const T &k1,
                                                  const T &k2,
                                                  const T &k3,
                                                  const T &p1,
                                                  const T &p2,
                                                  const T &normalized_x,
                                                  const T &normalized_y,
                                                  T *image_x,
                                                  T *image_y) {
  T x = normalized_x;
  T y = normalized_y;

  // Apply distortion to the normalized points to get (xd, yd).
  T r2 = x*x + y*y;
  T r4 = r2 * r2;
  T r6 = r4 * r2;
  T r_coeff = (T(1) + k1*r2 + k2*r4 + k3*r6);
  T xd = x * r_coeff + T(2)*p1*x*y + p2*(r2 + T(2)*x*x);
  T yd = y * r_coeff + T(2)*p2*x*y + p1*(r2 + T(2)*y*y);

  // Apply focal length and principal point to get the final image coordinates.
  *image_x = focal_length_x * xd + principal_point_x;
  *image_y = focal_length_y * yd + principal_point_y;
}

/**
 * Apply camera intrinsics to the normalized point to get image coordinates.
 * Templated for use with autodifferentiation.
 */
template <typename T>
inline void applyCameraIntrinsics(const T &focal_length_x,
                                  const T &focal_length_y,
                                  const T &principal_point_x,
                                  const T &principal_point_y,
                                  const T &normalized_x,
                                  const T &normalized_y,
                                  T *image_x,
                                  T *image_y) {
  // Apply focal length and principal point to get the final image coordinates.
  *image_x = focal_length_x * normalized_x + principal_point_x;
  *image_y = focal_length_y * normalized_y + principal_point_y;
}

/**
 * Cost functor which computes reprojection error of 3D point X
 * on camera defined by angle-axis rotation and it's translation
 * (which are in the same block due to optimization reasons).
 * This functor uses a radial distortion model.
 */
struct RadialDistortionReprojectionError {
  RadialDistortionReprojectionError(const double &_observed_x, const double &_observed_y)
      : observed_x(_observed_x), observed_y(_observed_y) {}

  template <typename T>
  bool operator()(const T* const intrinsics,
                  const T* const R_t,  // Rotation denoted by angle axis
                                       // followed with translation
                  const T* const X,    // Point coordinates 3x1.
                  T* residuals) const {
    // Unpack the intrinsics.
    const T& focal_length_x    = intrinsics[0];
    const T& focal_length_y    = intrinsics[1];
    const T& principal_point_x = intrinsics[2];
    const T& principal_point_y = intrinsics[3];
    const T& k1                = intrinsics[4];
    const T& k2                = intrinsics[5];
    const T& k3                = intrinsics[6];
    const T& p1                = intrinsics[7];
    const T& p2                = intrinsics[8];

    // Compute projective coordinates: x = RX + t.
    T x[3];

    ceres::AngleAxisRotatePoint(R_t, X, x);
    x[0] += R_t[3];
    x[1] += R_t[4];
    x[2] += R_t[5];

    // Compute normalized coordinates: x /= x[2].
    T xn = x[0] / x[2];
    T yn = x[1] / x[2];

    T predicted_x, predicted_y;

    // Apply distortion to the normalized points to get (xd, yd).
    applyRadialDistortionCameraIntrinsics(focal_length_x,
                                          focal_length_y,
                                          principal_point_x,
                                          principal_point_y,
                                          k1, k2, k3,
                                          p1, p2,
                                          xn, yn,
                                          &predicted_x,
                                          &predicted_y);

    // The error is the difference between the predicted and observed position.
    residuals[0] = predicted_x - T(observed_x);
    residuals[1] = predicted_y - T(observed_y);

    return true;
  }

  const double observed_x;
  const double observed_y;
};

/**
 * Cost functor which computes reprojection error of 3D point X
 * on camera defined by angle-axis rotation and it's translation
 * (which are in the same block due to optimization reasons).
 */
struct ReprojectionError {
  ReprojectionError(const double &_observed_x, const double &_observed_y)
      : observed_x(_observed_x), observed_y(_observed_y) {}

  template <typename T>
  bool operator()(const T* const intrinsics,
                  const T* const R_t,  // Rotation denoted by angle axis
                                       // followed with translation
                  const T* const X,    // Point coordinates 3x1.
                  T* residuals) const {
    // Unpack the intrinsics.
    const T& focal_length_x    = intrinsics[0];
    const T& focal_length_y    = intrinsics[1];
    const T& principal_point_x = intrinsics[2];
    const T& principal_point_y = intrinsics[3];

    // Compute projective coordinates: x = RX + t.
    T x[3];

    ceres::AngleAxisRotatePoint(R_t, X, x);
    x[0] += R_t[3];
    x[1] += R_t[4];
    x[2] += R_t[5];

    // Compute normalized coordinates: x /= x[2].
    T xn = x[0] / x[2];
    T yn = x[1] / x[2];

    T predicted_x, predicted_y;

    // Apply distortion to the normalized points to get (xd, yd).
    applyCameraIntrinsics(focal_length_x, focal_length_y, principal_point_x, principal_point_y,
                          xn, yn, &predicted_x, &predicted_y);

    // The error is the difference between the predicted and observed position.
    residuals[0] = predicted_x - T(observed_x);
    residuals[1] = predicted_y - T(observed_y);

    return true;
  }

  const double observed_x;
  const double observed_y;
};

/**
 * This functor uses a radial distortion model.
 */
struct RadialDistortionReprojectionAndDepthError {
  RadialDistortionReprojectionAndDepthError(const double &_observed_x, const double &_observed_y, const double &_inv_depth, const double &_depth_err_weight)
      : observed_x(_observed_x), observed_y(_observed_y), inv_depth(_inv_depth),
        depth_err_weight(_depth_err_weight) {}

  template <typename T>
  bool operator()(const T* const intrinsics,
                  const T* const R_t,  // Rotation denoted by angle axis
                                       // followed with translation
                  const T* const X,    // Point coordinates 3x1.
                  T* residuals) const {
    // Unpack the intrinsics.
    const T& focal_length_x    = intrinsics[0];
    const T& focal_length_y    = intrinsics[1];
    const T& principal_point_x = intrinsics[2];
    const T& principal_point_y = intrinsics[3];
    const T& k1                = intrinsics[4];
    const T& k2                = intrinsics[5];
    const T& k3                = intrinsics[6];
    const T& p1                = intrinsics[7];
    const T& p2                = intrinsics[8];

    // Compute projective coordinates: x = RX + t.
    T x[3];

    ceres::AngleAxisRotatePoint(R_t, X, x);
    x[0] += R_t[3];
    x[1] += R_t[4];
    x[2] += R_t[5];

    // Compute normalized coordinates: x /= x[2].
    T xn = x[0] / x[2];
    T yn = x[1] / x[2];

    T predicted_x, predicted_y;

    // Apply distortion to the normalized points to get (xd, yd).
    applyRadialDistortionCameraIntrinsics(focal_length_x,
                                          focal_length_y,
                                          principal_point_x,
                                          principal_point_y,
                                          k1, k2, k3,
                                          p1, p2,
                                          xn, yn,
                                          &predicted_x,
                                          &predicted_y);

    // The error is the difference between the predicted and observed position.
    residuals[0] = predicted_x - T(observed_x);
    residuals[1] = predicted_y - T(observed_y);
    residuals[2] = T(depth_err_weight)*(T(1.)/x[2] - T(inv_depth));

    return true;
  }

  const double observed_x;
  const double observed_y;
  const double inv_depth;
  const double depth_err_weight;
};

/**
 * Cost functor which computes reprojection and a RGBD-depth error of 3D point X
 * on camera defined by angle-axis rotation and it's translation
 * (which are in the same block due to optimization reasons).
 */
struct ReprojectionAndDepthError {
  ReprojectionAndDepthError(const double &_observed_x, const double &_observed_y, const double &_inv_depth,
        const double &_depth_err_weight)
      : observed_x(_observed_x), observed_y(_observed_y), inv_depth(_inv_depth),
        depth_err_weight(_depth_err_weight) {}

  template <typename T>
  bool operator()(const T* const intrinsics,
                  const T* const R_t,  // Rotation denoted by angle axis
                                       // followed with translation
                  const T* const X,    // Point coordinates 3x1.
                  T* residuals) const {
    // Unpack the intrinsics.
    const T& focal_length_x    = intrinsics[0];
    const T& focal_length_y    = intrinsics[1];
    const T& principal_point_x = intrinsics[2];
    const T& principal_point_y = intrinsics[3];

    // Compute projective coordinates: x = RX + t.
    T x[3];

    ceres::AngleAxisRotatePoint(R_t, X, x);
    x[0] += R_t[3];
    x[1] += R_t[4];
    x[2] += R_t[5];

    // Compute normalized coordinates: x /= x[2].
    T xn = x[0] / x[2];
    T yn = x[1] / x[2];

    T predicted_x, predicted_y;

    // Apply distortion to the normalized points to get (xd, yd).
    applyCameraIntrinsics(focal_length_x, focal_length_y, principal_point_x, principal_point_y,
                          xn, yn, &predicted_x, &predicted_y);

    // The error is the difference between the predicted and observed position.
    residuals[0] = predicted_x - T(observed_x);
    residuals[1] = predicted_y - T(observed_y);
    residuals[2] = T(depth_err_weight)*(T(1.)/x[2] - T(inv_depth));

//    print_values("dx,dy,dd: ",residuals[0],residuals[1],residuals[2],T(0));

    return true;
  }

  const double observed_x;
  const double observed_y;
  const double inv_depth;
  const double depth_err_weight;
};

/**
 * Cost functor which computes reprojection and a RGBD-depth error of 3D point X
 * on camera defined by angle-axis rotation and it's translation
 * (which are in the same block due to optimization reasons).
 */
struct ReprojectionErrorGlobalPoseCamViewData {
  ReprojectionErrorGlobalPoseCamViewData(const double &_observed_x, const double &_observed_y)
      : observed_x(_observed_x), observed_y(_observed_y) {}

  template <typename T>
  bool operator()(const T* const intrinsics,
                  const T* const Rt0,   // Rotation denoted by angle axis followed with translation (...to keyframe)
                  const T* const Rt1,   // ..to proj. frame
                  const T* const X0,    // Point coordinates 3x1. (in keyframe coordinates)
                  T* residuals) const {
    // Unpack the intrinsics.
    const T& focal_length_x    = intrinsics[0];
    const T& focal_length_y    = intrinsics[1];
    const T& principal_point_x = intrinsics[2];
    const T& principal_point_y = intrinsics[3];

    // transform point to global coordinates
    T xg[3];
    T invRt0[6];
    v4r::invPose6(Rt0, &Rt0[3], invRt0, &invRt0[3]);
    ceres::AngleAxisRotatePoint(invRt0, X0, xg);
    xg[0] += invRt0[3];
    xg[1] += invRt0[4];
    xg[2] += invRt0[5];

    // Compute projective coordinates: x = RX + t.
    T x[3];

    ceres::AngleAxisRotatePoint(Rt1, xg, x);
    x[0] += Rt1[3];
    x[1] += Rt1[4];
    x[2] += Rt1[5];

    // Compute normalized coordinates: x /= x[2].
    T xn = x[0] / x[2];
    T yn = x[1] / x[2];

    T predicted_x, predicted_y;

    // Apply distortion to the normalized points to get (xd, yd).
    applyCameraIntrinsics(focal_length_x, focal_length_y, principal_point_x, principal_point_y, xn, yn, &predicted_x, &predicted_y);

    // The error is the difference between the predicted and observed position.
    residuals[0] = predicted_x - T(observed_x);
    residuals[1] = predicted_y - T(observed_y);

    //print_values("dx,dy,dd: ",residuals[0],residuals[1],residuals[2],T(0));

    return true;
  }

  const double observed_x;
  const double observed_y;
};

/**
 * Cost functor which computes reprojection and a RGBD-depth error of 3D point X
 * on camera defined by angle-axis rotation and it's translation
 * (which are in the same block due to optimization reasons).
 */
struct ReprojectionErrorGlobalPoseDeltaPoseCamViewData {
  ReprojectionErrorGlobalPoseDeltaPoseCamViewData(const double &_observed_x, const double &_observed_y)
      : observed_x(_observed_x), observed_y(_observed_y) {}

  template <typename T>
  bool operator()(const T* const intrinsics,
                  const T* const Rt0,   // Rotation denoted by angle axis followed with translation (...to keyframe)
                  const T* const Rt1,   // ..to proj. frame
                  const T* const Rtd,   // ..to proj. frame
                  const T* const X0,    // Point coordinates 3x1. (in keyframe coordinates)
                  T* residuals) const {

    // Unpack the intrinsics.
    const T& focal_length_x    = intrinsics[0];
    const T& focal_length_y    = intrinsics[1];
    const T& principal_point_x = intrinsics[2];
    const T& principal_point_y = intrinsics[3];

    // transform point to cloud pose
    T xd[3];
    T invRtd[6];
    v4r::invPose6(Rtd, &Rtd[3], invRtd, &invRtd[3]);
    ceres::AngleAxisRotatePoint(invRtd, X0, xd);
    xd[0] += invRtd[3];
    xd[1] += invRtd[4];
    xd[2] += invRtd[5];

    // transform point to global coordinates
    T xg[3];
    T invRt0[6];
    v4r::invPose6(Rt0, &Rt0[3], invRt0, &invRt0[3]);
    ceres::AngleAxisRotatePoint(invRt0, xd, xg);
    xd[0] += invRt0[3];
    xd[1] += invRt0[4];
    xd[2] += invRt0[5];

    // Compute projective coordinates: x = RX + t.
    T x[3];
    ceres::AngleAxisRotatePoint(Rt1, xg, x);
    x[0] += Rt1[3];
    x[1] += Rt1[4];
    x[2] += Rt1[5];

    // Compute projective coordinates: x = RX + t.
    T xrgb[3];
    ceres::AngleAxisRotatePoint(Rtd, xg, xrgb);
    xrgb[0] += Rtd[3];
    xrgb[1] += Rtd[4];
    xrgb[2] += Rtd[5];

    // Compute normalized coordinates: x /= x[2].
    T xn = xrgb[0] / xrgb[2];
    T yn = xrgb[1] / xrgb[2];

    T predicted_x, predicted_y;

    // Apply distortion to the normalized points to get (xd, yd).
    applyCameraIntrinsics(focal_length_x, focal_length_y, principal_point_x, principal_point_y, xn, yn, &predicted_x, &predicted_y);

    // The error is the difference between the predicted and observed position.
    residuals[0] = predicted_x - T(observed_x);
    residuals[1] = predicted_y - T(observed_y);

    //print_values("dx,dy,dd: ",residuals[0],residuals[1],residuals[2],T(0));

    return true;
  }

  const double observed_x;
  const double observed_y;
};

/**
 * This functor uses a radial distortion model.
 */
struct RadialDistortionReprojectionErrorGlobalPoseCamViewData {
  RadialDistortionReprojectionErrorGlobalPoseCamViewData(const double &_observed_x, const double &_observed_y)
      : observed_x(_observed_x), observed_y(_observed_y) {}

  template <typename T>
  bool operator()(const T* const intrinsics,
                  const T* const Rt0,   // Rotation denoted by angle axis followed with translation (...to keyframe)
                  const T* const Rt1,   // ..to proj. frame
                  const T* const X0,    // Point coordinates 3x1.
                  T* residuals) const {
    // Unpack the intrinsics.
    const T& focal_length_x    = intrinsics[0];
    const T& focal_length_y    = intrinsics[1];
    const T& principal_point_x = intrinsics[2];
    const T& principal_point_y = intrinsics[3];
    const T& k1                = intrinsics[4];
    const T& k2                = intrinsics[5];
    const T& k3                = intrinsics[6];
    const T& p1                = intrinsics[7];
    const T& p2                = intrinsics[8];

    // transform point to global coordinates
    T xg[3];
    T invRt0[6];
    v4r::invPose6(Rt0, &Rt0[3], invRt0, &invRt0[3]);
    ceres::AngleAxisRotatePoint(invRt0, X0, xg);
    xg[0] += invRt0[3];
    xg[1] += invRt0[4];
    xg[2] += invRt0[5];

    // Compute projective coordinates: x = RX + t.
    T x[3];

    ceres::AngleAxisRotatePoint(Rt1, xg, x);
    x[0] += Rt1[3];
    x[1] += Rt1[4];
    x[2] += Rt1[5];

    // Compute normalized coordinates: x /= x[2].
    T xn = x[0] / x[2];
    T yn = x[1] / x[2];

    T predicted_x, predicted_y;

    // Apply distortion to the normalized points to get (xd, yd).
    applyRadialDistortionCameraIntrinsics(focal_length_x, focal_length_y, principal_point_x, principal_point_y, k1, k2, k3, p1, p2, xn, yn, &predicted_x, &predicted_y);

    // The error is the difference between the predicted and observed position.
    residuals[0] = predicted_x - T(observed_x);
    residuals[1] = predicted_y - T(observed_y);

    return true;
  }

  const double observed_x;
  const double observed_y;
};

/**
 * This functor uses a radial distortion model.
 */
struct RadialDistortionReprojectionErrorGlobalPoseDeltaPoseCamViewData {
  RadialDistortionReprojectionErrorGlobalPoseDeltaPoseCamViewData(const double &_observed_x, const double &_observed_y)
      : observed_x(_observed_x), observed_y(_observed_y) {}

  template <typename T>
  bool operator()(const T* const intrinsics,
                  const T* const Rt0,   // Rotation denoted by angle axis followed with translation (...to keyframe)
                  const T* const Rt1,   // ..to proj. frame
                  const T* const Rtd,   // ..cloud->RGB delta pose
                  const T* const X0,    // Point coordinates 3x1.
                  T* residuals) const {
    // Unpack the intrinsics.
    const T& focal_length_x    = intrinsics[0];
    const T& focal_length_y    = intrinsics[1];
    const T& principal_point_x = intrinsics[2];
    const T& principal_point_y = intrinsics[3];
    const T& k1                = intrinsics[4];
    const T& k2                = intrinsics[5];
    const T& k3                = intrinsics[6];
    const T& p1                = intrinsics[7];
    const T& p2                = intrinsics[8];

    // transform point to cloud pose
    T xd[3];
    T invRtd[6];
    v4r::invPose6(Rtd, &Rtd[3], invRtd, &invRtd[3]);
    ceres::AngleAxisRotatePoint(invRtd, X0, xd);
    xd[0] += invRtd[3];
    xd[1] += invRtd[4];
    xd[2] += invRtd[5];

    // transform point to global coordinates
    T xg[3];
    T invRt0[6];
    v4r::invPose6(Rt0, &Rt0[3], invRt0, &invRt0[3]);
    ceres::AngleAxisRotatePoint(invRt0, xd, xg);
    xg[0] += invRt0[3];
    xg[1] += invRt0[4];
    xg[2] += invRt0[5];

    // Compute projective coordinates: x = RX + t.
    T x[3];
    ceres::AngleAxisRotatePoint(Rt1, xg, x);
    x[0] += Rt1[3];
    x[1] += Rt1[4];
    x[2] += Rt1[5];

    // Compute projective coordinates: x = RX + t.
    T xrgb[3];
    ceres::AngleAxisRotatePoint(Rtd, x, xrgb);
    xrgb[0] += Rtd[3];
    xrgb[1] += Rtd[4];
    xrgb[2] += Rtd[5];

    // Compute normalized coordinates: x /= x[2].
    T xn = xrgb[0] / xrgb[2];
    T yn = xrgb[1] / xrgb[2];

    T predicted_x, predicted_y;

    // Apply distortion to the normalized points to get (xd, yd).
    applyRadialDistortionCameraIntrinsics(focal_length_x, focal_length_y, principal_point_x, principal_point_y, k1, k2, k3, p1, p2, xn, yn, &predicted_x, &predicted_y);

    // The error is the difference between the predicted and observed position.
    residuals[0] = predicted_x - T(observed_x);
    residuals[1] = predicted_y - T(observed_y);

    return true;
  }

  const double observed_x;
  const double observed_y;
};

/**
 * This functor uses a point to plane error model
 */
struct PointToPlaneErrorGlobalPoseCamViewData {
  PointToPlaneErrorGlobalPoseCamViewData(const Eigen::Vector3d &_pt0, const Eigen::Vector3d &_n0, const  Eigen::Vector3d &_pt1, const double &_w)
      : pt0(_pt0), n0(_n0), pt1(_pt1), error_weight(_w) {}

  template <typename T>
  bool operator()(const T* const Rt0,   // Rotation denoted by angle axis followed with translation (...to keyframe)
                  const T* const Rt1,   // ..to proj. frame
                  T* residuals) const {

    Eigen::Matrix<T,3,1> pt0e, n0e, pt0g, n0g, pt1t, n1t, pt1e;

    pt0e[0] = T(pt0[0]);
    pt0e[1] = T(pt0[1]);
    pt0e[2] = T(pt0[2]);
    n0e[0] = T(n0[0]);
    n0e[1] = T(n0[1]);
    n0e[2] = T(n0[2]);
    pt1e[0] = T(pt1[0]);
    pt1e[1] = T(pt1[1]);
    pt1e[2] = T(pt1[2]);

    // transform point to global coordinates
    T invRt0[6];
    v4r::invPose6(Rt0, &Rt0[3], invRt0, &invRt0[3]);
    ceres::AngleAxisRotatePoint(invRt0, &n0e[0], &n0g[0]);
    ceres::AngleAxisRotatePoint(invRt0, &pt0e[0], &pt0g[0]);
    pt0g[0] += invRt0[3];
    pt0g[1] += invRt0[4];
    pt0g[2] += invRt0[5];

    // transform to view 1: x = RX + t.
    ceres::AngleAxisRotatePoint(Rt1, &n0g[0], &n1t[0]);
    ceres::AngleAxisRotatePoint(Rt1, &pt0g[0], &pt1t[0]);
    pt1t[0] += Rt1[3];
    pt1t[1] += Rt1[4];
    pt1t[2] += Rt1[5];

    // compute the point to plane distance
    //residuals[0] = T(error_weight) * /*(T(2.)/(pt1e[2]+pt1t[2])) **/ (pt1e-pt1t).dot(n1t);
    T weight = T(error_weight) * (T(2.)/(v4r::sqr(pt1e[2])+v4r::sqr(pt1t[2])));
    Eigen::Matrix<T,3,1> diff;
    diff = pt1e-pt1t;
    residuals[0] = weight * diff[0]*n1t[0];
    residuals[1] = weight * diff[1]*n1t[1];
    residuals[2] = weight * diff[2]*n1t[2];

    return true;
  }

  const Eigen::Vector3d pt0;
  const Eigen::Vector3d n0;
  const Eigen::Vector3d pt1;
  const double error_weight;
};

/**
 * This functor uses a point to plane error model
 */
struct PointToPlaneErrorGlobalPoseCamViewDataOptiPt1 {
  PointToPlaneErrorGlobalPoseCamViewDataOptiPt1(const Eigen::Vector3d &_pt0, const Eigen::Vector3d &_n0, const double &_w)
      : pt0(_pt0), n0(_n0), error_weight(_w) {}

  template <typename T>
  bool operator()(const T* const Rt0,   // Rotation denoted by angle axis followed with translation (...to keyframe)
                  const T* const Rt1,   // ..to proj. frame
                  const T* const X1,
                  T* residuals) const {

    Eigen::Matrix<T,3,1> pt0e, n0e, pt0g, n0g, pt1t, n1t, pt1e;

    pt0e[0] = T(pt0[0]);
    pt0e[1] = T(pt0[1]);
    pt0e[2] = T(pt0[2]);
    n0e[0] = T(n0[0]);
    n0e[1] = T(n0[1]);
    n0e[2] = T(n0[2]);
    pt1e[0] = X1[0];
    pt1e[1] = X1[1];
    pt1e[2] = X1[2];

    // transform point to global coordinates
    T invRt0[6];
    v4r::invPose6(Rt0, &Rt0[3], invRt0, &invRt0[3]);
    ceres::AngleAxisRotatePoint(invRt0, &n0e[0], &n0g[0]);
    ceres::AngleAxisRotatePoint(invRt0, &pt0e[0], &pt0g[0]);
    pt0g[0] += invRt0[3];
    pt0g[1] += invRt0[4];
    pt0g[2] += invRt0[5];

    // transform to view 1: x = RX + t.
    ceres::AngleAxisRotatePoint(Rt1, &n0g[0], &n1t[0]);
    ceres::AngleAxisRotatePoint(Rt1, &pt0g[0], &pt1t[0]);
    pt1t[0] += Rt1[3];
    pt1t[1] += Rt1[4];
    pt1t[2] += Rt1[5];

    // compute the point to plane distance
    //residuals[0] = T(error_weight) * /*(T(2.)/(pt1e[2]+pt1t[2])) **/ (pt1e-pt1t).dot(n1t);
    T weight = T(error_weight) * (T(2.)/(v4r::sqr(pt1e[2])+v4r::sqr(pt1t[2])));
    Eigen::Matrix<T,3,1> diff;
    diff = pt1e-pt1t;
    residuals[0] = weight * diff[0]*n1t[0];
    residuals[1] = weight * diff[1]*n1t[1];
    residuals[2] = weight * diff[2]*n1t[2];

    return true;
  }

  const Eigen::Vector3d pt0;
  const Eigen::Vector3d n0;
  const double error_weight;
};

/**
 * This functor uses a point to plane error model
 */
struct PointToPlaneErrorGlobalPoseDeltaPoseCamViewDataOptiPt1 {
  PointToPlaneErrorGlobalPoseDeltaPoseCamViewDataOptiPt1(const Eigen::Vector3d &_pt0, const Eigen::Vector3d &_n0, const double &_w)
      : pt0(_pt0), n0(_n0), error_weight(_w) {}

  template <typename T>
  bool operator()(const T* const Rtd,   // ..to rgb. frame
                  const T* const X1,
                  T* residuals) const {

    Eigen::Matrix<T,3,1> pt0e, n0e, pt1t, n1t, pt1e;

    pt0e[0] = T(pt0[0]);
    pt0e[1] = T(pt0[1]);
    pt0e[2] = T(pt0[2]);
    n0e[0] = T(n0[0]);
    n0e[1] = T(n0[1]);
    n0e[2] = T(n0[2]);
    pt1e[0] = X1[0];
    pt1e[1] = X1[1];
    pt1e[2] = X1[2];

    // transform to view 1: x = RX + t.
    ceres::AngleAxisRotatePoint(Rtd, &n0e[0], &n1t[0]);
    ceres::AngleAxisRotatePoint(Rtd, &pt0e[0], &pt1t[0]);
    pt1t[0] += Rtd[3];
    pt1t[1] += Rtd[4];
    pt1t[2] += Rtd[5];

    // compute the point to plane distance
    //residuals[0] = T(error_weight) * /*(T(2.)/(pt1e[2]+pt1t[2])) **/ (pt1e-pt1t).dot(n1t);
    T weight = T(error_weight) * (T(2.)/(v4r::sqr(pt1e[2])+v4r::sqr(pt1t[2])));
    Eigen::Matrix<T,3,1> diff;
    diff = pt1e-pt1t;
    residuals[0] = weight * diff[0]*n1t[0];
    residuals[1] = weight * diff[1]*n1t[1];
    residuals[2] = weight * diff[2]*n1t[2];

    return true;
  }

  const Eigen::Vector3d pt0;
  const Eigen::Vector3d n0;
  const double error_weight;
};

}

#endif


/*
 This program is free software; you can redistribute it and/or modify it under
 the terms of the European Union Public Licence - EUPL v.1.1 as published by
 the European Commission.

 This program is distributed in the hope that it will be useful, but WITHOUT
 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 FOR A PARTICULAR PURPOSE. See the European Union Public Licence - EUPL v.1.1
 for more details.

 You should have received a copy of the European Union Public Licence - EUPL v.1.1
 along with this program.

 Further information about the European Union Public Licence - EUPL v.1.1 can
 also be found on the world wide web at http://ec.europa.eu/idabc/eupl

*/

/*
 ------ Copyright (C) 2010 STA Steering Board (space.trajectory.analysis AT gmail.com) ----
*/

/*
 ------------------ Author: Chris Laurel  -------------------------------------------------
 ------------------ E-mail: (claurel@gmail.com) ----------------------------

 Based on code by Chris Laurel from Celestia (http://www.shatters.net/celestia)
 */

#include <Eigen/Core>
//#include "Eigen/src/Core/Map.h"
#include <QtCore>
#include <QDataStream>
#include <QIODevice>
#include "date.h"
#include "jplephemeris.h"

using namespace Eigen;
using namespace std;

typedef Eigen::Matrix< double, 3, 3 > 	MyMatrix3d;
typedef Eigen::Matrix< double, 3, 1 > 	MyVector3d;

static const unsigned int DE200RecordSize    =  826;
static const unsigned int DE405RecordSize    = 1018;
static const unsigned int DE406RecordSize    =  728;

static const unsigned int NConstants         =  400;
static const unsigned int ConstantNameLength =  6;

static const unsigned int MaxChebyshevCoeffs = 32;

static const int LabelSize = 84;

sta::JPLEphemeris::JPLEphemeris()
{
}

sta::JPLEphemeris::~JPLEphemeris()
{
}

unsigned int
sta::JPLEphemeris::getDENumber() const
{
    return DENum;
}

double
sta::JPLEphemeris::getStartDate() const
{
    return startDate;
}

double
sta::JPLEphemeris::getEndDate() const
{
    return endDate;
}

sta::Ephemeris::TimeInterval
sta::JPLEphemeris::validTimeInterval(const StaBody* /* body */) const
{
    return TimeInterval(sta::JdToMjd(startDate), sta::JdToMjd(endDate));
}

// Internal method to map an STA body id to a JPL ephemeris item
sta::JPLEphemItem
sta::JPLEphemeris::mapStaId(StaBodyId id) const
{
    switch (id)
    {
    case STA_SOLAR_SYSTEM_BARYCENTER:
        return JPLEph_SSB;
    case STA_EARTH_BARYCENTER:
        return JPLEph_EarthMoonBary;
    case STA_MERCURY:
        return JPLEph_Mercury;
    case STA_VENUS:
        return JPLEph_Venus;
    case STA_EARTH:
        return JPLEph_Earth;
    case STA_MARS:
        return JPLEph_Mars;
    case STA_JUPITER:
        return JPLEph_Jupiter;
    case STA_SATURN:
        return JPLEph_Saturn;
    case STA_URANUS:
        return JPLEph_Uranus;
    case STA_NEPTUNE:
        return JPLEph_Neptune;
    case STA_PLUTO:
        return JPLEph_Pluto;
    case STA_SUN:
        return JPLEph_Sun;
    case STA_MOON:
        return JPLEph_Moon;
    default:
        return JPLEph_Invalid;
    }
}

const QList<StaBodyId>&
sta::JPLEphemeris::bodyList() const
{
    return m_bodies;
}

const StaBody*
sta::JPLEphemeris::parentBody(const StaBody* body) const
{
    JPLEphemItem id = mapStaId(body->id());
    if (id == JPLEph_Invalid)
    {
        return NULL;
    }
    else
    {
        if (id == JPLEph_SSB)
            return NULL;
        else if (id == JPLEph_Moon || id == JPLEph_Earth)
            return STA_SOLAR_SYSTEM->lookup(STA_EARTH_BARYCENTER);
        else
            return STA_SOLAR_SYSTEM->lookup(STA_SOLAR_SYSTEM_BARYCENTER);
    }
}

const sta::CoordinateSystemType
sta::JPLEphemeris::coordinateSystem(const StaBody* /* body */) const
{
    return sta::COORDSYS_EME_J2000;
}

sta::StateVector
sta::JPLEphemeris::stateVector(const StaBody* body,
                               double mjd,
                               const StaBody* center,
                               sta::CoordinateSystemType coordSys) const
{
    JPLEphemItem id = mapStaId(body->id());
    if (id == JPLEph_Invalid || body == center)
    {
        return StateVector::zero();
    }

    StateVector state = getPlanetStateVector(id, MjdToJd(mjd));
    const StaBody* parent = parentBody(body);

    // If the parent object in the ephemeris is the same as the center, then no
    // extra work is necessary. Otherwise, we also need compute the position of
    // the parent with respect to the center.
    if (parent == center)
    {
        return state;
    }
    else
    {
        const StaBody* ssb = STA_SOLAR_SYSTEM->lookup(STA_SOLAR_SYSTEM_BARYCENTER);
        StateVector parentState = parent->stateVector(mjd, ssb, coordSys);
        StateVector centerState = center->stateVector(mjd, ssb, coordSys);

        return StateVector(parentState.position - centerState.position + state.position,
                           parentState.velocity - centerState.velocity + state.velocity);
    }
}

// Return the position of an object relative to the solar system barycenter
// or the Earth (in the case of the Moon) at a specified TDB Julian date tjd.
// If tjd is outside the span covered by the ephemeris it is clamped to a
// valid time.
sta::StateVector
sta::JPLEphemeris::getPlanetStateVector(JPLEphemItem planet, double tjd) const
{
    // Solar system barycenter is the origin
    if (planet == JPLEph_SSB)
    {
        return StateVector(MyVector3d(0.0, 0.0, 0.0), MyVector3d(0.0, 0.0, 0.0));
    }

    // The position of the Earth must be computed from the positions of the
    // Earth-Moon barycenter and Moon. This is handled automatically because
    // Earth's 'parent' in the ephemeris is the Earth-Moon barycenter.
    if (planet == JPLEph_Earth)
    {
        // Get the geocentric position of the Moon
        sta::StateVector moonState = getPlanetStateVector(JPLEph_Moon, tjd);

        double f = 1.0 / (earthMoonMassRatio + 1.0);
        return StateVector(-moonState.position * f, -moonState.velocity * f);
    }

    // Clamp time to [ startDate, endDate ]
    if (tjd < startDate)
    {
        tjd = startDate;
    }
    else if (tjd > endDate)
    {
        tjd = endDate;
    }

    // recordIndex is always >= 0:
    unsigned int recordIndex = (unsigned int) ((tjd - startDate) / daysPerInterval);
    // Make sure we don't go past the end of the array if t == endDate
    if (recordIndex >= records.size())
    {
        recordIndex = records.size() - 1;
    }
    const JPLEphRecord& record = records[recordIndex];

    Q_ASSERT(coeffInfo[planet].nGranules >= 1);
    Q_ASSERT(coeffInfo[planet].nGranules <= 32);
    Q_ASSERT(coeffInfo[planet].nCoeffs <= MaxChebyshevCoeffs);

    // u is the normalized time (in [-1, 1]) for interpolating
    // coeffs is a pointer to the Chebyshev coefficients
    double u = 0.0;
    const double* coeffs = NULL;
    unsigned int nCoeffs = coeffInfo[planet].nCoeffs;
    double velocityScale = 1.0;

    if (coeffInfo[planet].nGranules == 0xffffffff)
    {
        coeffs = record.coeffs + coeffInfo[planet].offset;
        u = 2.0 * (tjd - record.t0) / daysPerInterval - 1.0;
        velocityScale = 2.0 / daysPerInterval;
    }
    else
    {
        // This interval is subdivided into shorter subintervals (called granules)
        // Adjust the coefficient pointer accordingly.
        double daysPerGranule = daysPerInterval / coeffInfo[planet].nGranules;
        int granule = (int) ((tjd - record.t0) / daysPerGranule);
        double granuleStartDate = record.t0 + daysPerGranule * (double) granule;
        coeffs = record.coeffs + coeffInfo[planet].offset +
                 granule * coeffInfo[planet].nCoeffs * 3;
        u = 2.0 * (tjd - granuleStartDate) / daysPerGranule - 1.0;
        velocityScale = 2.0 / daysPerGranule;
    }

    // Calculate the Chebyshev polynomial terms and their derivatives
    //Matrix<double, MaxChebyshevCoeffs, 1> positionTerms;
    double positionTerms[MaxChebyshevCoeffs];
    positionTerms[0] = 1.0;
    positionTerms[1] = u;

    //Matrix<double, MaxChebyshevCoeffs, 1> velocityTerms;
    double velocityTerms[MaxChebyshevCoeffs];
    velocityTerms[0] = 0.0;
    velocityTerms[1] = 1.0;
    for (unsigned int i = 2; i < nCoeffs; i++)
    {
        positionTerms[i] = 2.0 * u * positionTerms[i - 1] - positionTerms[i - 2];
        velocityTerms[i] = 2.0 * u * velocityTerms[i - 1] - velocityTerms[i - 2] + 2.0 * positionTerms[i - 1];
    }

    // Note that Eigen uses column-major storage, so we need to transpose the coefficient
    // matrix.
    MyVector3d position = Map<MatrixXd>(coeffs, nCoeffs, 3).transpose() * Map<MatrixXd>(positionTerms, nCoeffs, 1);
    MyVector3d velocity = Map<MatrixXd>(coeffs, nCoeffs, 3).transpose() * Map<MatrixXd>(velocityTerms, nCoeffs, 1);

    return StateVector(position, velocity * (velocityScale / 86400.0));
}

sta::JPLEphemeris*
sta::JPLEphemeris::load(QIODevice* device)
{
    QDataStream in(device);
    in.setByteOrder(QDataStream::BigEndian);

    JPLEphemeris* eph = NULL;

    // Skip past three header labels
    in.skipRawData(LabelSize * 3);
    if (in.status() != QDataStream::Ok)
        return NULL;

    // Skip past the constant names
    in.skipRawData(NConstants * ConstantNameLength);
    if (in.status() != QDataStream::Ok)
        return NULL;

    eph = new JPLEphemeris();
    if (eph == NULL)
        return NULL;

    // Read the start time, end time, and time interval
    in >> eph->startDate;
    in >> eph->endDate;
    in >> eph->daysPerInterval;
    if (in.status() != QDataStream::Ok)
    {
        delete eph;
        return NULL;
    }

    // Number of constants with valid values; not useful for us
    quint32 nConstants;
    in >> nConstants;

    in >> eph->au;     // kilometers per astronomical unit
    in >> eph->earthMoonMassRatio;

    // Read the coefficient information for each item in the ephemeris
    for (unsigned int i = 0; i < JPLEph_NItems; i++)
    {
        in >> eph->coeffInfo[i].offset;
        in >> eph->coeffInfo[i].nCoeffs;
        in >> eph->coeffInfo[i].nGranules;

        eph->coeffInfo[i].offset -= 3;
    }

    if (in.status() != QDataStream::Ok)
    {
        delete eph;
        return NULL;
    }

    in >> eph->DENum;

    switch (eph->DENum)
    {
    case 200:
        eph->recordSize = DE200RecordSize;
        break;
    case 405:
        eph->recordSize = DE405RecordSize;
        break;
    case 406:
        eph->recordSize = DE406RecordSize;
        break;
    default:
        delete eph;
        return NULL;
    }

    in >> eph->librationCoeffInfo.offset;
    in >> eph->librationCoeffInfo.nCoeffs;
    in >> eph->librationCoeffInfo.nGranules;
    if (in.status() != QDataStream::Ok)
    {
        delete eph;
        return NULL;
    }

    // Skip past the rest of the record
    in.skipRawData(eph->recordSize * 8 - 2856);
    // The next record contains constant values (which we don't need)
    in.skipRawData(eph->recordSize * 8);
    if (in.status() != QDataStream::Ok)
    {
        delete eph;
        return NULL;
    }

    unsigned int nRecords = (unsigned int) ((eph->endDate - eph->startDate) / eph->daysPerInterval);
    eph->records.resize(nRecords);

    for (unsigned int i = 0; i < nRecords; i++)
    {
        in >> eph->records[i].t0;
        in >> eph->records[i].t1;

    	// Allocate coefficient array for this record; the first two
    	// 'coefficients' are actually the start and end time (t0 and t1)
        eph->records[i].coeffs = new double[eph->recordSize - 2];
        for (unsigned int j = 0; j < eph->recordSize - 2; j++)
        {
            in >> eph->records[i].coeffs[j];
        }

    	// Make sure that we read this record successfully
        if (in.status() != QDataStream::Ok)
    	{
    	    delete eph;
    	    return NULL;
    	}
    }

    eph->m_bodies << STA_SUN << STA_MERCURY << STA_VENUS << STA_EARTH << STA_MARS
                  << STA_JUPITER << STA_SATURN << STA_URANUS << STA_NEPTUNE << STA_PLUTO
                  << STA_EARTH_BARYCENTER << STA_MOON;

    return eph;
}


#include <ciphey/freq.hpp>

#include <boost/math/special_functions/gamma.hpp>

#include <future>
#include <optional>
#include <random>
#include <set>
#include <thread>

namespace ciphey {
  static void merge_last(assoc_table& assoc) {
    // Take the last two elements
    auto to_merge = std::move(assoc.back());
    assoc.pop_back();
    auto target = assoc.back();
    assoc.pop_back();

    // Combine them
    target.expected += to_merge.expected;
    target.observed += to_merge.observed;

    // Insert them into the correct position
    auto new_pos = std::partition_point(assoc.begin(), assoc.end(),
                                        [&](auto& a) { return a.expected > target.expected; });
    assoc.insert(new_pos, std::move(target));

  }

  void prepare_chisq(assoc_table& assoc, freq_t count) {
    // IIRC copying is much, MUCH faster than mallocing, so a list would be slower in this case
    //
    // Benchmarks seem to confirm this

    // If there's nothing we can do, give up
    if (count < 10)
      return;

    std::sort(assoc.begin(), assoc.end(), [](auto& a, auto& b) { return a.expected > b.expected; });

    // We now perform a *really* simple bin packing to make sure that no expected frequency is lower that 1
    auto exp_1_target = 1. / count;

    while (assoc.size() >= 2 && assoc.back().expected < exp_1_target)
      merge_last(assoc);

    auto exp_5_target = 5. / count;
    if (exp_5_target > 0.2)
      return;
    while (assoc.size() >= 2) {
      auto upper_bound = std::partition_point(assoc.begin(), assoc.end(),
                                              [&](auto& a) { return a.expected > exp_5_target; });
      auto n_less_than_fifth = assoc.end() - upper_bound;
      if (n_less_than_fifth <= assoc.size() / 5.)
        break;
      merge_last(assoc);
    }
  }

  float_t run_chisq(const assoc_table& assoc, freq_t count) {
    float_t chisq = 0;

    for (auto const& elem : assoc) {
      if (elem.expected == 0 && elem.observed != 0)
        return std::numeric_limits<float_t>::infinity();

      float_t contrib = elem.expected - elem.observed;
      contrib *= contrib;
      contrib /= elem.expected;

      chisq += contrib;
    }

    // (f_e/n - f_o/n)^2/(f_e/n) = (1/n) (f_e - f_o)/(f_e)
    //
    // We need to normalise this for our stat to be interpretable
    chisq *= count;
    return chisq;
  }

  float_t run_g(const assoc_table& assoc, freq_t count) {
    float_t g = 0;

    for (auto const& elem : assoc) {
      if (elem.observed == 0)
        continue;

      if (elem.expected == 0)
        return std::numeric_limits<float_t>::infinity();

      float_t contrib = ::log(elem.observed / elem.expected); // `count` cancels out here
      contrib *= elem.observed;

      g += contrib;
    }

    // (f_e/n - f_o/n)^2/(f_e/n) = (1/n) (f_e - f_o)/(f_e)
    //
    // We need to normalise this for our stat to be interpretable
    g *= 2 * count;
    return g;
  }

  prob_t chisq_cdf(freq_t dof, float_t up_to) {
    // Handle the asymptopic value first, to save time
    if (up_to == std::numeric_limits<float_t>::infinity())
      return 1;
    else if (up_to == 0)
      return 0;
    else if (dof == 0) // https://en.wikipedia.org/wiki/Zero_degrees_of_freedom
      return 0;
    return boost::math::gamma_p(static_cast<float>(dof) / 2, up_to / 2);
  }

  prob_t gof_test(assoc_table assoc, freq_t count) {
//    prepare_chisq(assoc, count);
    auto stat = run_g(assoc, count);
    // We want the upper tail
    auto p_value = 1 - chisq_cdf(assoc.size() - 1, stat);

    return p_value;
//    return std::max(p_value, orig_p_value);
  }

  assoc_table create_assoc_table(prob_table const& observed, prob_table const& expected) {
    // TODO: optimise
    //
    // First, let's get all the keys
    struct key_elem {
      std::optional<float_t> obs_val, exp_val;
    };

    std::map<char_t, key_elem> keys;
    // We can use index notation for this, as that gets default constructed
    for (auto& i : expected)
      keys[i.first].exp_val = i.second;
    for (auto& i : observed) {
      if (i.second == 0.)
        continue;
      keys[i.first].obs_val = i.second;
    }

    // We can now fill in all the values, with non-existent expected values being defined as zero
    assoc_table ret;
//    ret.reserve(keys.size());

    for (auto& i : keys)
      ret.emplace_back(assoc_table_elem{
        .observed = i.second.obs_val.value_or(0),
        .expected = i.second.exp_val.value_or(0)});

    return ret;
  }

  prob_table freq_conv(freq_table const& freqs, freq_t total_len) {
    prob_table ret;
    for (auto& i : freqs)
      ret[i.first] = static_cast<prob_t>(i.second) / total_len;
    return ret;
  }
  prob_table freq_conv(freq_table const& freqs) {
    freq_t total_len = 0;
    for (auto& i : freqs)
      total_len += i.second;
    return freq_conv(freqs, total_len);
  }
  windowed_prob_table freq_conv(windowed_freq_table& freqs, freq_t total_len) {
    windowed_prob_table ret;
    ret.reserve(total_len);
    for (auto& i : freqs)
      // TODO: work the length out here, rather than do slow counting
      ret.emplace_back(freq_conv(i));
    return ret;
  }
  windowed_prob_table freq_conv(windowed_freq_table& freqs) {
    freq_t total_len = 0;
    for (auto& i : freqs)
      for (auto& entry: i)
        total_len += entry.second;
    return freq_conv(freqs, total_len);
  }

  size_t filter_missing(freq_table& target, prob_table const& lookup) {
    std::vector<char_t> to_remove;
    size_t ret = 0;
    for (auto& i : target) {
      if (auto iter = lookup.find(i.first); iter == lookup.end() || iter->second == 0) {
        to_remove.push_back(i.first);
        ret += i.second;
      }
    }

    for (auto i : to_remove)
      target.erase(i);

    return ret;
  }

  size_t filter_missing(freq_table& target, domain_t const& tab) {
    std::vector<char_t> to_remove;
    size_t ret = 0;
    for (auto& i : target) {
      if (!tab.count(i.first)) {
        to_remove.push_back(i.first);
        ret += i.second;
      }
    }

    for (auto i : to_remove)
      target.erase(i);

    return ret;
  }

//  size_t filter_missing(windowed_freq_table& target, prob_table const& lookup) {
//    size_t acc = 0;
//    for (auto& i : target)
//      acc += filter_missing(i, lookup);
//    return acc;
//  }

  void freq_analysis(freq_table& tab, string_const_ref_t const& str) {
    for (auto& i : str)
      ++tab[i];
  }
  size_t freq_analysis(freq_table& tab, string_const_ref_t const& str, domain_t const& domain) {
    size_t n = 0;
    for (auto& i : str) {
      if (domain.count(i)) {
        ++tab[i];
        ++n;
      }
    }
    return n;
  }
  void freq_analysis(windowed_freq_table& tabs, string_const_ref_t const& str, size_t offset) {
    for (size_t i = 0; i < str.size(); ++i)
      ++tabs[(offset + i) % tabs.size()][str[i]];
  }

  size_t freq_analysis(windowed_freq_table& tabs, string_const_ref_t const& str, domain_t const& domain, size_t offset) {
    for (auto& c : str) {
      if (domain.count(c)) {
        ++tabs[offset % tabs.size()][c];
        ++offset;
      }
    }
    return offset;
  }

  string_t generate_fuzz(prob_table const& tab, size_t len) {
    string_t ret;
    ret.resize(len);
    std::mt19937 rng;

    {
      thread_local std::random_device seed_rng;
      thread_local std::uniform_int_distribution<decltype(rng)::result_type> seed_dist;
      rng.seed(seed_dist(seed_rng));
    }

    thread_local std::uniform_real_distribution<float_t> dist{0, 1};

    for (auto& rand_char : ret) {
    restat_char:
      float_t stat = dist(rng);
      // Iterate through the table, removing the probabilities until we are within a bracket
      for (auto const& i : tab) {
        if ((stat -= i.second) <= 0) {
          rand_char = i.first;
          goto next_char;
        }
      }
      // This should not happen in normal usage!
      //
      // However, floats do weird things with rounding, so we will be leniant
      goto restat_char;
      next_char: {}
    }

    return ret;
  }

  assoc_table closeness_assoc(prob_table const& observed, prob_table const& expected) {
    // What can we possibly do with this?
    if (observed.size() == 0)
      return {{.observed = 0, .expected = 1}};

    assoc_table assoc;

    // Sort the expected values
    std::vector<prob_t> expected_sorted;
    expected_sorted.reserve(expected_sorted.size());
    for (auto& i : expected)
      expected_sorted.emplace_back(i.second);
    std::sort(expected_sorted.rbegin(), expected_sorted.rend());

    // Sort the observed values
    std::vector<prob_t> observed_sorted;
    observed_sorted.reserve(observed.size());
    for (auto& i : observed)
      observed_sorted.emplace_back(i.second);
    std::sort(observed_sorted.rbegin(), observed_sorted.rend());

    // Fill table with observed values, or zeroes where appropriate
    size_t i;
    if (observed.size() > expected.size()) {
      for (i = 0; i < expected.size(); ++i)
        assoc.emplace_back(assoc_table_elem{.observed = observed_sorted[i], .expected = expected_sorted[i]});
      for (; i < observed.size(); ++i)
        assoc.emplace_back(assoc_table_elem{.observed = observed_sorted[i], .expected = 0.});
    }
    else {
      for (i = 0; i < observed.size(); ++i)
        assoc.emplace_back(assoc_table_elem{.observed = observed_sorted[i], .expected = expected_sorted[i]});
      for (; i < expected.size(); ++i)
        assoc.emplace_back(assoc_table_elem{.observed = 0, .expected = expected_sorted[i]});
    }

    return assoc;
  }

  prob_t closeness_test(prob_table const& observed, prob_table const& expected, freq_t count) {
    // Quick bypass to avoid more filling than we have to
    if (observed.size() > expected.size())
      return 0;

    auto assoc = closeness_assoc(observed, expected);
    prepare_chisq(assoc, count);
    auto stat = run_chisq(assoc, count);

    return 1 - chisq_cdf(assoc.size() - 1, stat);
  }

  prob_t closeness_test(windowed_prob_table const& observed, prob_table const& expected, freq_t count) {
    struct imdt_res_t {
      size_t tab_size;
      float_t chi_sq;
    };

    std::vector<std::future<imdt_res_t>> asyncs(observed.size());

    for (size_t i = 0; i < observed.size(); ++i) {
      asyncs[i] = std::async(std::launch::async, [&, i]() -> imdt_res_t {
        auto assoc = closeness_assoc(observed[i], expected);
        prepare_chisq(assoc, count);

        // We can normalise at the end for efficiency
        return {.tab_size = assoc.size(), .chi_sq = run_chisq(assoc)};
      });
    }

    float_t stat = 0;
    size_t total_tab_len = 0;
    for (auto& i : asyncs) {
      auto res = i.get();
      total_tab_len += res.tab_size;
      stat += res.chi_sq;
    }

    // Finally, we can normalise
    stat *= count;

    return 1 - chisq_cdf(total_tab_len - 1, stat);
  }

  float_t calculate_entropy(std::map<uint8_t, freq_t> const& freqs, size_t len) {
    float_t ret = 0;
    for (auto& i : freqs)  {
      float_t observed_prob = static_cast<float_t>(i.second) / len;
      ret -= observed_prob * ::log2(observed_prob);
      // Some nice printing, kept for future debug
//      std::cout << "  " << std::hex << (int)i.first << ": " << i.second << std::endl;
    }

//    std::cout << ret << std::endl;
    return ret;
  }

  float_t information_content(bytes_const_ref_t b) {
    // TODO:
    // I would prefer something more empirical,
    // rather than just brute forcing the token size

    std::map<uint8_t, freq_t> byte_freqs;
    std::map<uint8_t, freq_t> nybble_freqs;
    std::map<uint8_t, freq_t> pair_freqs;
    std::map<uint8_t, freq_t> bit_freqs;
    for (auto i : b) {
      ++byte_freqs[i];
      ++nybble_freqs[i>>4];
      ++nybble_freqs[i&0xf];
      ++pair_freqs[i&0x3];
      ++pair_freqs[(i>>2)&0x3];
      ++pair_freqs[(i>>4)&0x3];
      ++pair_freqs[(i>>6)&0x3];
      for (int bit = 0; bit < 8; ++bit)
        ++bit_freqs[(i >> bit) & 1];
    }

    return std::min({
                      calculate_entropy(byte_freqs,   b.size()    ) / 8,
                      calculate_entropy(nybble_freqs, b.size() * 2) / 4,
                      calculate_entropy(pair_freqs,   b.size() * 4) / 2,
                      calculate_entropy(bit_freqs,    b.size() * 8)
                    }) * b.size();
  }

  freq_t hamming_weight(uint8_t byte) {
    /// Precomputed hamming weight table for each byte
    static std::array<freq_t, 256> weights {
      0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4,
      1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
      1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
      2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
      1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
      2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
      2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
      3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
      1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
      2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
      2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
      3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
      2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
      3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
      3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
      4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8,
    };

    return weights[byte];
  }

  freq_t hamming_weight(bytes_const_ref_t b) {
    freq_t ret = 0;
    for (auto i : b)
      ret += hamming_weight(i);
    return ret;
  }

  freq_t hamming_distance(bytes_const_ref_t x, bytes_const_ref_t y) {
    if (x.size() != y.size())
      throw std::invalid_argument("Lengths must be the same for hamming distance");
    freq_t ret = 0;
    for (size_t i = 0; i < x.size(); ++i)
      ret += hamming_weight(x[i] ^ y[i]);
    return ret;
  }
}


#include <polyfem/RhsAssembler.hpp>

#include <polyfem/BoundarySampler.hpp>
#include <polyfem/LinearSolver.hpp>
// #include <polyfem/UIState.hpp>

#include <polyfem/Logger.hpp>

#include <Eigen/Sparse>

#ifdef USE_TBB
#include <tbb/parallel_for.h>
#include <tbb/task_scheduler_init.h>
#include <tbb/enumerable_thread_specific.h>
#endif

#include <iostream>
#include <map>
#include <memory>

namespace polyfem
{
	namespace
	{
		class LocalThreadScalarStorage
		{
		public:
			double val;
            ElementAssemblyValues vals;

			LocalThreadScalarStorage()
			{
				val = 0;
			}
		};
	}

	RhsAssembler::RhsAssembler(const Mesh &mesh, const int n_basis, const int size, const std::vector< ElementBases > &bases, const std::vector< ElementBases > &gbases, const std::string &formulation, const Problem &problem)
	: mesh_(mesh), n_basis_(n_basis), size_(size), bases_(bases), gbases_(gbases), formulation_(formulation), problem_(problem)
	{ }

	void RhsAssembler::assemble(Eigen::MatrixXd &rhs, const double t) const
	{
		rhs = Eigen::MatrixXd::Zero(n_basis_ * size_, 1);
		if(!problem_.is_rhs_zero())
		{
			Eigen::MatrixXd rhs_fun;

			const int n_elements = int(bases_.size());
			ElementAssemblyValues vals;
			for(int e = 0; e < n_elements; ++e)
			{
				vals.compute(e, mesh_.is_volume(), bases_[e], gbases_[e]);

				const Quadrature &quadrature = vals.quadrature;

				problem_.rhs(formulation_, vals.val, t, rhs_fun);

				for(int d = 0; d < size_; ++d)
					rhs_fun.col(d) = rhs_fun.col(d).array() * vals.det.array() * quadrature.weights.array();

				const int n_loc_bases_ = int(vals.basis_values.size());
				for(int i = 0; i < n_loc_bases_; ++i)
				{
					const AssemblyValues &v = vals.basis_values[i];

					for(int d = 0; d < size_; ++d)
					{
						const double rhs_value = (rhs_fun.col(d).array() * v.val.array()).sum();
						for(std::size_t ii = 0; ii < v.global.size(); ++ii)
							rhs(v.global[ii].index*size_+d) +=  rhs_value * v.global[ii].val;
					}
				}
			}
		}
	}

	void RhsAssembler::initial_solution(Eigen::MatrixXd &sol) const
	{
		time_bc([&](const Eigen::MatrixXd&pts, Eigen::MatrixXd&val){ problem_.initial_solution(pts, val);}, sol);
	}

	void RhsAssembler::initial_velocity(Eigen::MatrixXd &sol) const
	{
		time_bc([&](const Eigen::MatrixXd&pts, Eigen::MatrixXd&val){ problem_.initial_velocity(pts, val);}, sol);
	}

	void RhsAssembler::initial_acceleration(Eigen::MatrixXd &sol) const
	{
		time_bc([&](const Eigen::MatrixXd&pts, Eigen::MatrixXd&val){ problem_.initial_acceleration(pts, val);}, sol);
	}

	void RhsAssembler::time_bc(const std::function<void(const Eigen::MatrixXd&, Eigen::MatrixXd&)> &fun,Eigen::MatrixXd &sol) const
	{
		sol = Eigen::MatrixXd::Zero(n_basis_ * size_, 1);
		Eigen::MatrixXd loc_sol;

		const int n_elements = int(bases_.size());
        ElementAssemblyValues vals;
		for(int e = 0; e < n_elements; ++e)
		{
			vals.compute(e, mesh_.is_volume(), bases_[e], gbases_[e]);

			const Quadrature &quadrature = vals.quadrature;
			//problem_.initial_solution(vals.val, loc_sol);
			fun(vals.val, loc_sol);

			for(int d = 0; d < size_; ++d)
				loc_sol.col(d) = loc_sol.col(d).array() * vals.det.array() * quadrature.weights.array();

			const int n_loc_bases_ = int(vals.basis_values.size());
			for(int i = 0; i < n_loc_bases_; ++i)
			{
				const AssemblyValues &v = vals.basis_values[i];

				for(int d = 0; d < size_; ++d)
				{
					const double sol_value = (loc_sol.col(d).array() * v.val.array()).sum();
					for(std::size_t ii = 0; ii < v.global.size(); ++ii)
						sol(v.global[ii].index*size_+d) +=  sol_value * v.global[ii].val;
				}
			}
		}
	}

	void RhsAssembler::set_bc(
			const std::function<void(const Eigen::MatrixXi&, const Eigen::MatrixXd&, const Eigen::MatrixXd&, Eigen::MatrixXd &)> &df,
			const std::function<void(const Eigen::MatrixXi&, const Eigen::MatrixXd&, const Eigen::MatrixXd&, Eigen::MatrixXd &)> &nf,
			const std::vector< LocalBoundary > &local_boundary, const std::vector<int> &bounday_nodes, const int resolution, const std::vector< LocalBoundary > &local_neumann_boundary, Eigen::MatrixXd &rhs) const
	{
		const int n_el=int(bases_.size());

		Eigen::MatrixXd uv, samples, gtmp, rhs_fun;
		Eigen::VectorXi global_primitive_ids;

		int index = 0;
		std::vector<int> indices; indices.reserve(n_el*10);
		// std::map<int, int> global_index_to_col;

		long total_size = 0;

		Eigen::Matrix<bool, Eigen::Dynamic, 1> is_boundary(n_basis_); is_boundary.setConstant(false);
		Eigen::VectorXi global_index_to_col(n_basis_); global_index_to_col.setConstant(-1);

		const int actual_dim = problem_.is_scalar() ? 1 : mesh_.dimension();

		// assert((bounday_nodes.size()/actual_dim)*actual_dim == bounday_nodes.size());

		for(int b : bounday_nodes)
			is_boundary[b/actual_dim] = true;

		for(const auto &lb : local_boundary)
		{
			const int e = lb.element_id();
			bool has_samples = sample_boundary(lb, resolution, true, uv, samples, global_primitive_ids);

			if(!has_samples)
				continue;

			const ElementBases &bs = bases_[e];
			const int n_local_bases = int(bs.bases.size());

			total_size += samples.rows();

			for(int j = 0; j < n_local_bases; ++j)
			{
				const Basis &b=bs.bases[j];

				for(std::size_t ii = 0; ii < b.global().size(); ++ii)
				{
					//pt found
					// if(std::find(bounday_nodes.begin(), bounday_nodes.end(), size_ * b.global()[ii].index) != bounday_nodes.end())
					if(is_boundary[b.global()[ii].index])
					{
						//if(global_index_to_col.find( b.global()[ii].index ) == global_index_to_col.end())
						if(global_index_to_col(b.global()[ii].index) == -1)
						{
							// global_index_to_col[b.global()[ii].index] = index++;
							global_index_to_col(b.global()[ii].index) = index++;
							indices.push_back(b.global()[ii].index);
							assert(indices.size() == size_t(index));
						}
					}
				}
			}
		}

		// Eigen::MatrixXd global_mat = Eigen::MatrixXd::Zero(total_size, indices.size());
		Eigen::MatrixXd global_rhs = Eigen::MatrixXd::Zero(total_size, size_);

		const long buffer_size = total_size * long(indices.size());
		std::vector< Eigen::Triplet<double> > entries, entries_t;
		// entries.reserve(buffer_size);
		// entries_t.reserve(buffer_size);

		index = 0;

		int global_counter = 0;
		Eigen::MatrixXd mapped;

		std::vector<AssemblyValues> tmp_val;

		for(const auto &lb : local_boundary)
		{
			const int e = lb.element_id();
			bool has_samples = sample_boundary(lb, resolution, false, uv, samples, global_primitive_ids);

			if(!has_samples)
				continue;

			const ElementBases &bs = bases_[e];
			const ElementBases &gbs = gbases_[e];
			const int n_local_bases = int(bs.bases.size());

			gbs.eval_geom_mapping(samples, mapped);

			bs.evaluate_bases(samples, tmp_val);
			for(int j = 0; j < n_local_bases; ++j)
			{
				const Basis &b=bs.bases[j];
				const auto &tmp = tmp_val[j].val;

				for(std::size_t ii = 0; ii < b.global().size(); ++ii)
				{
					// auto item = global_index_to_col.find(b.global()[ii].index);
					// if(item != global_index_to_col.end()){
					auto item = global_index_to_col(b.global()[ii].index);
					if(item != -1){
						for(int k = 0; k < int(tmp.size()); ++k)
						{
							// entries.push_back(Eigen::Triplet<double>(global_counter+k, item->second, tmp(k, j) * b.global()[ii].val));
							// entries_t.push_back(Eigen::Triplet<double>(item->second, global_counter+k, tmp(k, j) * b.global()[ii].val));
							entries.push_back(Eigen::Triplet<double>(global_counter+k, item, tmp(k) * b.global()[ii].val));
							entries_t.push_back(Eigen::Triplet<double>(item, global_counter+k, tmp(k) * b.global()[ii].val));
						}
						// global_mat.block(global_counter, item->second, tmp.size(), 1) = tmp;
					}
				}
			}

			// problem_.bc(mesh_, global_primitive_ids, mapped, t, rhs_fun);
			df(global_primitive_ids, uv, mapped, rhs_fun);
			global_rhs.block(global_counter, 0, rhs_fun.rows(), rhs_fun.cols()) = rhs_fun;
			global_counter += rhs_fun.rows();

			//UIState::ui_state().debug_data().add_points(mapped, Eigen::MatrixXd::Constant(1, 3, 0));

			//Eigen::MatrixXd asd(mapped.rows(), 3);
			//asd.col(0)=mapped.col(0);
			//asd.col(1)=mapped.col(1);
			//asd.col(2)=rhs_fun;
			//UIState::ui_state().debug_data().add_points(asd, Eigen::MatrixXd::Constant(1, 3, 0));
		}

		assert(global_counter == total_size);

		if(total_size > 0)
		{
			const double mmin = global_rhs.minCoeff();
			const double mmax = global_rhs.maxCoeff();

			if(fabs(mmin) < 1e-8 && fabs(mmax) < 1e-8)
			{
				// std::cout<<"is all zero, skipping"<<std::endl;
				for(size_t i = 0; i < indices.size(); ++i){
					for(int d = 0; d < size_; ++d){
						if(problem_.all_dimentions_dirichelt() || std::find(bounday_nodes.begin(), bounday_nodes.end(), indices[i]*size_+d) != bounday_nodes.end())
							rhs(indices[i]*size_+d) = 0;
					}
				}
			}
			else
			{
				StiffnessMatrix mat(int(total_size), int(indices.size()));
				mat.setFromTriplets(entries.begin(), entries.end());

				StiffnessMatrix mat_t(int(indices.size()), int(total_size));
				mat_t.setFromTriplets(entries_t.begin(), entries_t.end());

				StiffnessMatrix A = mat_t * mat;
				Eigen::MatrixXd b = mat_t * global_rhs;

				Eigen::MatrixXd coeffs(b.rows(), b.cols());

				json params = {
				{"mtype", -2}, // matrix type for Pardiso (2 = SPD)
				// {"max_iter", 0}, // for iterative solvers
				// {"tolerance", 1e-9}, // for iterative solvers
				};

				// auto solver = LinearSolver::create("", "");
				auto solver = LinearSolver::create(LinearSolver::defaultSolver(), LinearSolver::defaultPrecond());
				solver->setParameters(params);
				solver->analyzePattern(A);
				solver->factorize(A);
				for(long i = 0; i < b.cols(); ++i){
					solver->solve(b.col(i), coeffs.col(i));
				}
				logger().trace("RHS solve error {}", (A*coeffs-b).norm());

				for(long i = 0; i < coeffs.rows(); ++i){
					for(int d = 0; d < size_; ++d){
						if(problem_.all_dimentions_dirichelt() || std::find(bounday_nodes.begin(), bounday_nodes.end(), indices[i]*size_+d) != bounday_nodes.end())
							rhs(indices[i]*size_+d) = coeffs(i, d);
					}
				}
			}
		}

		//Neumann
		Eigen::MatrixXd points;
		Eigen::VectorXd weights;

        ElementAssemblyValues vals;
		for(const auto &lb : local_neumann_boundary)
		{
			const int e = lb.element_id();
			bool has_samples = boundary_quadrature(lb, resolution, false, uv, points, weights, global_primitive_ids);

			if(!has_samples)
				continue;

			const ElementBases &gbs = gbases_[e];
			const ElementBases &bs = bases_[e];

			vals.compute(e, mesh_.is_volume(), points, bs, gbs);
			// problem_.neumann_bc(mesh_, global_primitive_ids, vals.val, t, rhs_fun);
			nf(global_primitive_ids, uv, vals.val, rhs_fun);

			// UIState::ui_state().debug_data().add_points(vals.val, Eigen::RowVector3d(0,1,0));

			for(int d = 0; d < size_; ++d)
				rhs_fun.col(d) = rhs_fun.col(d).array() * weights.array();

			for(int i = 0; i < lb.size(); ++i)
			{
				const int primitive_global_id = lb.global_primitive_id(i);
				const auto nodes = bs.local_nodes_for_primitive(primitive_global_id, mesh_);

				for(long n = 0; n < nodes.size(); ++n)
				{
					// const auto &b = bs.bases[nodes(n)];
					const AssemblyValues &v = vals.basis_values[nodes(n)];
					for(int d = 0; d < size_; ++d)
					{
						const double rhs_value = (rhs_fun.col(d).array() * v.val.array()).sum();

						for(size_t g = 0; g < v.global.size(); ++g)
						{
							const int g_index = v.global[g].index*size_+d;
							const bool is_neumann = std::find(bounday_nodes.begin(), bounday_nodes.end(), g_index ) == bounday_nodes.end();

							if(is_neumann){
								rhs(g_index) += rhs_value * v.global[g].val;
								// UIState::ui_state().debug_data().add_points(v.global[g].node, Eigen::RowVector3d(1,0,0));
							}
							// else
								// std::cout<<"skipping "<<g_index<<" "<<rhs_value * v.global[g].val<<std::endl;
						}
					}
				}
			}
		}
	}

	void RhsAssembler::set_bc(const std::vector< LocalBoundary > &local_boundary, const std::vector<int> &bounday_nodes, const int resolution, const std::vector< LocalBoundary > &local_neumann_boundary, Eigen::MatrixXd &rhs, const double t) const
	{
		set_bc(
			[&](const Eigen::MatrixXi &global_ids, const Eigen::MatrixXd &uv, const Eigen::MatrixXd &pts, Eigen::MatrixXd &val){ problem_.bc(mesh_, global_ids, uv, pts, t, val);},
			[&](const Eigen::MatrixXi &global_ids, const Eigen::MatrixXd &uv, const Eigen::MatrixXd &pts, Eigen::MatrixXd &val){ problem_.neumann_bc(mesh_, global_ids, uv, pts, t, val);},
			local_boundary, bounday_nodes, resolution, local_neumann_boundary, rhs);
	}

	void RhsAssembler::set_velocity_bc(const std::vector< LocalBoundary > &local_boundary, const std::vector<int> &bounday_nodes, const int resolution, const std::vector< LocalBoundary > &local_neumann_boundary, Eigen::MatrixXd &rhs, const double t) const
	{
		set_bc(
			[&](const Eigen::MatrixXi &global_ids, const Eigen::MatrixXd &uv, const Eigen::MatrixXd &pts, Eigen::MatrixXd &val){ problem_.velocity_bc(mesh_, global_ids, uv, pts, t, val);},
			[&](const Eigen::MatrixXi &global_ids, const Eigen::MatrixXd &uv, const Eigen::MatrixXd &pts, Eigen::MatrixXd &val){ problem_.neumann_velocity_bc(mesh_, global_ids, uv, pts, t, val);},
			local_boundary, bounday_nodes, resolution, local_neumann_boundary, rhs);
	}

	void RhsAssembler::set_acceleration_bc(const std::vector< LocalBoundary > &local_boundary, const std::vector<int> &bounday_nodes, const int resolution, const std::vector< LocalBoundary > &local_neumann_boundary, Eigen::MatrixXd &rhs, const double t) const
	{
		set_bc(
			[&](const Eigen::MatrixXi &global_ids, const Eigen::MatrixXd &uv, const Eigen::MatrixXd &pts, Eigen::MatrixXd &val){ problem_.acceleration_bc(mesh_, global_ids, uv, pts, t, val);},
			[&](const Eigen::MatrixXi &global_ids, const Eigen::MatrixXd &uv, const Eigen::MatrixXd &pts, Eigen::MatrixXd &val){ problem_.neumann_acceleration_bc(mesh_, global_ids, uv, pts, t, val);},
			local_boundary, bounday_nodes, resolution, local_neumann_boundary, rhs);
	}

	void RhsAssembler::compute_energy_grad(const std::vector< LocalBoundary > &local_boundary, const std::vector<int> &bounday_nodes, const int resolution, const std::vector< LocalBoundary > &local_neumann_boundary, const Eigen::MatrixXd &final_rhs, const double t, Eigen::MatrixXd &rhs) const
	{
		if(problem_.is_linear_in_time()){
			if(problem_.is_time_dependent())
				rhs = final_rhs;
			else
				rhs = final_rhs * t;
		}
		else
		{
			assemble(rhs, t);
			rhs *= -1;
			set_bc(local_boundary, bounday_nodes, resolution, local_neumann_boundary, rhs, t);
		}
	}

	double RhsAssembler::compute_energy(const Eigen::MatrixXd &displacement, const std::vector< LocalBoundary > &local_neumann_boundary, const int resolution, const double t) const
	{
		Eigen::Matrix<double, Eigen::Dynamic, 1, 0, 3, 1> local_displacement(size_);

		double res = 0;
		Eigen::MatrixXd forces;

		if(!problem_.is_rhs_zero())
		{
#ifdef USE_TBB
		typedef tbb::enumerable_thread_specific< LocalThreadScalarStorage > LocalStorage;
		LocalStorage storages((LocalThreadScalarStorage()));
#else
		LocalThreadScalarStorage loc_storage;
#endif

		const int n_bases = int(bases_.size());

#ifdef USE_TBB
		tbb::parallel_for( tbb::blocked_range<int>(0, n_bases), [&](const tbb::blocked_range<int> &r) {
		LocalStorage::reference loc_storage = storages.local();
		for (int e = r.begin(); e != r.end(); ++e) {
#else
		for(int e=0; e < n_bases; ++e) {
#endif
			ElementAssemblyValues &vals = loc_storage.vals;
			vals.compute(e, mesh_.is_volume(), bases_[e], gbases_[e]);

			const Quadrature &quadrature = vals.quadrature;
			const Eigen::VectorXd da = vals.det.array() * quadrature.weights.array();

			problem_.rhs(formulation_, vals.val, t, forces);
			assert(forces.rows() == da.size());
			assert(forces.cols() == size_);

			for(long p = 0; p < da.size(); ++p)
			{
				local_displacement.setZero();

				for(size_t i = 0; i < vals.basis_values.size(); ++i)
				{
					const auto &bs = vals.basis_values[i];
					assert(bs.val.size() == da.size());
					const double b_val = bs.val(p);

					for(int d = 0; d < size_; ++d)
					{
						for(std::size_t ii = 0; ii < bs.global.size(); ++ii)
						{
							local_displacement(d) += (bs.global[ii].val * b_val) * displacement(bs.global[ii].index*size_ + d);
						}
					}
				}

				for(int d = 0; d < size_; ++d)
					loc_storage.val += forces(p, d) * local_displacement(d) * da(p);
					// res += forces(p, d) * local_displacement(d) * da(p);
			}
#ifdef USE_TBB
		}});
#else
		}
#endif

#ifdef USE_TBB
		for (LocalStorage::iterator i = storages.begin(); i != storages.end();  ++i)
		{
			res += i->val;
		}
#else
		res = loc_storage.val;
#endif
		}

		ElementAssemblyValues vals;
		//Neumann
		Eigen::MatrixXd points, uv;
		Eigen::VectorXd weights;
		Eigen::VectorXi global_primitive_ids;
		for(const auto &lb : local_neumann_boundary)
		{
			const int e = lb.element_id();
			bool has_samples = boundary_quadrature(lb, resolution, false, uv, points, weights, global_primitive_ids);

			if(!has_samples)
				continue;

			const ElementBases &gbs = gbases_[e];
			const ElementBases &bs = bases_[e];

			vals.compute(e, mesh_.is_volume(), points, bs, gbs);
			problem_.neumann_bc(mesh_, global_primitive_ids, uv, vals.val, t, forces);

			// UIState::ui_state().debug_data().add_points(vals.val, Eigen::RowVector3d(1,0,0));

			for(long p = 0; p < weights.size(); ++p)
			{
				local_displacement.setZero();

				for(size_t i = 0; i < vals.basis_values.size(); ++i)
				{
					const auto &vv = vals.basis_values[i];
					assert(vv.val.size() == weights.size());
					const double b_val = vv.val(p);

					for(int d = 0; d < size_; ++d)
					{
						for(std::size_t ii = 0; ii < vv.global.size(); ++ii)
						{
							local_displacement(d) += (vv.global[ii].val * b_val) * displacement(vv.global[ii].index*size_ + d);
						}
					}
				}

				for(int d = 0; d < size_; ++d)
					res -= forces(p, d) * local_displacement(d) * weights(p);
			}
		}

		return res;
	}

	bool RhsAssembler::boundary_quadrature(const LocalBoundary &local_boundary, const int order, const bool skip_computation, Eigen::MatrixXd &uv, Eigen::MatrixXd &points, Eigen::VectorXd &weights, Eigen::VectorXi &global_primitive_ids) const
	{
		uv.resize(0, 0);
		points.resize(0, 0);
		weights.resize(0);
		global_primitive_ids.resize(0);

		for(int i = 0; i < local_boundary.size(); ++i)
		{
			const int gid = local_boundary.global_primitive_id(i);
			Eigen::MatrixXd tmp_p, tmp_uv;
			Eigen::VectorXd tmp_w;
			switch(local_boundary.type())
			{
				case BoundaryType::TriLine:	 BoundarySampler::quadrature_for_tri_edge(local_boundary[i], order, gid, mesh_, tmp_uv, tmp_p, tmp_w); break;
				case BoundaryType::QuadLine: BoundarySampler::quadrature_for_quad_edge(local_boundary[i], order, gid, mesh_, tmp_uv, tmp_p, tmp_w); break;
				case BoundaryType::Quad: 	 BoundarySampler::quadrature_for_quad_face(local_boundary[i], order, gid, mesh_, tmp_uv, tmp_p, tmp_w); break;
				case BoundaryType::Tri: 	 BoundarySampler::quadrature_for_tri_face(local_boundary[i], order, gid, mesh_, tmp_uv, tmp_p, tmp_w); break;
				case BoundaryType::Invalid:  assert(false); break;
				default: assert(false);
			}

			uv.conservativeResize(uv.rows() + tmp_uv.rows(), tmp_uv.cols());
			uv.bottomRows(tmp_uv.rows()) = tmp_uv;

			points.conservativeResize(points.rows() + tmp_p.rows(), tmp_p.cols());
			points.bottomRows(tmp_p.rows()) = tmp_p;

			weights.conservativeResize(weights.rows() + tmp_w.rows(), tmp_w.cols());
			weights.bottomRows(tmp_w.rows()) = tmp_w;

			global_primitive_ids.conservativeResize(global_primitive_ids.rows() + tmp_p.rows());
			global_primitive_ids.bottomRows(tmp_p.rows()).setConstant(gid);
		}

		assert(uv.rows() == global_primitive_ids.size());
		assert(points.rows() == global_primitive_ids.size());
		assert(weights.size() == global_primitive_ids.size());

		return true;
	}

	bool RhsAssembler::sample_boundary(const LocalBoundary &local_boundary, const int n_samples, const bool skip_computation, Eigen::MatrixXd &uv, Eigen::MatrixXd &samples, Eigen::VectorXi &global_primitive_ids) const
	{
		uv.resize(0, 0);
		samples.resize(0, 0);
		global_primitive_ids.resize(0);

		for(int i = 0; i < local_boundary.size(); ++i)
		{
			Eigen::MatrixXd tmp, tmp_uv;
			switch(local_boundary.type())
			{
				case BoundaryType::TriLine:	 BoundarySampler::sample_parametric_tri_edge(local_boundary[i], n_samples, tmp_uv, tmp); break;
				case BoundaryType::QuadLine: BoundarySampler::sample_parametric_quad_edge(local_boundary[i], n_samples, tmp_uv, tmp); break;
				case BoundaryType::Quad: 	 BoundarySampler::sample_parametric_quad_face(local_boundary[i], n_samples, tmp_uv, tmp); break;
				case BoundaryType::Tri: 	 BoundarySampler::sample_parametric_tri_face(local_boundary[i], n_samples, tmp_uv, tmp); break;
				case BoundaryType::Invalid:  assert(false); break;
				default: assert(false);
			}

			uv.conservativeResize(uv.rows() + tmp_uv.rows(), tmp_uv.cols());
			uv.bottomRows(tmp_uv.rows()) = tmp_uv;

			samples.conservativeResize(samples.rows() + tmp.rows(), tmp.cols());
			samples.bottomRows(tmp.rows()) = tmp;

			global_primitive_ids.conservativeResize(global_primitive_ids.rows() + tmp.rows());
			global_primitive_ids.bottomRows(tmp.rows()).setConstant(local_boundary.global_primitive_id(i));
		}

		assert(uv.rows() == global_primitive_ids.size());
		assert(samples.rows() == global_primitive_ids.size());

		return true;
	}

}