#ifndef PROMP_EMISSION_H
#define PROMP_EMISSION_H

#include <armadillo>
#include <emissions.hpp>
#include <iostream>
#include <json.hpp>
#include <ForwardBackward.hpp>
#include <map>
#include <memory>
#include <robotics.hpp>

using namespace arma;
using namespace hsmm;
using namespace robotics;
using namespace std;

namespace hsmm {

    class NormalInverseWishart {
        public:
            NormalInverseWishart(vec mu_0, double lambda, mat Phi, int dof) :
                    Phi_(Phi), dof_(dof), mu_0_(mu_0), lambda_(lambda) {
                assert(dof > Phi.n_rows + 1);
                vec eigenvalues = eig_sym(Phi);
                assert(eigenvalues(0) > 0);
                assert(!(lambda < 0));
                assert(mu_0.n_rows == Phi.n_rows);
            }

            // This constructor assumes there is no prior for the mean.
            // (i.e. only the inverse-Wishart part).
            NormalInverseWishart(mat Phi, int dof) : NormalInverseWishart(
                    zeros<vec>(Phi.n_rows), 0.0, Phi, dof) {}

            vec getMu0() {
                return mu_0_;
            }

            double getLambda() {
                return lambda_;
            }

            mat getPhi() {
                return Phi_;
            }

            int getDof() {
                return dof_;
            }

        private:
            vec mu_0_;
            double lambda_;
            mat Phi_;
            int dof_;
    };

    double zeroMeanGaussianLogLikelihood(const vec& x, const mat &precision) {
        return -0.5 * as_scalar(x.t()*precision*x - log(det(precision)) +
                x.n_elem * log(2*datum::pi));
    }

    class ProMPsEmission : public AbstractEmissionOnlineSetting {
        public:
            ProMPsEmission(vector<FullProMP> promps) :
                    AbstractEmissionOnlineSetting(promps.size(),
                    promps.at(0).get_num_joints()), promps_(promps),
                    diagonal_sigma_y_(true) {
                for(int i = 0; i < getNumberStates(); i++)
                    assert(promps_.at(i).get_num_joints() == getDimension());
            }

            ProMPsEmission* clone() const {
                return new ProMPsEmission(*this);
            }

            void resetMemCaches() {
                cachePhis_.clear();
                cacheInvS_.clear();
                cacheK_.clear();
                cachePosteriorSigma_.clear();
            }

            // TODO: remove this once this behavior is moved to a subclass.
            void setDelta(double delta) {
                //assert(delta > 0);
                sample_locations_delta_ = delta;
            }

            // Derived classes could implement a different kinf of dependence
            // on the segment duration. Default implementation normalizes
            // the input to be between 0 and 1.
            // TODO: Implement the segment independent in a subclass.
            virtual vec getSampleLocations(int length) const {
                if (sample_locations_delta_ < 0)
                    return linspace<vec>(0, 1.0, length);
                else
                    return linspace<vec>(0, (length-1)*sample_locations_delta_,
                            length);
            }

            // This initialization mechanism assumes all the hidden states
            // share the same basis functions and their hyperparameters.
            void init_params_from_data(int min_duration, int ndurations,
                    const arma::field<arma::field<arma::mat>>& mobs) {
                vector<pair<double, vec>> pairs;
                vector<double> noise_vars;
                for(auto &obs : mobs) {
                    for(int t = 0; t < obs.n_elem; t++) {
                        for(int d = 0; d < ndurations; d++) {
                            if (t + min_duration + d > obs.n_elem)
                                break;
                            int end_idx = t + min_duration + d - 1;
                            auto &segment = obs.rows(t, end_idx);
                            vec lsq_omega = least_squares_omega(0, segment);
                            double var = var_isotropic_gaussian_given_omega(0,
                                    lsq_omega, segment);
                            pairs.push_back(make_pair(var, lsq_omega));
                            noise_vars.push_back(var);
                        }
                    }
                }
                sort(pairs.begin(), pairs.end(), [](const pair<double, vec>& a,
                        const pair<double, vec>& b)
                        {return a.first < b.first;});
                int cutoff_index = (int)((pairs.size() - 1) *
                        init_fraction_);
                double threshold = pairs.at(cutoff_index).first;
                cout << "Cuttoff index for init: " << cutoff_index <<
                        " var: " << threshold << endl;

                // Showing a histogram of the distribution of noise vars.
                vec vars(noise_vars);
                vec histogram_cutoffs = linspace<vec>(0, max(vars), 20);
                uvec histogram = histc(vars, histogram_cutoffs);
                cout << "Hist. of noise vars:" << endl << histogram << endl;

                mat remaining_w(promps_.at(0).get_model().get_mu_w().n_rows,
                        cutoff_index + 1);
                for(int i = 0; i <= cutoff_index; i++)
                    remaining_w.col(i) = pairs.at(i).second;
                mat means;
                kmeans(means, remaining_w, getNumberStates(), static_subset,
                        10, false);

                // Updating the means of the ProMPs based on the k-means.
                for(int i = 0; i < getNumberStates(); i++) {
                    ProMP promp = promps_.at(i).get_model();
                    promp.set_mu_w(means.col(i));
                    promps_.at(i).set_model(promp);
                }
            }

            // Note that the method is returning the posterior covariance for
            // the given state and duration which is independent of obs.
            mat generateCachedMatrices(const pair<int, int> &p) const {
                if (cacheInvS_.find(p) == cacheInvS_.end() ||
                        cacheK_.find(p) == cacheK_.end()) {
                    int state = p.first;
                    int dur = p.second;
                    const FullProMP& promp = promps_.at(state);
                    const cube& Phis = getPhiCube(state, dur);
                    mat Sigma(promp.get_model().get_Sigma_w());
                    mat Sigma_y(promp.get_model().get_Sigma_y());
                    field<mat> invS(dur);
                    field<mat> K(dur);
                    for(int i = 0; i < dur; i++) {
                        const mat& Phi = Phis.slice(i);
                        mat S = Phi * Sigma * Phi.t() + Sigma_y;
                        invS(i) = inv_sympd(S);

                        // Kalman updating (correcting) step.
                        K(i) = Sigma * Phi.t() * inv(S);
                        Sigma = Sigma - K(i) * S * K(i).t();
                    }

                    // Caching the matrices for faster likelihood evaluation.
                    cacheInvS_[p] = invS;
                    cacheK_[p] = K;
                    cachePosteriorSigma_[p] = Sigma;
                }
                return cachePosteriorSigma_.at(p);
            }

            double loglikelihood(int state, const field<mat>& obs) const {

                // Making sure all the required matrices are already
                // precomputed.
                pair<int, int> p = make_pair(state, obs.n_elem);
                generateCachedMatrices(p);

                const FullProMP& promp = promps_.at(state);
                const cube& Phis = getPhiCube(state, obs.n_elem);
                vec mu(promp.get_model().get_mu_w());
                const field<mat>& invS = cacheInvS_[p];
                const field<mat>& K = cacheK_[p];
                double ret = 0;
                for(int i = 0; i < obs.n_elem; i++) {
                    const mat& Phi = Phis.slice(i);
                    if (obs(i).is_empty()) {

                        // Making sure all the missing obs are at the end.
                        // Other missing obs patterns are not supported yet.
                        for(int j = i; j < obs.n_elem; j++)
                            assert(obs(j).is_empty());
                        break;
                    }
                    vec diff = obs(i) - Phi * mu;

                    // p(y_t | y_1, ..., y_{t-1}).
                    ret += zeroMeanGaussianLogLikelihood(diff, invS(i));
                    mu = mu + K(i) * diff;
                }
                return ret;
            }

            double informationFilterLoglikelihood(int state,
                    const field<mat>& obs) const {
                const FullProMP& promp = promps_.at(state);

                // The samples are assumed to be equally spaced.
                vec sample_locations = getSampleLocations(obs.n_elem);

                // Moment based parameterization.
                vec mu(promp.get_model().get_mu_w());
                mat Sigma(promp.get_model().get_Sigma_w());
                mat Sigma_y(promp.get_model().get_Sigma_y());

                // Canonical parameterization.
                mat information_matrix = inv_sympd(Sigma);
                vec information_state = information_matrix * mu;
                mat inv_obs_noise = inv_sympd(Sigma_y);

                double ret = 0;
                for(int i = 0; i < obs.n_elem; i++) {
                    mat Phi = promp.get_phi_t(sample_locations(i));

                    // Computing the likelihood under the current filtering
                    // distribution.
                    mat filtered_Sigma = inv_sympd(information_matrix);
                    vec filtered_mu = filtered_Sigma * information_state;

                    // p(y_t | y_{1:t-1}).
                    random::NormalDist dist = random::NormalDist(
                            Phi * filtered_mu,
                            Phi * filtered_Sigma * Phi.t() + Sigma_y);
                    ret = ret + log_normal_density(dist, obs(i));

                    mat aux = Phi.t() * inv_obs_noise;
                    mat I_k = aux * Phi;
                    mat i_k = aux * obs(i);
                    information_matrix = information_matrix + I_k;
                    information_state = information_state + i_k;
                }
                return ret;
            }

            double loglikelihoodBatchVersion(int state, const arma::mat& obs) {
                int dimension = obs.n_rows;
                assert(dimension == getDimension());
                random::NormalDist dist = getNormalDistForMultipleTimeSteps(state,
                        obs.n_cols);
                vec stacked_obs = vectorise(obs);
                return random::log_normal_density(dist, stacked_obs);
            }

            void reestimate(int min_duration,
                    const arma::field<arma::cube>& meta,
                    const arma::field<arma::field<arma::mat>>& mobs) {
                int nseq = mobs.n_elem;
                for(int i = 0; i < getNumberStates(); i++) {
                    ProMP promp = promps_.at(i).get_model();
                    const mat inv_Sigma_w = inv_sympd(promp.get_Sigma_w());
                    const mat inv_Sigma_y = inv_sympd(promp.get_Sigma_y());
                    const vec mu_w = promp.get_mu_w();

                    vector<double> mult_c;
                    vector<double> denominator_Sigma_y;
                    for(int s = 0; s < nseq; s++) {
                        const cube& eta = meta(s);
                        int nobs = mobs(s).n_elem;
                        int ndurations = eta.n_cols;
                        for(int t = min_duration - 1; t < nobs; t++) {
                            for(int d = 0; d < ndurations; d++) {
                                int first_idx_seg = t - min_duration - d + 1;
                                if (first_idx_seg < 0)
                                    break;
                                mult_c.push_back(eta(i, d, t));
                                denominator_Sigma_y.push_back(eta(i, d, t) +
                                        log(min_duration + d));
                            }
                        }
                    }

                    // Computing the multiplicative constants for mu_w and
                    // Sigma_w.
                    vec mult_c_normalized(mult_c);
                    mult_c_normalized -= logsumexp(mult_c_normalized);
                    mult_c_normalized = exp(mult_c_normalized);

                    // Computing the multiplicative constants for Sigma_y since
                    // they have a different denominator.
                    vec mult_c_Sigma_y_normalized(mult_c);
                    vec den_Sigma_y(denominator_Sigma_y);
                    mult_c_Sigma_y_normalized -= logsumexp(den_Sigma_y);
                    mult_c_Sigma_y_normalized = exp(mult_c_Sigma_y_normalized);

                    mat new_Sigma_y(size(promp.get_Sigma_y()), fill::zeros);

                    // EM for ProMPs.
                    vec weighted_sum_post_mean(size(mu_w), fill::zeros);
                    mat weighted_sum_post_cov(size(inv_Sigma_w), fill::zeros);
                    mat weighted_sum_post_mean_mean_T(size(inv_Sigma_w),
                            fill::zeros);
                    int idx_mult_c = 0;
                    for(int s = 0; s < nseq; s++) {
                        auto& obs = mobs(s);
                        int nobs = obs.n_elem;
                        int ndurations = meta(s).n_cols;
                        for(int t = min_duration - 1; t < nobs; t++) {
                            for(int d = 0; d < ndurations; d++) {
                                int first_idx_seg = t - min_duration - d + 1;
                                if (first_idx_seg < 0)
                                    break;
                                const int current_duration = min_duration + d;
                                const cube& Phis = getPhiCube(i, current_duration);

                                // E step for the emission hidden variables (Ws).
                                // Computing the posterior of W given Y and Theta.

                                // Computing the posterior covariance of the hidden
                                // variable w for this segment.
                                mat posterior_cov(size(inv_Sigma_w), fill::zeros);
                                for(int step = 0; step < current_duration; step++) {
                                    posterior_cov += Phis.slice(step).t() *
                                        inv_Sigma_y * Phis.slice(step);
                                }
                                posterior_cov = (posterior_cov+posterior_cov.t())/2.0;
                                posterior_cov = posterior_cov + inv_Sigma_w;
                                posterior_cov = (posterior_cov+posterior_cov.t())/2.0;
                                posterior_cov = inv_sympd(posterior_cov);
                                posterior_cov = (posterior_cov+posterior_cov.t())/2.0;

                                // Computing the posterior mean of the hidden
                                // variable w for this segment.
                                vec posterior_mean(size(mu_w), fill::zeros);
                                for(int step = 0; step < current_duration; step++) {
                                    const vec& ob = obs(first_idx_seg + step);
                                    posterior_mean += Phis.slice(step).t() *
                                        inv_Sigma_y * ob;
                                }
                                posterior_mean = inv_Sigma_w * mu_w + posterior_mean;
                                posterior_mean = posterior_cov * posterior_mean;

                                // Getting the multiplicative constants.
                                double mult_constant = mult_c_normalized(idx_mult_c);
                                double mult_constant_Sigma_y =
                                        mult_c_Sigma_y_normalized(idx_mult_c);
                                idx_mult_c++;

                                // Statistics required for updating mu_w & Sigma_w.
                                weighted_sum_post_mean += mult_constant *
                                        posterior_mean;
                                weighted_sum_post_cov += mult_constant *
                                        posterior_cov;
                                weighted_sum_post_mean_mean_T += mult_constant *
                                        posterior_mean * posterior_mean.t();

                                // Computing the new output noise covariance: Sigma_y.
                                mat Sigma_y_term(size(new_Sigma_y), fill::zeros);
                                for(int step = 0; step < current_duration; step++) {
                                    const mat& phi = Phis.slice(step);
                                    const vec& diff_y = obs(first_idx_seg + step) -
                                        phi * posterior_mean;
                                    Sigma_y_term += diff_y * diff_y.t() +
                                        phi * posterior_cov * phi.t();
                                }
                                new_Sigma_y += mult_constant_Sigma_y * Sigma_y_term;
                            }
                        }
                    }

                    // Expected number of segments generated by the i-th state.
                    double mle_den = exp(logsumexp(mult_c));

                    // M step for the emission variables.
                    vec new_mu_w_MLE(weighted_sum_post_mean);
                    mat new_Sigma_w_MLE = weighted_sum_post_cov +
                        weighted_sum_post_mean_mean_T - new_mu_w_MLE *
                        new_mu_w_MLE.t();

                    // If there is a prior then we do MAP instead.
                    mat new_Sigma_w;
                    mat new_mu_w;
                    if (normal_inverse_prior_) {
                        double v_0 = normal_inverse_prior_->getDof();
                        double D = mu_w.n_rows;
                        mat S_0 = normal_inverse_prior_->getPhi();
                        new_Sigma_w = (S_0 + mle_den * new_Sigma_w_MLE) /
                                (v_0 + mle_den + D + 2);

                        double k_0 = normal_inverse_prior_->getLambda();
                        vec m_0 = normal_inverse_prior_->getMu0();
                        new_mu_w = (k_0 * m_0 + mle_den * new_mu_w_MLE) /
                                (mle_den + k_0);
                    }
                    else {
                        new_Sigma_w = new_Sigma_w_MLE;
                        new_mu_w = new_mu_w_MLE;
                    }

                    cout << "State " << i << " MLE Den: " << mle_den << " ";
                    if (mle_den > epsilon_) {

                        // Making sure the noise covariance is diagonal.
                        if (diagonal_sigma_y_)
                            new_Sigma_y = diagmat(new_Sigma_y.diag());

                        // Checking that the new Sigma_w is a covariance matrix.
                        vec eigenvalues_map = eig_sym(new_Sigma_w);
                        assert(eigenvalues_map(0) > 0);

                        // Setting the new parameters.
                        promp.set_mu_w(new_mu_w);
                        promp.set_Sigma_w(new_Sigma_w);
                        promp.set_Sigma_y(new_Sigma_y);
                        promps_.at(i).set_model(promp);
                        resetMemCaches();
                        cout << ". Updated." << endl;
                    }
                    else
                        cout << ". Not updated." << endl;
                }
            }

            // Unshadowing this method from the AbstractEmission class.
            using AbstractEmission::sampleFromState;

            field<mat> sampleFromState(int state, int size,
                    mt19937 &rand_generator) const {
                return sampleFromProMP(promps_.at(state), size, rand_generator);
            }

            field<mat> sampleFromProMP(const FullProMP& fpromp, int size,
                    mt19937 &rand_generator) const {
                const ProMP& model = fpromp.get_model();
                assert(model.get_mu_w().n_rows == model.get_Sigma_w().n_rows);
                vector<vec> w_samples = random::sample_multivariate_normal(
                        rand_generator,
                        {model.get_mu_w(), model.get_Sigma_w()}, 1);
                vec w = w_samples.back();

                vec noise_mean = zeros<vec>(getDimension());
                vector<vec> output_noise = random::sample_multivariate_normal(
                        rand_generator, {noise_mean, model.get_Sigma_y()}, size);

                field<mat> ret(size);

                // The samples are assumed to be equally spaced.
                vec sample_locations = getSampleLocations(size);
                for(int i = 0; i < size; i++) {
                    double z = sample_locations(i);
                    mat phi_z =fpromp.get_phi_t(z);
                    ret(i) = phi_z * w + output_noise.at(i);
                }
                return ret;
            }

            // Equivalent to sampleFromState but uses conditioning.
            field<mat> sampleFromState2(int state, int size,
                    mt19937 &rand_generator) const {
                field<mat> no_obs;
                return sampleNextObsGivenPastObs(state, size, no_obs,
                        rand_generator);
            }

            field<mat> sampleNextObsGivenPastObs(int state, int seg_dur,
                    const field<mat>& past_obs, std::mt19937 &rng) const {
                assert(past_obs.n_elem < seg_dur);

                // Making sure all the required matrices are already computed.
                pair<int, int> p = make_pair(state, seg_dur);
                generateCachedMatrices(p);

                field<mat> ret(seg_dur);
                for(int i = 0; i < past_obs.n_elem; i++)
                    ret(i) = past_obs(i);

                const FullProMP& promp = promps_.at(state);
                const cube& Phis = getPhiCube(state, seg_dur);
                vec mu(promp.get_model().get_mu_w());

                const field<mat>& invS = cacheInvS_.at(p);
                const field<mat>& K = cacheK_.at(p);
                int i;
                for(i = 0; i < past_obs.n_elem; i++)
                    mu = mu + K(i) * (past_obs(i) - Phis.slice(i) * mu);

                // Now i indexes the offset we want to sample from.
                for(; i < seg_dur; i++) {
                    vector<vec> sample = random::sample_multivariate_normal(
                            rng, {Phis.slice(i) * mu, inv(invS(i))}, 1);
                    ret(i) = sample.at(0);
                    mu = mu + K(i) * (ret(i) - Phis.slice(i) * mu);
                }
                return ret;
            }

            field<mat> sampleFirstSegmentObsGivenLastSegment(int curr_state,
                    int curr_seg_dur, const field<mat> &last_segment,
                    int last_state, std::mt19937 &rng) const {
                mat last_obs = last_segment(last_segment.n_elem - 1);

                // Making sure all the required matrices are already computed.
                pair<int, int> p = make_pair(curr_state, curr_seg_dur);
                pair<int, int> last_p = make_pair(last_state,
                        last_segment.n_elem);
                generateCachedMatrices(p);
                mat last_p_Sigma = generateCachedMatrices(last_p);

                // Finding the posterior omega mean for the last segment.
                const cube& LPhis = getPhiCube(last_state, last_segment.n_elem);
                vec last_p_mu(promps_.at(last_state).get_model().get_mu_w());
                const field<mat>& last_K = cacheK_[last_p];
                for(int i = 0; i < last_segment.n_elem; i++) {
                    const mat& Phi = LPhis.slice(i);
                    vec diff = last_segment(i) - Phi * last_p_mu;
                    last_p_mu = last_p_mu + last_K(i) * diff;
                }
                FullProMP last_full_promp(promps_.at(last_state));
                mat zeros_Sigma_y = zeros<mat>(getDimension(), getDimension());
                ProMP last_p_promp(last_p_mu, last_p_Sigma, zeros_Sigma_y);
                last_full_promp.set_model(last_p_promp);

                // This gives the posterior distribution over q = (y, v) at the
                // last time step of the last segment.
                random::NormalDist last_pos_vel = last_full_promp.joint_dist(
                        1.0, true, true, false);
                vec last_p_pos = last_pos_vel.mean().head(getDimension());
                vec last_p_vel = last_pos_vel.mean().tail(getDimension());

                FullProMP promp = promps_.at(curr_state);
                promp = promp.condition_current_position(0, 1.0, last_p_pos);
                //promp = promp.condition_current_state(0, 1.0, last_p_pos,
                //        last_p_vel);
                return sampleFromProMP(promp, curr_seg_dur, rng);
            }

            nlohmann::json to_stream() const {
                vector<nlohmann::json> array_emission_params;
                for(int i = 0; i < getNumberStates(); i++) {
                    const ProMP& promp = promps_.at(i).get_model();
                    nlohmann::json whole_thing;
                    nlohmann::json promp_params;
                    promp_params["mu_w"] = vec2json(promp.get_mu_w());
                    promp_params["Sigma_w"] = mat2json(promp.get_Sigma_w());
                    promp_params["Sigma_y"] = mat2json(promp.get_Sigma_y());
                    whole_thing["model"] = promp_params;
                    whole_thing["num_joints"] = getDimension();

                    // Note that the information about the basis functions is not
                    // serialized.
                    array_emission_params.push_back(whole_thing);
                }
                nlohmann::json ret = array_emission_params;
                return ret;
            }

            void from_stream(const nlohmann::json& emission_params) {
                // Note that the given parameters need to be consistent with some
                // of the preexisting settings. Moreover, part of the structure of
                // this emission process is not read from the input json.
                // (e.g. the used basis functions).
                assert(emission_params.size() == getNumberStates());
                for(int i = 0; i < getNumberStates(); i++) {
                    const nlohmann::json& params = emission_params.at(i);
                    promps_.at(i).set_model(json2basic_promp(params.at("model")));
                    assert(params.at("num_joints") == getDimension());
                }
                resetMemCaches();
                return;
            }

            // TODO: change the name of this method.
            void set_Sigma_w_Prior(NormalInverseWishart prior) {
                normal_inverse_prior_ = std::make_shared<NormalInverseWishart>(
                        std::move(prior));
                int size_cov = promps_.at(0).get_model().get_Sigma_w().n_rows;
                assert(size_cov == normal_inverse_prior_->getPhi().n_rows);
            }

            void setParamsForInitialization(double fraction) {
                assert(fraction > 0 && fraction < 1.0);
                init_fraction_ = fraction;
            }

        protected:

            // Returns the marginal distribution of a particular state (ProMP) and
            // duration. Keep in mind that the covariance matrix grows quadratically
            // with respect to the duration. This could be cached for efficiency
            // but is mostly intended for debugging.
            random::NormalDist getNormalDistForMultipleTimeSteps(int state, int duration) {
                int nrows = getDimension();
                const FullProMP& promp = promps_.at(state);
                const mat stacked_Phi = getPhiStacked(state, duration);
                vec mean = stacked_Phi * promp.get_model().get_mu_w();
                mat cov = stacked_Phi * promp.get_model().get_Sigma_w() *
                    stacked_Phi.t();

                // Adding the noise variance.
                mat noise_cov(size(cov), fill::zeros);
                for(int i = 0; i < duration; i++)
                    noise_cov.submat(i * nrows, i * nrows, (i + 1) * nrows - 1,
                            (i + 1) * nrows - 1) = promp.get_model().get_Sigma_y();

                cov = cov + noise_cov;
                return random::NormalDist(mean, cov);
            }

            // Takes the covariance matrix for omega and returns it as if the
            // joints were independent (blockdiag operator in Sebastian's paper).
            mat getCovarianceIndependentJoints(const mat& cov) const {
                mat ret(size(cov), fill::zeros);
                assert(cov.n_rows == cov.n_cols);
                int dim = cov.n_rows;
                assert(dim % getDimension() == 0);
                int blocklen = dim / getDimension();
                for(int i = 0; i < dim; i++)
                    for(int j = 0; j < dim; j++)
                        if ((i/blocklen) == (j/blocklen))
                            ret(i, j) = cov(i, j);
                mat extra_term = eye<mat>(size(ret)) * 1e-4;
                return ret + extra_term;
            }

            mat getPhiStacked(int state, int duration) {
                cube Phis = getPhiCube(state, duration);
                mat PhiStacked(Phis.n_rows * Phis.n_slices, Phis.n_cols);
                for(int d = 0; d < duration; d++)
                    PhiStacked.rows(d * Phis.n_rows, (d + 1) * Phis.n_rows - 1) =
                        Phis.slice(d);
                return PhiStacked;
            }

            cube getPhiCube(int state, int duration) const {
                pair<int, int> p = make_pair(state, duration);
                if (cachePhis_.find(p) != cachePhis_.end())
                    return cachePhis_[p];
                const FullProMP& promp = promps_.at(state);

                // The samples are assumed to be equally spaced.
                vec sample_locations = getSampleLocations(duration);
                mat tmp = promp.get_phi_t(0);
                int ncols = tmp.n_cols;
                int nrows = tmp.n_rows;
                cube stacked_Phi(nrows, ncols, duration, fill::zeros);
                for(int i = 0; i < duration; i++)
                    stacked_Phi.slice(i) = promp.get_phi_t(sample_locations(i));
                cachePhis_[p] = stacked_Phi;
                return stacked_Phi;
            }

            // Returns the least squares solution to the problem:
            // y(t) = Phi(t) * w for a bunch of i.i.d. y's. Note that omega is
            // not assumed to be a random variable.
            vec least_squares_omega(int state,
                    const field<mat>& segment) const {
                int nobs = segment.n_elem;
                const cube& Phis = getPhiCube(state, nobs);
                mat A(size(Phis.slice(0).t() * Phis.slice(0)),
                        fill::zeros);
                vec b(A.n_rows, fill::zeros);
                for(int t = 0; t < nobs; t++) {
                    const mat& Phi = Phis.slice(t);
                    A = A + Phi.t() * Phi;
                    b = b + Phi.t() * segment(t);
                }

                // Solving A * w = b
                vec lsq_omega = solve(A, b);
                return lsq_omega;
            }

            // Compute the variance of an isotropic Gaussian noise model
            // given omega.
            double var_isotropic_gaussian_given_omega(int state, vec w,
                    const field<mat>& segment) const {
                int nobs = segment.n_elem;
                const cube& Phis = getPhiCube(state, nobs);
                double var = 0;
                for(int t = 0; t < nobs; t++) {
                    const mat& Phi = Phis.slice(t);
                    vec diff = segment(t) - Phi * w;
                    var += dot(diff, diff);
                }
                return var / (nobs * segment(0).n_rows);
            }

            vector<FullProMP> promps_;
            std::shared_ptr<NormalInverseWishart> normal_inverse_prior_;
            double epsilon_ = 1e-15;
            bool diagonal_sigma_y_;

            // Delta for emissions which are not dependent on the total
            // duration.
            double sample_locations_delta_ = -1;

            // Fraction of the total least squares omega estimates that will be
            // used for initialization.
            double init_fraction_ = 0.1;

            // Members for caching.
            mutable map<pair<int, int>, cube> cachePhis_;
            mutable map<pair<int, int>, field<mat>> cacheInvS_;
            mutable map<pair<int, int>, field<mat>> cacheK_;
            mutable map<pair<int, int>, mat> cachePosteriorSigma_;
    };

    // This version of the ProMP is specifically designed for HMMs where a
    // single observation is a time series itself
    class ProMPsEmissionHMM : public ProMPsEmission {
        public:
            ProMPsEmissionHMM(vector<FullProMP> promps) :
                    ProMPsEmission(promps) {}

            ProMPsEmission* clone() const {
                return new ProMPsEmission(*this);
            }

            double loglikelihood(int state, const field<mat>& obs) const {

                // Making sure that the duration is one.
                assert(obs.n_elem == 1);
                auto& sequence = obs(0);
                const FullProMP& promp = promps_.at(state);

                // The samples are assumed to be equally spaced.
                vec sample_locations = getSampleLocations(sequence.n_cols);

                vec mu(promp.get_model().get_mu_w());
                mat Sigma(promp.get_model().get_Sigma_w());
                mat Sigma_y(promp.get_model().get_Sigma_y());
                double ret = 0;
                for(int i = 0; i < sequence.n_cols; i++) {
                    mat Phi = promp.get_phi_t(sample_locations(i));
                    mat S = Phi * Sigma * Phi.t() + Sigma_y;

                    // Required for the marginal likelihood p(y_t | y_{1:t-1}).
                    random::NormalDist dist = random::NormalDist(Phi * mu, S);
                    ret = ret + log_normal_density(dist, sequence.col(i));

                    // Using the kalman updating step to compute this efficiently.
                    mat K = Sigma * Phi.t() * inv(S);
                    mu = mu + K * (sequence.col(i) - Phi * mu);
                    Sigma = Sigma - K * S * K.t();
                }
                return ret;
            }

            void reestimate(int min_duration,
                    const arma::field<arma::cube>& meta,
                    const arma::field<arma::field<arma::mat>>& mobs) {
                int nseq = mobs.n_elem;
                for(int i = 0; i < getNumberStates(); i++) {
                    ProMP promp = promps_.at(i).get_model();
                    const mat inv_Sigma_w = inv_sympd(promp.get_Sigma_w());
                    const mat inv_Sigma_y = inv_sympd(promp.get_Sigma_y());
                    const vec mu_w = promp.get_mu_w();

                    vector<double> mult_c;
                    vector<double> denominator_Sigma_y;
                    for(int s = 0; s < nseq; s++) {
                        const cube& eta = meta(s);
                        int nobs = mobs(s).n_elem;
                        int ndurations = eta.n_cols;
                        assert(ndurations == 1 && min_duration == 1);
                        for(int t = 0; t < nobs; t++) {
                            mult_c.push_back(eta(i, 0, t));
                            int current_duration = mobs(s)(t).n_cols;
                            denominator_Sigma_y.push_back(eta(i, 0, t) +
                                    log(current_duration));
                        }
                    }

                    // Computing the multiplicative constants for mu_w and
                    // Sigma_w.
                    vec mult_c_normalized(mult_c);
                    mult_c_normalized -= logsumexp(mult_c_normalized);
                    mult_c_normalized = exp(mult_c_normalized);

                    // Computing the multiplicative constants for Sigma_y since
                    // they have a different denominator.
                    vec mult_c_Sigma_y_normalized(mult_c);
                    vec den_Sigma_y(denominator_Sigma_y);
                    mult_c_Sigma_y_normalized -= logsumexp(den_Sigma_y);
                    mult_c_Sigma_y_normalized = exp(mult_c_Sigma_y_normalized);

                    mat new_Sigma_y(size(promp.get_Sigma_y()), fill::zeros);

                    // EM for ProMPs.
                    vec weighted_sum_post_mean(size(mu_w), fill::zeros);
                    mat weighted_sum_post_cov(size(inv_Sigma_w), fill::zeros);
                    mat weighted_sum_post_mean_mean_T(size(inv_Sigma_w),
                            fill::zeros);
                    int idx_mult_c = 0;
                    for(int s = 0; s < nseq; s++) {
                        auto& obs = mobs(s);
                        int nobs = obs.n_elem;
                        int ndurations = meta(s).n_cols;
                        for(int t = 0; t < nobs; t++) {

                                // Length of the current observation
                                const int current_duration = obs(t).n_cols;
                                const cube& Phis = getPhiCube(i, current_duration);

                                // E step for the emission hidden variables (Ws).
                                // Computing the posterior of W given Y and Theta.

                                // Computing the posterior covariance of the hidden
                                // variable w for this segment.
                                mat posterior_cov(size(inv_Sigma_w), fill::zeros);
                                for(int step = 0; step < current_duration; step++) {
                                    posterior_cov += Phis.slice(step).t() *
                                        inv_Sigma_y * Phis.slice(step);
                                }
                                posterior_cov = (posterior_cov+posterior_cov.t())/2.0;
                                posterior_cov = posterior_cov + inv_Sigma_w;
                                posterior_cov = (posterior_cov+posterior_cov.t())/2.0;
                                posterior_cov = inv_sympd(posterior_cov);
                                posterior_cov = (posterior_cov+posterior_cov.t())/2.0;

                                // Computing the posterior mean of the hidden
                                // variable w for this segment.
                                vec posterior_mean(size(mu_w), fill::zeros);
                                for(int step = 0; step < current_duration; step++) {
                                    const vec& ob = obs(t).col(step);
                                    posterior_mean += Phis.slice(step).t() *
                                        inv_Sigma_y * ob;
                                }
                                posterior_mean = inv_Sigma_w * mu_w + posterior_mean;
                                posterior_mean = posterior_cov * posterior_mean;

                                // Getting the multiplicative constants.
                                double mult_constant = mult_c_normalized(idx_mult_c);
                                double mult_constant_Sigma_y =
                                        mult_c_Sigma_y_normalized(idx_mult_c);
                                idx_mult_c++;

                                // Statistics required for updating mu_w & Sigma_w.
                                weighted_sum_post_mean += mult_constant *
                                        posterior_mean;
                                weighted_sum_post_cov += mult_constant *
                                        posterior_cov;
                                weighted_sum_post_mean_mean_T += mult_constant *
                                        posterior_mean * posterior_mean.t();

                                // Computing the new output noise covariance: Sigma_y.
                                mat Sigma_y_term(size(new_Sigma_y), fill::zeros);
                                for(int step = 0; step < current_duration; step++) {
                                    const mat& phi = Phis.slice(step);
                                    const vec& diff_y = obs(t).col(step) -
                                        phi * posterior_mean;
                                    Sigma_y_term += diff_y * diff_y.t() +
                                        phi * posterior_cov * phi.t();
                                }
                                new_Sigma_y += mult_constant_Sigma_y * Sigma_y_term;
                        }
                    }

                    // Expected number of segments generated by the i-th state.
                    double mle_den = exp(logsumexp(mult_c));

                    // M step for the emission variables.
                    vec new_mu_w_MLE(weighted_sum_post_mean);
                    mat new_Sigma_w_MLE = weighted_sum_post_cov +
                        weighted_sum_post_mean_mean_T - new_mu_w_MLE *
                        new_mu_w_MLE.t();

                    // If there is a prior then we do MAP instead.
                    mat new_Sigma_w;
                    mat new_mu_w;
                    if (normal_inverse_prior_) {
                        double v_0 = normal_inverse_prior_->getDof();
                        double D = mu_w.n_rows;
                        mat S_0 = normal_inverse_prior_->getPhi();
                        new_Sigma_w = (S_0 + mle_den * new_Sigma_w_MLE) /
                                (v_0 + mle_den + D + 2);

                        double k_0 = normal_inverse_prior_->getLambda();
                        vec m_0 = normal_inverse_prior_->getMu0();
                        new_mu_w = (k_0 * m_0 + mle_den * new_mu_w_MLE) /
                                (mle_den + k_0);
                    }
                    else {
                        new_Sigma_w = new_Sigma_w_MLE;
                        new_mu_w = new_mu_w_MLE;
                    }

                    cout << "State " << i << " MLE Den: " << mle_den << " ";
                    if (mle_den > epsilon_) {

                        // Making sure the noise covariance is diagonal.
                        if (diagonal_sigma_y_)
                            new_Sigma_y = diagmat(new_Sigma_y.diag());

                        // Checking that the new Sigma_w is a covariance matrix.
                        vec eigenvalues_map = eig_sym(new_Sigma_w);
                        assert(eigenvalues_map(0) > 0);

                        // Setting the new parameters.
                        promp.set_mu_w(new_mu_w);
                        promp.set_Sigma_w(new_Sigma_w);
                        promp.set_Sigma_y(new_Sigma_y);
                        promps_.at(i).set_model(promp);
                        cout << ". Updated." << endl;
                    }
                    else
                        cout << ". Not updated." << endl;
                }
            }

            field<mat> sampleFromState(int state, int size,
                    mt19937 &rand_generator) const {

                // Since the segments are modeled as single observations.
                assert(size == 1);
                const ProMP& model = promps_.at(state).get_model();
                vector<vec> w_samples = random::sample_multivariate_normal(
                        rand_generator, {model.get_mu_w(), model.get_Sigma_w()}, 1);
                vec w = w_samples.back();

                // Getting the actual size of the segment.
                size = getDurationForEachSegment(rand_generator);

                vec noise_mean = zeros<vec>(getDimension());
                vector<vec> output_noise = random::sample_multivariate_normal(
                        rand_generator, {noise_mean, model.get_Sigma_y()}, size);

                mat joint_sample(getDimension(), size);

                // The samples are assumed to be equally spaced.
                vec sample_locations = getSampleLocations(size);
                for(int i = 0; i < size; i++) {
                    double z = sample_locations(i);
                    mat phi_z = promps_.at(state).get_phi_t(z);
                    joint_sample.col(i) = phi_z * w + output_noise.at(i);
                }
                field<mat> ret = {joint_sample};
                return ret;
            }

        private:

            int getDurationForEachSegment(mt19937 &rand_generator) const {
                return 20;
            }

    };

};

#endif


/*
 * Mipmaps.hpp
 *
 *  Created on: Feb 4, 2012
 *      Author: david
 */

#ifndef INCLUDED_DENSITY_SCALEPYRAMID_HPP_
#define INCLUDED_DENSITY_SCALEPYRAMID_HPP_

#include <Slimage/Slimage.hpp>
#include <Eigen/Dense>
#include <stdexcept>
#include <vector>
#include <tuple>
#include <cassert>

namespace density {

Eigen::MatrixXf SumMipMapWithBlackBorder(const Eigen::MatrixXf& img_big);

template<unsigned int Q>
Eigen::MatrixXf SumMipMap(const Eigen::MatrixXf& img_big)
{
	// size of original image
	const unsigned int w_big = img_big.rows();
	const unsigned int h_big = img_big.cols();
	// size of reduced image
	const unsigned int w_sma = w_big / Q;
	const unsigned int h_sma = h_big / Q;
	// the computed mipmap will have 2^i size
	if(Q*w_sma != w_big || Q*h_sma != h_big) {
		throw std::runtime_error("ERROR: Q and size does not match in function SumMipMap!");
	}
	Eigen::MatrixXf img_small(w_sma, h_sma);
	for(unsigned int y=0; y<h_sma; ++y) {
		const unsigned int y_big = Q*y;
		for(unsigned int x=0; x<w_sma; ++x) {
			const unsigned int x_big = Q*x;
			float sum = 0.0f;
			for(unsigned int i=0; i<Q; ++i) {
				for(unsigned int j=0; j<Q; ++j) {
					sum += img_big(x_big+j, y_big+i);
				}
			}
			img_small(x, y) = sum;
		}
	}
	return img_small;
}

Eigen::MatrixXf ScaleUp(const Eigen::MatrixXf& img_small, unsigned int S);

std::vector<Eigen::MatrixXf> ComputeMipmaps(const Eigen::MatrixXf& img, unsigned int min_size);

std::vector<Eigen::MatrixXf> ComputeMipmapsLevels(const Eigen::MatrixXf& img, unsigned int levels);

std::vector<Eigen::MatrixXf> ComputeMipmaps640x480(const Eigen::MatrixXf& img);

std::vector<std::pair<Eigen::MatrixXf,Eigen::MatrixXf>> ComputeMipmapsWithAbs(const Eigen::MatrixXf& img, unsigned int min_size);

}

#endif


#include <boost/test/unit_test.hpp>
#include "functions/full_function_defs.hh"
#include "functions/addition.hh"
#include "functions/std_functions.hh"
#include "functions/operators.hh"
#include "functions/polynomial.hh"
#include "functions/all_simplifications.hh"

BOOST_AUTO_TEST_CASE(simplify_test) {
  using namespace manifolds;
  typedef Addition<Sin, Cos> Type;
  static_assert(std::is_same<SimplifiedType<Type>, Type>::value,
                "Failed to not simplify!");
  static_assert(std::is_same<SimplifiedType<Addition<Addition<Sin, Cos>, Tan> >,
                             Addition<Sin, Cos, Tan> >::value,
                "Failed to collapse nested additions!");
  Addition<Addition<Sin, Cos>, Tan> t;

  static_assert(
      std::is_same<decltype(Sin() + Sin()),
                   Composition<IntegralPolynomial<0, 2>, Sin> >::value,
      "Failed to make sin+sin == 2 * sin");

  auto ps = Sin() + Sin();
  BOOST_CHECK_EQUAL(std::get<0>(ps.GetFunctions()).GetCoeffs()[0], 0);
  BOOST_CHECK_EQUAL(std::get<0>(ps.GetFunctions()).GetCoeffs()[1], 2);

  auto ms = Sin() * Sin();
  BOOST_CHECK_EQUAL(std::get<0>(ms.GetFunctions()).GetCoeffs()[0], 0);
  BOOST_CHECK_EQUAL(std::get<0>(ms.GetFunctions()).GetCoeffs()[1], 0);
  BOOST_CHECK_EQUAL(std::get<0>(ms.GetFunctions()).GetCoeffs()[2], 1);

  auto ss = Sin()(Sin());
  auto mss = ss * ss;
  static_assert(
      std::is_same<decltype(mss),
                   Composition<IntegralPolynomial<0, 0, 1>, Sin, Sin> >::value,
      "Failed to simplify multiplication of "
      "composition of functions.");
}


#ifndef NAV_TYPES_HPP
# define NAV_TYPES_HPP

#include <Eigen/Dense>

using Eigen::Matrix;
using Eigen::DiagonalMatrix;

#define NN 30

typedef Eigen::Matrix<double,NN,NN> Cov;
typedef Eigen::Matrix<double,NN,1>  State;
typedef Eigen::Matrix<double,1,NN>  StateT;
typedef Eigen::Matrix<double,9,9>   Matrix9d;

const double FILTER_SMALL = 1e-12;  /* (--) small number used for underflow and div by zero */
const double PROP_DT      = 0.0125; /* (s)  propagation time step */
const double TAU          = 600.0;  /* (s)  time constant for ECRVs */
const double Q_ACCEL_PSD  = 1e-7;   /* accelerometer noise power spectral density */
const double Q_GYRO_PSD   = 1e-12;  /* gyroscope noise power spectral density */

#endif


// Copyright 2008 Gautam Sewani
// Copyright 2008 John Maddock
// Copyright 2021 Paul A. Bristow
//
// Use, modification and distribution are subject to the
// Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt
// or copy at http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_MATH_DISTRIBUTIONS_HYPERGEOMETRIC_HPP
#define BOOST_MATH_DISTRIBUTIONS_HYPERGEOMETRIC_HPP

#include <boost/math/distributions/detail/common_error_handling.hpp>
#include <boost/math/distributions/complement.hpp>
#include <boost/math/distributions/detail/hypergeometric_pdf.hpp>
#include <boost/math/distributions/detail/hypergeometric_cdf.hpp>
#include <boost/math/distributions/detail/hypergeometric_quantile.hpp>
#include <boost/math/special_functions/fpclassify.hpp>

namespace boost { namespace math {

   template <class RealType = double, class Policy = policies::policy<> >
   class hypergeometric_distribution
   {
   public:
      typedef RealType value_type;
      typedef Policy policy_type;

      hypergeometric_distribution(unsigned r, unsigned n, unsigned N) // Constructor. r=defective/failures/success, n=trials/draws, N=total population.
         : m_n(n), m_N(N), m_r(r)
      {
         static const char* function = "boost::math::hypergeometric_distribution<%1%>::hypergeometric_distribution";
         RealType ret;
         check_params(function, &ret);
      }
      // Accessor functions.
      unsigned total()const
      {
         return m_N;
      }

      unsigned defective()const // successes/failures/events
      {
         return m_r;
      }

      unsigned sample_count()const
      {
         return m_n;
      }

      bool check_params(const char* function, RealType* result)const
      {
         if(m_r > m_N)
         {
            *result = boost::math::policies::raise_domain_error<RealType>(
               function, "Parameter r out of range: must be <= N but got %1%", static_cast<RealType>(m_r), Policy());
            return false;
         }
         if(m_n > m_N)
         {
            *result = boost::math::policies::raise_domain_error<RealType>(
               function, "Parameter n out of range: must be <= N but got %1%", static_cast<RealType>(m_n), Policy());
            return false;
         }
         return true;
      }
      bool check_x(unsigned x, const char* function, RealType* result)const
      {
         if(x < static_cast<unsigned>((std::max)(0, (int)(m_n + m_r) - (int)(m_N))))
         {
            *result = boost::math::policies::raise_domain_error<RealType>(
               function, "Random variable out of range: must be > 0 and > m + r - N but got %1%", static_cast<RealType>(x), Policy());
            return false;
         }
         if(x > (std::min)(m_r, m_n))
         {
            *result = boost::math::policies::raise_domain_error<RealType>(
               function, "Random variable out of range: must be less than both n and r but got %1%", static_cast<RealType>(x), Policy());
            return false;
         }
         return true;
      }

   private:
      // Data members:
      unsigned m_n;  // number of items picked or drawn.
      unsigned m_N; // number of "total" items.
      unsigned m_r; // number of "defective/successes/failures/events items.

   }; // class hypergeometric_distribution

   typedef hypergeometric_distribution<double> hypergeometric;

   template <class RealType, class Policy>
   inline const std::pair<unsigned, unsigned> range(const hypergeometric_distribution<RealType, Policy>& dist)
   { // Range of permissible values for random variable x.
#ifdef _MSC_VER
#  pragma warning(push)
#  pragma warning(disable:4267)
#endif
      unsigned r = dist.defective();
      unsigned n = dist.sample_count();
      unsigned N = dist.total();
      unsigned l = static_cast<unsigned>((std::max)(0, (int)(n + r) - (int)(N)));
      unsigned u = (std::min)(r, n);
      return std::pair<unsigned, unsigned>(l, u);
#ifdef _MSC_VER
#  pragma warning(pop)
#endif
   }

   template <class RealType, class Policy>
   inline const std::pair<unsigned, unsigned> support(const hypergeometric_distribution<RealType, Policy>& d)
   {
      return range(d);
   }

   template <class RealType, class Policy>
   inline RealType pdf(const hypergeometric_distribution<RealType, Policy>& dist, const unsigned& x)
   {
      static const char* function = "boost::math::pdf(const hypergeometric_distribution<%1%>&, const %1%&)";
      RealType result = 0;
      if(!dist.check_params(function, &result))
         return result;
      if(!dist.check_x(x, function, &result))
         return result;

      return boost::math::detail::hypergeometric_pdf<RealType>(
         x, dist.defective(), dist.sample_count(), dist.total(), Policy());
   }

   template <class RealType, class Policy, class U>
   inline RealType pdf(const hypergeometric_distribution<RealType, Policy>& dist, const U& x)
   {
      BOOST_MATH_STD_USING
      static const char* function = "boost::math::pdf(const hypergeometric_distribution<%1%>&, const %1%&)";
      RealType r = static_cast<RealType>(x);
      unsigned u = itrunc(r, typename policies::normalise<Policy, policies::rounding_error<policies::ignore_error> >::type());
      if(u != r)
      {
         return boost::math::policies::raise_domain_error<RealType>(
            function, "Random variable out of range: must be an integer but got %1%", r, Policy());
      }
      return pdf(dist, u);
   }

   template <class RealType, class Policy>
   inline RealType cdf(const hypergeometric_distribution<RealType, Policy>& dist, const unsigned& x)
   {
      static const char* function = "boost::math::cdf(const hypergeometric_distribution<%1%>&, const %1%&)";
      RealType result = 0;
      if(!dist.check_params(function, &result))
         return result;
      if(!dist.check_x(x, function, &result))
         return result;

      return boost::math::detail::hypergeometric_cdf<RealType>(
         x, dist.defective(), dist.sample_count(), dist.total(), false, Policy());
   }

   template <class RealType, class Policy, class U>
   inline RealType cdf(const hypergeometric_distribution<RealType, Policy>& dist, const U& x)
   {
      BOOST_MATH_STD_USING
      static const char* function = "boost::math::cdf(const hypergeometric_distribution<%1%>&, const %1%&)";
      RealType r = static_cast<RealType>(x);
      unsigned u = itrunc(r, typename policies::normalise<Policy, policies::rounding_error<policies::ignore_error> >::type());
      if(u != r)
      {
         return boost::math::policies::raise_domain_error<RealType>(
            function, "Random variable out of range: must be an integer but got %1%", r, Policy());
      }
      return cdf(dist, u);
   }

   template <class RealType, class Policy>
   inline RealType cdf(const complemented2_type<hypergeometric_distribution<RealType, Policy>, unsigned>& c)
   {
      static const char* function = "boost::math::cdf(const hypergeometric_distribution<%1%>&, const %1%&)";
      RealType result = 0;
      if(!c.dist.check_params(function, &result))
         return result;
      if(!c.dist.check_x(c.param, function, &result))
         return result;

      return boost::math::detail::hypergeometric_cdf<RealType>(
         c.param, c.dist.defective(), c.dist.sample_count(), c.dist.total(), true, Policy());
   }

   template <class RealType, class Policy, class U>
   inline RealType cdf(const complemented2_type<hypergeometric_distribution<RealType, Policy>, U>& c)
   {
      BOOST_MATH_STD_USING
      static const char* function = "boost::math::cdf(const hypergeometric_distribution<%1%>&, const %1%&)";
      RealType r = static_cast<RealType>(c.param);
      unsigned u = itrunc(r, typename policies::normalise<Policy, policies::rounding_error<policies::ignore_error> >::type());
      if(u != r)
      {
         return boost::math::policies::raise_domain_error<RealType>(
            function, "Random variable out of range: must be an integer but got %1%", r, Policy());
      }
      return cdf(complement(c.dist, u));
   }

   template <class RealType, class Policy>
   inline RealType quantile(const hypergeometric_distribution<RealType, Policy>& dist, const RealType& p)
   {
      BOOST_MATH_STD_USING // for ADL of std functions

         // Checking function argument
         RealType result = 0;
      const char* function = "boost::math::quantile(const hypergeometric_distribution<%1%>&, %1%)";
      if (false == dist.check_params(function, &result)) return result;
      if(false == detail::check_probability(function, p, &result, Policy())) return result;

      return static_cast<RealType>(detail::hypergeometric_quantile(p, RealType(1 - p), dist.defective(), dist.sample_count(), dist.total(), Policy()));
   } // quantile

   template <class RealType, class Policy>
   inline RealType quantile(const complemented2_type<hypergeometric_distribution<RealType, Policy>, RealType>& c)
   {
      BOOST_MATH_STD_USING // for ADL of std functions

      // Checking function argument
      RealType result = 0;
      const char* function = "quantile(const complemented2_type<hypergeometric_distribution<%1%>, %1%>&)";
      if (false == c.dist.check_params(function, &result)) return result;
      if(false == detail::check_probability(function, c.param, &result, Policy())) return result;

      return static_cast<RealType>(detail::hypergeometric_quantile(RealType(1 - c.param), c.param, c.dist.defective(), c.dist.sample_count(), c.dist.total(), Policy()));
   } // quantile

   // https://www.wolframalpha.com/input/?i=kurtosis+hypergeometric+distribution

   template <class RealType, class Policy>
   inline RealType mean(const hypergeometric_distribution<RealType, Policy>& dist)
   {
      return static_cast<RealType>(dist.defective() * dist.sample_count()) / dist.total();
   } // RealType mean(const hypergeometric_distribution<RealType, Policy>& dist)

   template <class RealType, class Policy>
   inline RealType variance(const hypergeometric_distribution<RealType, Policy>& dist)
   {
      RealType r = static_cast<RealType>(dist.defective());
      RealType n = static_cast<RealType>(dist.sample_count());
      RealType N = static_cast<RealType>(dist.total());
      return n * r  * (N - r) * (N - n) / (N * N * (N - 1));
   } // RealType variance(const hypergeometric_distribution<RealType, Policy>& dist)

   template <class RealType, class Policy>
   inline RealType mode(const hypergeometric_distribution<RealType, Policy>& dist)
   {
      BOOST_MATH_STD_USING
      RealType r = static_cast<RealType>(dist.defective());
      RealType n = static_cast<RealType>(dist.sample_count());
      RealType N = static_cast<RealType>(dist.total());
      return floor((r + 1) * (n + 1) / (N + 2));
   }

   template <class RealType, class Policy>
   inline RealType skewness(const hypergeometric_distribution<RealType, Policy>& dist)
   {
      BOOST_MATH_STD_USING
      RealType r = static_cast<RealType>(dist.defective());
      RealType n = static_cast<RealType>(dist.sample_count());
      RealType N = static_cast<RealType>(dist.total());
      return (N - 2 * r) * sqrt(N - 1) * (N - 2 * n) / (sqrt(n * r * (N - r) * (N - n)) * (N - 2));
   } // RealType skewness(const hypergeometric_distribution<RealType, Policy>& dist)

   template <class RealType, class Policy>
   inline RealType kurtosis_excess(const hypergeometric_distribution<RealType, Policy>& dist)
   {
      // https://www.wolframalpha.com/input/?i=kurtosis+hypergeometric+distribution shown as plain text:
      //  mean | (m n)/N
      //  standard deviation | sqrt((m n(N - m) (N - n))/(N - 1))/N
      //  variance | (m n(1 - m/N) (N - n))/((N - 1) N)
      //  skewness | (sqrt(N - 1) (N - 2 m) (N - 2 n))/((N - 2) sqrt(m n(N - m) (N - n)))
      //  kurtosis | ((N - 1) N^2 ((3 m(N - m) (n^2 (-N) + (n - 2) N^2 + 6 n(N - n)))/N^2 - 6 n(N - n) + N(N + 1)))/(m n(N - 3) (N - 2) (N - m) (N - n))
     // Kurtosis[HypergeometricDistribution[n, m, N]]
      RealType m = static_cast<RealType>(dist.defective()); // Failures or success events. (Also symbols K or M are used).
      RealType n = static_cast<RealType>(dist.sample_count()); // draws or trials.
      RealType n2 = n * n; // n^2
      RealType N = static_cast<RealType>(dist.total()); // Total population from which n draws or trials are made.
      RealType N2 = N * N; // N^2
      // result = ((N - 1) N^2 ((3 m(N - m) (n^2 (-N) + (n - 2) N^2 + 6 n(N - n)))/N^2 - 6 n(N - n) + N(N + 1)))/(m n(N - 3) (N - 2) (N - m) (N - n));
      RealType result = ((N-1)*N2*((3*m*(N-m)*(n2*(-N)+(n-2)*N2+6*n*(N-n)))/N2-6*n*(N-n)+N*(N+1)))/(m*n*(N-3)*(N-2)*(N-m)*(N-n));
      // Agrees with kurtosis hypergeometric distribution(50,200,500) kurtosis = 2.96917
      // N[kurtosis[hypergeometricdistribution(50,200,500)], 55]  2.969174035736058474901169623721804275002985337280263464
      return result;
   } // RealType kurtosis_excess(const hypergeometric_distribution<RealType, Policy>& dist)

   template <class RealType, class Policy>
   inline RealType kurtosis(const hypergeometric_distribution<RealType, Policy>& dist)
   {
      return kurtosis_excess(dist) + 3;
   } // RealType kurtosis_excess(const hypergeometric_distribution<RealType, Policy>& dist)
}} // namespaces

// This include must be at the end, *after* the accessors
// for this distribution have been defined, in order to
// keep compilers that support two-phase lookup happy.
#include <boost/math/distributions/detail/derived_accessors.hpp>

#endif // include guard


/* vim: set tabstop=4 expandtab shiftwidth=4 softtabstop=4: */

/**
 * \file boost/numeric/ublasx/operation/lu.hpp
 *
 * \brief LU decomposition and solver.
 *
 * Computes an LU factorization of a general m-by-n matrix \f$A\f$ optionally
 * using partial pivoting with row interchanges.
 * The factorization has the form
 * \f[
 *  A = L U
 * \f]
 * or, if partial pivoting is used:
 * \f[
 *  A = P L U
 * \f]
 * where \f$P\f$ is a permutation matrix, \f$L\f$ is lower triangular with unit
 * diagonal elements (lower trapezoidal if \f$m > n\f$), and \f$U\f$ is upper
 * triangular (upper trapezoidal if \f$m < n\f$).
 * If matrix \f$A\f$ is rectangular \f$L and \f$P are square matrices each
 * having the same number of rows as \f$A\f$, while \f$U\f$ is exactly the same
 * shape as \f$A\f$.
 *
 * \author Marco Guazzone (marco.guazzone@gmail.com)
 *
 * <hr/>
 *
 * Copyright (c) 2010, Marco Guazzone
 *
 * Distributed under the Boost Software License, Version 1.0. (See
 * accompanying file LICENSE_1_0.txt or copy at
 * http://www.boost.org/LICENSE_1_0.txt)
 */

#ifndef BOOST_NUMERIC_UBLASX_OPERATION_LU_HPP
#define BOOST_NUMERIC_UBLASX_OPERATION_LU_HPP

//TODO: Use LAPACK functions in order to handle different types of matrices.
//TODO: How about full pivoting?
//TODO: Create a \c lu_decomposition class (e.g., \sa qr.hpp).

#include <boost/numeric/ublas/detail/temporary.hpp>
#include <boost/numeric/ublas/exception.hpp>
#include <boost/numeric/ublas/expression_types.hpp>
#include <boost/numeric/ublas/lu.hpp>
#include <boost/numeric/ublas/traits.hpp>
#include <boost/numeric/ublasx/operation/num_columns.hpp>
#include <boost/numeric/ublasx/operation/num_rows.hpp>
#include <boost/numeric/ublasx/operation/size.hpp>
#include <boost/numeric/ublasx/traits/layout_type.hpp>

namespace boost { namespace numeric { namespace ublasx {

using namespace ::boost::numeric::ublas;

/**
 * \brief LU decomposition without pivoting of the given matrix \a A.
 * \param A The matrix to be decomposed.
 * \return Zero if decomposition succeed; non-zero if decompositon fails (the
 *  value is 1 + the numer of the failing row). The input matrix \a A is
 *  modified in order to contain the L*U matrix.
 *
 * Perform LU decomposition of matrix \a A and replaces the strict lower
 * triangular part of \a m with the computed matrix L and the upper triangular
 * part of \a m is replaced by the computed matrix U.
 * For obtaining the single matrices L and U proceed as follows:
 * - for L: extract the strict lower-triangular part (i.e., without the main
 *   diagonal) from the computed matrix \a m and add to it the identity matrix
 *   of the same order of \a m;
 * - for U: extract the upper-triangular part (with the main diagonal) from the
 *   computer matrix \a m.
 */
template <typename MatrixT>
BOOST_UBLAS_INLINE
typename matrix_traits<MatrixT>::size_type lu_decompose_inplace(matrix_container<MatrixT>& A)
{
    return lu_factorize(A());
}

/**
 * \brief LU decomposition with partial pivoting of the matrix \a A.
 *
 * \param A The matrix to be decomposed.
 * \param P The permutation matrix reporting permutated rows of \a A after
 *  the decomposition.
 * \return Zero if decomposition succeed; non-zero if decompositon fails (the
 *  value is 1 + the numer of the failing row). The input matrix \a A is
 *  modified in order to contain the L*U matrix.
 *
 * Perform LUP decomposition of matrix \a A and replaces the strict lower
 * triangular part of \a A with the computed matrix L and the upper triangular
 * part of \a A is replaced by the computed matrix U.
 * For obtaining the single matrices L and U proceed as follows:
 * - for L: extract the strict lower-triangular part (i.e., without the main
 *   diagonal) from the computed matrix \a m and add to it the identity matrix
 *   of the same order of \a A;
 * - for U: extract the upper-triangular part (with the main diagonal) from the
 *   computer matrix \a A.
 * .
 */
template <typename MatrixT, typename PermutationMatrixT>
BOOST_UBLAS_INLINE
typename matrix_traits<MatrixT>::size_type lu_decompose_inplace(matrix_container<MatrixT>& A, PermutationMatrixT& P)
{
    typedef typename matrix_traits<MatrixT>::size_type size_type;

    // Safety check: P is squared && size(P) == num_rows(A)
    //BOOST_UBLAS_CHECK( size(P) == num_columns(P) && size(P) == num_rows(A), bad_size() );
    size_type nr_A = num_rows(A);
    if (size(P) != nr_A)
    {
        P.resize(nr_A, false);
    }

    return lu_factorize(A(), P);
//
//  // postcondition: P is squared && size(P) == num_rows(A)
//  BOOST_UBLAS_CHECK( size(P) == num_columns(P) && size(P) == num_rows(A), bad_size() );
}

/**
 * \brief LU decomposition without pivoting of the matrix \a A.
 *
 * \tparam MatrixT The type of the matrix to be decomposed.
 * \param A The matrix to be decomposed.
 * \return Zero if decomposition succeed; non-zero if decompositon fails (the
 *  value is 1 + the numer of the failing row). The input matrix \a A is
 *  modified in order to contain the L*U matrix.
 *
 * Perform LU decomposition of matrix \a A and replaces the strict lower
 * triangular part of \a A with the computed matrix L and the upper triangular
 * part of \a A is replaced by the computed matrix U.
 * For obtaining the single matrices L and U proceed as follows:
 * - for L: extract the strict lower-triangular part (i.e., without the main
 *   diagonal) from the computed matrix \a A and add to it the identity matrix
 *   of the same order of \a A;
 * - for U: extract the upper-triangular part (with the main diagonal) from the
 *   computer matrix \a A.
 * .
 */
template <typename AMatrixExprT, typename LUMatrixT>
BOOST_UBLAS_INLINE
typename matrix_traits<AMatrixExprT>::size_type lu_decompose(matrix_expression<AMatrixExprT> const& A, matrix_container<LUMatrixT>& LU)
{
    LU = A;

    return lu_decompose_inplace(LU);
}

/**
 * \brief LU decomposition with partial pivoting of the matrix \a A.
 *
 * \pre Matrix \a A must be squared.
 * \tparam MatrixT The type of the matrix to be decomposed.
 * \param A The matrix to be decomposed.
 * \param P The permutation matrix reporting permutated rows of \a A after
 *  the decomposition.
 * \return The LU matrix.
 *
 * Perform LUP decomposition of matrix \a A and return the LU matrix.
 * For obtaining the single matrices L and U proceed as follows:
 * - for L: extract the strict lower-triangular part (i.e., without the main
 *   diagonal) from the computed matrix \c LU and add to it the identity matrix
 *   of the same order of \c LU;
 * - for U: extract the upper-triangular part (with the main diagonal) from the
 *   computer matrix \c LU.
 * .
 */
template <typename AMatrixExprT, typename PermutationMatrixT, typename LUMatrixT>
BOOST_UBLAS_INLINE
typename matrix_traits<AMatrixExprT>::size_type lu_decompose(matrix_expression<AMatrixExprT> const& A, PermutationMatrixT& P, matrix_container<LUMatrixT>& LU)
{
    typedef typename matrix_traits<AMatrixExprT>::size_type size_type;

    // Safety check: size(P) == num_rows(A)
    //BOOST_UBLAS_CHECK( size(P) == num_rows(A), bad_size() );
    size_type nr_A = num_rows(A);
    if (size(P) != nr_A)
    {
        P.resize(nr_A, false);
    }

    LU = A;

    return lu_decompose_inplace(LU, P);
}

/**
 * \brief Complete the LU forward/backward substitution for solving the system
 *  \f$LUx=b\f$.
 *
 * \pre The size of \a b must be the same as the number of rows of \a LU.
 * \tparam MatrixExprT The type of the matrix obtained by the LU decomposition.
 * \tparam VectorT The type of the constant terms vector.
 * \param LU A matrix representing an LU decomposition.
 * \param b The vector of coefficients.
 * \return Nothing. However, the vector \a b is replaced with the value of
 *  unknowns \f$x_i\f$ satisfying the system \f$LUx=b\f$.
 *
 * An \f$n \times n\f$ linear system \f$Ax=b\f$ can often be solved efficiently
 * by \f$LU\f$ decomposition (that is by decomposing matrix \f$A\f$ into a
 * product \f$LU\f$, where \f$L\f$ is a lower-triangular matrix with ones on the
 * diagonal, and \f$U\f$ is an upper-triangular matrix). Solving the system
 * \f$Ax=b\f$ is then equivalent to solving two simpler systems \f$Ly=b\f$ and
 * \f$Ux=y\f$. Since \f$L\f$ is lower-triangular, the system \f$Ly=b\f$ can be
 * solved by forward substitution. Moreover, since \f$U\f$ is upper-triangular,
 * the system \f$Ux=y\f$ can be solved by backward substitution.
 */
template <typename MatrixExprT, typename VectorT>
BOOST_UBLAS_INLINE
void lu_apply_inplace(matrix_expression<MatrixExprT> const& LU, vector_container<VectorT>& b)
{
    // pre: size(b) == size(P)
    BOOST_UBLAS_CHECK( size(b) == num_rows(LU), bad_size() );

    lu_substitute(LU(), b());
}

/**
 * \brief Complete the LU forward/backward substitution for solving the system
 *  \f$LUX=B\f$.
 *
 * \pre The number of rows of \a LU and \a b must be the same.
 * \tparam LUMatrixExprT The type of the matrix obtained by the LU decomposition.
 * \tparam BMatrixT The type of the constant terms vector.
 * \param LU A matrix representing an LU decomposition.
 * \param B The matrix of coefficients.
 * \return Nothing. However, the matrix \a B is replaced with the value of
 *  unknowns \f$X_{ij}\f$ satisfying the system \f$LUX=B\f$.
 *
 * An \f$n \times n\f$ linear system \f$Ax=B\f$ can often be solved efficiently
 * by \f$LU\f$ decomposition (that is by decomposing matrix \f$A\f$ into a
 * product \f$LU\f$, where \f$L\f$ is a lower-triangular matrix with ones on the
 * diagonal, and \f$U\f$ is an upper-triangular matrix). Solving the system
 * \f$AX=B\f$ is then equivalent to solving two simpler systems \f$LY=B\f$ and
 * \f$UX=Y\f$. Since \f$L\f$ is lower-triangular, the system \f$LY=B\f$ can be
 * solved by forward substitution. Moreover, since \f$U\f$ is upper-triangular,
 * the system \f$UX=Y\f$ can be solved by backward substitution.
 */
template <typename LUMatrixExprT, typename BMatrixT>
BOOST_UBLAS_INLINE
void lu_apply_inplace(matrix_expression<LUMatrixExprT> const& LU, matrix_container<BMatrixT>& B)
{
    // pre: num_rows(B) == num_rows(LU)
    BOOST_UBLAS_CHECK( num_rows(B) == num_rows(LU), bad_size() );

    lu_substitute(LU(), B());
}

/**
 * \brief Complete the LUP forward/backward substitution for solving the system
 *  \f$LU*X=P*b\f$.
 *
 * \pre The size of \a P must be the same as the number of rows of \a LU.
 * \pre The size of \a b must be the same as the number of rows of \a LU.
 * \param LU A matrix representing an LU decomposition.
 * \param P A permutation matrix.
 * \param b The vector constant terms.
 * \return Nothing. However, the vector \a b is replaced with the value
 *  of unknowns \f$x_i\f$ satisfying the system \f$LU*x=P*b\f$.
 *
 * An \f$n \times n\f$ linear system \f$Ax=b\f$ can often be solved efficiently
 * by \f$LU\f$ decomposition (that is by decomposing matrix \f$A\f$ into a
 * product \f$LUP\f$, where \f$L\f$ is a lower-triangular matrix with ones on
 * the diagonal, \f$U\f$ is an upper-triangular matrix, and \f$P\f$ is a row
 * permutation matrix that is used to rearrange the rows of \f$A\f$ before so
 * that it can be decomposed). Solving the system \f$Ax=b\f$ is then equivalent
 * to solving two simpler systems \f$Ly=Pb\f$ and \f$Ux=y\f$. Since \f$L\f$ is
 * lower-triangular, the system \f$Ly=Pb\f$ can be solved by forward
 * substitution. Moreover, since \f$U\f$ is upper-triangular, the system
 * \f$Ux=y\f$ can be solved by backward substitution.
 */
template <typename LUMatrixExprT, typename PermutationMatrixT, typename BVectorT>
BOOST_UBLAS_INLINE
void lu_apply_inplace(matrix_expression<LUMatrixExprT> const& LU, PermutationMatrixT const& P, vector_container<BVectorT>& b)
{
    // precondition: size(P) == num_rows(LU)
    BOOST_UBLAS_CHECK( size(P) == num_rows(LU), bad_size() );
    // precondition: size(b) == num_rows(LU)
    BOOST_UBLAS_CHECK( size(b) == num_rows(LU), bad_size() );

    lu_substitute(LU(), P, b());
}

/**
 * \brief Complete the LUP forward/backward substitution for solving the system
 *  \f$LU*X=P*B\f$.
 *
 * \pre The size of \a P must be the same as the number of rows of \a LU.
 * \pre The number of rows of \a LU and of \a B must be the same.
 * \param LU A matrix representing an LU decomposition.
 * \param P A permutation matrix.
 * \param B The matrix of constant terms.
 * \return Nothing. However, the matrix \a B is replaced with the value
 *  of unknowns \f$X_{ij}\f$ satisfying the system \f$LU*X=P*B\f$.
 *
 * An \f$n \times n\f$ linear system \f$AX=B\f$ can often be solved efficiently
 * by \f$LU\f$ decomposition (that is by decomposing matrix \f$A\f$ into a
 * product \f$LUP\f$, where \f$L\f$ is a lower-triangular matrix with ones on
 * the diagonal, \f$U\f$ is an upper-triangular matrix, and \f$P\f$ is a row
 * permutation matrix that is used to rearrange the rows of \f$A\f$ before so
 * that it can be decomposed). Solving the system \f$AX=B\f$ is then equivalent
 * to solving two simpler systems \f$LY=Pb\f$ and \f$UX=y\f$. Since \f$L\f$ is
 * lower-triangular, the system \f$LY=PB\f$ can be solved by forward
 * substitution. Moreover, since \f$U\f$ is upper-triangular, the system
 * \f$UX=Y\f$ can be solved by backward substitution.
 */
template <typename LUMatrixExprT, typename PermutationMatrixT, typename BMatrixT>
BOOST_UBLAS_INLINE
void lu_apply_inplace(matrix_expression<LUMatrixExprT> const& LU, PermutationMatrixT const& P, matrix_container<BMatrixT>& B)
{
    // pre: size(P) == num_rows(LU)
    BOOST_UBLAS_CHECK( size(P) == num_rows(LU), bad_size() );
    // pre: num_rows(b) == num_rows(LU)
    BOOST_UBLAS_CHECK( num_rows(B) == num_rows(LU), bad_size() );

    lu_substitute(LU(), P, B());
}

/**
 * \brief Complete the LU forward/backward substitution for solving the system
 *  \f$LUx=b\f$.
 *
 * \param LU A matrix representing an LU decomposition.
 * \param b The vector of coefficients.
 * \return The vector of unknowns \f$x\f$ satisfying the system \f$LUPx=b\f$.
 *
 * An \f$n \times n\f$ linear system \f$Ax=b\f$ can often be solved efficiently
 * by \f$LU\f$ decomposition (that is by decomposing matrix \f$A\f$ into a
 * product \f$LU\f$, where \f$L\f$ is a lower-triangular matrix with ones on the
 * diagonal, and \f$U\f$ is an upper-triangular matrix). Solving the system
 * \f$Ax=b\f$ is then equivalent to solving two simpler systems \f$Ly=b\f$ and
 * \f$Ux=y\f$. Since \f$L\f$ is lower-triangular, the system \f$Ly=b\f$ can be
 * solved by forward substitution. Moreover, since \f$U\f$ is upper-triangular,
 * the system \f$Ux=y\f$ can be solved by backward substitution.
 */
template <typename LUMatrixExprT, typename BVectorExprT>
BOOST_UBLAS_INLINE
typename vector_temporary_traits<BVectorExprT>::type lu_apply(matrix_expression<LUMatrixExprT> const& LU, vector_expression<BVectorExprT> const& b)
{
    //  preconditions check delegated to lu_apply_inplace

    typedef typename vector_temporary_traits<BVectorExprT>::type out_vector_type;

    out_vector_type x(b);

    lu_apply_inplace(LU, x);

    return x;
}

/**
 * \brief Complete the LU forward/backward substitution for solving the system
 *  \f$LU*X=B\f$.
 *
 * \param LU A matrix representing an LU decomposition.
 * \param B The matrix of constant terms.
 * \return The matrix of unknowns \f$X\f$ satisfying the system \f$LU*X=B\f$.
 */
template <typename LUMatrixExprT, typename BMatrixExprT>
BOOST_UBLAS_INLINE
typename matrix_temporary_traits<BMatrixExprT>::type lu_apply(matrix_expression<LUMatrixExprT> const& LU, matrix_expression<BMatrixExprT> const& B)
{
    //  preconditions check delegated to lu_apply_inplace

    typedef typename matrix_temporary_traits<BMatrixExprT>::type out_matrix_type;

    out_matrix_type X(B);

    lu_apply_inplace(LU, X);

    return X;
}

/**
 * \brief Complete the LUP forward/backward substitution for solving the system
 *  \f$LUx=b\f$.
 *
 * \param LU A matrix representing an LU decomposition.
 * \param P A permutation matrix.
 * \param b The vector of coefficients.
 * \return The vector of unknowns \f$x\f$ satisfying the system \f$LUPx=b\f$.
 *
 * An \f$n \times n\f$ linear system \f$Ax=b\f$ can often be solved efficiently
 * by \f$LU\f$ decomposition (that is by decomposing matrix \f$A\f$ into a
 * product \f$LUP\f$, where \f$L\f$ is a lower-triangular matrix with ones on
 * the diagonal, \f$U\f$ is an upper-triangular matrix, and \f$P\f$ is a row
 * permutation matrix that is used to rearrange the rows of A before so that it
 * can be decomposed). Solving the system \f$Ax=b\f$ is then equivalent to
 * solving two simpler systems \f$Ly=Pb\f$ and \f$Ux=y\f$. Since \f$L\f$ is
 * lower-triangular, the system \f$Ly=Pb\f$ can be solved by forward
 * substitution. Moreover, since \f$U\f$ is upper-triangular, the system
 * \f$Ux=y\f$ can be solved by backward substitution.
 */
template <typename LUMatrixExprT, typename PermutationMatrixT, typename BVectorExprT>
BOOST_UBLAS_INLINE
typename vector_temporary_traits<BVectorExprT>::type lu_apply(matrix_expression<LUMatrixExprT> const& LU, PermutationMatrixT const& P, vector_expression<BVectorExprT> const& b)
{
    //  preconditions check delegated to lu_apply_inplace

    typedef typename vector_temporary_traits<BVectorExprT>::type out_vector_type;

    out_vector_type x(b);

    lu_apply_inplace(LU, P, x);

    return x;
}

/**
 * \brief Complete the LUP forward/backward substitution for solving the system
 *  \f$LUP*X=B\f$.
 *
 * \param LU A matrix representing an LU decomposition.
 * \param P A permutation matrix.
 * \param B The matrix of constant terms.
 * \return The matrix of unknowns \f$X\f$ satisfying the system \f$LUP*X=B\f$.
 */
template <typename LUMatrixExprT, typename PermutationMatrixT, typename BMatrixExprT>
BOOST_UBLAS_INLINE
typename matrix_temporary_traits<BMatrixExprT>::type lu_apply(matrix_expression<LUMatrixExprT> const& LU, PermutationMatrixT const& P, matrix_expression<BMatrixExprT> const& B)
{
    //  preconditions check delegated to lu_apply_inplace

    typedef typename matrix_temporary_traits<BMatrixExprT>::type out_matrix_type;

    out_matrix_type X(B);

    lu_apply_inplace(LU, P, X);

    return X;
}

/**
 * \brief Solve the linear system \f$Ax=b\f$ by LUP decomposition.
 *
 * \pre The size of \a b must be the same as the number of rows of \a A.
 * \param A The input matrix representing the coefficients of the unknowns.
 * \param b The vector of coefficients.
 * \return A zero value if the system is solvable; a number greater than zero if
 *  it is not. Moreover, the vector \a b is replaced with the value of unknowns
 *  \f$x_i\f$ satisfying the system \f$Ax=b\f$.
 *
 * An \f$n \times n\f$ linear system \f$Ax=b\f$ can often be solved efficiently
 * by \f$LU\f$ decomposition (that is by decomposing matrix \f$A\f$ into a
 * product \f$LUP\f$, where \f$L\f$ is a lower-triangular matrix with ones on
 * the diagonal, \f$U\f$ is an upper-triangular matrix, and \f$P\f$ is a row
 * permutation matrix that is used to rearrange the rows of \f$A\f$ before so
 * that it can be decomposed). Solving the system \f$Ax=b\f$ is then equivalent
 * to solving two simpler systems \f$Ly=Pb\f$ and \f$Ux=y\f$. Since \f$L\f$ is
 * lower-triangular, the system \f$Ly=Pb\f$ can be solved by forward
 * substitution. Moreover, since \f$U\f$ is upper-triangular, the system
 * \f$Ux=y\f$ can be solved by backward substitution.
 */
template <typename MatrixExprT, typename VectorT>
BOOST_UBLAS_INLINE
typename matrix_traits<MatrixExprT>::size_type lu_solve_inplace(matrix_expression<MatrixExprT> const& A, vector_container<VectorT>& b)
{
    // precondition: size(b) == num_rows(A)
    BOOST_UBLAS_CHECK( size(b) == num_rows(A), bad_size() );

    typedef typename matrix_traits<MatrixExprT>::size_type size_type;
    typedef typename matrix_traits<MatrixExprT>::value_type value_type;
    typedef typename layout_type<MatrixExprT>::type layout_type;

    // Ax=b ==> LUx=b ==> Ly=b AND Ux=y

    matrix<value_type, layout_type> LU(A);
    permutation_matrix<size_type> P(num_rows(LU));

    size_type singular;
    singular = lu_decompose_inplace(LU, P);

    if (!singular)
    {
        lu_apply_inplace(LU, P, b());
    }

    return singular;
}

/**
 * \brief Solve the linear system \f$AX=B\f$ by LUP decomposition.
 *
 * \pre The number of rows of \a A and \a B must be the same.
 * \param A The input matrix representing the coefficients of the unknowns.
 * \param B The matrix of coefficients.
 * \return A zero value if the system is solvable; a number greater than zero if
 *  it is not. Moreover, the matrix \a B is replaced with the value of unknowns
 *  \f$X_{ij}\f$ satisfying the system \f$AX=B\f$.
 *
 * An \f$n \times n\f$ linear system \f$AX=B\f$ can often be solved efficiently
 * by \f$LU\f$ decomposition (that is by decomposing matrix \f$A\f$ into a
 * product \f$LUP\f$, where \f$L\f$ is a lower-triangular matrix with ones on
 * the diagonal, \f$U\f$ is an upper-triangular matrix, and \f$P\f$ is a row
 * permutation matrix that is used to rearrange the rows of \f$A\f$ before so
 * that it can be decomposed). Solving the system \f$AX=B\f$ is then equivalent
 * to solving two simpler systems \f$LY=PB\f$ and \f$UX=Y\f$. Since \f$L\f$ is
 * lower-triangular, the system \f$LY=PB\f$ can be solved by forward
 * substitution. Moreover, since \f$U\f$ is upper-triangular, the system
 * \f$UX=Y\f$ can be solved by backward substitution.
 */
template <typename AMatrixExprT, typename BMatrixExprT>
BOOST_UBLAS_INLINE
typename matrix_traits<AMatrixExprT>::size_type lu_solve_inplace(matrix_expression<AMatrixExprT> const& A, matrix_container<BMatrixExprT>& B)
{
    // pre: num_rows(B) == num_rows(A)
    BOOST_UBLAS_CHECK( num_rows(B) == num_rows(A), bad_size() );

    typedef typename matrix_traits<AMatrixExprT>::size_type size_type;
    typedef typename matrix_traits<AMatrixExprT>::value_type value_type;
    typedef typename layout_type<AMatrixExprT>::type layout_type;

    // Ax=B ==> LUx=B ==> Ly=B AND Ux=y

    matrix<value_type, layout_type> LU(A());
    permutation_matrix<size_type> P(num_rows(LU));

    size_type singular;
    singular = lu_decompose_inplace(LU, P);

    if (!singular)
    {
        lu_apply_inplace(LU, P, B());
    }

    return singular;
}

/**
 * \brief Solve the linear system \f$Ax=b\f$ by LUP decomposition.
 *
 * \pre The size of \a b must be the same as the number of rows of \a A.
 * \param A The input matrix representing the coefficients of the unknowns..
 * \param b The vector of coefficients.
 * \return The vector of unknowns \f$x\f$ satisfying the system \f$Ax=b\f$.
 *
 * An \f$n \times n\f$ linear system \f$Ax=b\f$ can often be solved efficiently
 * by \f$LU\f$ decomposition (that is by decomposing matrix \f$A\f$ into a
 * product \f$LUP\f$, where \f$L\f$ is a lower-triangular matrix with ones on
 * the diagonal, \f$U\f$ is an upper-triangular matrix, and \f$P\f$ is a row
 * permutation matrix that is used to rearrange the rows of A before so that it
 * can be decomposed). Solving the system \f$Ax=b\f$ is then equivalent to
 * solving two simpler systems \f$Ly=Pb\f$ and \f$Ux=y\f$. Since \f$L\f$ is
 * lower-triangular, the system \f$Ly=Pb\f$ can be solved by forward
 * substitution. Moreover, since \f$U\f$ is upper-triangular, the system
 * \f$Ux=y\f$ can be solved by backward substitution.
 */
template <typename MatrixExprT, typename VectorExprT, typename OutVectorT>
BOOST_UBLAS_INLINE
typename matrix_traits<MatrixExprT>::size_type lu_solve(matrix_expression<MatrixExprT> const& A, vector_expression<VectorExprT> const& b, OutVectorT& x)
{
    // precondition: size(b) == num_rows(A)
    BOOST_UBLAS_CHECK( size(b) == num_rows(A), bad_size() );

    // Ax=b ==> LUx=b ==> Ly=b AND Ux=y

    typedef typename matrix_traits<MatrixExprT>::value_type size_type;
    typedef typename matrix_traits<MatrixExprT>::value_type value_type;
    typedef typename layout_type<MatrixExprT>::type layout_type;

    matrix<value_type,layout_type> LU(A);
    permutation_matrix<size_type> P(num_rows(LU));

    size_type singular;
    singular = lu_decompose_inplace(LU, P);

    if (!singular)
    {
        x = b;

        lu_apply_inplace(LU, P, x);
    }

    return singular;
}

/**
 * \brief Solve the linear system \f$Ax=b\f$ by LUP decomposition.
 *
 * \pre The size of \a b must be the same as the number of rows of \a A.
 * \param A The input matrix representing the coefficients of the unknowns..
 * \param b The vector of coefficients.
 * \return The vector of unknowns \f$x\f$ satisfying the system \f$Ax=b\f$.
 *
 * An \f$n \times n\f$ linear system \f$Ax=b\f$ can often be solved efficiently
 * by \f$LU\f$ decomposition (that is by decomposing matrix \f$A\f$ into a
 * product \f$LUP\f$, where \f$L\f$ is a lower-triangular matrix with ones on
 * the diagonal, \f$U\f$ is an upper-triangular matrix, and \f$P\f$ is a row
 * permutation matrix that is used to rearrange the rows of A before so that it
 * can be decomposed). Solving the system \f$Ax=b\f$ is then equivalent to
 * solving two simpler systems \f$Ly=Pb\f$ and \f$Ux=y\f$. Since \f$L\f$ is
 * lower-triangular, the system \f$Ly=Pb\f$ can be solved by forward
 * substitution. Moreover, since \f$U\f$ is upper-triangular, the system
 * \f$Ux=y\f$ can be solved by backward substitution.
 */
template <typename AMatrixExprT, typename BMatrixExprT, typename OutMatrixT>
BOOST_UBLAS_INLINE
typename matrix_traits<AMatrixExprT>::size_type lu_solve(matrix_expression<AMatrixExprT> const& A, matrix_expression<BMatrixExprT> const& B, OutMatrixT& X)
{
    // precondition: size(b) == num_rows(A)
    BOOST_UBLAS_CHECK( num_rows(B) == num_rows(A), bad_size() );

    // Ax=b ==> LUx=b ==> Ly=b AND Ux=y

    typedef typename matrix_traits<AMatrixExprT>::value_type size_type;
    typedef typename matrix_traits<AMatrixExprT>::value_type value_type;
    typedef typename layout_type<AMatrixExprT>::type layout_type;

    matrix<value_type,layout_type> LU(A);
    permutation_matrix<size_type> P(num_rows(LU));

    size_type singular;
    singular = lu_decompose_inplace(LU, P);

    if (!singular)
    {
        X = B;

        lu_apply_inplace(LU, P, X);
    }

    return singular;
}

}}} // Namespace boost::numeric::ublasx

#endif // BOOST_NUMERIC_UBLASX_OPERATION_LU_HPP


//[hello

#include <boost/simd/sdk/simd/pack.hpp>
#include <boost/simd/sdk/simd/io.hpp>
#include <boost/simd/include/functions/splat.hpp>
#include <boost/simd/include/functions/plus.hpp>
#include <boost/simd/include/functions/multiplies.hpp>
#include <iostream>

int main()
{
  typedef boost::simd::pack<float> p_t;

  p_t res;
  p_t u(10);
  p_t r = boost::simd::splat<p_t>(11);

  res = (u + r) * 2.f;

  std::cout << res << std::endl;

  return 0;
}
//]


#pragma once
#include <cmath>
#include <omp.h>
#include <vector>
#include <string>
#include <cstdio>
#include <random>
#include <fstream>
#include <algorithm>
#include <unordered_set>
#include <parallel/algorithm>

#include <Eigen/Core>
#include <Eigen/Dense>

#include "n2/hnsw.h"

using namespace std;
using namespace Eigen;

namespace parallel
{

struct nn_t
{
    int key;
    float val;
    nn_t() : key(-1), val(-987654321.0) {}
    bool operator < (const float& that){
        return this->val > that;
    }

    bool operator < (const nn_t& that){
        return this->val > that.val;
    }
};

struct topn_t
{
    nn_t* nns;
    int K;
    topn_t() : K(0) {}
    void alloc(int _K){
        K = _K;
        nns = new nn_t[K];
    }
    void free() {
        if(K){
            delete[] nns;
        }
    }
    void update(int& key, float& val){
        nn_t* ptr = lower_bound(nns, nns + K, val);
        int idx = (int)(ptr - nns);
        if(idx >= K)
            return;
        if(idx + 1 == K){
            nns[idx].key = key;
            nns[idx].val = val;
        }
        else{
            memmove((void*)&nns[idx + 1], (void*)&nns[idx], sizeof(nn_t) * (K - idx - 1));
            nns[idx].key = key;
            nns[idx].val = val;
        }
    }
};

void quickselect(float* scores, int rows, int cols, int32_t* result, int k, bool sorted, int num_threads){
    Map<Matrix<float, Dynamic, Dynamic, RowMajor>> _scores(scores, rows, cols);
    Map<Matrix<int, Dynamic, Dynamic, RowMajor>> _result(result, rows, k);
    omp_set_num_threads(num_threads);
    #pragma omp parallel for schedule(dynamic, 4)
    for (int i=0; i<rows; ++i){
        vector<int> ranks(cols);
        iota(ranks.begin(), ranks.end(), 0);  // initialize ranks as {0, 1, 2, ...}
        // high score has a priority
        nth_element(ranks.begin(), ranks.begin() + k - 1, ranks.end(),
                    [&](int lhs, int rhs){return _scores(i, lhs) > _scores(i, rhs);});
        // sort ranks[: k], since it is not guaranteed to be sorted
        if (sorted)
            sort(ranks.begin(), ranks.begin() + k,
                    [&](int lhs, int rhs){return _scores(i, lhs) > _scores(i, rhs);});
        copy(ranks.begin(), ranks.begin() + k, &_result(i, 0));
    }
}

void dot_topn(
        int32_t* indexes, int num_queries,
        float* _P, int p_rows, int p_cols,
        float* _Q, int q_rows, int q_cols,
        float* _Qb, int qb_rows,
        int32_t* _out_keys, float* _out_scores,
        int32_t* _pool, int pool_size,
        int k, int num_threads)
{
    bool is_same = _P == _Q;
    Map<Matrix<float, Dynamic, Dynamic, RowMajor>> P(_P, p_rows, p_cols);
    Map<Matrix<float, Dynamic, Dynamic, RowMajor>> Q(_Q, q_rows, q_cols);
    Map<Matrix<float, Dynamic, Dynamic, RowMajor>> Qb(_Qb, qb_rows, 1);
    Map<Matrix<int, Dynamic, Dynamic, RowMajor>> out_keys(_out_keys, num_queries, k);
    Map<Matrix<float , Dynamic, Dynamic, RowMajor>> out_scores(_out_scores, num_queries, k);

    unordered_set<int32_t> pool;
    for (int i=0; i < pool_size; ++i)
        pool.insert(_pool[i]);

    int correct_k = min(q_rows, k);
    if (pool_size)
        correct_k = min(pool_size, correct_k);

    omp_set_num_threads(num_threads);
    #pragma omp parallel for schedule(guided)
    for (int i=0; i < num_queries; ++i){
        topn_t topn;
        topn.alloc(correct_k);
        float last_one = -987654321.0f;
        int q = indexes[i];
        for (int j=0; j < q_rows; ++j) {
            if (is_same and q == j)
                continue;
            if (pool_size and pool.find(j) == pool.end())
                continue;
            float score = P.row(q).dot(Q.row(j));
            if (qb_rows)
                score += Qb(j);
            if (score > last_one) {
                topn.update(j, score);
                last_one = topn.nns[correct_k - 1].val;
            }
        }
        for (int j=0; j < correct_k; ++j) {
            out_keys(i, j) = topn.nns[j].key;
            out_scores(i, j) = topn.nns[j].val;
        }
        for (int j=correct_k; j < k; ++j) {
            out_keys(i, j) = -1;
            out_scores(i, j) = 0.0;
        }
        topn.free();
    }
}

void ann_search(
        string index_path,
        int ef_search,
        bool use_mmap,
        int32_t* indexes, int num_queries,
        float* _P, int p_rows, int p_cols,
        float* _Q, int q_rows, int q_cols,
        float* _Qb, int qb_rows,
        int32_t* _out_keys, float* _out_scores,
        int32_t* _pool, int pool_size,
        int k, int num_threads)
{
    Map<Matrix<float, Dynamic, Dynamic, RowMajor>> P(_P, p_rows, p_cols);
    Map<Matrix<float, Dynamic, Dynamic, RowMajor>> Q(_Q, q_rows, q_cols);
    Map<Matrix<float, Dynamic, Dynamic, RowMajor>> Qb(_Qb, qb_rows, 1);
    Map<Matrix<int, Dynamic, Dynamic, RowMajor>> out_keys(_out_keys, num_queries, k);
    Map<Matrix<float , Dynamic, Dynamic, RowMajor>> out_scores(_out_scores, num_queries, k);

    unordered_set<int32_t> pool;
    for (int i=0; i < pool_size; ++i)
        pool.insert(_pool[i]);

    int correct_k = min(q_rows, k);
    if (pool_size)
        correct_k = min(pool_size, correct_k);

    omp_set_num_threads(num_threads);
    n2::Hnsw hnsws[num_threads];
    for (int i=0; i < num_threads; ++i) {
        hnsws[i].LoadModel(index_path.c_str(), use_mmap);
    }

    #pragma omp parallel
    {
        int worker_id = omp_get_thread_num();
        #pragma omp for schedule(guided)
        for (int i=0; i < num_queries; ++i)
        {
            auto& hnsw = hnsws[worker_id];
            int q = indexes[i];
            vector<pair<int, float>> result;
            hnsw.SearchById(q, correct_k, ef_search, result);
            for (int j=0; j < (int)result.size(); ++j) {
                out_keys(i, j) = result[j].first;
                out_scores(i, j) = result[j].second;
            }
            for (int j=(int)result.size(); j < k; ++j) {
                out_keys(i, j) = -1;
                out_scores(i, j) = 0.0;
            }
        }
    }
    for (int i=0; i < num_threads; ++i) {
        hnsws[i].UnloadModel();
    }
}

}


#include <boost/numeric/odeint/stepper/adams_bashforth.hpp>