#include "MCHMM.h"
#include <cmath>
#include "Sampler.h"
#include <random>
#include <chrono>
#include <glog/logging.h>
#include <cassert>
#include <boost/math/distributions/students_t.hpp>
using namespace std;
using namespace google;

MCHMM::MCHMM(){
    pi = new vector<Sample>();
    m = new vector<Sample>();
    v = new vector<Sample>();
}

/**
 * @brief MCHMM::forward Goes one step forward in time according to the HMM distributions
 * @param observations Observations needed to do the reasoning based on the HMM
 * @param N Number of samples used in the resampling step of any sampling inside this method
 * @return Alpha DETree for the next time step distribution
 */
DETree * MCHMM::forward(vector<Observation> *observations, size_t N){
    vector<Sample> alpha_samples[2];
    Sampler sampler;
    size_t T = observations->size();

    alpha_samples[0] = sampler.resample_from(pi_tree, N);

    // STEP 2
    size_t t = 0;
    for (t = 1; t < T; t++){
        // STEP 2(a)
        vector<Sample> temp = sampler.likelihood_weighted_resampler(alpha_samples[(t - 1) % 2], N);
        double sum_densities = 0.0;

        for (size_t i = 0; i < temp.size(); i++){
            // STEP 2(b)
            Sample x = sampler.sample_given(m_tree, temp[i]);

            for (size_t i = 0; i < temp[i].size(); i++){
                x.values.pop_back();
            }

            // STEP 2(c)
            Sample v_temp = (*observations)[t].combine(x);
            double density = v_tree->density_value(v_temp, rho);
            x.p = density;

            sum_densities += density;
            temp[i] = x;
        }

        // Normalizing the probabilities
        for (size_t i = 0; i < temp.size(); i++){
            temp[i].p = temp[i].p / sum_densities;
        }

        // STEP 2(d)
        alpha_samples[t % 2] = temp;
    }

    vector<Sample> temp = sampler.likelihood_weighted_resampler(alpha_samples[(t - 1) % 2], N);
    for (size_t i = 0; i < temp.size(); i++){
        // STEP 2(b)
        Sample x = sampler.sample_given(m_tree, temp[i]);
        x.p = 1.0 / temp.size();
        temp[i] = x;
    }

    return new DETree(temp, pi_low_limit, pi_high_limit);
}

/**
 * @brief MCHMM::gamma It perform the forward-backward algorithms and combine the results to get the
 * gamma distribution, it is useful mostly for getting the most probable states
 * @param observations Observations needed to do the reasoning based on the HMM
 * @param N Number of samples used in the resampling step of any sampling inside this method
 * @return Gamma DETrees for each observation
 */
vector<DETree *> MCHMM::gamma(vector<Observation> *observations, size_t N){
    Sampler sampler;
    size_t T = observations->size();

    vector<Sample> alpha_samples[2];
    vector<Sample> beta_samples[2];

    vector<DETree*> alpha_trees;
    vector<DETree*> beta_trees;
    vector<DETree*> gamma_trees;

    // STEP 1
    alpha_samples[0] = sampler.resample_from(pi_tree, N);
    alpha_trees.push_back(new DETree(alpha_samples[0], pi_low_limit, pi_high_limit));

    // STEP 2
    for (size_t t = 1; t < T; t++){
        // STEP 2(a)
        vector<Sample> temp = sampler.likelihood_weighted_resampler(alpha_samples[(t - 1) % 2], N);
        double sum_densities = 0.0;

        for (size_t i = 0; i < temp.size(); i++){
            // STEP 2(b)
            Sample x = sampler.sample_given(m_tree, temp[i]);

            for (size_t i = 0; i < temp[i].size(); i++){
                x.values.pop_back();
            }

            // STEP 2(c)
            Sample v_temp = (*observations)[t].combine(x);
            double density = v_tree->density_value(v_temp, rho);
            x.p = density;

            sum_densities += density;
            temp[i] = x;
        }

        // Normalizing the probabilities
        for (size_t i = 0; i < temp.size(); i++){
            temp[i].p = temp[i].p / sum_densities;
        }

        // STEP 2(d)
        alpha_samples[t % 2] = temp;
        alpha_trees.push_back(new DETree(temp, pi_low_limit, pi_high_limit));
    }

    // STEP 3
    beta_samples[0] = sampler.uniform_sampling(pi_low_limit, pi_high_limit, N);
    beta_trees.push_back(new DETree(beta_samples[0], pi_low_limit, pi_high_limit));

    // STEP 4
    for (size_t t = T - 1; t >= 1; t--){
        // STEP 4(a)
        int index_t = ((T) - (t + 1)) % 2;
        vector<Sample> temp = sampler.likelihood_weighted_resampler(beta_samples[index_t], N);
        double sum_densities = 0.0;

        for (size_t i = 0; i < temp.size(); i++){
            // STEP 4(b)
            Sample x = sampler.sample_given(m_tree, temp[i]);

            for (size_t i = 0; i < temp[i].size(); i++){
                x.values.pop_back();
            }

            // STEP 4(c)
            Sample v_temp = (*observations)[t].combine(x);
            double density = v_tree->density_value(v_temp, rho);
            x.p = density;

            sum_densities += density;
            temp[i] = x;
        }

        // Normalizing the probabilities
        for (size_t i = 0; i < temp.size(); i++){
            temp[i].p = temp[i].p / sum_densities;
        }

        // STEP 4(d)
        beta_samples[(index_t + 1) % 2] = temp;
        beta_trees.push_back(new DETree(temp, pi_low_limit, pi_high_limit));
    }

    // STEP 5
    for (size_t t = 0; t < T; t++){
        vector<Sample> temp;
        double sum_density = 0.0;

        int index_t = (T) - (t + 1);

        // STEP 5(a)
        for (size_t j = 0; j < N / 2; j++){
            Sample sample = sampler.sample(alpha_trees[t]);
            sample.p = (*beta_trees[index_t]).density_value(sample, rho);
            sum_density += sample.p;
            temp.push_back(sample);
        }

        // STEP 5(b)
        for (size_t j = 0; j < N - (N / 2); j++){
            Sample sample = sampler.sample(beta_trees[index_t]);
            sample.p = (*alpha_trees[t]).density_value(sample, rho);
            sum_density += sample.p;
            temp.push_back(sample);
        }

        // Normalizing the probabilities
        for (size_t i = 0; i < temp.size(); i++){
            temp[i].p = temp[i].p / sum_density;
        }

        gamma_trees.push_back(new DETree(temp, pi_low_limit, pi_high_limit));
    }

    alpha_samples[0].clear();
    alpha_samples[1].clear();
    beta_samples[0].clear();
    beta_samples[1].clear();

    for (size_t i = 0; i < alpha_trees.size(); i++){
        delete alpha_trees[i];
    }

    for (size_t i = 0; i < beta_trees.size(); i++){
        delete beta_trees[i];
    }

    assert(gamma_trees.size() == T);

    return gamma_trees;
}

/**
 * @brief MCHMM::learn_hmm it takes in some observations and perform the EM as many iterations as the max_iteration arg
 * @param observations Observations needed for learning the HMM distributions
 * @param max_iteration Maximum number of iterations performed for the EM
 * @param N Number of samples used in the resampling step of any sampling inside this method
 */
void MCHMM::learn_hmm_KL(vector<Observation> *observations, double threshold, size_t max_iteration, int N){
    if (observations->size() < 2){
        LOG(ERROR) << "Not enough observation data!";
        return;
    }

    Sampler sampler;
    size_t T = observations->size();

    if (pi->size() < 1 || v->size() < 1 || m->size() < 1){
        LOG(INFO) << "Init HMM Randomly!";
        init_hmm_randomly(N, N, N);
    }

    bool cond = true;
    size_t iteration = 0;
    DETree* old_gamma_tree = NULL;

    vector<Sample> test_samples = sampler.uniform_sampling(pi_low_limit, pi_high_limit, 1000);

    while (cond){
        vector<Sample> alpha_samples[2];
        vector<Sample> beta_samples[2];

        vector<DETree*> alpha_trees;
        vector<DETree*> beta_trees;
        vector<DETree*> gamma_trees;

        /////////////////E STEP/////////////////
        {
            // STEP 1
            alpha_samples[0] = sampler.resample_from(pi_tree, N);
            alpha_trees.push_back(new DETree(alpha_samples[0], pi_low_limit, pi_high_limit));

            // STEP 2
            for (size_t t = 1; t < T; t++){
                // STEP 2(a)
                vector<Sample> temp = sampler.likelihood_weighted_resampler(alpha_samples[(t - 1) % 2], N);
                double sum_densities = 0.0;

                for (size_t i = 0; i < temp.size(); i++){
                    // STEP 2(b)
                    Sample x = sampler.sample_given(m_tree, temp[i]);

                    for (size_t i = 0; i < temp[i].size(); i++){
                        x.values.pop_back();
                    }

                    // STEP 2(c)
                    Sample v_temp = (*observations)[t].combine(x);
                    double density = v_tree->density_value(v_temp, rho);
                    x.p = density;

                    sum_densities += density;
                    temp[i] = x;
                }

                // Normalizing the probabilities
                for (size_t i = 0; i < temp.size(); i++){
                    temp[i].p = temp[i].p / sum_densities;
                }

                // STEP 2(d)
                alpha_samples[t % 2] = temp;
                alpha_trees.push_back(new DETree(temp, pi_low_limit, pi_high_limit));
            }

            // STEP 3
            beta_samples[0] = sampler.uniform_sampling(pi_low_limit, pi_high_limit, N);
            beta_trees.push_back(new DETree(beta_samples[0], pi_low_limit, pi_high_limit));

            // STEP 4
            for (size_t t = T - 1; t >= 1; t--){
                // STEP 4(a)
                int index_t = ((T) - (t + 1)) % 2;
                vector<Sample> temp = sampler.likelihood_weighted_resampler(beta_samples[index_t], N);
                double sum_densities = 0.0;

                for (size_t i = 0; i < temp.size(); i++){
                    // STEP 4(b)
                    Sample x = sampler.sample_given(m_tree, temp[i]);

                    for (size_t i = 0; i < temp[i].size(); i++){
                        x.values.pop_back();
                    }

                    // STEP 4(c)
                    Sample v_temp = (*observations)[t].combine(x);
                    double density = v_tree->density_value(v_temp, rho);
                    x.p = density;

                    sum_densities += density;
                    temp[i] = x;
                }

                // Normalizing the probabilities
                for (size_t i = 0; i < temp.size(); i++){
                    temp[i].p = temp[i].p / sum_densities;
                }

                // STEP 4(d)
                beta_samples[(index_t + 1) % 2] = temp;
                beta_trees.push_back(new DETree(temp, pi_low_limit, pi_high_limit));
            }

            // STEP 5
            for (size_t t = 1; t < T; t++){
                vector<Sample> temp;
                double sum_density = 0.0;

                int index_t = (T) - (t);

                // STEP 5(a)
                for (int j = 0; j < N / 2; j++){
                    Sample sample = sampler.sample(alpha_trees[t]);
                    sample.p = (*beta_trees[index_t]).density_value(sample, rho);
                    sum_density += sample.p;
                    temp.push_back(sample);
                }

                // STEP 5(b)
                for (int j = 0; j < N - (N / 2); j++){
                    Sample sample = sampler.sample(beta_trees[index_t]);
                    sample.p = (*alpha_trees[t]).density_value(sample, rho);
                    sum_density += sample.p;
                    temp.push_back(sample);
                }

                // Normalizing the probabilities
                for (size_t i = 0; i < temp.size(); i++){
                    temp[i].p = temp[i].p / sum_density;
                }

                gamma_trees.push_back(new DETree(temp, pi_low_limit, pi_high_limit));
            }

            alpha_samples[0].clear();
            alpha_samples[1].clear();
            beta_samples[0].clear();
            beta_samples[1].clear();

            LOG(INFO) << "End of E Step at iteration: " << iteration;
        }

        /////////////////M STEP/////////////////
        {
            unsigned seed = std::chrono::system_clock::now().time_since_epoch().count();
            std::default_random_engine gen(seed);

            vector<Sample> * temp_m = new vector<Sample>();
            vector<Sample> * temp_v = new vector<Sample>();
            vector<Sample> * temp_pi = new vector<Sample>();

            // STEP 1
            for (int i = 0; i < N; i++){
                uniform_real_distribution<double> dist(1, T - 2);
                int t = dist(gen);

                Sample x = sampler.sample(gamma_trees[t]);
                Sample x_prime = sampler.sample(gamma_trees[t + 1]);

                Sample temp = x.combine(x_prime.values);
                temp.p = 1.0 / N;
                temp_m->push_back(temp);
            }

            // STEP 2
            for (int i = 0; i < N; i++){
                uniform_real_distribution<double> dist(1, T - 1);
                int t = dist(gen);

                Sample x = sampler.sample(gamma_trees[t]);

                Sample temp = (*observations)[t].combine(x);
                temp.p = 1.0 / N;
                temp_v->push_back(temp);
            }

            // STEP 3
            int pi_size = pi->size();
            for (int i = 0; i < pi_size; i++){
                temp_pi->push_back(sampler.sample(gamma_trees[0]));
            }

            delete pi;
            pi = temp_pi;
            pi_tree->create_tree(*pi, pi_low_limit, pi_high_limit);

            delete m;
            m =  temp_m;
            m_tree->create_tree(*m, m_low_limit, m_high_limit);

            delete v;
            v = temp_v;
            v_tree->create_tree(*v, v_low_limit, v_high_limit);

            LOG(INFO) << "End of M Step at iteration: " << iteration;

        }

        /////////////////ANNEALING/////////////////
        if (rho > 0.01)
            rho = rho * rho_bar;

        /////////////////SAMPLE SET SIZE/////////////////
        if (N < (int)max_sample_size)
            N = N; // * eta;

        /////////////////STOP CONDITION/////////////////
        {
            LOG(INFO) << "Iteration " << iteration + 1 << " Finished!" << "\n";
            iteration++;
            if (iteration >= max_iteration){
                cond = false;
            }

            if (iteration > 1){ // Do the KL if we have at least on previous HMM parameters set!!!
                // Generate a lot of samples uniformly

                // Find the density estimations for each generated sample
                vector<double> estimates_old;
                vector<double> estimates_new;
                double sum_old = 0.0;
                double sum_new = 0.0;
                for (size_t r = 0; r < test_samples.size(); r++){
                    estimates_old.push_back(old_gamma_tree->density_value(test_samples[r], 0.5));
                    estimates_new.push_back(gamma_trees.back()->density_value(test_samples[r], 0.5));

                    sum_old += estimates_old.back();
                    sum_new += estimates_new.back();
                }

                // Normalize the density values
                for (size_t r = 0; r < test_samples.size(); r++){
                    estimates_old[r] = estimates_old[r] / sum_old;
                    estimates_new[r] = estimates_new[r] / sum_new;
                }

                // Compute the KL divergence factor
                double KLD = KLD_compute(estimates_old, estimates_new);

                LOG(ERROR) << "KLD: " << KLD;

                // If KLD < threshold --> STOP
                if (KLD < threshold){
                    cond = false;
                }
            }
        }

        for (size_t i = 0; i < alpha_trees.size(); i++){
            delete alpha_trees[i];
        }

        for (size_t i = 0; i < beta_trees.size(); i++){
            delete beta_trees[i];
        }

        for (size_t i = 0; i < gamma_trees.size() - 1; i++){
            delete gamma_trees[i];
        }

        if (old_gamma_tree){
            delete old_gamma_tree;
        }

        old_gamma_tree = gamma_trees.back();
    }

    initialized = true;
}

/**
 * @brief MCHMM::KLD_compute computes the Kullback-Leibler divergence of two distributions
 * @param P The true distribution (in this application, the true is our old distribution)
 * @param Q The estimated distribution (in this application, the estimated is our new distribution)
 * @return KLD value
 */
double MCHMM::  KLD_compute(vector<double> P, vector<double> Q){
    double KLD = 0.0;
    for (size_t i = 0; i < P.size(); i++){
        KLD += P[i] * std::log(P[i] / Q[i]);
    }
    return KLD;
}

vector<Sample> MCHMM::get_uniform_samples_from_pi(size_t N){
    Sampler sampler;
    return sampler.uniform_sampling(pi_low_limit, pi_high_limit, N);
}

/**
 * @brief MCHMM::learn_hmm it takes in some observations and perform the EM as many iterations as the max_iteration arg
 * @param observations Observations needed for learning the HMM distributions
 * @param max_iteration Maximum number of iterations performed for the EM
 * @param N Number of samples used in the resampling step of any sampling inside this method
 */
void MCHMM::learn_hmm(vector<Observation> *observations, size_t max_iteration, int N){

    if (observations->size() < 2){
        LOG(ERROR) << "Not enough observation data!";
        return;
    }

    Sampler sampler;
    size_t T = observations->size();

    if (pi->size() < 1 || v->size() < 1 || m->size() < 1){
        LOG(INFO) << "Init HMM Randomly!";
        init_hmm_randomly(N, N, N);
    }

    bool cond = true;
    size_t iteration = 0;

    while (cond){
        vector<Sample> alpha_samples[2];
        vector<Sample> beta_samples[2];

        vector<DETree*> alpha_trees;
        vector<DETree*> beta_trees;
        vector<DETree*> gamma_trees;

        /////////////////E STEP/////////////////
        {
            // STEP 1
            alpha_samples[0] = sampler.resample_from(pi_tree, N);
            alpha_trees.push_back(new DETree(alpha_samples[0], pi_low_limit, pi_high_limit));

            // STEP 2
            for (size_t t = 1; t < T; t++){
                // STEP 2(a)
                vector<Sample> temp = sampler.likelihood_weighted_resampler(alpha_samples[(t - 1) % 2], N);
                double sum_densities = 0.0;

                for (size_t i = 0; i < temp.size(); i++){
                    // STEP 2(b)
                    Sample x = sampler.sample_given(m_tree, temp[i]);

                    for (size_t i = 0; i < temp[i].size(); i++){
                        x.values.pop_back();
                    }

                    // STEP 2(c)
                    Sample v_temp = (*observations)[t].combine(x);
                    double density = v_tree->density_value(v_temp, rho);
                    x.p = density;

                    sum_densities += density;
                    temp[i] = x;
                }

                // Normalizing the probabilities
                for (size_t i = 0; i < temp.size(); i++){
                    temp[i].p = temp[i].p / sum_densities;
                }

                // STEP 2(d)
                alpha_samples[t % 2] = temp;
                alpha_trees.push_back(new DETree(temp, pi_low_limit, pi_high_limit));
            }

            // STEP 3
            beta_samples[0] = sampler.uniform_sampling(pi_low_limit, pi_high_limit, N);
            beta_trees.push_back(new DETree(beta_samples[0], pi_low_limit, pi_high_limit));

            // STEP 4
            for (size_t t = T - 1; t >= 1; t--){
                // STEP 4(a)
                int index_t = ((T) - (t + 1)) % 2;
                vector<Sample> temp = sampler.likelihood_weighted_resampler(beta_samples[index_t], N);
                double sum_densities = 0.0;

                for (size_t i = 0; i < temp.size(); i++){
                    // STEP 4(b)
                    Sample x = sampler.sample_given(m_tree, temp[i]);

                    for (size_t i = 0; i < temp[i].size(); i++){
                        x.values.pop_back();
                    }

                    // STEP 4(c)
                    Sample v_temp = (*observations)[t].combine(x);
                    double density = v_tree->density_value(v_temp, rho);
                    x.p = density;

                    sum_densities += density;
                    temp[i] = x;
                }

                // Normalizing the probabilities
                for (size_t i = 0; i < temp.size(); i++){
                    temp[i].p = temp[i].p / sum_densities;
                }

                // STEP 4(d)
                beta_samples[(index_t + 1) % 2] = temp;
                beta_trees.push_back(new DETree(temp, pi_low_limit, pi_high_limit));
            }

            // STEP 5
            for (size_t t = 1; t < T; t++){
                vector<Sample> temp;
                double sum_density = 0.0;

                int index_t = (T) - (t);

                // STEP 5(a)
                for (int j = 0; j < N / 2; j++){
                    Sample sample = sampler.sample(alpha_trees[t]);
                    sample.p = (*beta_trees[index_t]).density_value(sample, rho);
                    sum_density += sample.p;
                    temp.push_back(sample);
                }

                // STEP 5(b)
                for (int j = 0; j < N - (N / 2); j++){
                    Sample sample = sampler.sample(beta_trees[index_t]);
                    sample.p = (*alpha_trees[t]).density_value(sample, rho);
                    sum_density += sample.p;
                    temp.push_back(sample);
                }

                // Normalizing the probabilities
                for (size_t i = 0; i < temp.size(); i++){
                    temp[i].p = temp[i].p / sum_density;
                }

                gamma_trees.push_back(new DETree(temp, pi_low_limit, pi_high_limit));
            }

            alpha_samples[0].clear();
            alpha_samples[1].clear();
            beta_samples[0].clear();
            beta_samples[1].clear();

            LOG(INFO) << "End of E Step at iteration: " << iteration;
        }

        /////////////////M STEP/////////////////
        {
            unsigned seed = std::chrono::system_clock::now().time_since_epoch().count();
            std::default_random_engine gen(seed);

            vector<Sample> * temp_m = new vector<Sample>();
            vector<Sample> * temp_v = new vector<Sample>();
            vector<Sample> * temp_pi = new vector<Sample>();

            // STEP 1
            for (int i = 0; i < N; i++){
                uniform_real_distribution<double> dist(1, T - 2);
                int t = dist(gen);

                Sample x = sampler.sample(gamma_trees[t]);
                Sample x_prime = sampler.sample(gamma_trees[t + 1]);

                Sample temp = x.combine(x_prime.values);
                temp.p = 1.0 / N;
                temp_m->push_back(temp);
            }

            // STEP 2
            for (int i = 0; i < N; i++){
                uniform_real_distribution<double> dist(1, T - 1);
                int t = dist(gen);

                Sample x = sampler.sample(gamma_trees[t]);

                Sample temp = (*observations)[t].combine(x);
                temp.p = 1.0 / N;
                temp_v->push_back(temp);
            }

            // STEP 3
            int pi_size = pi->size();
            for (int i = 0; i < pi_size; i++){
                temp_pi->push_back(sampler.sample(gamma_trees[0]));
            }

            delete pi;
            pi = temp_pi;
            pi_tree->create_tree(*pi, pi_low_limit, pi_high_limit);

            delete m;
            m =  temp_m;
            m_tree->create_tree(*m, m_low_limit, m_high_limit);

            delete v;
            v = temp_v;
            v_tree->create_tree(*v, v_low_limit, v_high_limit);

            LOG(INFO) << "End of M Step at iteration: " << iteration;

        }

        /////////////////ANNEALING/////////////////
        if (rho > 0.01)
            rho = rho * rho_bar;

        /////////////////SAMPLE SET SIZE/////////////////
        if (N < (int)max_sample_size)
            N = N; // * eta;

        /////////////////STOP CONDITION/////////////////
        LOG(INFO) << "Iteration " << iteration + 1 << " Finished!" << "\n";
        iteration++;
        if (iteration >= max_iteration){
            cond = false;
        }

        for (size_t i = 0; i < alpha_trees.size(); i++){
            delete alpha_trees[i];
        }

        for (size_t i = 0; i < beta_trees.size(); i++){
            delete beta_trees[i];
        }

        for (size_t i = 0; i < gamma_trees.size(); i++){
            delete gamma_trees[i];
        }
    }

    initialized = true;
}

/**
 * @brief MCHMM::set_distributions Instead of learning the distributions one can use this to prime the HMM with pre-collected samples
 * @param pi Samples collected for the initial state distribution PI
 * @param m Samples collected for the transition distribution M
 * @param v Samples collected for the observation distribution NU
 * @param rho Amount of effect that different levels of the DETree have on the computed density value (default: 0.5)
 */
void MCHMM::set_distributions(vector<Sample> *pi, vector<Sample> *m, vector<Sample> *v, double rho){
    this->pi = new vector<Sample>();
    this->m = new vector<Sample>();
    this->v = new vector<Sample>();

    for (size_t i = 0; i < pi->size(); i++){
        this->pi->push_back((*pi)[i]);
    }

    for (size_t i = 0; i < m->size(); i++){
        this->m->push_back((*m)[i]);
    }

    for (size_t i = 0; i < v->size(); i++){
        this->v->push_back((*v)[i]);
    }

    this->rho = rho;

    pi_tree = new DETree(*pi, pi_low_limit, pi_high_limit);
    m_tree = new DETree(*m, m_low_limit, m_high_limit);
    v_tree = new DETree(*v, v_low_limit, v_high_limit);

    initialized = true;
}

void MCHMM::set_limits(vector<double> *pi_low_limit, vector<double> *pi_high_limit,
                       vector<double> *m_low_limit, vector<double> *m_high_limit,
                       vector<double> *v_low_limit, vector<double> *v_high_limit
                       )
{
    this->pi_low_limit = pi_low_limit;
    this->pi_high_limit = pi_high_limit;

    this->m_low_limit = m_low_limit;
    this->m_high_limit = m_high_limit;

    this->v_low_limit = v_low_limit;
    this->v_high_limit = v_high_limit;
}

void MCHMM::init_hmm_randomly(int sample_size_pi, int sample_size_m, int sample_size_v){

    if (pi_low_limit == NULL){
        LOG(FATAL) << "Please set the limits first and then run this method!!!";
    }

    for (int i = 0; i < sample_size_pi; i++){
        Sample sample;
        sample.init_rand(pi_low_limit, pi_high_limit);
        sample.p = 1.0 / sample_size_pi;
        pi->push_back(sample);
    }

    for (int i = 0; i < sample_size_m; i++){
        Sample sample;
        sample.init_rand(m_low_limit, m_high_limit);
        sample.p = 1.0 / sample_size_m;
        m->push_back(sample);
    }

    for (int i = 0; i < sample_size_v; i++){
        Sample sample;
        sample.init_rand(v_low_limit, v_high_limit);
        sample.p = 1.0 / sample_size_v;
        v->push_back(sample);
    }

    pi_tree = new DETree(*pi, pi_low_limit, pi_high_limit);
    m_tree = new DETree(*m, m_low_limit, m_high_limit);
    v_tree = new DETree(*v, v_low_limit, v_high_limit);

    initialized = true;
}

double MCHMM::_rho(){
    return this->rho;
}

bool MCHMM::initialized_(){
    return initialized;
}

DETree* MCHMM::pi_tree_(){
    return pi_tree;
}


#ifdef _WIN32
#pragma warning(disable:4503)
#pragma warning(push)
#pragma warning(disable:4996 4251 4275 4800 4190 4244)
#endif
#include <vector>
#include "opencv2/imgproc/imgproc.hpp"
#include "opencv2/imgcodecs.hpp"

#ifdef _MSC_VER
#include <intrin.h>
#define __builtin_popcount __popcnt
#endif

#ifdef _WIN32
#pragma warning(pop)
#endif

#include <boost/math/special_functions/erf.hpp>
#include "Block.h"
#include "ParamValidator.h"

using std::vector;
using std::string;
using cv::Mat;

namespace charliesoft
{
  BLOCK_BEGIN_INSTANTIATION(CORE_filter);
  inline double hamming_distance(char* vals1, char* vals2, int nbVals)
  {
    int output = 0;
    int i;
    int nbValsPacked = nbVals - 3;
    for (i = 0; i < nbValsPacked; i += 4) {
      output += __builtin_popcount(
        *((uint32_t *)(&vals1[i])) ^ *((uint32_t *)(&vals2[i])));
    }
    //if nbVals is not 4bytes packed
    for (; i < nbVals; i++)
      output += __builtin_popcount((unsigned char)vals1[i] ^ (unsigned char)vals2[i]);
    return output;
  };

  template<typename T>
  inline double euclidean2_distance(T* vals1, T* vals2, int nbVals)
  {
    double output = 0;
    int i;
    for (i = 0; i < nbVals; i ++) {
      T tmp = vals1[i] - vals2[i];
      output += static_cast<double>(tmp*tmp);
    }
    return output;
  };

  template<typename T>
  cv::Mat get_distances_float_opt(cv::Mat features, float sigma, cv::Mat points, double boostThreshold)
  {
    static cv::Mat output(cv::Size(features.rows, features.rows), CV_64FC1);
    if (output.rows != features.rows)
      output = Mat(cv::Size(features.rows, features.rows), CV_64FC1);

    double* tmp = new double[features.cols];
    float* tmpP = new float[points.cols];
    float* ptrP = points.ptr<float>();

    double sigma_carre = 2 * sigma*sigma;

    double* ptr = output.ptr<double>();
    T* val = features.ptr<T>();
    int nbFeatures = features.cols;
    int totalSwitch = 0;
    for (int i = 0; i < output.rows; i++)
    {
      int realAdress = i*output.cols;
      ptr[realAdress + i] = 1;

      T* val1 = &val[i*features.cols];
      int nbCopy = 1;
      for (int j = i + 1; j < output.cols; j++)
      {
        double test = ptr[realAdress + j] = exp(
          -euclidean2_distance(val1, &val[j*nbFeatures], nbFeatures)
          / sigma_carre);
        if (test > boostThreshold)//points are the same!
        {
          if (i + nbCopy < j)
          {
            //switch (i+nbCopy)th line with jth feature:
            memcpy(tmp, &val[j*nbFeatures], sizeof(T)*nbFeatures);
            memcpy(&val[j*nbFeatures], &val[(i + nbCopy)*nbFeatures], sizeof(T)*nbFeatures);
            memcpy(&val[(i + nbCopy)*nbFeatures], tmp, sizeof(T)*nbFeatures);

            //switch also the points
            float* pt1 = &ptrP[j*points.cols];
            float* pt2 = &ptrP[(i + nbCopy)*points.cols];
            memcpy(tmpP, pt1, sizeof(float)*points.cols);
            memcpy(pt1, pt2, sizeof(float)*points.cols);
            memcpy(pt2, tmpP, sizeof(float)*points.cols);

            //also switch the value of previously computed distance:
            for (int cpt = 0; cpt <= i; cpt++)
              std::swap(ptr[cpt*output.cols + j], ptr[cpt*output.cols + i + nbCopy]);

          }
          nbCopy++;
        }
      }
      //copy the distances of ith point:
      for (int cpt = 1; cpt < nbCopy; cpt++)
        memcpy(&ptr[realAdress + cpt*output.cols + i + cpt], &ptr[realAdress + i + cpt], sizeof(double)*(output.cols - (i + cpt)));

      nbCopy--;
      i += nbCopy;//skip the same points...
      totalSwitch += nbCopy;
    }
    delete[] tmp;
    delete[] tmpP;
    cv::Mat cleanImg;
    output.convertTo(cleanImg, CV_8UC1, 512);
    cv::imwrite("avecOpt.bmp", cleanImg);
    return output;
  }

  cv::Mat get_distances_binary_opt(cv::Mat features, double mu, int nbBits, cv::Mat points, double boostThreshold)
  {
    static cv::Mat output(cv::Size(features.rows, features.rows), CV_64FC1);
    if (output.rows != features.rows)
      output = Mat(cv::Size(features.rows, features.rows), CV_64FC1);

    double* tmp = new double[features.cols];
    float* tmpP = new float[points.cols];
    float* ptrP = points.ptr<float>();

    double* ptr = output.ptr<double>();
    char* val = features.ptr<char>();

    int nbFeatures = features.cols;
    int sizeFeatures = features.cols * 8;
    int totalSwitch = 0;
    for (int i = 0; i < output.rows; i++)
    {
      int realAdress = i*output.cols;
      ptr[realAdress + i] = 1;

      char* val1 = &val[i*features.cols];
      int nbCopy = 1;
      for (int j = i + 1; j < output.cols; j++)
      {
        double distTmp = hamming_distance(
          val1, &val[j*features.cols], nbFeatures);

        double test = ptr[realAdress + j] = (pow(mu, (int)distTmp) * pow(1. - mu, (int)(nbBits - distTmp)));

        if (distTmp / nbBits > boostThreshold)//points are the same!
        {
          if (i + nbCopy < j)
          {
            //switch (i+nbCopy)th line with jth feature:
            memcpy(tmp, &val[j*nbFeatures], nbFeatures);
            memcpy(&val[j*nbFeatures], &val[(i + nbCopy)*nbFeatures], nbFeatures);
            memcpy(&val[(i + nbCopy)*nbFeatures], tmp, nbFeatures);

            //switch also the points
            float* pt1 = &ptrP[j*points.cols];
            float* pt2 = &ptrP[(i + nbCopy)*points.cols];
            memcpy(tmpP, pt1, sizeof(float)*points.cols);
            memcpy(pt1, pt2, sizeof(float)*points.cols);
            memcpy(pt2, tmpP, sizeof(float)*points.cols);

            //also switch the value of previously computed distance:
            for (int cpt = 0; cpt <= i; cpt++)
              std::swap(ptr[cpt*output.cols + j], ptr[cpt*output.cols + i + nbCopy]);
          }
          nbCopy++;
        }
      }
      //copy the distances of ith point:
      for (int cpt = 1; cpt < nbCopy; cpt++)
        memcpy(&ptr[realAdress + cpt*output.cols + i + cpt], &ptr[realAdress + i + cpt], sizeof(double)*(output.cols - (i + cpt)));

      nbCopy--;
      i += nbCopy;//skip the same points...
      totalSwitch += nbCopy;
    }
    delete[] tmp;
    delete[] tmpP;
    return output;
  }

  template<typename T>
  cv::Mat get_distances_float(cv::Mat features, float sigma)
  {
    static cv::Mat output(cv::Size(features.rows, features.rows), CV_64FC1);
    if (output.rows != features.rows)
      output = Mat(cv::Size(features.rows, features.rows), CV_64FC1);

    float sigma_carre = 2 * sigma*sigma;

    double* ptr = output.ptr<double>();
    T* val = features.ptr<T>();
    int nbFeatures = features.cols;
    for (int i = 0; i < features.rows; i++)
    {
      int realAdress = i*features.rows;
      ptr[realAdress + i] = 0;//diagonal distance is null...
      T* val1 = &val[i*features.cols];
      for (int j = i + 1; j < features.rows; j++)
      {
        ptr[realAdress + j] = exp(
          -euclidean2_distance(val1, &val[j*features.cols], nbFeatures)
          / sigma_carre);
      }
    }
    return output;
  }

  cv::Mat get_distances_binary(cv::Mat features, double mu, int nbBits)
  {
    static cv::Mat output(cv::Size(features.rows, features.rows), CV_64FC1);
    if (output.rows != features.rows)
      output = Mat(cv::Size(features.rows, features.rows), CV_64FC1);
    output.setTo(0);

    double* ptr = output.ptr<double>();
    char* val = features.ptr<char>();
    int nbFeatures = features.cols;
    for (int i = 0; i < features.rows; i++)
    {
      int realAdress = i*features.rows;
      char* val1 = &val[i*features.cols];
      for (int j = i + 1; j < features.rows; j++)
      {
        double dist = hamming_distance(
          val1, &val[j*features.cols], nbFeatures);

        ptr[realAdress + j] = (pow(mu, (int)dist)) * (pow(1. - mu, (int)(nbBits - dist)));
      }
    }
    return output;
  }

  std::vector<double> get_criterions(cv::Mat distances)
  {
    std::vector<double> output;
    for (int i = 0; i < distances.cols; i++)
      output.push_back(0);
    //double divisor = 1.;// (1. / ((distances.rows - 1.) * pow((sigma*sqrt(2 * CV_PI)), D)));

    for (int i = 0; i < distances.cols; i++)
    {
      double* dist_i = distances.ptr<double>(i);
      double crit_i = 0;
      for (int j = i + 1; j < distances.rows; j++)
      {
        double crit_i = dist_i[j];
        output[i] += crit_i;
        output[j] += crit_i;
      }
    }

    return output;
  };

  double get_threshold_float(double p, double D, double sigma, double N)
  {
    double gamma = 2.*pow(boost::math::erf_inv(2 * p - 1), 2);
    double v = sigma * sigma * (D + 2. * sqrt(gamma*(D - gamma))) / (D - 2 * gamma);
    double Ci = pow(2. * CV_PI * v, -D / 2.);

    return Ci*(N - 1)*pow(sqrt(CV_PI*2.) * sigma, D);
  }

  double get_threshold_binary(double p, double D, double mu, double N)
  {
    double v = 0;
    double gamma = 2.*pow(boost::math::erf_inv(2 * p - 1), 2);

    if (p < 0.5)
      v = (2 * mu*D + gamma + sqrt(gamma * (8 * mu * D + gamma))) / (2.*D);
    if (p > 0.5)
      v = (2 * mu*D + gamma - sqrt(gamma * (8 * mu * D + gamma))) / (2.*D);
    if (p == 0.5)
      v = mu;

    double Ci = pow(1. - v, D);

    return Ci;
  }

  BLOCK_END_INSTANTIATION(CORE_filter, AlgoType::imgProcess, BLOCK__CORE_NAME);

  BEGIN_BLOCK_INPUT_PARAMS(CORE_filter);
  //Add parameters, with following parameters:
  ADD_PARAMETER(toBeLinked, Matrix, "BLOCK__CORE_IN_POINTS", "BLOCK__CORE_IN_POINTS_HELP");
  ADD_PARAMETER(toBeLinked, Matrix, "BLOCK__CORE_IN_DESC", "BLOCK__CORE_IN_DESC_HELP");
  ADD_PARAMETER_FULL(false, Float, "BLOCK__CORE_IN_THRESHOLD", "BLOCK__CORE_IN_THRESHOLD_HELP", 90.f);
  ADD_PARAMETER_FULL(false, Float, "BLOCK__CORE_IN_OPTIM_THRESHOLD", "BLOCK__CORE_IN_OPTIM_THRESHOLD_HELP", .75f);
  END_BLOCK_PARAMS();

  BEGIN_BLOCK_OUTPUT_PARAMS(CORE_filter);
  ADD_PARAMETER(toBeLinked, Matrix, "BLOCK__CORE_OUT_POINTS", "BLOCK__CORE_OUT_POINTS_HELP");
  ADD_PARAMETER(toBeLinked, Matrix, "BLOCK__CORE_OUT_DESC", "BLOCK__CORE_OUT_DESC_HELP");
  END_BLOCK_PARAMS();

  BEGIN_BLOCK_SUBPARAMS_DEF(CORE_filter);
  END_BLOCK_PARAMS();

  CORE_filter::CORE_filter() :Block("BLOCK__CORE_NAME", true){
    _myInputs["BLOCK__CORE_IN_POINTS"].addValidator({ new ValNeeded() });
    _myInputs["BLOCK__CORE_IN_THRESHOLD"].addValidator({ new ValRange(0,100) });
  };

  template <typename T>
  vector<size_t> sort_indexes(const vector<T> &v) {

    // initialize original index locations
    vector<size_t> idx(v.size());
    for (size_t i = 0; i != idx.size(); ++i) idx[i] = i;

    // sort indexes based on comparing values in v
    std::sort(idx.begin(), idx.end(),
      [&v](size_t i1, size_t i2) {return v[i1] < v[i2]; });

    return idx;
  }

  bool CORE_filter::run(bool oneShot){
    double percent = _myInputs["BLOCK__CORE_IN_THRESHOLD"].get<double>() / 100.;
    double percentOpt = _myInputs["BLOCK__CORE_IN_OPTIM_THRESHOLD"].get<double>();
    cv::Mat points = _myInputs["BLOCK__CORE_IN_POINTS"].get<cv::Mat>().clone();
    cv::Mat desc = _myInputs["BLOCK__CORE_IN_DESC"].get<cv::Mat>().clone();
    if (points.empty() || desc.empty())
    {
      _myOutputs["BLOCK__CORE_OUT_POINTS"] = points;
      _myOutputs["BLOCK__CORE_OUT_DESC"] = desc;
      return true;//not a problem, just nothing to produce...
    }
    int nbChanels = points.channels();
    if (nbChanels != 1)
      points = points.reshape(1, points.rows);
    if (points.depth() != CV_32F)
      points.convertTo(points, CV_32F);

    float sigma = 32.135f;
    double mu = 0.1;

    int dataSize = 8 * desc.cols;
    bool isBinary = true, isFloat = false;
    if (desc.depth() == CV_16U || desc.depth() == CV_16S)
      dataSize = 16 * desc.cols;
    if (desc.depth() == CV_32S)
      dataSize = 32 * desc.cols;
    if (desc.depth() == CV_32F)
    {
      dataSize = desc.cols;//we just need the number of values, not the number of bytes...
      isBinary = false;
      isFloat = true;
    }
    if (desc.depth() == CV_64F)
    {
      dataSize = desc.cols;//we just need the number of values, not the number of bytes...
      isBinary = false;
      isFloat = false;
    }
    std::vector<double> crit;
    cv::Mat distances;
    if (isBinary)
    {
      if (percentOpt >= 1)
        distances = get_distances_binary(desc, mu, dataSize);
      else
      {
        auto tick = boost::posix_time::microsec_clock::local_time();/*
        distances = get_distances_binary_opt(desc, 0.1, dataSize, points, percentOpt);
        distances = get_distances_binary_opt(desc, 0.1, dataSize, points, percentOpt);
        distances = get_distances_binary_opt(desc, 0.1, dataSize, points, percentOpt);
        distances = get_distances_binary_opt(desc, 0.1, dataSize, points, percentOpt);
        distances = get_distances_binary_opt(desc, 0.1, dataSize, points, percentOpt);*/
        distances = get_distances_binary_opt(desc, mu, dataSize, points, percentOpt);
        auto now = boost::posix_time::microsec_clock::local_time();
        static long long elapsed = 0;
        static int nbEchantillons = 0;
        nbEchantillons++;
        elapsed += (now - tick).total_microseconds();

        tick = boost::posix_time::microsec_clock::local_time();
        cv::Mat distances1 = get_distances_binary(desc, mu, dataSize);/*
        distances1 = get_distances_binary(desc, 0.1, dataSize);
        distances1 = get_distances_binary(desc, 0.1, dataSize);
        distances1 = get_distances_binary(desc, 0.1, dataSize);
        distances1 = get_distances_binary(desc, 0.1, dataSize);
        distances1 = get_distances_binary(desc, 0.1, dataSize);*/
        now = boost::posix_time::microsec_clock::local_time();
        static long long elapsed1 = 0;
        elapsed1 += (now - tick).total_microseconds();
        std::cout << elapsed / nbEchantillons << " sans optim : " << elapsed1 / nbEchantillons << " gain : " << elapsed * 100 / elapsed1 << std::endl;
        distances1.ptr<double>();
      }
    }
    else
    {
      if (percentOpt >= 1)
      {
        if (isFloat)
          distances = get_distances_float<float>(desc, sigma);
        else
          distances = get_distances_float<double>(desc, sigma);
      }
      else
      {
        if (isFloat)
        {
          auto tick = boost::posix_time::microsec_clock::local_time();
          distances = get_distances_float_opt<float>(desc, sigma, points, percentOpt);/*
          distances = get_distances_float_opt<float>(desc, sigma, points, percentOpt);
          distances = get_distances_float_opt<float>(desc, sigma, points, percentOpt);
          distances = get_distances_float_opt<float>(desc, sigma, points, percentOpt);
          distances = get_distances_float_opt<float>(desc, sigma, points, percentOpt);
          distances = get_distances_float_opt<float>(desc, sigma, points, percentOpt);*/
          auto now = boost::posix_time::microsec_clock::local_time();
          auto elapsed = (now - tick).total_microseconds();

          tick = boost::posix_time::microsec_clock::local_time();
          cv::Mat distances1 = get_distances_float<float>(desc, sigma);/*
          distances1 = get_distances_float<float>(desc, sigma);
          distances1 = get_distances_float<float>(desc, sigma);
          distances1 = get_distances_float<float>(desc, sigma);
          distances1 = get_distances_float<float>(desc, sigma);
          distances1 = get_distances_float<float>(desc, sigma);*/
          now = boost::posix_time::microsec_clock::local_time();
          auto elapsed1 = (now - tick).total_microseconds();
          std::cout << elapsed << " sans optim : " << elapsed1 << " gain : " << elapsed * 100 / elapsed1 << std::endl;
          distances1.ptr<double>();
        }
        else
          distances = get_distances_float_opt<double>(desc, sigma, points, percentOpt);
      }
    }
    crit = get_criterions(distances);

    std::vector<size_t> indices = sort_indexes(crit);

    int nbVals = 0;
    if (!isBinary)
    {
      //first count the correct values:
      double threshold = get_threshold_float(percent, desc.cols, sigma, crit.size());
      for (size_t i = 0; i < crit.size(); i++) {
        if (crit[i] < threshold)
          nbVals++;
        else
          break;
      }
    }
    else
    {
      //first count the correct values:
      double threshold = get_threshold_binary(percent, desc.cols, mu, crit.size());
      for (size_t i = 0; i < crit.size(); i++) {
        if (crit[i] < threshold)
          nbVals++;
        else
          break;
      }
    }
    //nbVals = (int)(points.rows*percent);

    cv::Mat outPoints(cv::Size(points.cols, nbVals), points.type());
    cv::Mat outDesc(cv::Size(desc.cols, nbVals), desc.type());
    for (int i = 0; i < nbVals; i++) {
      int bestIndice = indices[i];
      memcpy(outPoints.ptr<float>(i), points.ptr<float>(bestIndice), points.cols*sizeof(float));

      if (desc.depth() == CV_8U || desc.depth() == CV_8S)
        memcpy(outDesc.ptr<char>(i), desc.ptr<char>(bestIndice), desc.cols*sizeof(char));

      if (desc.depth() == CV_16U || desc.depth() == CV_16S)
        memcpy(outDesc.ptr<short>(i), desc.ptr<short>(bestIndice), desc.cols*sizeof(short));

      if (desc.depth() == CV_32S || desc.depth() == CV_32F)
        memcpy(outDesc.ptr<float>(i), desc.ptr<float>(bestIndice), desc.cols*sizeof(float));

      if (desc.depth() == CV_64F)
        memcpy(outDesc.ptr<double>(i), desc.ptr<double>(bestIndice), desc.cols*sizeof(double));
    }
    _myOutputs["BLOCK__CORE_OUT_POINTS"] = outPoints;
    _myOutputs["BLOCK__CORE_OUT_DESC"] = outDesc;

    return true;
  };
};


/*
 * The MIT License (MIT)
 *
 * Copyright (c) 2017 Hidekazu Ikeno
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */

///
/// \file fast_esprit.hpp
///
#ifndef MXPFIT_ESPRIT_HPP
#define MXPFIT_ESPRIT_HPP

#include <algorithm>
#include <type_traits>

#include <Eigen/Core>
#include <Eigen/Eigenvalues>
#include <Eigen/SVD>

#include <mxpfit/exponential_sum.hpp>
#include <mxpfit/prony_like_method_common.hpp>

namespace mxpfit
{
///
/// ### ESPRIT
///
/// \brief ESPRIT method for finding parameters of exponential sum
/// approximation from sampled data on uniform grid.
///
/// \tparam T  Scalar type of function values.
///
/// #### Description
///
/// This class implements the Estimation of Signal Parameters via Rotational
/// Invariance Techniques (ESPRIT) method for parameter estimation of decaying
/// exponential sum funcitons.
///
/// This class is implemented only for the benchmarking purpose to compare with
/// the fast ESPRIT algorithm. DO NOT use the class for practical applications,
/// use `FastESPRIT` class instead, which is much faster than ESPRIT.
///
///
/// #### References
///
/// 1. D. Potts and M. Tasche,"Parameter estimation for nonincreasing
///    exponential sums by Prony-like methods", Linear Algebra Appl. **439**
///    (2013) 1024-1039.
///    [DOI: https://doi.org/10.1016/j.laa.2012.10.036]
///
template <typename T>
class ESPRIT
{
public:
    using Scalar        = T;
    using RealScalar    = typename Eigen::NumTraits<Scalar>::Real;
    using ComplexScalar = std::complex<RealScalar>;
    using Index         = Eigen::Index;

    using Vector        = Eigen::Matrix<Scalar, Eigen::Dynamic, 1>;
    using RealVector    = Eigen::Matrix<RealScalar, Eigen::Dynamic, 1>;
    using ComplexVector = Eigen::Matrix<ComplexScalar, Eigen::Dynamic, 1>;

    using Matrix = Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>;
    using RealMatrix =
        Eigen::Matrix<RealScalar, Eigen::Dynamic, Eigen::Dynamic>;
    using ComplexMatrix =
        Eigen::Matrix<ComplexScalar, Eigen::Dynamic, Eigen::Dynamic>;

    using ResultType =
        typename detail::gen_prony_like_method_result<T>::ResultType;

private:
    enum
    {
        IsComplex = Eigen::NumTraits<Scalar>::IsComplex,
        Alignment = Eigen::internal::traits<Matrix>::Alignment
    };

    Index m_rows;
    Index m_cols;
    Index m_max_terms;
    Matrix m_matH;

public:
    ///
    /// Default constructor
    ///
    ESPRIT() = default;

    ///
    /// Constructor with memory preallocation
    ///
    /// \param[in] N  Number of sampling points
    /// \param[in] L  Window size. This is equals to the number of rows of
    ///               generalized Hankel matrix.
    /// \param[in] M  Maxumum number of terms used for the exponential sum.
    /// \pre  `N >= M >= 1` and `N - L + 1 >= M >= 1`.
    ///
    ESPRIT(Index N, Index L, Index M)
        : m_rows(L), m_cols(N - L + 1), m_max_terms(M), m_matH(m_rows, m_cols)
    {
        assert(m_rows >= M && m_cols >= M && M >= 1);
    }

    ///
    /// Destructor
    ///
    ~ESPRIT()
    {
    }

    ///
    /// Memory reallocation
    ///
    /// \param[in] N  Number of sampling points
    /// \param[in] L  Window size. This is equals to the number of rows of
    ///               generalized Hankel matrix.
    /// \param[in] M  Maxumum number of terms used for the exponential sum.
    /// \pre  `N >= M >= 1` and `N - L + 1 >= M >= 1`.
    ///
    void resize(Index N, Index L, Index M)
    {
        m_rows = L;
        m_cols = N - L + 1;
        assert(m_rows >= M && m_cols >= M && M >= 1);
        m_max_terms = M;
        m_matH.resize(m_rows, m_cols);
    }

    ///
    /// \return Number of sampling points.
    ///
    Index size() const
    {
        return m_rows + m_cols - 1;
    }

    ///
    /// Fit signals by a exponential sum
    ///
    /// \param[in] f The array of signals sampled on the equispaced grid. The
    ///    first `size()` elemnets of `nterms` are used as a sampled data. In
    ///    case `f.size() < size()` then, last `size() - f.size()` elements are
    ///    padded by zeros.
    /// \param[in] eps  Small positive number `(0 < eps < 1)` that
    ///    controlls the accuracy of the fit.
    /// \param[in] x0  Argument of first sampling point
    /// \param[in] delta Spacing between neighboring sample points.
    ///
    template <typename VectorT>
    ResultType compute(const Eigen::MatrixBase<VectorT>& h, RealScalar x0,
                       RealScalar delta, RealScalar eps);
};

template <typename T>
template <typename VectorT>
typename ESPRIT<T>::ResultType
ESPRIT<T>::compute(const Eigen::MatrixBase<VectorT>& h, RealScalar x0,
                   RealScalar delta, RealScalar eps)
{
    assert(h.size() == size() && "Number of data points mismatch.");
    //
    // Form rectangular Hankel matrix from sequance h.
    //
    for (Index j = 0; j < m_cols; ++j)
    {
        for (Index i = 0; i < m_rows; ++i)
        {
            m_matH.coeffRef(i, j) = h(i + j);
        }
    }
    // const Index nr = m_matH.rows();
    const Index nc = m_matH.cols();

    //-------------------------------------------------------------------------
    // Compute roots of Prony polynomials H = U * S * W^H
    //-------------------------------------------------------------------------
    Eigen::BDCSVD<Matrix> svd(m_matH, Eigen::ComputeThinV);

    const auto& sigma = svd.singularValues();
    Index nterms      = 0;
    // extract rank of matrix H from singular values
    while (nterms < m_max_terms)
    {
        if (sigma(nterms) < eps)
        {
            break;
        }
        ++nterms;
    }

    if (nterms == 0)
    {
        return ResultType();
    }

    ComplexVector roots(nterms);
    ComplexVector weights(nterms);

    {
        // --- Form the views of matrix W
        // Matrix W excluding the last row
        auto W0 = svd.matrixV().block(0, 0, nc - 1, nterms);
        // Matrix W excluding the first row
        auto W1 = svd.matrixV().block(1, 0, nc - 1, nterms);
        // adjoint of the last row of matrix W
        auto nu = svd.matrixV().block(nc - 1, 0, 1, nterms).adjoint();
        //
        // Compute the spectral matrix G = pinv(W0) * W1, where pinv indicate
        // the Moore-Penrose pseudo-inverse. The computation of the
        // pseudo-inverse of W0 can be avoided.
        //
        Matrix G(W0.adjoint() * W1);
        Vector phi(G.adjoint() * nu);
        auto scal = RealScalar(1) / (RealScalar(1) - nu.squaredNorm());
        G += scal * nu * phi.adjoint();
        //
        // Prony roots \f$\{z_i\{\}\f$ are the eigenvalues of matrix G. The
        // exponents for approximation are obtained as \f$ \log z_i \f$
        //
        roots = G.eigenvalues();
    }

    //----------------------------------------------------------------------
    // Solve overdetermined Vandermonde system to obtain the weights
    //----------------------------------------------------------------------
    {
        // Create Vandermonde matrix from prony roots
        ComplexMatrix matV(h.size(), nterms);
        for (Index j = 0; j < matV.cols(); ++j)
        {
            auto x     = roots(j);
            auto v     = ComplexScalar(1);
            matV(0, j) = v;
            for (Index i = 1; i < matV.rows(); ++i)
            {
                v *= x;
                matV(i, j) = v;
            }
        }
        // Solve least-squares problem V x = h
        weights = matV.jacobiSvd(Eigen::ComputeThinU | Eigen::ComputeThinV)
                      .solve(h.template cast<ComplexScalar>());
    }

    // Rescale the exponents and weights, and return the result
    return detail::gen_prony_like_method_result<T>::create(
        roots.array(), weights.array(), x0, delta);
}

} // namespace mxpfit

#endif /* MXPFIT_ESPRIT_HPP */


#pragma once

#include <deal.II/base/quadrature_lib.h>
#include <deal.II/base/utilities.h>
#include <deal.II/fe/fe_values.h>
#include <deal.II/numerics/vector_tools.h>
#include <iostream>

namespace boltzmann {

template <int dim, typename InVector>
double
compute_mean_value(const dealii::DoFHandler<dim>& dh, const InVector& v)
{
  dealii::QGauss<dim> quad(2);
  return dealii::VectorTools::compute_mean_value(dh, quad, v, 0);
}

namespace mpi {

template <int dim>
std::vector<double>
compute_mean_value(const dealii::DoFHandler<dim>& dh, const Epetra_MultiVector& src)
{
  unsigned int pid = dealii::Utilities::MPI::this_mpi_process(MPI_COMM_WORLD);
  dealii::QGauss<dim> quad(2);
  dealii::UpdateFlags flags =
      dealii::update_values | dealii::update_JxW_values | dealii::update_quadrature_points;
  dealii::FEValues<dim> fe_values(dh.get_fe(), quad, flags);

  int nvec = src.NumVectors();

  std::vector<double> lsum(nvec, 0.0);

  unsigned int dofs_per_cell = fe_values.dofs_per_cell;
  std::vector<unsigned int> local_dof_indices(dofs_per_cell);
  for (auto cell = dh.begin_active(); cell != dh.end(); ++cell) {
    if (cell->subdomain_id() == pid) {
      fe_values.reinit(cell);
      cell->get_dof_indices(local_dof_indices);
      for (int k = 0; k < nvec; ++k) {
        for (unsigned int ix = 0; ix < dofs_per_cell; ++ix) {
          int lid = src.Map().LID(int(local_dof_indices[ix]));
          BOOST_ASSERT(lid >= 0);
          // if (lid < 0) std::cout << "lid " << lid << " gid " << local_dof_indices[ix] << "\n";
          for (unsigned int q = 0; q < quad.size(); ++q) {
            lsum[k] += src[k][lid] * fe_values.shape_value(ix, q) * fe_values.JxW(q);
          }
        }
      }
    }
  }

  std::vector<double> sum(nvec, 0.0);
  MPI_Allreduce(lsum.data(), sum.data(), nvec, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
  return sum;
}
}  // namespace mpi

}  // namespace boltzmann


/*!
  @file main.cpp
  @author Klaus K. Holst
  @copyright 2018-2021, Klaus Kähler Holst

  @brief Example file

*/

#include <armadillo>
#include <cstdio>  // remove
#include <algorithm>  // max
#include <target/target.hpp>
#include <target/utils.hpp>

using namespace arma;

int main(int argc, char **argv) {
  vec y, w, p;
  mat a, x1, x2;
  if (argc == 1) {
    std::cout << "* Simulating data...\n\n";
    arma_rng::set_seed(1);
    unsigned n = 5;
    y = vec(n, 1);
    y.randn();
    w = mat(n, 1);
    w.fill(1);
    a = randi<mat>(n, 1, distr_param(0, 1));
    x2 = mat(n, 2);
    x2.randn();
    x1 = mat(n, 1);
    x1.fill(1);
    p = vec(4);
    p.fill(0.5);
  } else {
    const char *infile = argv[1];
    std::cout << "* Reading from '" << infile << "'\n\n";
  }

  unsigned n = std::min(3, (int)y.n_elem);
  std::cout << " Response:\n"
            << target::BLUE << y.rows(0, n - 1)
            << "\t..." << target::COL_RESET
            << std::endl;
  std::cout << " Exposure:\n"
            << target::CYAN << a.rows(0, n - 1)
            << "\t..." << target::COL_RESET
            << std::endl;
  std::cout << " X1:\n"
            << target::YELLOW << x1.rows(0, n - 1) << "\t..."
            << target::COL_RESET << std::endl;
  std::cout << " X2:\n"
            << target::YELLOW << x2.rows(0, n - 1) << "\t..."
            << target::COL_RESET << std::endl;
  std::cout << "\n parameter: " << target::GREEN << p.t()
            << target::COL_RESET
            << std::endl;

  target::RD<double> model(y, a, x1, x2, x2, p, w);
  vec res = model.loglik();
  std::cout << " loglik=\n"
            << target::RED << res << std::endl
            << target::COL_RESET;

  mat U = model.score(true);
  std::cout << " score=\n"
            << target::RED << U << std::endl
            << target::COL_RESET;

  arma::vec alpha2(1); alpha2.fill(1);
  U = model.est(alpha2);
  std::cout << " U=\n" << target::RED << U <<
    std::endl << target::COL_RESET;

  //arma::mat pp = model.target::TargetBinary<double>::pa();
  //std::cout << "pp=\n" << pp << std::endl;
  // const char *filen = "tmp/a.h5";
  // std::remove(filen);
  // y.save(hdf5_name(filen, "y", hdf5_opts::append+hdf5_opts::trans));
  // a.save(hdf5_name(filen, "a", hdf5_opts::append+hdf5_opts::trans));
  // w.save(hdf5_name(filen, "w", hdf5_opts::append+hdf5_opts::trans));
  // x1.save(hdf5_name(filen, "x1", hdf5_opts::append+hdf5_opts::trans));
  // x2.save(hdf5_name(filen, "x2", hdf5_opts::append+hdf5_opts::trans));
  // p.save(hdf5_name(filen, "p", hdf5_opts::append+hdf5_opts::trans));

  return 0;
}


#include <stdio.h>
#include <stdlib.h>
#include <cstdlib>
#include <time.h>
#include <math.h>
#include <dirent.h>
#include <iostream>
#include <iomanip>
#include <fstream>
#include <iomanip>
#include <limits>
#include <string>
#include <algorithm>
#include <boost/filesystem.hpp>
#include <boost/algorithm/string/predicate.hpp>
#include <boost/algorithm/string/replace.hpp>
#include <boost/math/distributions/normal.hpp>

#include "opencv2/opencv.hpp"
#include "opencv2/core/core.hpp"
#include "opencv2/features2d/features2d.hpp"
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/imgproc/imgproc.hpp"
#include "opencv2/calib3d/calib3d.hpp"
#include "opencv2/nonfree/features2d.hpp"
#include "opencv2/flann/flann.hpp"

#include "FastEMD/emd_hat.hpp"
#include "FastEMD/emd_hat_signatures_interface.hpp"
#include "RubnerEMD/emd.hpp"
#include "FastEMD/tictoc.hpp"

using namespace cv;
using namespace std;

namespace fs = boost::filesystem;

#define IMAGE_WIDTH 1000

// TODO: config
/* CONFIGURATIONS */
int keypoint_filter_threshold = 1;
bool filter_kp = 0;
int sort_type = -1;
bool show_images = 0;
/* END CONFIGURATIONS */

void help();
void find_images(const char *path);
void print_progress(time_t start, int total, int& completed, int per_cout);
void print_keypoints(const vector<KeyPoint>& kp);

bool extract_descriptors(int index, int fc, vector<KeyPoint>& kps, Mat& descs);
void select_keypoints(vector<KeyPoint>& k, vector<KeyPoint>& s, int fc);
void filter_keypoints(vector<KeyPoint>& k, vector<KeyPoint>& f);
void calc_response_density_img(const Mat& img, vector<KeyPoint> kp, Mat& ds);
void calc_reponse_response_density_space(const vector<KeyPoint>& kp);
void calc_scale_density_img(const Mat& img, vector<KeyPoint> kp, Mat& ds);
void calc_scale_response_density_space(const vector<KeyPoint>& kp);
float keypoint_dist(const KeyPoint& k1, const KeyPoint& k2);
float point_dist(const Point2f& p1, const Point2f& p2);
bool response_comparator(const KeyPoint& p1, const KeyPoint& p2);
float get_density(const KeyPoint& k);
bool density_comparator(const KeyPoint& p1, const KeyPoint& p2);
float get_scale_density(const KeyPoint& k);
bool scale_density_comparator(const KeyPoint& p1, const KeyPoint& p2);
bool scale_comparator(const KeyPoint& p1, const KeyPoint& p2);

vector<fs::path> img_paths;

float response_density_space[IMAGE_WIDTH][IMAGE_WIDTH];
float scale_response_density_space[IMAGE_WIDTH][IMAGE_WIDTH];
Mat current_img;

vector<KeyPoint> filtered_keypoints;
struct response_density_classcomp {
	bool operator()(const int& lhs, const int& rhs) const {
		KeyPoint kp1, kp2;
		kp1 = filtered_keypoints[lhs];
		kp2 = filtered_keypoints[rhs];
		return density_comparator(kp1, kp2);
	}
};
struct scale_response_density_classcomp {
	bool operator()(const int& lhs, const int& rhs) const {
		KeyPoint kp1, kp2;
		kp1 = filtered_keypoints[lhs];
		kp2 = filtered_keypoints[rhs];
		return scale_density_comparator(kp1, kp2);
	}
};

int main(int argc, char *argv[]) {
	if (argc < 4) {
		help();
		return -1;
	}

	int fc = atoi(argv[1]);
	sort_type = atoi(argv[2]);
	const char *imgs_path = argv[3];
	const char *index_path = argv[4];

	cout << "Reading image paths from " << imgs_path << endl;
	string img_path;
	int img_count;
	ifstream imgs_file(imgs_path);
	imgs_file >> img_count;
	for (int i = 0; i < img_count; ++i) {
		imgs_file >> img_path;
		img_paths.push_back(img_path);
	}
	cout << img_count << " images found" << endl;

	int total_process = img_count;
	int completed_process = 0;
	time_t start = time(0);

	cout << "Building index..." << endl;

	tictoc timer;
	timer.tic();

	vector<KeyPoint> kps;
	Mat descs;
	FileStorage fs(index_path, FileStorage::WRITE);
	fs << "index" << "[";
	for (int i = 0; i < img_count; ++i) {
		if (extract_descriptors(i, fc, kps, descs)) {
			fs << "{";
			fs << "path" << img_paths[i].string();
			fs << "keypoints" << kps;
			fs << "descriptors" << descs;
			fs << "}";
		}
	}
	fs << "]";
	fs.release();

	timer.toc();
	cout << "Time in seconds: " << timer.totalTimeSec() << endl;

	return 0;
}

void help() {
	cout << "Usage: ./build_index <fc> <sort_type> <imgs_path> <index_path>" << endl;
}

void print_progress(time_t start, int total, int& completed, int per_count) {
	++completed;
	if (completed % per_count != 0)
		return;
	int remaining_process = total - completed;
	time_t current = time(0);
	float elapsed_time = difftime(current, start);
	float remaining_time = elapsed_time * remaining_process / completed;
	cout << fixed << setprecision(3);
	cout << "\rRemaining time: ";
	cout << remaining_time / 60 << "min";
}

void print_keypoints(const vector<KeyPoint>& kp) {
	KeyPoint k;
	cout << setprecision(3) << fixed;
	for (int i = 0; i < (int) kp.size(); ++i) {
		k = kp[i];
		cout << "(" << k.pt.x << "\t" << k.pt.y << "):\t" << k.response << "\t" << k.size << "\t" << get_scale_density(k) << endl;
	}
}

void find_images(const char *path) {
	fs::path dir_path(path);
	fs::recursive_directory_iterator end_iter;

	for (fs::recursive_directory_iterator iter(dir_path); iter != end_iter; ++iter) {
		if (fs::is_regular_file(iter->status())) {
			if (!boost::algorithm::ends_with(iter->path().c_str(), ".png"))
				continue;
			img_paths.push_back(iter->path());
		}
	}

	sort(img_paths.begin(), img_paths.end());
}

bool extract_descriptors(int index, int fc, vector<KeyPoint>& kps, Mat& descs) {
	string img_path = img_paths[index].string();
	fs::path yaml = img_paths[index];
	string yaml_path = yaml.replace_extension(".yml").string();

	// cout << "Calc img matrix: " << img_path << endl;
	current_img = imread(img_path.c_str(), CV_LOAD_IMAGE_GRAYSCALE);

	if (!current_img.data) {
		cout << "Error reading image: " << img_path << endl;
		return 0;
	}

	vector<KeyPoint> keypoints;
	if (fs::exists(yaml_path)) {
		FileStorage fs(yaml_path.c_str(), FileStorage::READ);
		read(fs["keypoints"], keypoints);
	} else {
		SiftFeatureDetector detector;
		detector.detect(current_img, keypoints);
		FileStorage fs(yaml_path.c_str(), FileStorage::WRITE);
		fs << "keypoints" << keypoints;
	}
	// print_keypoints(keypoints);

	// int fc = keypoints.size() * fc / 100;
	// fc = fc > 10 ? fc : 10;

	// not enough feature, dont use them
	if (keypoints.size() == 0) {
		cout << "Image " << img_path << " has no feature" << endl;
		return 0;
	}

	vector<KeyPoint> selected_keypoints;
	if (fc == 0) {
		selected_keypoints = keypoints;
	} else {
		select_keypoints(keypoints, selected_keypoints, fc);
	}

	cout << setw(50) << left << img_path
			<< setw(10) << right << keypoints.size()
			<< setw(10) << selected_keypoints.size() << endl;

	// extract descriptors
	Mat descriptors;
	SiftDescriptorExtractor extractor;
	extractor.compute(current_img, selected_keypoints, descriptors);

	// cout << img_path << ":" << endl;
	// print_keypoints(selected_keypoints);
	// cout << descriptors << endl;

	if (show_images) {
		Mat white = Mat::ones(current_img.rows, current_img.cols, CV_8UC1) * 255;
		Mat black = Mat::zeros(current_img.rows, current_img.cols, CV_8UC1);
		Mat img_density = Mat::zeros(current_img.rows, current_img.cols, CV_8UC1);
		Mat img(current_img);

		drawKeypoints(current_img, keypoints, img, Scalar::all(-1), DrawMatchesFlags::DRAW_RICH_KEYPOINTS);
		imshow("kp_on_img", img);

		drawKeypoints(black, keypoints, img, Scalar::all(-1), DrawMatchesFlags::DRAW_RICH_KEYPOINTS);
		imshow("kp_on_black", img);

		drawKeypoints(black, selected_keypoints, img, Scalar::all(-1), DrawMatchesFlags::DRAW_RICH_KEYPOINTS);
		imshow("selected_kp_on_white", img);

		calc_scale_density_img(current_img, keypoints, img_density);
		imshow("density", img_density);

		drawKeypoints(img_density, selected_keypoints, img, Scalar::all(-1), DrawMatchesFlags::DRAW_RICH_KEYPOINTS);
		imshow("selected_kp_on_density", img);

		imshow("img", current_img);
		waitKey(0);
	}

	descs = descriptors;
	kps = selected_keypoints;
	return 1;
}

void select_keypoints(vector<KeyPoint>& k, vector<KeyPoint>& s, int fc) {
	// filter keypoints
	filtered_keypoints.clear();
	if (filter_kp) {
		filter_keypoints(k, filtered_keypoints);
	} else {
		filtered_keypoints = k;
	}
	int filtered_kp_size = filtered_keypoints.size();

	if (sort_type == -5) { // scale + scale density
		sort(filtered_keypoints.begin(), filtered_keypoints.end(), scale_comparator);
		vector<KeyPoint> scale_selected;
		for (int i = 0; i < 2 * fc && i < filtered_kp_size; ++i) {
			scale_selected.push_back(filtered_keypoints[i]);
		}
		calc_scale_response_density_space(scale_selected);
		sort(scale_selected.begin(), scale_selected.end(), scale_density_comparator);
		for (int i = 0; i < fc && i < filtered_kp_size; ++i) {
			s.push_back(scale_selected[i]);
		}
	} else if (sort_type == -4) { // scale density with substrating
		calc_scale_response_density_space(filtered_keypoints);
		sort(filtered_keypoints.begin(), filtered_keypoints.end(), scale_density_comparator);

		for (int i = 0; i < fc && i < filtered_kp_size; ++i) {
			KeyPoint keypoint = filtered_keypoints[0];
			s.push_back(keypoint);

			// subtract the contribution of this keypoint
			for (int j = 0; j < filtered_kp_size; j++) {
				KeyPoint kp = filtered_keypoints[i];
				float dist = keypoint_dist(kp, keypoint);
				boost::math::normal_distribution<> norm_dist(0.0, 1 / keypoint.size);
				float density = boost::math::pdf(norm_dist, dist);
				int x = kp.pt.x;
				int y = kp.pt.y;
				scale_response_density_space[x][y] -= density;
			}

			// delete the selected keypoint
			filtered_keypoints.erase(filtered_keypoints.begin());
			// resort
			sort(filtered_keypoints.begin(), filtered_keypoints.end(), scale_density_comparator);
		}
	} else if (sort_type == -3) { // scale density with kdtree
		calc_scale_response_density_space(filtered_keypoints);

		// create keypoint map for easy deletion
		map<int, KeyPoint, scale_response_density_classcomp> mk;
		map<int, KeyPoint>::iterator it;

		// create kd-tree index
		Mat coord = Mat::zeros(filtered_kp_size, 2, CV_32F);
		for (int i = 0; i < filtered_kp_size; ++i) {
			mk[i] = filtered_keypoints[i];
			coord.ptr<float>(i)[0] = filtered_keypoints[i].pt.x;
			coord.ptr<float>(i)[1] = filtered_keypoints[i].pt.y;
		}

		flann::Index kdtree_index(coord, flann::KDTreeIndexParams(4));

		int max_count = 100;
		Mat indices = Mat::zeros(1, max_count, CV_32F);
		Mat dists = Mat::zeros(1, max_count, CV_32F);

		for (int i = 0; i < fc && i < filtered_kp_size; ++i) {
			if (mk.size() == 0)
				break;

			KeyPoint query_keypoint = mk.begin()->second;

			s.push_back(query_keypoint);

			Mat query = Mat::zeros(1, 2, CV_32F);
			query.ptr<float>(0)[0] = query_keypoint.pt.x;
			query.ptr<float>(0)[1] = query_keypoint.pt.y;
			float radius = query_keypoint.size * query_keypoint.size / 4;
			int found_count = kdtree_index.radiusSearch(query, indices, dists, radius, max_count, flann::SearchParams());
			for (int j = 0; j < found_count; ++j) {
				mk.erase(indices.at<int>(j));
			}
		}
	} else if (sort_type == -2) { // scale density
		calc_scale_response_density_space(filtered_keypoints);
		sort(filtered_keypoints.begin(), filtered_keypoints.end(), scale_density_comparator);
		for (int i = 0; i < fc && i < filtered_kp_size; ++i) {
			s.push_back(filtered_keypoints[i]);
		}
	} else if (sort_type == -1) { // scale
		sort(filtered_keypoints.begin(), filtered_keypoints.end(), scale_comparator);
		for (int i = 0; i < fc && i < filtered_kp_size; ++i) {
			s.push_back(filtered_keypoints[i]);
		}
	} else if (sort_type == 1) { // response
		sort(filtered_keypoints.begin(), filtered_keypoints.end(), response_comparator);
		for (int i = 0; i < fc && i < filtered_kp_size; ++i) {
			s.push_back(filtered_keypoints[i]);
		}
	} else if (sort_type == 2) { // response density
		calc_reponse_response_density_space(filtered_keypoints);
		sort(filtered_keypoints.begin(), filtered_keypoints.end(), density_comparator);
		for (int i = 0; i < fc && i < filtered_kp_size; ++i) {
			s.push_back(filtered_keypoints[i]);
		}
	} else if (sort_type == 3) { // response density with kdtree
		calc_reponse_response_density_space(filtered_keypoints);

		// create keypoint map for easy deletion
		map<int, KeyPoint, response_density_classcomp> mk;
		map<int, KeyPoint>::iterator it;

		// create kd-tree index
		Mat coord = Mat::zeros(filtered_kp_size, 2, CV_32F);
		for (int i = 0; i < filtered_kp_size; ++i) {
			mk[i] = filtered_keypoints[i];
			coord.ptr<float>(i)[0] = filtered_keypoints[i].pt.x;
			coord.ptr<float>(i)[1] = filtered_keypoints[i].pt.y;
		}

		flann::Index kdtree_index(coord, flann::KDTreeIndexParams(4));

		int max_count = 100;
		Mat indices = Mat::zeros(1, max_count, CV_32F);
		Mat dists = Mat::zeros(1, max_count, CV_32F);

		for (int i = 0; i < fc && i < filtered_kp_size; ++i) {
			if (mk.size() == 0)
				break;

			KeyPoint query_keypoint = mk.begin()->second;

			s.push_back(query_keypoint);

			Mat query = Mat::zeros(1, 2, CV_32F);
			query.ptr<float>(0)[0] = query_keypoint.pt.x;
			query.ptr<float>(0)[1] = query_keypoint.pt.y;
			float radius = query_keypoint.size * query_keypoint.size / 4;
			int found_count = kdtree_index.radiusSearch(query, indices, dists, radius, max_count, flann::SearchParams());
			for (int j = 0; j < found_count; ++j) {
				mk.erase(indices.at<int>(j));
			}
		}
	}
}

void filter_keypoints(vector<KeyPoint>& k, vector<KeyPoint>& f) {
	if (sort_type < 0)
		sort(k.begin(), k.end(), scale_comparator);
	else
		sort(k.begin(), k.end(), response_comparator);

	for (int i = k.size() - 1; i >= 0; --i) {
		KeyPoint k1 = k[i];
		bool add = 1;
		for (int j = 0; j < f.size(); ++j) {
			KeyPoint k2 = f[j];
			float dist = keypoint_dist(k1, k2);
			if (dist < keypoint_filter_threshold) {
				add = 0;
				break;
			}
		}
		if (add)
			f.push_back(k1);
	}
}

void calc_response_density_img(const Mat& img, vector<KeyPoint> kp, Mat& ds) {
	int rows = img.rows;
	int cols = img.cols;
	int size = kp.size();
	Mat tmp = Mat::zeros(rows, cols, CV_32F);

	float max_value = 0, min_value = 9999999;
	for (int i = 0; i < cols; i++) {
		for (int j = 0; j < rows; j++) {
			float sum = 0;
			for (int k = 0; k < size; k++) {
				KeyPoint p = kp[k];
				float dist = sqrt(pow(i - p.pt.x, 2) + pow(j - p.pt.y, 2));
				boost::math::normal_distribution<> norm_dist(0.0, 1 / p.response * 100);
				sum += boost::math::pdf(norm_dist, dist);
			}
			tmp.ptr<float>(j)[i] = sum;
			if (sum > max_value)
				max_value = sum;
			if (sum < min_value)
				min_value = sum;
		}
	}

	ds = Mat(rows, cols, CV_8UC1);
	for (int i = 0; i < cols; i++) {
		for (int j = 0; j < rows; j++) {
			ds.at<uchar>(j, i) = (tmp.at<float>(j, i) - min_value) / (max_value - min_value) * 255;
		}
	}

	equalizeHist(ds, ds);
}

void calc_reponse_response_density_space(const vector<KeyPoint>& kp) {
	int kp_size = kp.size();
	for (int i = 0; i < kp_size; i++) {
		KeyPoint ki = kp[i];
		float density = 0;
		for (int j = 0; j < kp_size; j++) {
			KeyPoint kj = kp[j];
			float dist = keypoint_dist(ki, kj);
			boost::math::normal_distribution<> norm_dist(0.0, 1 / kj.response * 100);
			density += boost::math::pdf(norm_dist, dist);
		}
		int x = ki.pt.x;
		int y = ki.pt.y;
		response_density_space[x][y] = density;
	}
}

void calc_scale_density_img(const Mat& img, vector<KeyPoint> kp, Mat& ds) {
	int rows = img.rows;
	int cols = img.cols;
	int size = kp.size();
	Mat tmp = Mat::zeros(rows, cols, CV_32F);

	float max_value = 0, min_value = 9999999;
	for (int i = 0; i < cols; i++) {
		for (int j = 0; j < rows; j++) {
			float sum = 0;
			for (int k = 0; k < size; k++) {
				KeyPoint p = kp[k];
				float dist = sqrt(pow(i - p.pt.x, 2) + pow(j - p.pt.y, 2));
				boost::math::normal_distribution<> norm_dist(0.0, 20 / p.size);
				sum += boost::math::pdf(norm_dist, dist);
			}
			tmp.ptr<float>(j)[i] = sum;
			if (sum > max_value)
				max_value = sum;
			if (sum < min_value)
				min_value = sum;
		}
	}

	ds = Mat(rows, cols, CV_8UC1);
	for (int i = 0; i < cols; i++) {
		for (int j = 0; j < rows; j++) {
			ds.at<uchar>(j, i) = (tmp.at<float>(j, i) - min_value) / (max_value - min_value) * 255;
		}
	}

	equalizeHist(ds, ds);
}

void calc_scale_response_density_space(const vector<KeyPoint>& kp) {
	int kp_size = kp.size();
	for (int i = 0; i < kp_size; i++) {
		KeyPoint ki = kp[i];
		float density = 0;
		for (int j = 0; j < kp_size; j++) {
			KeyPoint kj = kp[j];
			float dist = keypoint_dist(ki, kj);
			boost::math::normal_distribution<> norm_dist(0.0, 20 / kj.size);
			density += boost::math::pdf(norm_dist, dist);
		}
		int x = ki.pt.x;
		int y = ki.pt.y;
		scale_response_density_space[x][y] = density;
	}
}

float keypoint_dist(const KeyPoint& k1, const KeyPoint& k2) {
	float dist = point_dist(k1.pt, k2.pt);
	return dist;
}

float point_dist(const Point2f& p1, const Point2f& p2) {
	float x_dist = p1.x - p2.x;
	float y_dist = p1.y - p2.y;
	float dist = sqrt(x_dist*x_dist + y_dist*y_dist);
	return dist;
}

bool response_comparator(const KeyPoint& p1, const KeyPoint& p2) {
	return p1.response > p2.response;
}

float get_density(const KeyPoint& k) {
	int x = k.pt.x;
	int y = k.pt.y;
	return response_density_space[x][y];
}

bool density_comparator(const KeyPoint& p1, const KeyPoint& p2) {
	return get_density(p1) > get_density(p2);
}

float get_scale_density(const KeyPoint& k) {
	int x = k.pt.x;
	int y = k.pt.y;
	return scale_response_density_space[x][y];
}

bool scale_density_comparator(const KeyPoint& p1, const KeyPoint& p2) {
	return get_scale_density(p1) > get_scale_density(p2);
}

bool scale_comparator(const KeyPoint& p1, const KeyPoint& p2) {
	return p1.size > p2.size;
}


#include <Eigen/Core>
#include <filesystem/resolver.h>
#include <fstream>
#include <nori/integrator.h>
#include <nori/ray.h>
#include <nori/scene.h>
#include <random>
#include <sh/default_image.h>
#include <sh/spherical_harmonics.h>
#include <stb_image.h>

NORI_NAMESPACE_BEGIN

namespace ProjEnv
{
    std::vector<std::unique_ptr<float[]>>
    LoadCubemapImages(const std::string &cubemapDir, int &width, int &height,
                      int &channel)
    {
        std::vector<std::string> cubemapNames{"negx.jpg", "posx.jpg", "posy.jpg",
                                              "negy.jpg", "posz.jpg", "negz.jpg"};
        std::vector<std::unique_ptr<float[]>> images(6);
        for (int i = 0; i < 6; i++)
        {
            std::string filename = cubemapDir + "/" + cubemapNames[i];
            int w, h, c;
            float *image = stbi_loadf(filename.c_str(), &w, &h, &c, 3);
            if (!image)
            {
                std::cout << "Failed to load image: " << filename << std::endl;
                exit(-1);
            }
            if (i == 0)
            {
                width = w;
                height = h;
                channel = c;
            }
            else if (w != width || h != height || c != channel)
            {
                std::cout << "Dismatch resolution for 6 images in cubemap" << std::endl;
                exit(-1);
            }
            images[i] = std::unique_ptr<float[]>(image);
            int index = (0 * 128 + 0) * channel;
            // std::cout << images[i][index + 0] << "\t" << images[i][index + 1] << "\t"
            //           << images[i][index + 2] << std::endl;
        }
        return images;
    }

    const Eigen::Vector3f cubemapFaceDirections[6][3] = {
        {{0, 0, 1}, {0, -1, 0}, {-1, 0, 0}},  // negx
        {{0, 0, 1}, {0, -1, 0}, {1, 0, 0}},   // posx
        {{1, 0, 0}, {0, 0, -1}, {0, -1, 0}},  // negy
        {{1, 0, 0}, {0, 0, 1}, {0, 1, 0}},    // posy
        {{-1, 0, 0}, {0, -1, 0}, {0, 0, -1}}, // negz
        {{1, 0, 0}, {0, -1, 0}, {0, 0, 1}},   // posz
    };

    float CalcPreArea(const float &x, const float &y)
    {
        return std::atan2(x * y, std::sqrt(x * x + y * y + 1.0));
    }

    float CalcArea(const float &u_, const float &v_, const int &width,
                   const int &height)
    {
        // transform from [0..res - 1] to [- (1 - 1 / res) .. (1 - 1 / res)]
        // ( 0.5 is for texel center addressing)
        float u = (2.0 * (u_ + 0.5) / width) - 1.0;
        float v = (2.0 * (v_ + 0.5) / height) - 1.0;

        // shift from a demi texel, mean 1.0 / size  with u and v in [-1..1]
        float invResolutionW = 1.0 / width;
        float invResolutionH = 1.0 / height;

        // u and v are the -1..1 texture coordinate on the current face.
        // get projected area for this texel
        float x0 = u - invResolutionW;
        float y0 = v - invResolutionH;
        float x1 = u + invResolutionW;
        float y1 = v + invResolutionH;
        float angle = CalcPreArea(x0, y0) - CalcPreArea(x0, y1) -
                      CalcPreArea(x1, y0) + CalcPreArea(x1, y1);

        return angle;
    }

    // template <typename T> T ProjectSH() {}

    template <size_t SHOrder>
    std::vector<Eigen::Array3f> PrecomputeCubemapSH(const std::vector<std::unique_ptr<float[]>> &images,
                                                    const int &width, const int &height,
                                                    const int &channel)
    {
        /// cubemapDirs，保存了cubemap 6张贴图上每个像素对应的单位向量
        std::vector<Eigen::Vector3f> cubemapDirs;
        cubemapDirs.reserve(6 * width * height);

        /// 写入cubemapDirs
        for (int i = 0; i < 6; i++)
        {
            Eigen::Vector3f faceDirX = cubemapFaceDirections[i][0];
            Eigen::Vector3f faceDirY = cubemapFaceDirections[i][1];
            Eigen::Vector3f faceDirZ = cubemapFaceDirections[i][2];
            for (int y = 0; y < height; y++)
            {
                for (int x = 0; x < width; x++)
                {
                    float u = 2 * ((x + 0.5) / width) - 1;
                    float v = 2 * ((y + 0.5) / height) - 1;
                    // Eigen::Vector3f dir = (faceDirX * u + faceDirY * v + faceDirZ).normalized();
                    // cubemapDirs.push_back(dir);
                    cubemapDirs.emplace_back((faceDirX * u + faceDirY * v + faceDirZ).normalized());
                }
            }
        }

        constexpr int SHNum = (SHOrder + 1) * (SHOrder + 1);

        //        std::vector<Eigen::Array3f> SHCoeffiecents(SHNum);
        //        for (int i = 0; i < SHNum; i++)
        //            SHCoeffiecents[i] = Eigen::Array3f(0);
        std::vector<Eigen::Array3f> SHCoeffiecents(SHNum, Eigen::Array3f(0));
        /// fill up SHCoeffiecents

        /// unused?
        // float sumWeight = 0;
        for (int i = 0; i < 6; i++)
        {
            for (int y = 0; y < height; y++)
            {
                for (int x = 0; x < width; x++)
                {
                    // TODO: here you need to compute light sh of each face of cubemap of each pixel
                    // TODO: 此处你需要计算每个像素下cubemap某个面的球谐系数

                    /// Note: requiring double in EvalSH()
                    const Eigen::Vector3d dir = cubemapDirs[i * width * height + y * width + x].cast<double>().normalized();
                    int index = (y * width + x) * channel;
                    /// L_env, environment light
                    Eigen::Array3f Le(images[i][index + 0], images[i][index + 1], images[i][index + 2]);

                    /// finish code here
                    for (int level = 0; level <= SHOrder; ++level)
                    {
                        for (int m = -level; m <= level; ++m)
                        {
                            SHCoeffiecents[sh::GetIndex(level, m)] += Le * sh::EvalSH(level, m, dir) * CalcArea(x, y, width, height);
                        }
                    }
                }
            }
        }
        return SHCoeffiecents;
    }
} // namespace ProjEnv

class PRTIntegrator : public Integrator
{
public: // static
    static constexpr int SHOrder = 2;

    static constexpr int SHCoeffLength = (SHOrder + 1) * (SHOrder + 1);

public: // class definition
    enum class Type
    {
        Unshadowed = 0,
        Shadowed = 1,
        Interreflection = 2
    };

public: // constructor
    PRTIntegrator(const PropertyList &props)
    {
        /* No parameters this time */
        m_SampleCount = props.getInteger("PRTSampleCount", 100);
        m_CubemapPath = props.getString("cubemap");
        auto type = props.getString("type", "unshadowed");
        if (type == "unshadowed")
        {
            m_Type = Type::Unshadowed;
        }
        else if (type == "shadowed")
        {
            m_Type = Type::Shadowed;
        }
        else if (type == "interreflection")
        {
            m_Type = Type::Interreflection;
            m_Bounce = props.getInteger("bounce", 1);
        }
        else
        {
            throw NoriException("Unsupported type: %s.", type);
        }
    }

public: // virtual func
    void preprocess(const Scene *scene) override
    {
        // Here only compute one mesh
        const auto mesh = scene->getMeshes()[0];
        // Projection environment
        /// get cubePath (file path)
        auto cubePath = getFileResolver()->resolve(m_CubemapPath);

        int width, height, channel;
        std::vector<std::unique_ptr<float[]>> images = ProjEnv::LoadCubemapImages(cubePath.str(), width, height, channel);

        /// write
        auto lightPath = cubePath / "light.txt";
        auto transPath = cubePath / "transport.txt";

        std::ofstream lightFout(lightPath.str());
        std::ofstream fout(transPath.str());

        /// get pre-computed SH coeffs
        auto envCoeffs = ProjEnv::PrecomputeCubemapSH<SHOrder>(images, width, height, channel);

        m_LightCoeffs.resize(3, SHCoeffLength);
        for (int i = 0; i < envCoeffs.size(); i++)
        {
            lightFout << (envCoeffs)[i].x() << " " << (envCoeffs)[i].y() << " " << (envCoeffs)[i].z() << std::endl;
            m_LightCoeffs.col(i) = (envCoeffs)[i];
        }
        std::cout << "Computed light sh coeffs from: " << cubePath.str() << " to: " << lightPath.str() << std::endl;

        // Projection transport
        /// (SHOrder+1)^2 rows, VertexCount cols
        /// SHCoeffLength = (SHOrder+1)^2
        m_TransportSHCoeffs.resize(SHCoeffLength, mesh->getVertexCount());
        /// write in file
        fout << mesh->getVertexCount() << std::endl;

        for (int i = 0; i < mesh->getVertexCount(); i++)
        {
            const Point3f &v = mesh->getVertexPositions().col(i);
            const Normal3f &n = mesh->getVertexNormals().col(i);

            /// input: w_i, output: L(w_i) * max(0, cos(\theta)).
            auto shFunc = [&](double phi, double theta) -> double
            {
                Eigen::Array3d d = sh::ToVector(phi, theta);

                /// sample vector
                const auto wi = Vector3f(d.x(), d.y(), d.z());

                double H = wi.dot(n);

                if (H < 0.0)
                    return 0.0;
                if (m_Type == Type::Unshadowed)
                {
                    return H;
                    // TODO: here you need to calculate unshadowed transport term of a given direction
                    // TODO: 此处你需要计算给定方向下的unshadowed传输项球谐函数值
                    // return 0;
                }
                else
                {
                    // test if there should be shadow
                    return (scene->rayIntersect(Ray3f(v, wi))) ? 0.0 : H;

                    // TODO: here you need to calculate shadowed transport term of a given direction
                    // TODO: 此处你需要计算给定方向下的shadowed传输项球谐函数值
                    // return 0;
                }
            };

            /// lambda is passed to sh::ProjectFunction
            /// sh::ProjectFunction returns std::vector<double>
            auto shCoeff = sh::ProjectFunction(SHOrder, shFunc, m_SampleCount);
            for (int j = 0; j < shCoeff->size(); j++)
            {
                m_TransportSHCoeffs.col(i).coeffRef(j) = (*shCoeff)[j];
            }
        }
        if (m_Type == Type::Interreflection)
        {
            Eigen::MatrixXf interreflectionCoefs;
            interreflectionCoefs.resize(SHCoeffLength, mesh->getVertexCount());

            for (int i = 0; i < mesh->getVertexCount(); ++i)
            {
                bool vis = false;
                for (int j = 0; j < SHCoeffLength; ++j)
                {
                    if (m_TransportSHCoeffs.col(i).coeffRef(j) != 0)
                    {
                        vis = true;
                        break;
                    }
                }
                if (vis)
                    continue;
                const Point3f v = mesh->getVertexPositions().col(i);
                const Normal3f n = mesh->getVertexNormals().col(i);
                auto shFunc = [&](double phi, double theta) -> std::vector<double>
                {
                    Eigen::Array3d d = sh::ToVector(phi, theta);

                    /// sample vector
                    const auto wi = Vector3f(d.x(), d.y(), d.z());

                    double H = wi.dot(n);
                    std::vector<double> ans(SHCoeffLength, 0.0);
                    if (H < 0.0)
                        return ans;
                    Intersection it;
                    bool intersect = scene->rayIntersect(Ray3f(v, wi), it);
                    if (!intersect)
                        return ans;
                    auto a = m_TransportSHCoeffs.col(it.tri_index.x());
                    auto b = m_TransportSHCoeffs.col(it.tri_index.y());
                    auto c = m_TransportSHCoeffs.col(it.tri_index.z());
                    auto x = it.bary.x();
                    auto y = it.bary.y();
                    auto z = it.bary.z();
                    auto result = (a * x + b * y + c * z) * H;
                    for (int i = 0; i < ans.size(); ++i)
                    {
                        ans[i] = result[i];
                    }
                    return ans;
                };

                const int sample_side = static_cast<int>(floor(sqrt(m_SampleCount)));
                // auto shCoeff = sh::ProjectFunction(SHOrder, shFunc, m_SampleCount);

                std::random_device rd;
                std::mt19937 gen(rd());
                std::uniform_real_distribution<> rng(0.0, 1.0);
                for (int t = 0; t < sample_side; t++)
                {
                    for (int p = 0; p < sample_side; p++)
                    {
                        double alpha = (t + rng(gen)) / sample_side;
                        double beta = (p + rng(gen)) / sample_side;
                        // See http://www.bogotobogo.com/Algorithms/uniform_distribution_sphere.php
                        double phi = 2.0 * M_PI * beta;
                        double theta = acos(2.0 * alpha - 1.0);

                        // evaluate the analytic function for the current spherical coords
                        auto func_value = shFunc(phi, theta);

                        // evaluate the SH basis functions up to band O, scale them by the
                        // function's value and accumulate them over all generated samples
                        for (int j = 0; j < SHCoeffLength; ++j)
                        {
                            interreflectionCoefs.col(i).coeffRef(j) += func_value[j];
                        }
                    }
                }

                // scale by the probability of a particular sample, which is
                // 4pi/sample_side^2. 4pi for the surface area of a unit sphere, and
                // 1/sample_side^2 for the number of samples drawn uniformly.
                interreflectionCoefs.col(i) *= 4.0 * M_PI / (sample_side * sample_side);
            }
            // TODO: leave for bonus

            for (int i = 0; i < mesh->getVertexCount(); ++i)
            {
                m_TransportSHCoeffs.col(i) += interreflectionCoefs.col(i);
            }
        }

        // Save in face format
        for (int f = 0; f < mesh->getTriangleCount(); f++)
        {
            const MatrixXu &F = mesh->getIndices();
            uint32_t idx0 = F(0, f), idx1 = F(1, f), idx2 = F(2, f);
            for (int j = 0; j < SHCoeffLength; j++)
            {
                fout << m_TransportSHCoeffs.col(idx0).coeff(j) << " ";
            }
            fout << std::endl;
            for (int j = 0; j < SHCoeffLength; j++)
            {
                fout << m_TransportSHCoeffs.col(idx1).coeff(j) << " ";
            }
            fout << std::endl;
            for (int j = 0; j < SHCoeffLength; j++)
            {
                fout << m_TransportSHCoeffs.col(idx2).coeff(j) << " ";
            }
            fout << std::endl;
        }
        std::cout << "Computed SH coeffs"
                  << " to: " << transPath.str() << std::endl;
    }

public: // member func
    Color3f Li(const Scene *scene, Sampler *sampler, const Ray3f &ray) const
    {
        Intersection its;
        if (!scene->rayIntersect(ray, its))
            return Color3f(0.0f);

        const Eigen::Matrix<Vector3f::Scalar, SHCoeffLength, 1> sh0 = m_TransportSHCoeffs.col(its.tri_index.x()),
                                                                sh1 = m_TransportSHCoeffs.col(its.tri_index.y()),
                                                                sh2 = m_TransportSHCoeffs.col(its.tri_index.z());
        const Eigen::Matrix<Vector3f::Scalar, SHCoeffLength, 1> rL = m_LightCoeffs.row(0), gL = m_LightCoeffs.row(1), bL = m_LightCoeffs.row(2);

        Color3f c0 = Color3f(rL.dot(sh0), gL.dot(sh0), bL.dot(sh0)),
                c1 = Color3f(rL.dot(sh1), gL.dot(sh1), bL.dot(sh1)),
                c2 = Color3f(rL.dot(sh2), gL.dot(sh2), bL.dot(sh2));

        const Vector3f &bary = its.bary;
        Color3f c = bary.x() * c0 + bary.y() * c1 + bary.z() * c2;
        // TODO: you need to delete the following four line codes after finishing your calculation to SH,
        //       we use it to visualize the normals of model for debug.
        // TODO: 在完成了球谐系数计算后，你需要删除下列四行，这四行代码的作用是用来可视化模型法线
        // if (c.isZero()) {
        //     auto n_ = its.shFrame.n.cwiseAbs();
        //     return Color3f(n_.x(), n_.y(), n_.z());
        // }
        return c;
    }

    std::string toString() const
    {
        return "PRTIntegrator[]";
    }

private:
    Type m_Type;
    int m_Bounce = 1;
    int m_SampleCount = 100;
    std::string m_CubemapPath;
    Eigen::MatrixXf m_TransportSHCoeffs;
    Eigen::MatrixXf m_LightCoeffs;
};

NORI_REGISTER_CLASS(PRTIntegrator, "prt");
NORI_NAMESPACE_END