#pragma once
#include <tuple>
#include <fastad_bits/reverse/core/expr_base.hpp>
#include <fastad_bits/reverse/core/value_adj_view.hpp>
#include <fastad_bits/reverse/core/constant.hpp>
#include <fastad_bits/util/type_traits.hpp>
#include <fastad_bits/util/numeric.hpp>
#include <Eigen/Dense>

namespace ad {
namespace stat {
namespace details {

template <class XExprType
        , class MeanExprType
        , class SigmaExprType>
struct NormalBase:
    core::ValueAdjView<util::common_value_t<
                        XExprType,
                        MeanExprType,
                        SigmaExprType>, ad::scl>
{
    using x_t = XExprType;
    using mean_t = MeanExprType;
    using sigma_t = SigmaExprType;
    using common_value_t = util::common_value_t<
        x_t, mean_t, sigma_t>;
    using value_adj_view_t = core::ValueAdjView<common_value_t, ad::scl>;
    using typename value_adj_view_t::value_t;
    using typename value_adj_view_t::shape_t;
    using typename value_adj_view_t::var_t;
    using typename value_adj_view_t::ptr_pack_t;

    NormalBase(const x_t& x,
               const mean_t& mean,
               const sigma_t& sigma)
        : value_adj_view_t(nullptr, nullptr, 1, 1)
        , x_{x}
        , mean_{mean}
        , sigma_{sigma}
    {}

    ptr_pack_t bind_cache(ptr_pack_t begin)
    {
        begin = x_.bind_cache(begin);
        begin = mean_.bind_cache(begin);
        begin = sigma_.bind_cache(begin);
        auto adj = begin.adj;
        begin.adj = nullptr;
        begin = value_adj_view_t::bind(begin);
        begin.adj = adj;
        return begin;
    }

    util::SizePack bind_cache_size() const
    {
        return single_bind_cache_size() +
                x_.bind_cache_size() +
                mean_.bind_cache_size() +
                sigma_.bind_cache_size();
    }

    util::SizePack single_bind_cache_size() const
    {
        return {this->size(), 0};
    }

protected:
    x_t x_;
    mean_t mean_;
    sigma_t sigma_;
};

} // namespace details

/**
 * NormalAdjLogPDFNode represents the normal log pdf
 * adjusted to omit all fixed constants, i.e. omits -n/2*log(2*pi).
 *
 * It assumes the value type that is common to all three expressions.
 * Since it represents a log-pdf, it is always a scalar expression.
 *
 * The only possible shape combinations are as follows:
 * x -> scalar, mean -> scalar, sigma -> scalar
 * x -> vec, mean -> scalar | vector, sigma -> scalar | vector | self adjoint matrix
 *
 * No other shapes are permitted for this node.
 *
 * At construction, the actual sizes of the three expressions are checked -
 * specifically if x is a vector, and mean and sigma are not scalar,
 * then size of x must be the same as that of mean rows and sigma rows.
 * Additionally, we check that sigma is square if it is a matrix.
 *
 * @tparam  XExprType           type of x expression at which to evaluate log-pdf
 * @tparam  MeanExprType        type of mean expression
 * @tparam  SigmaExprType       type of sigma expression
 */
template <class XExprType
        , class MeanExprType
        , class SigmaExprType
        , class = std::tuple<
            typename util::shape_traits<XExprType>::shape_t,
            typename util::shape_traits<MeanExprType>::shape_t,
            typename util::shape_traits<SigmaExprType>::shape_t> >
struct NormalAdjLogPDFNode;

// Case 1: sss
template <class XExprType
        , class MeanExprType
        , class SigmaExprType>
struct NormalAdjLogPDFNode<XExprType, MeanExprType, SigmaExprType,
                           std::tuple<scl, scl, scl> >:
    details::NormalBase<XExprType, MeanExprType, SigmaExprType>,
    core::ExprBase<NormalAdjLogPDFNode<XExprType, MeanExprType, SigmaExprType>>
{
private:
    using base_t = details::NormalBase<
        XExprType, MeanExprType, SigmaExprType>;

public:
    using typename base_t::x_t;
    using typename base_t::mean_t;
    using typename base_t::sigma_t;
    using typename base_t::value_t;
    using typename base_t::var_t;
    using base_t::x_;
    using base_t::mean_;
    using base_t::sigma_;

    NormalAdjLogPDFNode(const x_t& x,
                        const mean_t& mean,
                        const sigma_t& sigma)
        : base_t(x, mean, sigma)
        , log_sigma_{0}
    {
        if constexpr (util::is_constant_v<sigma_t>) {
            this->update_cache();
        }
    }

    const var_t& feval()
    {
        auto&& x = x_.feval();
        auto&& m = mean_.feval();
        auto&& s = sigma_.feval();

        if (s <= 0) return this->get() = util::neg_inf<value_t>;

        if constexpr (!util::is_constant_v<sigma_t>) {
            this->update_cache();
        }

        auto z = (x - m) / s;

        return this->get() = -0.5 * z * z - log_sigma_;
    }

    void beval(value_t seed)
    {
        if (seed == 0 || sigma_.get() <= 0) return;

        value_t inv_s = 1./sigma_.get();
        value_t z = (x_.get() - mean_.get()) * inv_s;

        if constexpr (!util::is_constant_v<sigma_t>) {
            sigma_.beval(seed * (z*z - 1) * inv_s);
        }
        value_t adj = seed * z * inv_s;
        mean_.beval(adj);
        x_.beval(-adj);
    }

private:
    void update_cache() {
        log_sigma_ = std::log(sigma_.get());
    }

    value_t log_sigma_;
};

// Case 2: vss
template <class XExprType
        , class MeanExprType
        , class SigmaExprType>
struct NormalAdjLogPDFNode<XExprType, MeanExprType, SigmaExprType,
                           std::tuple<vec, scl, scl> >:
    details::NormalBase<XExprType, MeanExprType, SigmaExprType>,
    core::ExprBase<NormalAdjLogPDFNode<XExprType, MeanExprType, SigmaExprType>>
{
private:
    using base_t = details::NormalBase<
        XExprType, MeanExprType, SigmaExprType>;

public:
    using typename base_t::x_t;
    using typename base_t::mean_t;
    using typename base_t::sigma_t;
    using typename base_t::value_t;
    using typename base_t::var_t;
    using base_t::x_;
    using base_t::mean_;
    using base_t::sigma_;

    NormalAdjLogPDFNode(const x_t& x,
                        const mean_t& mean,
                        const sigma_t& sigma)
        : base_t(x, mean, sigma)
        , log_sigma_{0}
        , z_sq{0}
        , x_mean_{0}
        , x_var_{0}
    {
        if constexpr (util::is_constant_v<sigma_t>) {
            this->update_cache();
        }

        // optimization when x_ is constant
        // reduced exponential form
        if constexpr (util::is_constant_v<x_t>) {
            x_mean_ = x_.get().mean();
            x_var_ = (x_.get().array() - x_mean_).matrix().squaredNorm();
        }
    }

    const var_t& feval()
    {
        auto&& x = x_.feval().array();
        auto&& m = mean_.feval();
        auto&& s = sigma_.feval();

        if (s <= 0) return this->get() = util::neg_inf<value_t>;

        if constexpr (!util::is_constant_v<sigma_t>) {
            this->update_cache();
        }

        if constexpr (util::is_constant_v<x_t>) {
            value_t centered = (m - x_mean_);
            value_t inv_s_sq = 1./(s * s);
            return this->get() =
                -0.5 * inv_s_sq * (x_var_ + x_.rows() * centered * centered)
                        - x_.rows() * log_sigma_;
        } else {
            auto z = (x - m).matrix();
            z_sq = z.squaredNorm() / (s * s);
            return this->get() = -0.5 * z_sq - x_.rows() * log_sigma_;
        }
    }

    void beval(value_t seed)
    {
        if (seed == 0 || sigma_.get() <= 0) return;

        value_t inv_s = 1./sigma_.get();
        value_t inv_s_sq = inv_s * inv_s;

        auto&& x = x_.get().array();
        auto&& m = mean_.get();

        // if x is constant, more optimized beval
        if constexpr (util::is_constant_v<x_t>) {

            if constexpr (!util::is_constant_v<sigma_t>) {
                value_t c = (m - x_mean_);
                value_t sigma_adj = ((x_var_ + x_.rows() * c * c) * inv_s_sq - x_.rows()) * inv_s;
                sigma_.beval(seed * sigma_adj);
            }

            value_t mean_adj = x_.rows() * (x_mean_ - m) * inv_s_sq;
            mean_.beval(seed * mean_adj);

        } else {

            if constexpr (!util::is_constant_v<sigma_t>) {
                sigma_.beval(seed * (z_sq - x_.rows()) * inv_s);
            }

            value_t mean_adj = (x.array() - m).sum() * inv_s_sq;
            mean_.beval(seed * mean_adj);

            if constexpr (!util::is_constant_v<x_t>) {
                x_.beval((seed * inv_s_sq) * (m - x));
            }

        }
    }

private:
    void update_cache() {
        log_sigma_ = std::log(sigma_.get());
    }

    value_t log_sigma_;

    // only used when x is not constant
    value_t z_sq;

    // only used when x is constant
    value_t x_mean_;
    value_t x_var_;
};

// Case 3: vvs
template <class XExprType
        , class MeanExprType
        , class SigmaExprType>
struct NormalAdjLogPDFNode<XExprType, MeanExprType, SigmaExprType,
                           std::tuple<vec, vec, scl> >:
    details::NormalBase<XExprType, MeanExprType, SigmaExprType>,
    core::ExprBase<NormalAdjLogPDFNode<XExprType, MeanExprType, SigmaExprType>>
{
private:
    using base_t = details::NormalBase<
        XExprType, MeanExprType, SigmaExprType>;

public:
    using typename base_t::x_t;
    using typename base_t::mean_t;
    using typename base_t::sigma_t;
    using typename base_t::value_t;
    using typename base_t::var_t;
    using base_t::x_;
    using base_t::mean_;
    using base_t::sigma_;

    NormalAdjLogPDFNode(const x_t& x,
                        const mean_t& mean,
                        const sigma_t& sigma)
        : base_t(x, mean, sigma)
        , log_sigma_{0}
        , z_sq{0}
    {
        if constexpr (util::is_constant_v<sigma_t>) {
            this->update_cache();
        }
    }

    const var_t& feval()
    {
        auto&& x = x_.feval().array();
        auto&& m = mean_.feval().array();
        auto&& s = sigma_.feval();

        if (s <= 0) return this->get() = util::neg_inf<value_t>;

        if constexpr (!util::is_constant_v<sigma_t>) {
            this->update_cache();
        }

        auto z = (x - m).matrix();
        z_sq = z.squaredNorm() / (s * s);

        return this->get() = -0.5 * z_sq - x_.rows() * log_sigma_;
    }

    void beval(value_t seed)
    {
        if (seed == 0 || sigma_.get() <= 0) return;

        value_t inv_s = 1./sigma_.get();
        value_t inv_s_sq = inv_s * inv_s;

        auto&& x = x_.get().array();
        auto&& m = mean_.get().array();

        if constexpr (!util::is_constant_v<sigma_t>) {
            sigma_.beval(seed * (z_sq - x_.rows()) * inv_s);
        }

        mean_.beval((seed * inv_s_sq) * (x - m));
        x_.beval((seed * inv_s_sq) * (m - x));
    }

private:
    void update_cache() {
        log_sigma_ = std::log(sigma_.get());
    }

    value_t log_sigma_;
    value_t z_sq;
};

// Case 4: vsv
template <class XExprType
        , class MeanExprType
        , class SigmaExprType>
struct NormalAdjLogPDFNode<XExprType, MeanExprType, SigmaExprType,
                           std::tuple<vec, scl, vec> >:
    details::NormalBase<XExprType, MeanExprType, SigmaExprType>,
    core::ExprBase<NormalAdjLogPDFNode<XExprType, MeanExprType, SigmaExprType>>
{
private:
    using base_t = details::NormalBase<
        XExprType, MeanExprType, SigmaExprType>;

public:
    using typename base_t::x_t;
    using typename base_t::mean_t;
    using typename base_t::sigma_t;
    using typename base_t::value_t;
    using typename base_t::var_t;
    using base_t::x_;
    using base_t::mean_;
    using base_t::sigma_;

    NormalAdjLogPDFNode(const x_t& x,
                        const mean_t& mean,
                        const sigma_t& sigma)
        : base_t(x, mean, sigma)
        , log_sigma_{0}
        , is_pos_def_{false}
        , sq_term_{0}
        , lin_term_{0}
        , const_term_{0}
    {
        if constexpr (util::is_constant_v<sigma_t>) {
            this->update_cache();

            // if additionally x is constant, more optimized form
            if constexpr (util::is_constant_v<x_t>) {
                auto&& x = x_.get().array();
                auto&& s = sigma_.get().array();
                sq_term_ = (x/s).matrix().squaredNorm();
                lin_term_ = (x/(s * s)).sum();
                const_term_ = (1./s).matrix().squaredNorm();
            }
        }
    }

    const var_t& feval()
    {
        auto&& x = x_.feval().array();
        auto&& m = mean_.feval();
        auto&& s = sigma_.feval().array();

        if constexpr (!util::is_constant_v<sigma_t>) {
            this->update_cache();
        }

        if (!is_pos_def_) {
            return this->get() = util::neg_inf<value_t>;
        }

        if constexpr (util::is_constant_v<x_t> &&
                      util::is_constant_v<sigma_t>) {
            return this->get() =
                -0.5 * (sq_term_ - 2 * m * lin_term_ + m * m * const_term_)
                    - log_sigma_;
        } else {
            auto z = ((x - m) / s).matrix();
            return this->get() = -0.5 * z.squaredNorm() - log_sigma_;
        }
    }

    void beval(value_t seed)
    {
        if (seed == 0 || !is_pos_def_) return;

        auto&& x = x_.get().array();
        auto&& m = mean_.get();
        auto&& s = sigma_.get().array();

        if constexpr (util::is_constant_v<x_t> &&
                      util::is_constant_v<sigma_t>) {
            value_t mean_adj = lin_term_ - m * const_term_;
            mean_.beval(seed * mean_adj);
        } else {

            if constexpr (!util::is_constant_v<sigma_t>) {
                sigma_.beval((seed / s) * ( ((x - m)/s).square() - 1. ));
            }

            value_t mean_adj = ((x - m) / s.square()).sum();
            mean_.beval(seed * mean_adj);
            x_.beval((seed / s.square()) * (m - x));

        }
    }

private:
    void update_cache() {
        is_pos_def_ = (sigma_.get().array() > 0).all();
        if (is_pos_def_) {
            log_sigma_ = sigma_.get().array().log().sum();
        }
    }

    value_t log_sigma_;
    size_t is_pos_def_;

    // only used when x and sigma are both constant
    value_t sq_term_;
    value_t lin_term_;
    value_t const_term_;
};

// Case 5: vvv
template <class XExprType
        , class MeanExprType
        , class SigmaExprType>
struct NormalAdjLogPDFNode<XExprType, MeanExprType, SigmaExprType,
                           std::tuple<vec, vec, vec> >:
    details::NormalBase<XExprType, MeanExprType, SigmaExprType>,
    core::ExprBase<NormalAdjLogPDFNode<XExprType, MeanExprType, SigmaExprType>>
{
private:
    using base_t = details::NormalBase<
        XExprType, MeanExprType, SigmaExprType>;

public:
    using typename base_t::x_t;
    using typename base_t::mean_t;
    using typename base_t::sigma_t;
    using typename base_t::value_t;
    using typename base_t::var_t;
    using base_t::x_;
    using base_t::mean_;
    using base_t::sigma_;

    NormalAdjLogPDFNode(const x_t& x,
                        const mean_t& mean,
                        const sigma_t& sigma)
        : base_t(x, mean, sigma)
        , log_sigma_{0}
        , is_pos_def_{false}
    {
        if constexpr (util::is_constant_v<sigma_t>) {
            this->update_cache();
        }
    }

    const var_t& feval()
    {
        auto&& x = x_.feval().array();
        auto&& m = mean_.feval().array();
        auto&& s = sigma_.feval().array();

        if constexpr (!util::is_constant_v<sigma_t>) {
            this->update_cache();
        }

        if (!is_pos_def_) {
            return this->get() = util::neg_inf<value_t>;
        }

        auto z = ((x - m) / s).matrix();

        return this->get() = -0.5 * z.squaredNorm() - log_sigma_;
    }

    void beval(value_t seed)
    {
        if (seed == 0 || !is_pos_def_) return;

        auto&& x = x_.get().array();
        auto&& m = mean_.get().array();
        auto&& s = sigma_.get().array();

        if constexpr (!util::is_constant_v<sigma_t>) {
            sigma_.beval((seed / s) * ( ((x - m)/s).square() - 1. ));
        }

        mean_.beval(seed * (x - m) / s.square());
        x_.beval(seed * (m - x) / s.square());
    }

private:
    void update_cache()
    {
        is_pos_def_ = (sigma_.get().array() > 0).all();
        if (is_pos_def_) {
            log_sigma_ = sigma_.get().array().log().sum();
        }
    }

    value_t log_sigma_;
    size_t is_pos_def_;
};

// Case 6: vsm
template <class XExprType
        , class MeanExprType
        , class SigmaExprType>
struct NormalAdjLogPDFNode<XExprType, MeanExprType, SigmaExprType,
                           std::tuple<vec, scl,
                                std::enable_if_t<util::is_mat_v<SigmaExprType>,
                                    typename util::shape_traits<SigmaExprType>::shape_t>> >:
    details::NormalBase<XExprType, MeanExprType, SigmaExprType>,
    core::ExprBase<NormalAdjLogPDFNode<XExprType, MeanExprType, SigmaExprType>>
{
private:
    using base_t = details::NormalBase<
        XExprType, MeanExprType, SigmaExprType>;

public:
    using typename base_t::x_t;
    using typename base_t::mean_t;
    using typename base_t::sigma_t;
    using typename base_t::value_t;
    using typename base_t::var_t;
    using base_t::x_;
    using base_t::mean_;
    using base_t::sigma_;

    NormalAdjLogPDFNode(const x_t& x,
                        const mean_t& mean,
                        const sigma_t& sigma)
        : base_t(x, mean, sigma)
        , llt_(sigma.rows())
        , log_det_{0}
        , is_pos_def_{false}
        , inv_(sigma.rows(), sigma.cols())
        , z_(mean.cols())
    {
        // must be square matrix
        assert(sigma_.rows() == sigma_.cols());
        assert(x_.rows() == sigma_.rows());

        if constexpr (util::is_constant_v<sigma_t>) {
            this->update_cache();
        }
    }

    const var_t& feval()
    {
        auto&& x = x_.feval().array();
        auto&& m = mean_.feval();
        sigma_.feval();

        if constexpr (!util::is_constant_v<sigma_t>) {
            this->update_cache();
        }

        if (!is_pos_def_) {
            return this->get() = util::neg_inf<value_t>;
        }

        z_ = inv_ * (x - m).matrix();
        value_t sq_term = (x - m).matrix().transpose() * z_;

        return this->get() = -0.5 * sq_term - log_det_;
    }

    void beval(value_t seed)
    {
        if (seed == 0 || !is_pos_def_) return;

        if constexpr (!util::is_constant_v<sigma_t>) {
            auto adj = (-0.5 * seed) * (inv_ - z_ * z_.transpose());
            sigma_.beval(adj.array());
        }

        mean_.beval(seed * z_.sum());
        x_.beval((-seed) * z_.array());
    }

private:
    void update_cache() {
        llt_.compute(sigma_.get());
        is_pos_def_ = (llt_.info() == Eigen::Success);
        if (is_pos_def_) {
            log_det_ = std::log(llt_.matrixL().determinant());
            inv_ = llt_.solve(mat_t::Identity(sigma_.rows(), sigma_.cols()));
        }
    }

    using mat_t = Eigen::Matrix<value_t, Eigen::Dynamic, Eigen::Dynamic>;
    using vec_t = Eigen::Matrix<value_t, Eigen::Dynamic, 1>;

    Eigen::LLT<mat_t, Eigen::Lower> llt_;
    value_t log_det_;
    bool is_pos_def_;
    mat_t inv_;
    vec_t z_;
};

// Case 7: vvm
template <class XExprType
        , class MeanExprType
        , class SigmaExprType>
struct NormalAdjLogPDFNode<XExprType, MeanExprType, SigmaExprType,
                           std::tuple<vec, vec,
                                std::enable_if_t<util::is_mat_v<SigmaExprType>,
                                    typename util::shape_traits<SigmaExprType>::shape_t>> >:
    details::NormalBase<XExprType, MeanExprType, SigmaExprType>,
    core::ExprBase<NormalAdjLogPDFNode<XExprType, MeanExprType, SigmaExprType>>
{
private:
    using base_t = details::NormalBase<
        XExprType, MeanExprType, SigmaExprType>;

public:
    using typename base_t::x_t;
    using typename base_t::mean_t;
    using typename base_t::sigma_t;
    using typename base_t::value_t;
    using typename base_t::var_t;
    using base_t::x_;
    using base_t::mean_;
    using base_t::sigma_;

    NormalAdjLogPDFNode(const x_t& x,
                        const mean_t& mean,
                        const sigma_t& sigma)
        : base_t(x, mean, sigma)
        , llt_(sigma.rows())
        , log_det_{0}
        , is_pos_def_{false}
        , inv_(sigma.rows(), sigma.cols())
        , z_(mean.cols())
    {
        // must be square matrix
        assert(sigma_.rows() == sigma_.cols());
        assert(x_.rows() == mean_.rows());
        assert(x_.rows() == sigma_.rows());

        if constexpr (util::is_constant_v<sigma_t>) {
            this->update_cache();
        }
    }

    const var_t& feval()
    {
        auto&& x = x_.feval();
        auto&& m = mean_.feval();
        sigma_.feval();

        if constexpr (!util::is_constant_v<sigma_t>) {
            this->update_cache();
        }

        if (!is_pos_def_) {
            return this->get() = util::neg_inf<value_t>;
        }

        z_ = inv_ * (x - m);
        value_t sq_term = (x - m).transpose() * z_;

        return this->get() = -0.5 * sq_term - log_det_;
    }

    void beval(value_t seed)
    {
        if (seed == 0 || !is_pos_def_) return;

        if constexpr (!util::is_constant_v<sigma_t>) {
            auto adj = (-0.5 * seed) * (inv_ - z_ * z_.transpose());
            sigma_.beval(adj.array());
        }

        mean_.beval(seed * z_.array());
        x_.beval((-seed) * z_.array());
    }

private:
    void update_cache() {
        llt_.compute(sigma_.get());
        is_pos_def_ = (llt_.info() == Eigen::Success);
        if (is_pos_def_) {
            log_det_ = std::log(llt_.matrixL().determinant());
            inv_ = llt_.solve(mat_t::Identity(sigma_.rows(), sigma_.cols()));
        }
    }

    using mat_t = Eigen::Matrix<value_t, Eigen::Dynamic, Eigen::Dynamic>;
    using vec_t = Eigen::Matrix<value_t, Eigen::Dynamic, 1>;

    Eigen::LLT<mat_t, Eigen::Lower> llt_;
    value_t log_det_;
    bool is_pos_def_;
    mat_t inv_;
    vec_t z_;
};

} // namespace stat

template <class XType
        , class MeanType
        , class SigmaType
        , class = std::enable_if_t<
            util::is_convertible_to_ad_v<XType> &&
            util::is_convertible_to_ad_v<MeanType> &&
            util::is_convertible_to_ad_v<SigmaType> &&
            util::any_ad_v<XType, MeanType, SigmaType> > >
inline auto normal_adj_log_pdf(const XType& x,
                               const MeanType& mean,
                               const SigmaType& sigma)
{
    using x_expr_t = util::convert_to_ad_t<XType>;
    using mean_expr_t = util::convert_to_ad_t<MeanType>;
    using sigma_expr_t = util::convert_to_ad_t<SigmaType>;
    x_expr_t x_expr = x;
    mean_expr_t mean_expr = mean;
    sigma_expr_t sigma_expr = sigma;
    return stat::NormalAdjLogPDFNode<
        x_expr_t, mean_expr_t, sigma_expr_t>(x_expr, mean_expr, sigma_expr);
}

} // namespace ad


#include <string>
#include <stdexcept>

#include <cstdlib>

#include <armadillo>

#include "Conv2D/Mesh.hpp"
#include "Conv2D/EulerDefaultBase.hpp"

class Euler : public EulerDefaultBase
{
public:
    Euler() : EulerDefaultBase()
    {
    }

    ~Euler() = default;

    void computeBottomFlux(const arma::vec &UInt, const arma::rowvec &n,
                           arma::vec &U, arma::vec &F, double &s) const;

    void computeRightFlux(const arma::vec &UInt, const arma::rowvec &n,
                          arma::vec &U, arma::vec &F, double &s) const;

    void computeTopFlux(const arma::vec &UInt, const arma::rowvec &n,
                        arma::vec &U, arma::vec &F, double &s) const;

    void computeLeftFlux(const arma::vec &UInt, const arma::rowvec &n,
                         arma::vec &U, arma::vec &F, double &s) const;
};

void Euler::computeTopFlux(const arma::vec &UInt, const arma::rowvec &n,
                           arma::vec &U, arma::vec &F, double &s) const
{
    applyInvisidWallBC(UInt, n, U, F, s);
}

void Euler::computeBottomFlux(const arma::vec &UInt, const arma::rowvec &n,
                              arma::vec &U, arma::vec &F, double &s) const
{
    applyInvisidWallBC(UInt, n, U, F, s);
}

void Euler::computeLeftFlux(const arma::vec &UInt, const arma::rowvec &n,
                            arma::vec &U, arma::vec &F, double &s) const
{
    applyInflowBC(UInt, n, U, F, s);
}

void Euler::computeRightFlux(const arma::vec &UInt, const arma::rowvec &n,
                             arma::vec &U, arma::vec &F, double &s) const
{
    applyOutflowBC(UInt, n, U, F, s);
}

int main(int argc, char **argv)
{
    if (argc != 2)
    {
        throw std::runtime_error("No choice of input mesh given");
    }

    int choice = std::atoi(argv[1]);

    if (choice < 0 || choice > 4)
    {
        throw std::runtime_error("Invalid choice of input mesh");
    }

    const std::string choiceString = std::to_string(choice);

    const std::string meshFile       = "bump" + choiceString + ".gri";
    const std::string residualFile   = "FirstOrderSolverResidual" + choiceString + ".dat";
    const std::string validationFile = "FirstOrderSolverValidation" + choiceString + ".dat";
    const std::string pressureFile   = "FirstOrderSolverPressureCoefficients" + choiceString + ".dat";
    const std::string machFile       = "FirstOrderSolverMachNumbers" + choiceString + ".vtk";
    const std::string solutionFile   = "FirstOrderSolverSolution" + choiceString + ".dat";

    Mesh mesh;
    mesh.readFromFile(meshFile);
    mesh.computeMatrices();

    Euler problem;
    problem.setMesh(mesh);
    problem.setGasConstant(1.0);
    problem.setSpecificHeatRatio(1.4);
    problem.setFreeFlowMachNumber(0.5);
    problem.setFreeFlowStaticPressure(1.0);
    problem.setCFLNumber(0.5);

    problem.setInitialState();

    const double tolerance = 1.0e-07;
    problem.runFirstOrderSolver(tolerance, residualFile);

    problem.writeValidationValuesTofile(validationFile);
    problem.writePressureCoefficientsToFile(pressureFile);
    problem.writeMachNumbersToFile(machFile);

    problem.writeStateToFile(solutionFile);

    return 0;
}


/*
 * fusion_runge_kutta.hpp
 *
 * Copyright 2010-2011 Mario Mulansky
 * Copyright 2010-2012 Karsten Ahnert
 *
 * 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 FUSION_EXPLICIT_RK_HPP_
#define FUSION_EXPLICIT_RK_HPP_

#include <boost/mpl/vector.hpp>
#include <boost/mpl/push_back.hpp>
#include <boost/mpl/for_each.hpp>
#include <boost/mpl/range_c.hpp>
#include <boost/mpl/copy.hpp>
#include <boost/mpl/size_t.hpp>

#include <boost/fusion/container.hpp>
#include <boost/fusion/algorithm/iteration.hpp>

#include <boost/array.hpp>

#include "fusion_algebra.hpp"
//#include "fusion_foreach_performance.hpp"

namespace mpl = boost::mpl;
namespace fusion = boost::fusion;

using namespace std;

struct intermediate_stage {};
struct last_stage {};

template< class T , class Constant >
struct array_wrapper
{
    typedef const typename boost::array< T , Constant::value > type;
};

template< class T , size_t i , class StageCategory >
struct stage
{
    T c;
    boost::array< T , i > a;
    typedef StageCategory category;
};

template< class T , size_t i>
struct stage< T , i , last_stage >
{
    T c;
    boost::array< T , i > b;
    typedef last_stage category;
};

template< class T , class Constant , class StageCategory >
struct stage_wrapper
{
    typedef stage< T , Constant::value , StageCategory > type;
};

template< class StateType , size_t stage_count >
class explicit_rk
{

public:

    typedef StateType state_type;

    typedef mpl::range_c< size_t , 1 , stage_count > stage_indices;

    typedef typename fusion::result_of::as_vector
    <
        typename mpl::copy
        <
            stage_indices ,
            mpl::inserter
            <
                mpl::vector0< > ,
                mpl::push_back< mpl::_1 , array_wrapper< double , mpl::_2 > >
            >
        >::type
    >::type coef_a_type;

    typedef boost::array< double , stage_count > coef_b_type;
    typedef boost::array< double , stage_count > coef_c_type;

    typedef typename fusion::result_of::as_vector
    <
        typename mpl::push_back
        <
            typename mpl::copy
            <
                stage_indices,
                mpl::inserter
                <
                    mpl::vector0<> ,
                    mpl::push_back< mpl::_1 , stage_wrapper< double , mpl::_2 , intermediate_stage > >
                >
            >::type ,
            stage< double , stage_count , last_stage >
        >::type
    >::type stage_vector_base;

    struct stage_vector : public stage_vector_base
    {
        struct do_insertion
        {
            stage_vector_base &m_base;
            const coef_a_type &m_a;
            const coef_c_type &m_c;

            do_insertion( stage_vector_base &base , const coef_a_type &a , const coef_c_type &c )
            : m_base( base ) , m_a( a ) , m_c( c ) { }

            template< class Index >
            void operator()( Index ) const
            {
                //fusion::at< Index >( m_base ) = stage< double , Index::value+1 , intermediate_stage >( m_c[ Index::value ] , fusion::at< Index >( m_a ) );
                fusion::at< Index >( m_base ).c  = m_c[ Index::value ];
                fusion::at< Index >( m_base ).a = fusion::at< Index >( m_a );
            }
        };

        stage_vector( const coef_a_type &a , const coef_b_type &b , const coef_c_type &c )
        {
            typedef mpl::range_c< size_t , 0 , stage_count - 1 > indices;
            mpl::for_each< indices >( do_insertion( *this , a , c ) );
            //fusion::at_c< 0 >( fusion::at_c< stage_count - 1 >( *this ) ) = stage_count - 1 ;
            fusion::at_c< stage_count - 1 >( *this ).c = c[ stage_count - 1 ];
            fusion::at_c< stage_count - 1 >( *this ).b = b;
        }
    };

    template< class System >
    struct calculate_stage
    {
        System &system;
        state_type &x , &x_tmp;
        state_type *F;
        const double t;
        const double dt;

        calculate_stage( System &_system , state_type &_x , state_type &_x_tmp , state_type *_F ,
                            const double _t , const double _dt )
        : system( _system ) , x( _x ) , x_tmp( _x_tmp ) , F( _F ) , t( _t ) , dt( _dt )
        {}

        template< typename T , size_t stage_number >
        void inline operator()( stage< T , stage_number , intermediate_stage > const &stage ) const
        //typename stage_fusion_wrapper< T , mpl::size_t< stage_number > , intermediate_stage >::type const &stage ) const
        {
            if( stage_number == 1 )
                system( x , F[stage_number-1] , t + stage.c * dt );
            else
                system( x_tmp , F[stage_number-1] , t + stage.c * dt );

            fusion_algebra<stage_number>::foreach( x_tmp , x , stage.a , F , dt);
        }

        template< typename T , size_t stage_number >
        void inline operator()( stage< T , stage_number , last_stage > const &stage ) const
        //void operator()( typename stage_fusion_wrapper< T , mpl::size_t< stage_number > , last_stage >::type const &stage ) const
        {
            if( stage_number == 1 )
                system( x , F[stage_number-1] , t + stage.c * dt );
            else
                system( x_tmp , F[stage_number-1] , t + stage.c * dt );

            fusion_algebra<stage_number>::foreach( x , x , stage.b , F , dt);
        }

    };

public:

    explicit_rk( const coef_a_type &a ,
                  const coef_b_type &b ,
                  const coef_c_type &c )
    : m_stages( a , b , c )

    { }

    template< class System >
    void inline do_step( System system , state_type &x , const double t , const double dt )
    {
        fusion::for_each( m_stages , calculate_stage< System >( system , x , m_x_tmp , m_F , t , dt ) );
    }

private:

    stage_vector m_stages;
    state_type m_x_tmp;

protected:
    state_type m_F[stage_count];

};

#endif /* FUSION_EXPLICIT_RK_HPP_ */


#ifndef STAN_MATH_PRIM_PROB_VON_MISES_RNG_HPP
#define STAN_MATH_PRIM_PROB_VON_MISES_RNG_HPP

#include <stan/math/prim/meta.hpp>
#include <stan/math/prim/err.hpp>
#include <stan/math/prim/fun/constants.hpp>
#include <stan/math/prim/fun/max_size.hpp>
#include <boost/random/uniform_real_distribution.hpp>
#include <boost/random/variate_generator.hpp>
#include <cmath>

namespace stan {
namespace math {

/** \ingroup prob_dists
 * Return a von Mises random variate for the given location and concentration
 * using the specified random number generator.
 *
 * mu and kappa can each be a scalar or a vector. Any non-scalar
 * inputs must be the same length.
 *
 * The algorithm used in von_mises_rng is a modified version of the
 * algorithm in:
 *
 * Efficient Simulation of the von Mises Distribution
 * D. J. Best and N. I. Fisher
 * Journal of the Royal Statistical Society. Series C (Applied Statistics),
 * Vol. 28, No. 2 (1979), pp. 152-157
 *
 * For kappa < 1.4e-8, this reduces to a uniform distribution.
 *
 * @tparam T_loc type of location parameter
 * @tparam T_conc type of scale (concentration) parameter
 * @tparam RNG type of random number generator
 *
 * @param mu (Sequence of) location parameter(s)
 * @param kappa (Sequence of) non-negative scale (concentration) parameter(s)
 * @param rng random number generator
 * @return (Sequence of) von Mises random variate(s)
 * @throw std::domain_error if mu is infinite or kappa is negative
 * @throw std::invalid_argument if non-scalar arguments are of different
 * sizes
 */
template <typename T_loc, typename T_conc, class RNG>
inline typename VectorBuilder<true, double, T_loc, T_conc>::type von_mises_rng(
    const T_loc& mu, const T_conc& kappa, RNG& rng) {
  using boost::random::uniform_real_distribution;
  using boost::variate_generator;
  static const char* function = "von_mises_rng";

  check_finite(function, "Location parameter", mu);
  check_nonnegative(function, "Scale parameter", kappa);
  check_finite(function, "Scale parameter", kappa);
  check_consistent_sizes(function, "Location parameter", mu, "Scale parameter",
                         kappa);

  scalar_seq_view<T_loc> mu_vec(mu);
  scalar_seq_view<T_conc> kappa_vec(kappa);
  size_t N = max_size(mu, kappa);
  VectorBuilder<true, double, T_loc, T_conc> output(N);

  variate_generator<RNG&, uniform_real_distribution<> > uniform_rng(
      rng, uniform_real_distribution<>(0.0, 1.0));

  for (size_t n = 0; n < N; ++n) {
    // for kappa sufficiently close to zero, it reduces to a
    // circular uniform distribution centered at mu
    if (kappa_vec[n] < 1.4e-8) {
      output[n] = (uniform_rng() - 0.5) * TWO_PI
                  + std::fmod(std::fmod(mu_vec[n], TWO_PI) + TWO_PI, TWO_PI);
      continue;
    }

    double r = 1 + std::pow((1 + 4 * kappa_vec[n] * kappa_vec[n]), 0.5);
    double rho = 0.5 * (r - std::pow(2 * r, 0.5)) / kappa_vec[n];
    double s = 0.5 * (1 + rho * rho) / rho;

    bool done = false;
    double W;
    while (!done) {
      double Z = std::cos(pi() * uniform_rng());
      W = (1 + s * Z) / (s + Z);
      double Y = kappa_vec[n] * (s - W);
      double U2 = uniform_rng();
      done = Y * (2 - Y) - U2 > 0;

      if (!done) {
        done = std::log(Y / U2) + 1 - Y >= 0;
      }
    }

    double U3 = uniform_rng() - 0.5;
    double sign = ((U3 >= 0) - (U3 <= 0));

    //  it's really an fmod() with a positivity constraint
    output[n] = sign * std::acos(W)
                + std::fmod(std::fmod(mu_vec[n], TWO_PI) + TWO_PI, TWO_PI);
  }

  return output.data();
}

}  // namespace math
}  // namespace stan
#endif


// Copyright 2018 The Simons Foundation, Inc. - All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <Eigen/Dense>
#include <algorithm>
#include <cmath>
#include <iostream>
#include "Utils/random_utils.hpp"
#include <vector>
#include "Utils/json_utils.hpp"
#include "abstract_hilbert.hpp"

#ifndef NETKET_BOSONS_HPP
#define NETKET_BOSONS_HPP

namespace netket {

/**
  Hilbert space for integer or bosons.
  The hilbert space is truncated to some maximum occupation number.
*/

class Boson : public AbstractHilbert {
  int nsites_;

  std::vector<double> local_;

  // total number of bosons
  // if constraint is activated
  int nbosons_;

  bool constraintN_;

  // maximum local occupation number
  int nmax_;

  int nstates_;

 public:
  explicit Boson(const json &pars) {
    if (!FieldExists(pars["Hilbert"], "Nsites")) {
      std::cerr << "Nsites is not defined" << std::endl;
    }

    nsites_ = pars["Hilbert"]["Nsites"];

    if (!FieldExists(pars["Hilbert"], "Nmax")) {
      std::cerr << "Nmax is not defined" << std::endl;
    }

    nmax_ = pars["Hilbert"]["Nmax"];

    Init();

    if (FieldExists(pars["Hilbert"], "Nbosons")) {
      SetNbosons(pars["Hilbert"]["Nbosons"]);
    } else {
      constraintN_ = false;
    }
  }

  void Init() {
    if (nsites_ <= 0) {
      std::cerr << "Invalid number of sites" << std::endl;
      std::abort();
    }

    if (nmax_ <= 0) {
      std::cerr << "Invalid maximum occupation number" << std::endl;
      std::abort();
    }

    nstates_ = nmax_ + 1;

    local_.resize(nstates_);

    for (int i = 0; i < nstates_; i++) {
      local_[i] = i;
    }
  }

  void SetNbosons(int nbosons) {
    constraintN_ = true;
    nbosons_ = nbosons;

    if (nbosons_ > nsites_ * nmax_) {
      std::cerr << "Cannot set the desired number of bosons" << std::endl;
      std::abort();
    }
  }

  bool IsDiscrete() const override { return true; }

  int LocalSize() const override { return nstates_; }

  int Size() const override { return nsites_; }

  std::vector<double> LocalStates() const override { return local_; }

  void RandomVals(Eigen::VectorXd &state,
                  netket::default_random_engine &rgen) const override {
    assert(state.size() == nsites_);

    if (!constraintN_) {
      std::uniform_int_distribution<int> distribution(0, nstates_ - 1);
      // unconstrained random
      for (int i = 0; i < state.size(); i++) {
        state(i) = distribution(rgen);
      }
    } else {
      state.setZero();

      std::uniform_int_distribution<int> distribution(0, nsites_ - 1);
      for (int i = 0; i < nbosons_; i++) {
        int rsite = distribution(rgen);

        while (state(rsite) >= nmax_) {
          rsite = distribution(rgen);
        }

        state(rsite) += 1;
      }
    }
  }

  bool CheckConstraint(Eigen::VectorXd &v) const {
    int tot = 0;
    for (int i = 0; i < v.size(); i++) {
      tot += int(v(i));
    }

    return tot == nbosons_;
  }

  void UpdateConf(Eigen::VectorXd &v, const std::vector<int> &tochange,
                  const std::vector<double> &newconf) const override {
    assert(v.size() == nsites_);

    int i = 0;
    for (auto sf : tochange) {
      v(sf) = newconf[i];
      i++;
      assert(v(sf) <= nmax_);
    }

    if (constraintN_) {
      assert(CheckConstraint(v));
    }
  }
};

}  // namespace netket
#endif


#include "transport.hpp"

#include <assert.h>
#include <limits>
#include <vector>

#include <boost/intrusive/list.hpp>
#include <boost/intrusive/slist.hpp>

#include "union-find.hpp"

template <typename T>
class TransportGraph {
    public:
        TransportGraph(int sizeSupply, int sizeDemand)
            : _sizeSupply(sizeSupply)
            , _sizeDemand(sizeDemand)
            , _nodes(sizeSupply+sizeDemand)
            , _treeEdges(2*(sizeSupply+sizeDemand - 1))
        {
            for (int i = 0; i < _sizeSupply+_sizeDemand; ++i)
                _nodes[i].id = i;
        }

        typedef boost::intrusive::list_base_hook<
            boost::intrusive::link_mode<boost::intrusive::normal_link>
            > ListHook;
        struct Edge : public ListHook {
            int target;
            int edgeIdx;
            Edge* reverse;
        };

        typedef boost::intrusive::list<Edge,
                boost::intrusive::base_hook<ListHook>> EdgeList;

        typedef boost::intrusive::slist_member_hook<
            boost::intrusive::link_mode<boost::intrusive::normal_link>
            > StackHook;

        struct Node {
            int id = 0;
            T potential = 0;
            EdgeList outEdges;
            typename EdgeList::iterator dfsOutEdge;
            StackHook stackHook;
            bool inDfsStack = false;
        };
        typedef boost::intrusive::slist<Node,
                boost::intrusive::member_hook<Node,
                                              StackHook,
                                              &Node::stackHook>
                > NodeStack;

        void addTreeEdge(int supply, int demand) {
            assert(_numTreeEdges+2 <= static_cast<int>(_treeEdges.size()));
            Edge& e1 = _treeEdges[_numTreeEdges];
            Edge& e2 = _treeEdges[_numTreeEdges+1];
            _numTreeEdges += 2;

            int edgeIdx = supply*_sizeDemand+demand;
            e1.target = demand+_sizeSupply;
            e1.edgeIdx = edgeIdx;
            e1.reverse = &e2;
            e2.target = supply;
            e2.edgeIdx = edgeIdx;
            e2.reverse = &e1;

            _nodes[supply].outEdges.push_back(e1);
            _nodes[demand+_sizeSupply].outEdges.push_back(e2);
        }

        void augment(int supply, int demand, T* flows) {
            int edgeIdx = supply*_sizeDemand+demand;

            for (auto& n : _nodes)
                n.inDfsStack = false;

            NodeStack stack;
            {
                Node& startNode = _nodes.at(supply);
                stack.push_front(startNode);
                startNode.dfsOutEdge = startNode.outEdges.begin();
                startNode.inDfsStack = true;
            }
            while (!stack.empty() && stack.front().id != demand+_sizeSupply) {
                auto& n = stack.front();
                if (n.dfsOutEdge == n.outEdges.end()) {
                    stack.pop_front();
                    assert(!stack.empty());
                    auto& nextN = stack.front();
                    assert(nextN.dfsOutEdge != nextN.outEdges.end());
                    assert(nextN.dfsOutEdge->target == n.id);
                    nextN.dfsOutEdge++;
                    continue;
                } else {
                    int target = n.dfsOutEdge->target;
                    auto& nextN = _nodes.at(target);
                    if (nextN.inDfsStack) {
                        n.dfsOutEdge++;
                        continue;
                    } else {
                        stack.push_front(nextN);
                        nextN.dfsOutEdge = nextN.outEdges.begin();
                        nextN.inDfsStack = true;
                        continue;
                    }
                }
            }
            assert(!stack.empty());
            assert(stack.front().id == demand+_sizeSupply);
            assert((stack.size() % 2) == 0);
            stack.pop_front();

            T bottleneckFlow = std::numeric_limits<T>::max();
            Edge* bottleneckEdge = nullptr;
            bool parity = true;
            for (auto& n : stack) {
                if (parity) {
                    auto& e = *n.dfsOutEdge;
                    auto f = flows[e.edgeIdx];
                    if (f < bottleneckFlow) {
                        bottleneckFlow = f;
                        bottleneckEdge = &e;
                    }
                }
                parity = !parity;
            }
            assert(bottleneckFlow < std::numeric_limits<T>::max());
            assert(bottleneckEdge != nullptr);

            parity = true;
            for (auto& n : stack) {
                auto& e = *n.dfsOutEdge;
                if (parity)
                    flows[e.edgeIdx] -= bottleneckFlow;
                else
                    flows[e.edgeIdx] += bottleneckFlow;
                parity = !parity;
            }
            flows[edgeIdx] += bottleneckFlow;

            Edge& e1 = *bottleneckEdge;
            Edge& e2 = *e1.reverse;
            {
                Node& n1 = _nodes[e2.target];
                Node& n2 = _nodes[e1.target];
                n1.outEdges.erase(n1.outEdges.iterator_to(e1));
                n2.outEdges.erase(n2.outEdges.iterator_to(e2));
            }
            {
                Node& n1 = _nodes[supply];
                Node& n2 = _nodes[demand+_sizeSupply];
                e1.target = demand+_sizeSupply;
                e1.edgeIdx = edgeIdx;
                e1.reverse = &e2;
                e2.target = supply;
                e2.edgeIdx = edgeIdx;
                e2.reverse = &e1;
                n1.outEdges.push_back(e1);
                n2.outEdges.push_back(e2);
            }
        }

        void updatePotentials(const T* costs) {
            for (auto& n : _nodes) {
                n.potential = 0;
                n.inDfsStack = false;
            }

            NodeStack stack;
            {
                Node& startNode = _nodes.at(0);
                stack.push_front(startNode);
                startNode.dfsOutEdge = startNode.outEdges.begin();
                startNode.inDfsStack = true;
            }
            while (!stack.empty()) {
                auto& n = stack.front();
                if (n.dfsOutEdge == n.outEdges.end()) {
                    stack.pop_front();
                    if(!stack.empty()) {
                        auto& nextN = stack.front();
                        assert(nextN.dfsOutEdge != nextN.outEdges.end());
                        assert(nextN.dfsOutEdge->target == n.id);
                        nextN.dfsOutEdge++;
                    }
                    continue;
                } else {
                    int target = n.dfsOutEdge->target;
                    auto& nextN = _nodes.at(target);
                    if (nextN.inDfsStack) {
                        n.dfsOutEdge++;
                        continue;
                    } else {
                        int edgeIdx = n.dfsOutEdge->edgeIdx;
                        T c = costs[edgeIdx];
                        nextN.potential = c - n.potential;

                        stack.push_front(nextN);
                        nextN.dfsOutEdge = nextN.outEdges.begin();
                        nextN.inDfsStack = true;
                        continue;
                    }
                }
            }
        }

        std::tuple<int, int, T> findPivot(const T* costs) {
            std::tuple<int, int, T> minEdge
                = {0, 0, std::numeric_limits<T>::max()};
            for (int i = 0; i < _sizeSupply; ++i) {
                for (int j = 0; j < _sizeDemand; ++j) {
                    int edgeIdx = i*_sizeDemand+j;
                    T resCost = costs[edgeIdx]
                        - _nodes[i].potential
                        - _nodes[j+_sizeSupply].potential;
                    if (resCost < std::get<2>(minEdge))
                        minEdge = {i, j, resCost};
                }
            }
            return minEdge;
        }

        int _sizeSupply;
        int _sizeDemand;
        int _numTreeEdges = 0;
        std::vector<Node> _nodes;
        std::vector<Edge> _treeEdges;
};

template <typename T>
double solveTransport(int sizeSupply, int sizeDemand, const T* costs,
        const T* supply, const T* demand, T* flow) {
    T sumSupply = 0;
    for (int i = 0; i < sizeSupply; ++i)
        sumSupply += supply[i];
    T sumDemand = 0;
    for (int j = 0; j < sizeDemand; ++j)
        sumDemand += demand[j];
    assert(fabs(sumSupply - sumDemand) < 1e-5);

    for (int k = 0; k < sizeSupply*sizeDemand; ++k)
        flow[k] = 0;

    std::vector<T> resSupply(supply, supply+sizeSupply);
    std::vector<T> resDemand(demand, demand+sizeDemand);

    TransportGraph<T> graph{sizeSupply, sizeDemand};

    int numNodes = sizeSupply+sizeDemand;
    auto uf = UnionFind{numNodes};
    for (int k = 0; k < numNodes-1; ++k) {
        T minCost = std::numeric_limits<T>::max();
        int minI = 0;
        int minJ = 0;
        for (int i = 0; i < sizeSupply; ++i) {
            if (resSupply.at(i) <= 0)
                continue;
            auto comp_i = uf.Find(i);
            for (int j = 0; j < sizeDemand; ++j) {
                if (resDemand.at(j) <= 0)
                    continue;
                auto comp_j = uf.Find(j+sizeSupply);
                if (comp_i == comp_j)
                    continue;
                if (costs[i*sizeDemand+j] < minCost) {
                    minCost = costs[i*sizeDemand+j];
                    minI = i;
                    minJ = j;
                }
            }
        }
        if (minCost == std::numeric_limits<T>::max())
            break;
        assert(uf.Find(minI) != uf.Find(minJ+sizeSupply));
        T f = std::min(resSupply.at(minI), resDemand.at(minJ));
        resSupply.at(minI) -= f;
        resDemand.at(minJ) -= f;
        flow[minI*sizeDemand+minJ] += f;
        graph.addTreeEdge(minI, minJ);
        uf.Merge(minI, minJ+sizeSupply);
    }
    // Add any unconnected components to connect to 0
    for (int i = 0; i < sizeSupply; ++i) {
        if (uf.Find(i) != uf.Find(0+sizeSupply)) {
            graph.addTreeEdge(i, 0);
            uf.Merge(i, 0+sizeSupply);
        }
    }
    for (int j = 0; j < sizeDemand; ++j) {
        if (uf.Find(0) != uf.Find(j+sizeSupply)) {
            graph.addTreeEdge(0, j);
            uf.Merge(0, j+sizeSupply);
        }
    }

#ifndef NDEBUG
    for (auto s : resSupply) {
        assert(fabs(s) < 1e-5);
    }
    for (auto d : resDemand)
        assert(fabs(d) < 1e-5);
#endif

    while (true) {
        graph.updatePotentials(costs);
        std::tuple<int, int, T> minEdge = graph.findPivot(costs);
        if (std::get<2>(minEdge) < 0)
            graph.augment(std::get<0>(minEdge), std::get<1>(minEdge), flow);
        else
            break;
    }

    double cost = 0;
    for (int i = 0; i < sizeSupply; ++i) {
        for (int j = 0; j < sizeDemand; ++j) {
            int idx = i*sizeDemand + j;
            cost += flow[idx] * costs[idx];
        }
    }
    return cost;
}

#define INSTANTIATE_TRANSPORT(T) \
    template double solveTransport<T>(int sizeSupply, int sizeDemand, \
            const T* costs, const T* supply, const T* demand, T* flow);
INSTANTIATE_TRANSPORT(double)
INSTANTIATE_TRANSPORT(int32_t)
INSTANTIATE_TRANSPORT(int64_t)
#undef INSTANTIATE_TRANSPORT


// --------------------------------------------------------------------------
//                   OpenMS -- Open-Source Mass Spectrometry
// --------------------------------------------------------------------------
// Copyright The OpenMS Team -- Eberhard Karls University Tuebingen,
// ETH Zurich, and Freie Universitaet Berlin 2002-2016.
//
// This software is released under a three-clause BSD license:
//  * Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
//  * Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//  * Neither the name of any author or any participating institution
//    may be used to endorse or promote products derived from this software
//    without specific prior written permission.
// For a full list of authors, refer to the file AUTHORS.
// --------------------------------------------------------------------------
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL ANY OF THE AUTHORS OR THE CONTRIBUTING
// INSTITUTIONS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
// OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
// OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
// ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// --------------------------------------------------------------------------
// $Maintainer: Stephan Aiche $
// $Authors: Andreas Bertsch $
// --------------------------------------------------------------------------
//

#include <sstream>
#include <iostream>
#include <cmath>
#include <boost/math/special_functions/gamma.hpp>
#include <boost/math/special_functions/digamma.hpp>

#include <unsupported/Eigen/NonLinearOptimization>

#include <OpenMS/MATH/STATISTICS/GammaDistributionFitter.h>

// #define GAMMA_DISTRIBUTION_FITTER_VERBOSE
// #undef  GAMMA_DISTRIBUTION_FITTER_VERBOSE

namespace OpenMS
{
  namespace Math
  {
    GammaDistributionFitter::GammaDistributionFitter() :
      init_param_(1.0, 5.0)
    {
    }

    GammaDistributionFitter::~GammaDistributionFitter()
    {
    }

    void GammaDistributionFitter::setInitialParameters(const GammaDistributionFitResult& param)
    {
      init_param_ = param;
    }

    struct GammaFunctor
    {
      int inputs() const { return m_inputs; }
      int values() const { return m_values; }

      GammaFunctor(unsigned dimensions, const std::vector<DPosition<2> >* data) :
        m_inputs(dimensions), m_values(data->size()), m_data(data) {}

      int operator()(const Eigen::VectorXd& x, Eigen::VectorXd& fvec)
      {

        double b = x(0);
        double p = x(1);

        UInt i = 0;

        // gamma distribution is only defined for positive parameter values
        if (b > 0.0 && p > 0.0)
        {
          for (std::vector<DPosition<2> >::const_iterator it = m_data->begin(); it != m_data->end(); ++it, ++i)
          {
            double the_x = it->getX();
            fvec(i) =  std::pow(b, p) / boost::math::tgamma(p) * std::pow(the_x, p - 1) * std::exp(-b * the_x) - it->getY();
          }
        }
        else
        {
          for (std::vector<DPosition<2> >::const_iterator it = m_data->begin(); it != m_data->end(); ++it, ++i)
          {
            fvec(i) = -it->getY();
          }
        }

        return 0;
      }

      // compute Jacobian matrix for the different parameters
      int df(const Eigen::VectorXd& x, Eigen::MatrixXd& J)
      {

        double b = x(0);
        double p = x(1);

        UInt i = 0;
        // gamma distribution is only defined for positive parameter values
        if (b > 0.0 && p > 0.0)
        {
          for (std::vector<DPosition<2> >::const_iterator it = m_data->begin(); it != m_data->end(); ++it, ++i)
          {
            double the_x = it->getX();

            // partial deviation regarding b
            double part_dev_b = std::pow(the_x, p - 1) * std::exp(-the_x * b) / boost::math::tgamma(p) * (p * std::pow(b, p - 1) - the_x * std::pow(b, p));
            J(i, 0) = part_dev_b;

            // partial deviation regarding p
            double factor = std::exp(-b * the_x) * std::pow(the_x, p - 1) * std::pow(b, p) / std::pow(boost::math::tgamma(p), 2);
            double argument = (std::log(b) + std::log(the_x)) * boost::math::tgamma(p) - boost::math::tgamma(p) * boost::math::digamma(p);
            double part_dev_p = factor * argument;
            J(i, 1) = part_dev_p;
          }
        }
        else
        {
          for (std::vector<DPosition<2> >::const_iterator it = m_data->begin(); it != m_data->end(); ++it, ++i)
          {
            J(i, 0) = 0.0;
            J(i, 1) = 0.0;
          }
        }
        return 0;
      }

      const int m_inputs, m_values;
      const std::vector<DPosition<2> >* m_data;
    };

    GammaDistributionFitter::GammaDistributionFitResult GammaDistributionFitter::fit(const std::vector<DPosition<2> >& input)
    {
      Eigen::VectorXd x_init(2);
      x_init << init_param_.b, init_param_.p;
      GammaFunctor functor(2, &input);
      Eigen::LevenbergMarquardt<GammaFunctor> lmSolver(functor);
      Eigen::LevenbergMarquardtSpace::Status status = lmSolver.minimize(x_init);

      //the states are poorly documented. after checking the source, we believe that
      //all states except NotStarted, Running and ImproperInputParameters are good
      //termination states.
      if (status <= Eigen::LevenbergMarquardtSpace::ImproperInputParameters)
      {
        throw Exception::UnableToFit(__FILE__, __LINE__, __PRETTY_FUNCTION__, "UnableToFit-GammaDistributionFitter", "Could not fit the gamma distribution to the data");
      }

#ifdef GAMMA_DISTRIBUTION_FITTER_VERBOSE
      std::stringstream formula;
      formula << "f(x)=" << "(" << x_init(0) << " ** " << x_init(1) << ") / gamma(" << x_init(1) << ") * x ** (" << x_init(1) << " - 1) * exp(- " << x_init(0) << " * x)";
      std::cout << formula.str() << std::endl;
#endif

      return GammaDistributionFitResult(x_init(0), x_init(1));
    }

  } //namespace Math
} // namespace OpenMS


#ifndef STAN_MATH_PRIM_SCAL_PROB_NEG_BINOMIAL_LPMF_HPP
#define STAN_MATH_PRIM_SCAL_PROB_NEG_BINOMIAL_LPMF_HPP

#include <stan/math/prim/scal/meta/is_constant_struct.hpp>
#include <stan/math/prim/scal/meta/operands_and_partials.hpp>
#include <stan/math/prim/scal/err/check_consistent_sizes.hpp>
#include <stan/math/prim/scal/err/check_nonnegative.hpp>
#include <stan/math/prim/scal/err/check_positive_finite.hpp>
#include <stan/math/prim/scal/fun/size_zero.hpp>
#include <stan/math/prim/scal/fun/constants.hpp>
#include <stan/math/prim/scal/fun/value_of.hpp>
#include <stan/math/prim/scal/fun/binomial_coefficient_log.hpp>
#include <stan/math/prim/scal/fun/multiply_log.hpp>
#include <stan/math/prim/scal/fun/digamma.hpp>
#include <stan/math/prim/scal/fun/lgamma.hpp>
#include <stan/math/prim/scal/fun/lbeta.hpp>
#include <stan/math/prim/scal/meta/length.hpp>
#include <stan/math/prim/scal/meta/scalar_seq_view.hpp>
#include <stan/math/prim/scal/meta/VectorBuilder.hpp>
#include <stan/math/prim/scal/meta/partials_return_type.hpp>
#include <stan/math/prim/scal/meta/return_type.hpp>
#include <stan/math/prim/scal/meta/include_summand.hpp>
#include <stan/math/prim/scal/fun/grad_reg_inc_beta.hpp>
#include <stan/math/prim/scal/fun/inc_beta.hpp>
#include <stan/math/prim/scal/meta/max_size.hpp>
#include <boost/math/special_functions/digamma.hpp>
#include <boost/random/negative_binomial_distribution.hpp>
#include <boost/random/variate_generator.hpp>
#include <cmath>

namespace stan {
namespace math {

// NegBinomial(n|alpha, beta)  [alpha > 0;  beta > 0;  n >= 0]
template <bool propto, typename T_n, typename T_shape, typename T_inv_scale>
typename return_type<T_shape, T_inv_scale>::type neg_binomial_lpmf(
    const T_n& n, const T_shape& alpha, const T_inv_scale& beta) {
  typedef typename stan::partials_return_type<T_n, T_shape, T_inv_scale>::type
      T_partials_return;

  static const char* function = "neg_binomial_lpmf";

  if (size_zero(n, alpha, beta))
    return 0.0;

  T_partials_return logp(0.0);
  check_nonnegative(function, "Failures variable", n);
  check_positive_finite(function, "Shape parameter", alpha);
  check_positive_finite(function, "Inverse scale parameter", beta);
  check_consistent_sizes(function, "Failures variable", n, "Shape parameter",
                         alpha, "Inverse scale parameter", beta);

  if (!include_summand<propto, T_shape, T_inv_scale>::value)
    return 0.0;

  using std::log;
  using std::log;

  scalar_seq_view<T_n> n_vec(n);
  scalar_seq_view<T_shape> alpha_vec(alpha);
  scalar_seq_view<T_inv_scale> beta_vec(beta);
  size_t size = max_size(n, alpha, beta);

  operands_and_partials<T_shape, T_inv_scale> ops_partials(alpha, beta);

  size_t len_ab = max_size(alpha, beta);
  VectorBuilder<true, T_partials_return, T_shape, T_inv_scale> lambda(len_ab);
  for (size_t i = 0; i < len_ab; ++i)
    lambda[i] = value_of(alpha_vec[i]) / value_of(beta_vec[i]);

  VectorBuilder<true, T_partials_return, T_inv_scale> log1p_beta(length(beta));
  for (size_t i = 0; i < length(beta); ++i)
    log1p_beta[i] = log1p(value_of(beta_vec[i]));

  VectorBuilder<true, T_partials_return, T_inv_scale> log_beta_m_log1p_beta(
      length(beta));
  for (size_t i = 0; i < length(beta); ++i)
    log_beta_m_log1p_beta[i] = log(value_of(beta_vec[i])) - log1p_beta[i];

  VectorBuilder<true, T_partials_return, T_inv_scale, T_shape>
      alpha_times_log_beta_over_1p_beta(len_ab);
  for (size_t i = 0; i < len_ab; ++i)
    alpha_times_log_beta_over_1p_beta[i]
        = value_of(alpha_vec[i])
          * log(value_of(beta_vec[i]) / (1.0 + value_of(beta_vec[i])));

  VectorBuilder<!is_constant_struct<T_shape>::value, T_partials_return, T_shape>
      digamma_alpha(length(alpha));
  if (!is_constant_struct<T_shape>::value) {
    for (size_t i = 0; i < length(alpha); ++i)
      digamma_alpha[i] = digamma(value_of(alpha_vec[i]));
  }

  VectorBuilder<!is_constant_struct<T_shape>::value, T_partials_return,
                T_inv_scale>
      log_beta(length(beta));
  if (!is_constant_struct<T_shape>::value) {
    for (size_t i = 0; i < length(beta); ++i)
      log_beta[i] = log(value_of(beta_vec[i]));
  }

  VectorBuilder<!is_constant_struct<T_inv_scale>::value, T_partials_return,
                T_shape, T_inv_scale>
      lambda_m_alpha_over_1p_beta(len_ab);
  if (!is_constant_struct<T_inv_scale>::value) {
    for (size_t i = 0; i < len_ab; ++i)
      lambda_m_alpha_over_1p_beta[i]
          = lambda[i]
            - (value_of(alpha_vec[i]) / (1.0 + value_of(beta_vec[i])));
  }

  for (size_t i = 0; i < size; i++) {
    if (alpha_vec[i] > 1e10) {  // reduces numerically to Poisson
      if (include_summand<propto>::value)
        logp -= lgamma(n_vec[i] + 1.0);
      if (include_summand<propto, T_shape, T_inv_scale>::value)
        logp += multiply_log(n_vec[i], lambda[i]) - lambda[i];

      if (!is_constant_struct<T_shape>::value)
        ops_partials.edge1_.partials_[i]
            += n_vec[i] / value_of(alpha_vec[i]) - 1.0 / value_of(beta_vec[i]);
      if (!is_constant_struct<T_inv_scale>::value)
        ops_partials.edge2_.partials_[i]
            += (lambda[i] - n_vec[i]) / value_of(beta_vec[i]);
    } else {  // standard density definition
      if (include_summand<propto, T_shape>::value)
        if (n_vec[i] != 0)
          logp += binomial_coefficient_log(
              n_vec[i] + value_of(alpha_vec[i]) - 1.0, n_vec[i]);
      if (include_summand<propto, T_shape, T_inv_scale>::value)
        logp += alpha_times_log_beta_over_1p_beta[i] - n_vec[i] * log1p_beta[i];

      if (!is_constant_struct<T_shape>::value)
        ops_partials.edge1_.partials_[i]
            += digamma(value_of(alpha_vec[i]) + n_vec[i]) - digamma_alpha[i]
               + log_beta_m_log1p_beta[i];
      if (!is_constant_struct<T_inv_scale>::value)
        ops_partials.edge2_.partials_[i]
            += lambda_m_alpha_over_1p_beta[i]
               - n_vec[i] / (value_of(beta_vec[i]) + 1.0);
    }
  }
  return ops_partials.build(logp);
}

template <typename T_n, typename T_shape, typename T_inv_scale>
inline typename return_type<T_shape, T_inv_scale>::type neg_binomial_lpmf(
    const T_n& n, const T_shape& alpha, const T_inv_scale& beta) {
  return neg_binomial_lpmf<false>(n, alpha, beta);
}

}  // namespace math
}  // namespace stan
#endif


#include <iostream>
#include "specex_linalg.h"
#include <boost/archive/xml_oarchive.hpp>
#include <boost/archive/text_oarchive.hpp>
#include <boost/archive/basic_text_oarchive.hpp>
#include <time.h>
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/triangular.hpp>
#include <boost/random.hpp>

using namespace std;

int main() {

  harp::vector_double x(4);
  x(1)=1;
  x[2]=2;

  if(0) {
  boost::archive::xml_oarchive xml_oa ( cout );
  xml_oa << BOOST_SERIALIZATION_NVP(x);

  boost::archive::text_oarchive text_oa ( cout );
  text_oa << BOOST_SERIALIZATION_NVP(x);

  boost::archive::text_oarchive basic_text_oa ( cout );
  basic_text_oa << BOOST_SERIALIZATION_NVP(x);
  }

  {

    //ublas::mapped_vector<double> sv (4000);
    //ublas::compressed_vector<double> sv (4000);
     cout << "==================" << endl;
  cout << " SPARSE VECTORS   " << endl;
  cout << "==================" << endl;

  int n = 8000;
  //ublas::coordinate_vector<double> sv (n);
  //ublas::mapped_vector<double> sv (n);
  ublas::compressed_vector<double> sv (n);
  //ublas::vector<double> sv (n);
  harp::vector_double v (n);
  harp::vector_double v2 (n);
  double w = 0.2;
  v.clear();
  v2.clear();
  sv.clear();

  for (unsigned i = 0; i < sv.size (); i ++) {
    v2(i)=i;
  }
  for (unsigned i = 0; i < sv.size (); i += 10) {
    sv(i) = i;
    v(i)=i;
  }

  {
    harp::matrix_double a(10,10);
    a.clear();
    harp::vector_double h = ublas::project(v2,ublas::range(0,10));
    cout << h << endl;
    specex::syr(1,ublas::project(h,ublas::range(0,5)),a);
    cout << a << endl;
    a.clear();
    specex::syr(1,ublas::project(h,ublas::range(5,10)),a);
    cout << a << endl;

    exit(12);
  }

    // see http://www.boost.org/doc/libs/1_38_0/libs/numeric/ublas/doc/blas.htm

  harp::matrix_double m(n,n);

  int N=500;
  {
    m.clear();
    clock_t tstart = clock();
    for(int i=0;i<N;i++)
      specex::syr(w,v,m);
    clock_t tstop = clock();
    cout << "#1 n clocks = " << tstop-tstart << " " << float(tstop-tstart)/float(CLOCKS_PER_SEC) << endl;
  }

  //ublas::coordinate_matrix<double> sm (n,n);
  ublas::compressed_matrix<double> sm (n,n);
  //ublas::triangular_matrix<double,  ublas::lower> sm (n,n); // crashed
  //ublas::mapped_matrix<double> sm (n,n);
  //ublas::matrix<double> sm (n,n);
  {

    clock_t  tstart = clock();
    //ublas::outer_prod(sv,sv);
    clock_t tstop = clock();
    cout << "#1 n clocks = " << tstop-tstart << " " << float(tstop-tstart)/float(CLOCKS_PER_SEC) << endl;
  }
  {

    clock_t  tstart = clock();
    //ublas::outer_prod(sv,sv);

    //ublas::triangular_adaptor<ublas::compressed_matrix<double>,ublas::lower> tsm(sm);
    for(int i=0;i<N;i++)
      ublas::noalias(sm) += w*ublas::outer_prod(sv,sv);
    //ublas::noalias(m) += ublas::sparse_prod(sv,sv);

    //ublas::blas_3::srk(m,1,1,v);

    clock_t tstop = clock();
    cout << "#2 n clocks = " << tstop-tstart << " " << float(tstop-tstart)/float(CLOCKS_PER_SEC) << endl;
  }
  harp::matrix_double dm = m - sm;
  for(unsigned i = 0; i < 100;i++) {
    cout << "m dm ("<< i << ") =" << m(i,i) << " " << dm(i,i) << endl;
  }

  if(0) {
    harp::matrix_double m(n,n); m.clear();
    clock_t tstart = clock();
    for(int i=0;i<N;i++)
      specex::syr(w,v2,m);
    clock_t tstop = clock();
    cout << "#3 n clocks = " << tstop-tstart << " " << float(tstop-tstart)/float(CLOCKS_PER_SEC) << endl;
  }

    return 0;

}
  cout << "==================" << endl;
  cout << " LINEAR SYSTEMS   " << endl;
  cout << "==================" << endl;

  typedef boost::ecuyer1988 base_generator_type;
  base_generator_type generator(42u);
  boost::uniform_real < double > dist ( 0.0, 1.0 );
  boost::variate_generator < base_generator_type&, boost::uniform_real < double > > uni ( generator, dist );

  int ndata = 12;
  int nparams = 10;
  vector<harp::vector_double> h;
  for( int i=0; i<ndata; i++) {
    harp::vector_double hi(nparams);
    for(int p=0; p<nparams ;p++) {
      hi(p)=uni(); // stupid number
    }
    //cout << "data " << i << " h=" << hi << endl;
    h.push_back(hi);
  }

  harp::vector_double true_params(nparams);
  for(int p=0; p<nparams ;p++) {
    true_params(p)=p;
  }

  harp::matrix_double A(nparams,nparams);
  harp::vector_double B(nparams);
  A.clear();
  B.clear();

  harp::vector_double data(ndata);
  for( int i=0; i<ndata; i++) {
    data[i]=specex::dot(true_params,h[i]);
    double w=uni()+1; // stupid number
    specex::syr(w,h[i],A);
    specex::axpy(w*data(i),h[i],B);
  }

  //cout << "A=" << A << endl;
  //cout << "B=" << B << endl;

  int status = specex::cholesky_solve(A,B);
  cout << "cholesky_solve status = " << status << endl;
  harp::vector_double diff = B-true_params;
  cout << "parameter residuals = " << diff << endl;
  harp::vector_double res(ndata);
  for( int i=0; i<ndata; i++) {
    res(i) = data(i)-specex::dot(B,h[i]);
  }
  cout << "data residuals = " << res << endl;

  cout << "==================" << endl;
  cout << " PROJECTIONS   " << endl;
  cout << "==================" << endl;

  harp::vector_double one(16);
  harp::vector_double zero(16);
  zero.clear();

  for(size_t i=0;i<one.size();i++) one[i]=i+1;
  for(int j=0;j<4;j++)
    ublas::project(zero,ublas::range(4*j,4*(j+1))) += pow(10,j)* ublas::project(one,ublas::range(4*j,4*(j+1)));
  cout << zero << endl;

  return 0;
  harp::vector_double y(12);
  for(int i=7;i<7+4;i++) y[i]=1;

  try {
    //cout << specex::dot(x,y) << endl; // this throws an exception

    ublas::range r(0,1);
    cout << specex::dot(x,ublas::project(y,ublas::range(0,x.size()))) << endl;
    cout << specex::dot(x,ublas::project(y,ublas::range(7,7+x.size()))) << endl;

    y.clear();
    size_t x_size = x.size();
    for(size_t k=0;(k+1)*x_size<=y.size();k++) {
      ublas::project(y,ublas::range(k*x_size,(k+1)*x_size)) = (k+1)*x;
    }

    cout << y << endl;
    cout << ublas::inner_prod (y,y) << endl;
    cout << specex::dot(y,y) << endl;
    cout << ublas::outer_prod (y,y) << endl;

    cout << "==================" << endl;
    cout << " SPARSE VECTORS   " << endl;
    cout << "==================" << endl;

    //ublas::mapped_vector<double> sv (8000);
    ublas::compressed_vector<double> sv (8000);
    //ublas::coordinate_vector<double> sv (8000);
    harp::vector_double v (8000);

    if(0) {

      cout << sv << endl;
      // for coordinate vector
      //for(ublas::coordinate_vector<double>::iterator it=sv.begin();it!=sv.end();++it) {
      //cout << *it << " ";
      //}
      //cout << endl;
    }

    double res1=0;
    {
      clock_t tstart = clock();
      ublas::vector<double> v2(v.size());
      ublas::matrix<double> m = ublas::outer_prod(v,v); boost::numeric::bindings::blas::gemv(1,m,v,0,v2);
      //boost::numeric::bindings::blas::gemv(1,(ublas::matrix<double>&)ublas::outer_prod(v,v),v,0,v2);
      //cout << v2 << endl;
      clock_t tstop = clock();
      res1=ublas::inner_prod(v,v2);
      cout << res1 << endl;
      cout << "n clocks = " << tstop-tstart << " " << float(tstop-tstart)/float(CLOCKS_PER_SEC) << endl;
    }
    if(1){ // ok
      clock_t tstart = clock();
      ublas::vector<double> v2(sv.size());
      ublas::matrix<double> m = ublas::outer_prod(sv,sv); boost::numeric::bindings::blas::gemv(1,m,v,0,v2);
      double res2 = ublas::inner_prod(sv,v2);
      clock_t tstop = clock();
      cout << res2 << " " << res2/res1-1 << endl;
      cout << "n clocks = " << tstop-tstart << " " << float(tstop-tstart)/float(CLOCKS_PER_SEC) << endl;
    }
    if(1){ // ok
      clock_t tstart = clock();
      double res2 = ublas::inner_prod(sv,ublas::prod(ublas::outer_prod(sv,sv),sv));
      clock_t tstop = clock();
      cout << res2 << " " << res2/res1-1 << endl;
      cout << "n clocks = " << tstop-tstart << " " << float(tstop-tstart)/float(CLOCKS_PER_SEC) << endl;
    }

    // same with "pre-existing" matrix
    {
      ublas::matrix<double, ublas::column_major> m(v.size(),v.size());
      clock_t tstart = clock();
      //m += ublas::outer_prod(sv,sv); // terribly slow
      // ublas::add(m,ublas::outer_prod(sv,sv)); // doesn't compile
      boost::numeric::bindings::blas::syr(1.,v,boost::numeric::bindings::lower(m)); // compiles, and fast, but fills only part of the matrix !!
      //boost::numeric::bindings::blas::syr(1.,sv,boost::numeric::bindings::lower(m)); // doesn't compile

      //specex::syr(1.,v,m);
      double res2 = ublas::inner_prod(sv,ublas::prod(m,sv));
      clock_t tstop = clock();
      cout << res2 << " " << res2/res1-1 << endl;
      cout << "n clocks = " << tstop-tstart << " " << float(tstop-tstart)/float(CLOCKS_PER_SEC) << endl;
    }

  }catch(harp::exception e) {
    cout << "ohoh " << e.what() << endl;
    return EXIT_FAILURE;
  }
  return EXIT_SUCCESS;
}


/*=============================================================================
Copyright (c) 2011, The Trustees of Indiana University
All rights reserved.

Authors: Michael Hansen (mihansen@indiana.edu), Shinya Ito

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

  1. Redistributions of source code must retain the above copyright notice,
     this list of conditions and the following disclaimer.

  2. Redistributions in binary form must reproduce the above copyright notice,
     this list of conditions and the following disclaimer in the documentation
     and/or other materials provided with the distribution.

  3. Neither the name of Indiana University nor the names of its contributors
     may be used to endorse or promote products derived from this software
     without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
=============================================================================*/

#include <bitset>
#include <cmath>
#include <algorithm>
#include <numeric>
#include <cassert>

#include <boost/mpl/arithmetic.hpp>
#include <boost/static_assert.hpp>
#include <boost/mpl/plus.hpp>
#include <boost/limits.hpp>

#define MAX_XY_ORDER 64

namespace mpl = boost::mpl;
namespace boost { namespace mpl {
  template <std::size_t N, std::size_t Power>
  struct pow
  {
    static const std::size_t value = N * pow<N, Power - 1>::value;
  };

  template <std::size_t N>
  struct pow<N, 0>
  {
    static const std::size_t value = 1;
  };

} }

// Computes the higher-order transfer entropy matrix for all pairs.
// x and y orders must be known at compile time.
template <typename TimeSeriesCollection, typename ResultMatrix,
         std::size_t x_order, std::size_t y_order>
void transent_ho
(const TimeSeriesCollection& all_series,
 const typename TimeSeriesCollection::value_type::value_type y_delay,
 const typename TimeSeriesCollection::value_type::value_type duration,
 ResultMatrix& te_result,
 std::size_t row_start = 0, std::size_t rows = 0,
 std::size_t col_start = 0, std::size_t cols = 0) {

  // Typedefs
  typedef typename TimeSeriesCollection::value_type TimeSeries;
  typedef typename TimeSeries::value_type TimeType;
  typedef typename TimeSeries::const_iterator TimeSeriesIter;

  typedef std::pair<TimeSeriesIter, TimeType> IterShiftPair;
  typedef std::pair<TimeType, std::size_t> TimeIndexPair;

  // Constants
  const std::size_t num_series = 1 + y_order + x_order,
                    num_counts = mpl::pow<2, num_series>::value,
                    num_x = mpl::pow<2, mpl::plus<mpl::int_<x_order>, mpl::int_<1> >::value>::value,
                    num_y = mpl::pow<2, y_order>::value;

  BOOST_STATIC_ASSERT(x_order > 0);
  BOOST_STATIC_ASSERT(y_order > 0);
  assert(y_delay > 0);
  BOOST_STATIC_ASSERT(num_series <= MAX_XY_ORDER);

  if (rows == 0) {
    rows = all_series.size();
  }

  if (cols == 0) {
    cols = all_series.size();
  }

  // Locals
  std::vector<TimeType> counts(num_counts);
  std::bitset<MAX_XY_ORDER> code;
  std::size_t idx = 0;
  double te_final, prob_2, prob_3;

  IterShiftPair ord_iter[num_series];
  TimeType ord_times[num_series];
  TimeSeriesIter ord_end[num_series];

  const std::size_t window = std::max(y_order + y_delay, x_order + 1);
  TimeType cur_time, next_time, end_time = duration - window + 1, shift;

  // Calculate TE
  for (std::size_t i = row_start; i < (rows + row_start); ++i) {
    for (std::size_t j = col_start; j < (cols + col_start); ++j) {

      // NOTE: Time series are assumed to be 1-based, so everything is shifted by 1 too.
      // Order is x^(k+1), y^(l)
      idx = 0;

      // x^(k+1)
      for (std::size_t k = 0; k < (x_order + 1); ++k) {
        shift = (window - 1) - k;
        ord_end[idx] = all_series[i].end();
        ord_iter[idx] = std::make_pair(std::lower_bound(all_series[i].begin(), ord_end[idx], shift + 1), shift);
        ord_times[idx] = *(ord_iter[idx].first) - ord_iter[idx].second;
        ++idx;
      }

      // y^(l)
      for (std::size_t k = 0; k < y_order; ++k) {
        shift = (window - 1) - y_delay - k;
        ord_end[idx] = all_series[j].end();
        ord_iter[idx] = std::make_pair(std::lower_bound(all_series[j].begin(), ord_end[idx], shift + 1), shift);
        ord_times[idx] = *(ord_iter[idx].first) - ord_iter[idx].second;
        ++idx;
      }

      // Count spikes
      std::fill(counts.begin(), counts.end(), 0);
      cur_time = *(std::min_element(ord_times, ord_times + num_series));

      while (cur_time <= end_time) {

        code.reset();
        next_time = std::numeric_limits<TimeType>::max();

        // Calculate hash code for this time
        for (std::size_t k = 0; k < num_series; ++k) {
          if (ord_times[k] == cur_time) {
            code[k] = 1;

            // Next spike
            ++(ord_iter[k].first);

            if (ord_iter[k].first == ord_end[k]) {
              ord_times[k] = std::numeric_limits<TimeType>::max();
            }
            else {
              ord_times[k] = *(ord_iter[k].first) - ord_iter[k].second;
            }
          }

          if (ord_times[k] < next_time) {
            next_time = ord_times[k];
          }
        }

        ++(counts[code.to_ulong()]);
        cur_time = next_time;

      } // while spikes left

      // Fill in zero count
      counts[0] = end_time - std::accumulate(counts.begin() + 1, counts.end(), 0);

      // =====================================================================

      // Use counts to calculate TE
      te_final = 0;

      // Order is x^(k), y^(l), x(n+1)
      for (std::size_t k = 0; k < num_counts; ++k) {
        if (counts[k] == 0) {
          continue;
        }

        prob_2 = (double)counts[k] / (double)(counts[k] + counts[k ^ 1]);

        std::size_t c1 = 0, c2 = 0;
        for (std::size_t l = 0; l < num_y; ++l) {
          idx = (k & (num_x - 1)) + (l << (x_order + 1));
          c1 += counts[idx];
          c2 += (counts[idx] + counts[idx ^ 1]);
        }

        prob_3 = (double)c1 / (double)c2;

        te_final += ((double)counts[k] * (log2(prob_2) - log2(prob_3)));
      }

      te_result[i - row_start][j - col_start] = te_final / (double)end_time;

    } // for j

  } // for i

} //transent_ho

// Computes the 1st order transfer entropy matrix for all pairs.
template <typename TimeSeriesCollection, typename ResultMatrix>
void transent_1
(const TimeSeriesCollection& all_series,
 const typename TimeSeriesCollection::value_type::value_type y_delay,
 const typename TimeSeriesCollection::value_type::value_type duration,
 ResultMatrix& te_result,
 std::size_t row_start = 0, std::size_t rows = 0,
 std::size_t col_start = 0, std::size_t cols = 0) {

  return (transent_ho<TimeSeriesCollection, ResultMatrix, 1, 1>
          (all_series, y_delay, duration, te_result,
           row_start, rows, col_start, cols));

} // transent_1

// Computes the higher-order transfer entropy matrix for all pairs.
template <typename TimeSeriesCollection, typename ResultMatrix>
void transent_ho
(const TimeSeriesCollection& all_series,
 const typename std::size_t x_order, std::size_t y_order,
 const typename TimeSeriesCollection::value_type::value_type y_delay,
 const typename TimeSeriesCollection::value_type::value_type duration,
 ResultMatrix& te_result,
 std::size_t row_start = 0, std::size_t rows = 0,
 std::size_t col_start = 0, std::size_t cols = 0) {

  // Typedefs
  typedef typename TimeSeriesCollection::value_type TimeSeries;
  typedef typename TimeSeries::value_type TimeType;
  typedef typename TimeSeries::const_iterator TimeSeriesIter;

  typedef std::pair<TimeSeriesIter, TimeType> IterShiftPair;
  typedef std::pair<TimeType, std::size_t> TimeIndexPair;

  // Constants
  const std::size_t num_series = 1 + y_order + x_order,
                    num_counts = (std::size_t)pow(2, num_series),
                    num_x = (std::size_t)pow(2, x_order + 1),
                    num_y = (std::size_t)pow(2, y_order);

  assert(x_order > 0);
  assert(y_order > 0);
  assert(y_delay > 0);
  assert(num_series <= MAX_XY_ORDER);

  if (rows == 0) {
    rows = all_series.size();
  }

  if (cols == 0) {
    cols = all_series.size();
  }

  // Locals
  std::vector<TimeType> counts(num_counts);
  std::bitset<MAX_XY_ORDER> code;
  std::size_t idx = 0;
  double te_final, prob_2, prob_3;

  IterShiftPair ord_iter[num_series];
  TimeType ord_times[num_series];
  TimeSeriesIter ord_end[num_series];

  const std::size_t window = std::max(y_order + y_delay, x_order + 1);
  TimeType cur_time, next_time, end_time = duration - window + 1, shift;

  // Calculate TE
  for (std::size_t i = row_start; i < (rows + row_start); ++i) {
    for (std::size_t j = col_start; j < (cols + col_start); ++j) {

      // NOTE: Time series are assumed to be 1-based, so everything is shifted by 1 too.
      // Order is x^(k+1), y^(l)
      idx = 0;

      // x^(k+1)
      for (std::size_t k = 0; k < (x_order + 1); ++k) {
        shift = (window - 1) - k;
        ord_end[idx] = all_series[i].end();
        ord_iter[idx] = std::make_pair(std::lower_bound(all_series[i].begin(), ord_end[idx], shift + 1), shift);
        ord_times[idx] = *(ord_iter[idx].first) - ord_iter[idx].second;
        ++idx;
      }

      // y^(l)
      for (std::size_t k = 0; k < y_order; ++k) {
        ord_end[idx] = all_series[j].end();
        ord_iter[idx] = std::make_pair(all_series[j].begin(), -k);
        ord_times[idx] = *(ord_iter[idx].first) - ord_iter[idx].second;
        ++idx;
      }

      // Count spikes
      std::fill(counts.begin(), counts.end(), 0);
      cur_time = *(std::min_element(ord_times, ord_times + num_series));

      while (cur_time <= end_time) {

        code.reset();
        next_time = std::numeric_limits<TimeType>::max();

        // Calculate hash code for this time
        for (std::size_t k = 0; k < num_series; ++k) {
          if (ord_times[k] == cur_time) {
            code[k] = 1;

            // Next spike
            ++(ord_iter[k].first);

            if (ord_iter[k].first == ord_end[k]) {
              ord_times[k] = std::numeric_limits<TimeType>::max();
            }
            else {
              ord_times[k] = *(ord_iter[k].first) - ord_iter[k].second;
            }
          }

          if (ord_times[k] < next_time) {
            next_time = ord_times[k];
          }
        }

        ++(counts[code.to_ulong()]);
        cur_time = next_time;

      } // while spikes left

      // Fill in zero count
      counts[0] = end_time - std::accumulate(counts.begin() + 1, counts.end(), 0);

      // =====================================================================

      // Use counts to calculate TE
      te_final = 0;

      // Order is x^(k), y^(l), x(n+1)
      for (std::size_t k = 0; k < num_counts; ++k) {
        if (counts[k] == 0) {
          continue;
        }

        prob_2 = (double)counts[k] / (double)(counts[k] + counts[k ^ 1]);

        std::size_t c1 = 0, c2 = 0;
        for (std::size_t l = 0; l < num_y; ++l) {
          idx = (k & (num_x - 1)) + (l << (x_order + 1));
          c1 += counts[idx];
          c2 += (counts[idx] + counts[idx ^ 1]);
        }

        prob_3 = (double)c1 / (double)c2;

        te_final += ((double)counts[k] * (log2(prob_2) - log2(prob_3)));
      }

      te_result[i - row_start][j - col_start] = te_final / (double)end_time;

    } // for j

  } // for i

} // transent_ho