/*
 Copyright (C) 2016 Quaternion Risk Management Ltd
 All rights reserved.

 This file is part of ORE, a free-software/open-source library
 for transparent pricing and risk analysis - http://opensourcerisk.org

 ORE is free software: you can redistribute it and/or modify it
 under the terms of the Modified BSD License.  You should have received a
 copy of the license along with this program.
 The license is also available online at <http://opensourcerisk.org>

 This program is distributed on the basis that it will form a useful
 contribution to risk analytics and model standardisation, but WITHOUT
 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 FITNESS FOR A PARTICULAR PURPOSE. See the license for more details.
*/

/*! \file models/crossassetmodel.hpp
    \brief cross asset model
    \ingroup crossassetmodel
*/

#ifndef quantext_crossasset_model_hpp
#define quantext_crossasset_model_hpp

#include <qle/models/cirppparametrization.hpp>
#include <qle/models/crcirpp.hpp>
#include <qle/models/crlgm1fparametrization.hpp>
#include <qle/models/eqbsparametrization.hpp>
#include <qle/models/fxbsparametrization.hpp>
#include <qle/models/infdkparametrization.hpp>
#include <qle/models/infjyparameterization.hpp>
#include <qle/models/lgm.hpp>

#include <qle/processes/crossassetstateprocess.hpp>

#include <ql/math/distributions/normaldistribution.hpp>
#include <ql/math/integrals/integral.hpp>
#include <ql/math/matrix.hpp>
#include <ql/models/model.hpp>

#include <boost/bind.hpp>

namespace QuantExt {
using namespace QuantLib;

namespace CrossAssetModelTypes {
//! Cross Asset Type
//! \ingroup crossassetmodel
enum AssetType { IR, FX, INF, CR, EQ, AUX };
static constexpr Size crossAssetModelAssetTypes = 6;

std::ostream& operator<<(std::ostream& out, const AssetType& type);

/*! the model types supported by the CrossAssetModel or derived classes;
  a model type may applicable to several asset types (like BS for FX, EQ) */
enum ModelType { LGM1F, BS, DK, CIRPP, JY };

} // namespace CrossAssetModelTypes

//! Choice of measure
struct Measure {
    enum Type { LGM, BA };
};

using namespace CrossAssetModelTypes;

//! Cross Asset Model
/*! \ingroup crossassetmodel
 */
class CrossAssetModel : public LinkableCalibratedModel {
public:
    /*! Parametrizations must be given in the following order
        - IR  (first parametrization defines the domestic currency)
        - FX  (for all pairs domestic-ccy defined by the IR models)
        - INF (optionally, ccy must be a subset of the IR ccys)
        - CR  (optionally, ccy must be a subset of the IR ccys)
        - EQ  (for all names equity currency defined in Parametrization)
        If the correlation matrix is not given, it is initialized
        as the unit matrix (and can be customized after
        construction of the model).
    */
    CrossAssetModel(const std::vector<boost::shared_ptr<Parametrization>>& parametrizations,
                    const Matrix& correlation = Matrix(), SalvagingAlgorithm::Type salvaging = SalvagingAlgorithm::None,
                    Measure::Type measure = Measure::LGM);

    /*! IR-FX model based constructor */
    CrossAssetModel(const std::vector<boost::shared_ptr<LinearGaussMarkovModel>>& currencyModels,
                    const std::vector<boost::shared_ptr<FxBsParametrization>>& fxParametrizations,
                    const Matrix& correlation = Matrix(), SalvagingAlgorithm::Type salvaging = SalvagingAlgorithm::None,
                    Measure::Type measure = Measure::LGM);

    /*! returns the state process with a given discretization */
    const boost::shared_ptr<StochasticProcess>
    stateProcess(CrossAssetStateProcess::discretization disc = CrossAssetStateProcess::exact) const;

    /*! total dimension of model (sum of number of state variables) */
    Size dimension() const;

    /*! total number of Brownian motions (this is less or equal to dimension) */
    Size brownians() const;

    /*! total number of parameters that can be calibrated */
    Size totalNumberOfParameters() const;

    /*! number of components for an asset class */
    Size components(const AssetType t) const;

    /*! number of brownian motions for a component */
    Size brownians(const AssetType t, const Size i) const;

    /*! number of state variables for a component */
    Size stateVariables(const AssetType t, const Size i) const;

    /*! model type of a component */
    ModelType modelType(const AssetType t, const Size i) const;

    /*! Choice of probability measure */
    Measure::Type measure() const { return measure_; }

    /*! return index for currency (0 = domestic, 1 = first
      foreign currency and so on) */
    Size ccyIndex(const Currency& ccy) const;

    /*! return index for equity (0 = first equity) */
    Size eqIndex(const std::string& eqName) const;

    /*! return index for inflation (0 = first inflation index) */
    Size infIndex(const std::string& index) const;

    /*! return index for credit (0 = first credit name) */
    Size crName(const std::string& name) const;

    /*! observer and linked calibrated model interface */
    void update();
    void generateArguments();

    /*! the vector of parametrizations */
    const std::vector<boost::shared_ptr<Parametrization>>& parametrizations() const { return p_; }

    /*! components per asset class, see below for specific model type inspectors */
    const boost::shared_ptr<Parametrization> ir(const Size ccy) const;
    const boost::shared_ptr<Parametrization> fx(const Size ccy) const;
    const boost::shared_ptr<Parametrization> inf(const Size i) const;
    const boost::shared_ptr<Parametrization> cr(const Size i) const;
    const boost::shared_ptr<Parametrization> eq(const Size i) const;

    /*! LGM1F components, ccy=0 refers to the domestic currency */
    const boost::shared_ptr<LinearGaussMarkovModel> lgm(const Size ccy) const;

    const boost::shared_ptr<IrLgm1fParametrization> irlgm1f(const Size ccy) const;

    /*! LGM measure numeraire */
    Real numeraire(const Size ccy, const Time t, const Real x,
                   Handle<YieldTermStructure> discountCurve = Handle<YieldTermStructure>()) const;

    /*! Bank account measure numeraire B(t) as a function of drifted LGM state variable x and drift-free auxiliary state
     * variable y */
    Real bankAccountNumeraire(const Size ccy, const Time t, const Real x, const Real y,
                              Handle<YieldTermStructure> discountCurve = Handle<YieldTermStructure>()) const;

    Real discountBond(const Size ccy, const Time t, const Time T, const Real x,
                      Handle<YieldTermStructure> discountCurve = Handle<YieldTermStructure>()) const;

    Real reducedDiscountBond(const Size ccy, const Time t, const Time T, const Real x,
                             Handle<YieldTermStructure> discountCurve = Handle<YieldTermStructure>()) const;

    Real discountBondOption(const Size ccy, Option::Type type, const Real K, const Time t, const Time S, const Time T,
                            Handle<YieldTermStructure> discountCurve = Handle<YieldTermStructure>()) const;

    /*! FXBS components, ccy=0 referes to the first foreign currency,
        so it corresponds to ccy+1 if you want to get the corresponding
        irmgl1f component */
    const boost::shared_ptr<FxBsParametrization> fxbs(const Size ccy) const;

    /*! INF DK components */
    const boost::shared_ptr<InfDkParametrization> infdk(const Size i) const;

    //! Inflation JY component
    const boost::shared_ptr<InfJyParameterization> infjy(const Size i) const;

    /*! CR LGM 1F components */
    const boost::shared_ptr<CrLgm1fParametrization> crlgm1f(const Size i) const;

    /*! CR CIR++ components */
    const boost::shared_ptr<CrCirpp> crcirppModel(const Size i) const;
    const boost::shared_ptr<CrCirppParametrization> crcirpp(const Size i) const;

    /*! EQBS components */
    const boost::shared_ptr<EqBsParametrization> eqbs(const Size ccy) const;

    /* ... add more components here ...*/

    /*! correlation linking the different marginal models, note that
        the use of asset class pairs specific inspectors is
        recommended instead of the global matrix directly */
    const Matrix& correlation() const;

    /*! check if correlation matrix is valid */
    void checkCorrelationMatrix() const;

    /*! index of component in the parametrization vector */
    Size idx(const AssetType t, const Size i) const;

    /*! index of component in the correlation matrix, by offset */
    Size cIdx(const AssetType t, const Size i, const Size offset = 0) const;

    /*! index of component in the stochastic process array, by offset */
    Size pIdx(const AssetType t, const Size i, const Size offset = 0) const;

    /*! correlation between two components */
    const Real& correlation(const AssetType s, const Size i, const AssetType t, const Size j, const Size iOffset = 0,
                            const Size jOffset = 0) const;
    /*! set correlation */
    void correlation(const AssetType s, const Size i, const AssetType t, const Size j, const Real value,
                     const Size iOffset = 0, const Size jOffset = 0);

    /*! get salvaging algorithm */
    SalvagingAlgorithm::Type salvagingAlgorithm() const { return salvaging_; }

    /*! analytical moments require numerical integration,
      which can be customized here */
    void setIntegrationPolicy(const boost::shared_ptr<Integrator> integrator,
                              const bool usePiecewiseIntegration = true) const;
    const boost::shared_ptr<Integrator> integrator() const;

    /*! return (V(t), V^tilde(t,T)) in the notation of the book */
    std::pair<Real, Real> infdkV(const Size i, const Time t, const Time T);

    /*! return (I(t), I^tilde(t,T)) in the notation of the book, note that
        I(0) is normalized to 1 here, i.e. you have to multiply the result
        with the index value (as of the base date of the inflation ts) */
    std::pair<Real, Real> infdkI(const Size i, const Time t, const Time T, const Real z, const Real y);

    /*! return YoYIIS(t) in the notation of the book, the year on year
        swaplet price from S to T, at time t */
    Real infdkYY(const Size i, const Time t, const Time S, const Time T, const Real z, const Real y, const Real irz);

    /*! returns (S(t), S^tilde(t,T)) in the notation of the book */
    std::pair<Real, Real> crlgm1fS(const Size i, const Size ccy, const Time t, const Time T, const Real z,
                                   const Real y) const;

    /*! returns (S(t), S^tilde(t,T)) in the notation of the book */
    std::pair<Real, Real> crcirppS(const Size i, const Time t, const Time T, const Real y, const Real s) const;

    /*! tentative: more generic interface that is agnostic of the model type - so far only for CR */
    virtual Handle<DefaultProbabilityTermStructure> crTs(const Size i) const;
    virtual std::pair<Real, Real> crS(const Size i, const Size ccy, const Time t, const Time T, const Real z,
                                      const Real y) const;

    /*! calibration procedures */

    /*! calibrate irlgm1f volatilities to a sequence of ir options with
        expiry times equal to step times in the parametrization */
    void calibrateIrLgm1fVolatilitiesIterative(const Size ccy,
                                               const std::vector<boost::shared_ptr<BlackCalibrationHelper>>& helpers,
                                               OptimizationMethod& method, const EndCriteria& endCriteria,
                                               const Constraint& constraint = Constraint(),
                                               const std::vector<Real>& weights = std::vector<Real>());

    /*! calibrate irlgm1f reversion to a sequence of ir options with
        maturities equal to step times in the parametrization */
    void calibrateIrLgm1fReversionsIterative(const Size ccy,
                                             const std::vector<boost::shared_ptr<BlackCalibrationHelper>>& helpers,
                                             OptimizationMethod& method, const EndCriteria& endCriteria,
                                             const Constraint& constraint = Constraint(),
                                             const std::vector<Real>& weights = std::vector<Real>());

    /*! calibrate irlgm1f parameters for one ccy globally to a set
        of ir options */
    void calibrateIrLgm1fGlobal(const Size ccy, const std::vector<boost::shared_ptr<BlackCalibrationHelper>>& helpers,
                                OptimizationMethod& method, const EndCriteria& endCriteria,
                                const Constraint& constraint = Constraint(),
                                const std::vector<Real>& weights = std::vector<Real>());

    /*! calibrate eq or fx volatilities to a sequence of options with
            expiry times equal to step times in the parametrization */
    void calibrateBsVolatilitiesIterative(const AssetType& assetType, const Size aIdx,
                                          const std::vector<boost::shared_ptr<BlackCalibrationHelper>>& helpers,
                                          OptimizationMethod& method, const EndCriteria& endCriteria,
                                          const Constraint& constraint = Constraint(),
                                          const std::vector<Real>& weights = std::vector<Real>());

    /*! calibrate eq/fx volatilities globally to a set of fx options */
    void calibrateBsVolatilitiesGlobal(const AssetType& assetType, const Size aIdx,
                                       const std::vector<boost::shared_ptr<BlackCalibrationHelper>>& helpers,
                                       OptimizationMethod& method, const EndCriteria& endCriteria,
                                       const Constraint& constraint = Constraint(),
                                       const std::vector<Real>& weights = std::vector<Real>());

    /*! calibrate infdk volatilities to a sequence of cpi options with
        expiry times equal to step times in the parametrization */
    void calibrateInfDkVolatilitiesIterative(const Size index,
                                             const std::vector<boost::shared_ptr<BlackCalibrationHelper>>& helpers,
                                             OptimizationMethod& method, const EndCriteria& endCriteria,
                                             const Constraint& constraint = Constraint(),
                                             const std::vector<Real>& weights = std::vector<Real>());

    /*! calibrate infdk reversions to a sequence of cpi options with
        maturity times equal to step times in the parametrization */
    void calibrateInfDkReversionsIterative(const Size index,
                                           const std::vector<boost::shared_ptr<BlackCalibrationHelper>>& helpers,
                                           OptimizationMethod& method, const EndCriteria& endCriteria,
                                           const Constraint& constraint = Constraint(),
                                           const std::vector<Real>& weights = std::vector<Real>());

    /*! calibrate infdk volatilities globally to a sequence of cpi cap/floors */
    void calibrateInfDkVolatilitiesGlobal(const Size index,
                                          const std::vector<boost::shared_ptr<BlackCalibrationHelper>>& helpers,
                                          OptimizationMethod& method, const EndCriteria& endCriteria,
                                          const Constraint& constraint = Constraint(),
                                          const std::vector<Real>& weights = std::vector<Real>());

    /*! calibrate infdk reversions globally to a sequence of cpi cap/floors */
    void calibrateInfDkReversionsGlobal(const Size index,
                                        const std::vector<boost::shared_ptr<BlackCalibrationHelper>>& helpers,
                                        OptimizationMethod& method, const EndCriteria& endCriteria,
                                        const Constraint& constraint = Constraint(),
                                        const std::vector<Real>& weights = std::vector<Real>());

    /*! Calibrate JY inflation parameters globally.

        The parameter \p toCalibrate indicates which parameters of the JY inflation model that we want to calibrate.
        The map key should be in {0, 1, 2} where 0 indicates the real rate volatility, 1 indicates the real rate
        reversion and 2 indicates the inflation index volatility. The value is \c true if we wish to calibrate
        the parameter and \p false if we do not want to calibrate it.
    */
    void calibrateInfJyGlobal(QuantLib::Size index,
                              const std::vector<boost::shared_ptr<QuantLib::CalibrationHelper>>& helpers,
                              QuantLib::OptimizationMethod& method, const QuantLib::EndCriteria& endCriteria,
                              const std::map<QuantLib::Size, bool>& toCalibrate,
                              const QuantLib::Constraint& constraint = QuantLib::Constraint(),
                              const std::vector<QuantLib::Real>& weights = std::vector<QuantLib::Real>());

    /*! Calibrate a single JY inflation parameter iteratively.

        Calibrate one of real rate volatility, real rate reversion or inflation index volatility. The
        \p parameterIndex indicates the parameter that should be calibrated where 0 indicates the real rate
        volatility, 1 indicates the real rate reversion and 2 indicates the inflation index volatility.
    */
    void calibrateInfJyIterative(QuantLib::Size inflationModelIndex, QuantLib::Size parameterIndex,
                                 const std::vector<boost::shared_ptr<QuantLib::CalibrationHelper>>& helpers,
                                 QuantLib::OptimizationMethod& method, const QuantLib::EndCriteria& endCriteria,
                                 const QuantLib::Constraint& constraint = QuantLib::Constraint(),
                                 const std::vector<QuantLib::Real>& weights = std::vector<QuantLib::Real>());

    /*! calibrate crlgm1f volatilities to a sequence of cds options with
        expiry times equal to step times in the parametrization */
    void calibrateCrLgm1fVolatilitiesIterative(const Size index,
                                               const std::vector<boost::shared_ptr<BlackCalibrationHelper>>& helpers,
                                               OptimizationMethod& method, const EndCriteria& endCriteria,
                                               const Constraint& constraint = Constraint(),
                                               const std::vector<Real>& weights = std::vector<Real>());

    /*! calibrate crlgm1f reversions to a sequence of cds options with
        maturity times equal to step times in the parametrization */
    void calibrateCrLgm1fReversionsIterative(const Size index,
                                             const std::vector<boost::shared_ptr<BlackCalibrationHelper>>& helpers,
                                             OptimizationMethod& method, const EndCriteria& endCriteria,
                                             const Constraint& constraint = Constraint(),
                                             const std::vector<Real>& weights = std::vector<Real>());

    /* ... add more calibration procedures here ... */

protected:
    /* ctor to be used in extensions, initialize is not called */
    CrossAssetModel(const std::vector<boost::shared_ptr<Parametrization>>& parametrizations, const Matrix& correlation,
                    SalvagingAlgorithm::Type salvaging, Measure::Type measure, const bool)
        : LinkableCalibratedModel(), p_(parametrizations), rho_(correlation), salvaging_(salvaging), measure_(measure) {
    }

    /*! number of arguments for a component */
    Size arguments(const AssetType t, const Size i) const;

    /*! index of component in the arguments vector, by offset */
    Size aIdx(const AssetType t, const Size i, const Size offset = 0) const;

    /*! asset and model type for given parametrization */
    virtual std::pair<AssetType, ModelType> getComponentType(const Size i) const;
    /*! number of parameters for given parametrization */
    virtual Size getNumberOfParameters(const Size i) const;
    /*! number of brownians for given parametrization */
    virtual Size getNumberOfBrownians(const Size i) const;
    /*! number of state variables for given parametrization */
    virtual Size getNumberOfStateVariables(const Size i) const;
    /*! helper function to init component indices */
    void updateIndices(const AssetType& t, const Size i, const Size cIdx, const Size pIdx, const Size aIdx);

    /* init methods */
    virtual void initialize();
    virtual void initializeParametrizations();
    virtual void initializeCorrelation();
    virtual void initializeArguments();
    virtual void finalizeArguments();
    virtual void checkModelConsistency() const;
    virtual void initDefaultIntegrator();
    virtual void initStateProcess();

    /* helper function for infdkI, crlgm1fS */
    Real infV(const Size idx, const Size ccy, const Time t, const Time T) const;
    Real crV(const Size idx, const Size ccy, const Time t, const Time T) const;

    // cache for infdkI, crlgm1fS method
    struct cache_key {
        Size i, ccy;
        double t, T;
        bool operator==(const cache_key& o) const { return (i == o.i) && (ccy == o.ccy) && (t == o.t) && (T == o.T); }
    };

    struct cache_hasher : std::unary_function<cache_key, std::size_t> {
        std::size_t operator()(cache_key const& x) const {
            std::size_t seed = 0;
            boost::hash_combine(seed, x.i);
            boost::hash_combine(seed, x.ccy);
            boost::hash_combine(seed, x.t);
            boost::hash_combine(seed, x.T);
            return seed;
        }
    };

    mutable boost::unordered_map<cache_key, std::pair<Real, Real>, cache_hasher> cache_crlgm1fS_, cache_infdkI_;

    /* members */

    // components per asset type
    std::vector<Size> components_;
    // indices per asset type and component number within asset type
    std::vector<std::vector<Size>> idx_, cIdx_, pIdx_, aIdx_, brownians_, stateVariables_, numArguments_;
    // counter
    Size totalDimension_, totalNumberOfBrownians_, totalNumberOfParameters_;
    // model type per asset type and component number within asset type
    std::vector<std::vector<ModelType>> modelType_;
    // parametrizations, models
    std::vector<boost::shared_ptr<Parametrization>> p_;
    std::vector<boost::shared_ptr<LinearGaussMarkovModel>> lgm_;
    std::vector<boost::shared_ptr<CrCirpp>> crcirppModel_;
    Matrix rho_;
    SalvagingAlgorithm::Type salvaging_;
    Measure::Type measure_;
    mutable boost::shared_ptr<Integrator> integrator_;
    boost::shared_ptr<CrossAssetStateProcess> stateProcessExact_, stateProcessEuler_;

    /* calibration constraints */

    void appendToFixedParameterVector(const AssetType t, const AssetType v, const Size param, const Size index,
                                      const Size i, std::vector<bool>& res) {
        for (Size j = 0; j < components(t); ++j) {
            for (Size k = 0; k < arguments(t, j); ++k) {
                std::vector<bool> tmp1(p_[idx(t, j)]->parameter(k)->size(), true);
                if ((param == Null<Size>() || k == param) && t == v && index == j) {
                    for (Size ii = 0; ii < tmp1.size(); ++ii) {
                        if (i == Null<Size>() || i == ii) {
                            tmp1[ii] = false;
                        }
                    }
                }
                res.insert(res.end(), tmp1.begin(), tmp1.end());
            }
        }
    }

    // move parameter param (e.g. vol, reversion, or all if null) of asset type component t / index at step i (or at all
    // steps if i is null)
    Disposable<std::vector<bool>> MoveParameter(const AssetType t, const Size param, const Size index, const Size i) {
        QL_REQUIRE(param == Null<Size>() || param < arguments(t, index),
                   "parameter for " << t << " at " << index << " (" << param << ") out of bounds 0..."
                                    << arguments(t, index) - 1);
        std::vector<bool> res(0);
        appendToFixedParameterVector(IR, t, param, index, i, res);
        appendToFixedParameterVector(FX, t, param, index, i, res);
        appendToFixedParameterVector(INF, t, param, index, i, res);
        appendToFixedParameterVector(CR, t, param, index, i, res);
        appendToFixedParameterVector(EQ, t, param, index, i, res);
        if (measure_ == Measure::BA)
            appendToFixedParameterVector(AUX, t, param, index, i, res);
        return res;
    }
};

//! Utility function to return a handle to the inflation term structure given the inflation index.
QuantLib::Handle<QuantLib::ZeroInflationTermStructure>
inflationTermStructure(const boost::shared_ptr<CrossAssetModel>& model, QuantLib::Size index);

// inline

inline const boost::shared_ptr<StochasticProcess>
CrossAssetModel::stateProcess(CrossAssetStateProcess::discretization disc) const {
    return disc == CrossAssetStateProcess::exact ? stateProcessExact_ : stateProcessEuler_;
}

inline Size CrossAssetModel::dimension() const { return totalDimension_; }

inline Size CrossAssetModel::brownians() const { return totalNumberOfBrownians_; }

inline Size CrossAssetModel::totalNumberOfParameters() const { return totalNumberOfParameters_; }

inline const boost::shared_ptr<Parametrization> CrossAssetModel::ir(const Size ccy) const { return p_[idx(IR, ccy)]; }

inline const boost::shared_ptr<Parametrization> CrossAssetModel::fx(const Size ccy) const { return p_[idx(FX, ccy)]; }

inline const boost::shared_ptr<Parametrization> CrossAssetModel::inf(const Size i) const { return p_[idx(INF, i)]; }

inline const boost::shared_ptr<Parametrization> CrossAssetModel::cr(const Size i) const { return p_[idx(CR, i)]; }

inline const boost::shared_ptr<Parametrization> CrossAssetModel::eq(const Size i) const {
    return boost::static_pointer_cast<Parametrization>(p_[idx(EQ, i)]);
}

inline const boost::shared_ptr<LinearGaussMarkovModel> CrossAssetModel::lgm(const Size ccy) const {
    boost::shared_ptr<LinearGaussMarkovModel> tmp = lgm_[idx(IR, ccy)];
    QL_REQUIRE(tmp, "model at " << ccy << " is not IR-LGM1F");
    return tmp;
}

inline const boost::shared_ptr<IrLgm1fParametrization> CrossAssetModel::irlgm1f(const Size ccy) const {
    return lgm(ccy)->parametrization();
}

inline const boost::shared_ptr<InfDkParametrization> CrossAssetModel::infdk(const Size i) const {
    boost::shared_ptr<InfDkParametrization> tmp = boost::dynamic_pointer_cast<InfDkParametrization>(p_[idx(INF, i)]);
    QL_REQUIRE(tmp, "model at " << i << " is not INF-DK");
    return tmp;
}

inline const boost::shared_ptr<InfJyParameterization> CrossAssetModel::infjy(const Size i) const {
    auto tmp = boost::dynamic_pointer_cast<InfJyParameterization>(p_[idx(INF, i)]);
    QL_REQUIRE(tmp, "model at " << i << " is not INF-JY");
    return tmp;
}

inline const boost::shared_ptr<CrLgm1fParametrization> CrossAssetModel::crlgm1f(const Size i) const {
    boost::shared_ptr<CrLgm1fParametrization> tmp = boost::dynamic_pointer_cast<CrLgm1fParametrization>(p_[idx(CR, i)]);
    QL_REQUIRE(tmp, "model at " << i << " is not CR-LGM");
    return tmp;
}

inline const boost::shared_ptr<CrCirpp> CrossAssetModel::crcirppModel(const Size i) const {
    boost::shared_ptr<CrCirpp> tmp = crcirppModel_[i];
    QL_REQUIRE(tmp, "model at " << i << " is not CR-CIRPP");
    return tmp;
}

inline const boost::shared_ptr<CrCirppParametrization> CrossAssetModel::crcirpp(const Size i) const {
    boost::shared_ptr<CrCirppParametrization> tmp = boost::dynamic_pointer_cast<CrCirppParametrization>(p_[idx(CR, i)]);
    QL_REQUIRE(tmp, "model at " << i << " is not CR-CIRPP");
    return tmp;
}

inline const boost::shared_ptr<EqBsParametrization> CrossAssetModel::eqbs(const Size name) const {
    boost::shared_ptr<EqBsParametrization> tmp = boost::dynamic_pointer_cast<EqBsParametrization>(p_[idx(EQ, name)]);
    QL_REQUIRE(tmp, "model at " << name << " is not EQ-BS");
    return tmp;
}

inline Real CrossAssetModel::numeraire(const Size ccy, const Time t, const Real x,
                                       Handle<YieldTermStructure> discountCurve) const {
    return lgm(ccy)->numeraire(t, x, discountCurve);
}

inline Real CrossAssetModel::bankAccountNumeraire(const Size ccy, const Time t, const Real x, const Real y,
                                                  Handle<YieldTermStructure> discountCurve) const {
    return lgm(ccy)->bankAccountNumeraire(t, x, y, discountCurve);
}

inline Real CrossAssetModel::discountBond(const Size ccy, const Time t, const Time T, const Real x,
                                          Handle<YieldTermStructure> discountCurve) const {
    return lgm(ccy)->discountBond(t, T, x, discountCurve);
}

inline Real CrossAssetModel::reducedDiscountBond(const Size ccy, const Time t, const Time T, const Real x,
                                                 Handle<YieldTermStructure> discountCurve) const {
    return lgm(ccy)->reducedDiscountBond(t, T, x, discountCurve);
}

inline Real CrossAssetModel::discountBondOption(const Size ccy, Option::Type type, const Real K, const Time t,
                                                const Time S, const Time T,
                                                Handle<YieldTermStructure> discountCurve) const {
    return lgm(ccy)->discountBondOption(type, K, t, S, T, discountCurve);
}

inline const boost::shared_ptr<FxBsParametrization> CrossAssetModel::fxbs(const Size ccy) const {
    boost::shared_ptr<FxBsParametrization> tmp = boost::dynamic_pointer_cast<FxBsParametrization>(p_[idx(FX, ccy)]);
    QL_REQUIRE(tmp, "model at " << ccy << " is not FX-BS");
    return tmp;
}

inline const Matrix& CrossAssetModel::correlation() const { return rho_; }

inline const boost::shared_ptr<Integrator> CrossAssetModel::integrator() const { return integrator_; }

inline Handle<DefaultProbabilityTermStructure> CrossAssetModel::crTs(const Size i) const {
    if (modelType(CR, i) == LGM1F)
        return crlgm1f(i)->termStructure();
    if (modelType(CR, i) == CIRPP)
        return crcirpp(i)->termStructure();
    QL_FAIL("model at " << i << " is not CR-*");
}

inline std::pair<Real, Real> CrossAssetModel::crS(const Size i, const Size ccy, const Time t, const Time T,
                                                  const Real z, const Real y) const {
    if (modelType(CR, i) == LGM1F)
        return crlgm1fS(i, ccy, t, T, z, y);
    if (modelType(CR, i) == CIRPP) {
        QL_REQUIRE(ccy == 0, "CrossAssetModelPlus::crS() only implemented for ccy=0, got " << ccy);
        return crcirppS(i, t, T, z, y);
    }
    QL_FAIL("model at " << i << " is not CR-*");
}

} // namespace QuantExt

#endif


#include "galaxy.hpp"
#include <boost/python.hpp>
//#include "quadrature.hpp"

namespace gd {
using namespace boost::python;

void py_export_galaxy() {
	class_< Galaxy, boost::noncopyable >("Galaxy", no_init)
		//.def("densityr", pure_virtual(&Profile::densityr))
	;
	class_<Galaxy1C_constant_anisotropy, bases<Galaxy> >("Galaxy1C_constant_anisotropy", init<Profile*, double, double>())
		.def("findzmax", (&Galaxy1C_constant_anisotropy::findzmax))
		.def("rmax_at_E", (&Galaxy1C_constant_anisotropy::rmax_at_E))
		.def("Tr", (&Galaxy1C_constant_anisotropy::Tr))
	;

	class_<Galaxy2C_constant_anisotropy, bases<Galaxy> >("Galaxy2C_constant_anisotropy", init<Profile*, Profile*, double, double>())
		.def("potentialr", (&Galaxy2C_constant_anisotropy::potentialr))
		.def("findzmax", (&Galaxy2C_constant_anisotropy::findzmax))
		.def("rmax_at_E", (&Galaxy2C_constant_anisotropy::rmax_at_E))
		.def("Tr", (&Galaxy2C_constant_anisotropy::Tr))
	;

}

double Galaxy1C_constant_anisotropy::Tr(double E, double L, double ra, double rp) {
	double Lsq = L*L;
	auto f = [&](double r) {
		return 2/sqrt(-Lsq/r*r-2*this->potentialr(r)+2*E);
	};
	IntegratorGSL<> integratorGSL(f); // the integrator
	//return integratorGSL.integrate(ra, rp, 100000, 0, 1e-4);
	double center = (ra+rp)/2;
	return integratorGSL.integrate_to_sigularity(rp, center) + integratorGSL.integrate_to_sigularity(center, ra);

}

}


/* -*- c++ -*- */
/*
 * Copyright 2002,2008 Free Software Foundation, Inc.
 *
 * This file is part of GNU Radio
 *
 * GNU Radio is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3, or (at your option)
 * any later version.
 *
 * GNU Radio is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with GNU Radio; see the file COPYING.  If not, write to
 * the Free Software Foundation, Inc., 51 Franklin Street,
 * Boston, MA 02110-1301, USA.
 */

#include <gnuradio/atsc/sssr_impl.h>
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <assert.h>
#include <gnuradio/atsc/diag_output_impl.h>
#include <gnuradio/math.h>
#include <stdio.h>
#include <boost/math/special_functions/sign.hpp>
#include <iostream>

/*
 * ----------------------------------------------------------------
 * Segment Integrator Pre-Processor
 *
 * Compute weight passed to integrator based on
 * correlation result and ncorr_gain2
 */

inline static int
sipp (bool digcorr_output)
{
  if (digcorr_output)
    return +2;			// positive correlation
  else
    return -1;			// no correlation
}

/*
 * ----------------------------------------------------------------
 * Segment Sync Integrator...
 */

static const int SSI_MIN = -16;
static const int SSI_MAX =  15;

void
sssr::seg_sync_integrator::reset ()
{
  for (int i = 0; i < ATSC_DATA_SEGMENT_LENGTH; i++)
    d_integrator[i] = SSI_MIN;
}

int
sssr::seg_sync_integrator::update (int weight, int index)
{
  int	t = d_integrator[index] + weight;
  t = std::max (t, SSI_MIN);
  t = std::min (t, SSI_MAX);
  d_integrator[index] = t;

  return t;
}

int
sssr::seg_sync_integrator::find_max (int *v)
{
  int	best_value = d_integrator[0];
  int	best_index = 0;

  for (int i = 1; i < ATSC_DATA_SEGMENT_LENGTH; i++)
    if (d_integrator[i] > best_value){
      best_value = d_integrator[i];
      best_index = i;
    }

  *v = best_value;
  return best_index;
}

/*
 * ----------------------------------------------------------------
 * Segment Sync and Symbol recovery
 */

static const int	SYMBOL_INDEX_OFFSET = 3;
static const int	MIN_SEG_LOCK_CORRELATION_VALUE = 5;

atsci_sssr::atsci_sssr ()
  : d_debug_fp(0)
{
  reset ();

  if (_SSSR_DIAG_OUTPUT_){
    const char *file = "sssr.ssout";
    if ((d_debug_fp = fopen (file, "wb")) == 0){
      perror (file);
      exit (1);
    }
  }
}

atsci_sssr::~atsci_sssr ()
{
  if (d_debug_fp)
    fclose (d_debug_fp);
}

void
atsci_sssr::reset ()
{
  d_correlator.reset ();
  d_integrator.reset ();
  d_quad_filter.reset ();

  for (int i = 0; i < ATSC_DATA_SEGMENT_LENGTH; i++)
    d_quad_output[i] = 0;

  d_timing_adjust = 0;
  d_counter = 0;
  d_symbol_index = 0;
  d_seg_locked = false;
}

void
atsci_sssr::update (sssr::sample_t sample_in,	 // input
		   bool		 *seg_locked,	 // are we seeing segment syncs?
		   int		 *symbol_index,	 // 0..831
		   double	 *timing_adjust) // how much to adjust timing
{
  double qo = d_quad_filter.update (sample_in);
  d_quad_output[d_counter] = qo;

  int bit = boost::math::signbit (sample_in);
  if (bit != 0)
    bit = 0;
  else
    bit = 1;
  int corr_out = d_correlator.update (bit);
  int weight = sipp (corr_out);
  int corr_value = d_integrator.update (weight, d_counter);
  int best_correlation_index = -1;

  incr_symbol_index ();
  if (incr_counter ()){		// counter just wrapped...
    int	best_correlation_value;
    best_correlation_index = d_integrator.find_max (&best_correlation_value);
    d_seg_locked = best_correlation_value >= MIN_SEG_LOCK_CORRELATION_VALUE;
    //std::cout << "best = " << best_correlation_value << " min is " << MIN_SEG_LOCK_CORRELATION_VALUE << std::endl;
    d_timing_adjust = d_quad_output[best_correlation_index];

    d_symbol_index = SYMBOL_INDEX_OFFSET - 1 - best_correlation_index;
    if (d_symbol_index < 0)
      d_symbol_index += ATSC_DATA_SEGMENT_LENGTH;
  }

  *seg_locked = d_seg_locked;
  *symbol_index = d_symbol_index;
  *timing_adjust = d_timing_adjust;

  if (_SSSR_DIAG_OUTPUT_){
    float	iodata[7];
    iodata[0] = bit;
    iodata[1] = corr_value;
    iodata[2] = qo;
    iodata[3] = d_counter;
    iodata[4] = d_symbol_index;
    iodata[5] = best_correlation_index;
    iodata[6] = d_timing_adjust;

    if (fwrite (iodata, sizeof (iodata), 1, d_debug_fp) != 1){
      perror ("fwrite: sssr");
      exit (1);
    }
  }

}

/*
 * ----------------------------------------------------------------
 * Interpolator control for Seg & Symbol Sync Recovery
 */

static const double	LOOP_FILTER_TAP = 0.00025;	// 0.0005 works
static const double	ADJUSTMENT_GAIN = 1.0e-5 / (10 * ATSC_DATA_SEGMENT_LENGTH);

atsci_interpolator::atsci_interpolator (double nominal_ratio_of_rx_clock_to_symbol_freq)
  : d_debug_fp(0)
{
// Tweaked ratio from 1.8 to 1.78 to support input rate of 19.2MHz
  assert (nominal_ratio_of_rx_clock_to_symbol_freq >= 1.78);
  d_nominal_ratio_of_rx_clock_to_symbol_freq =
    nominal_ratio_of_rx_clock_to_symbol_freq;

  d_loop.set_taps (LOOP_FILTER_TAP);

  reset ();

  if (_SSSR_DIAG_OUTPUT_){
    const char *file = "interp.ssout";
    if ((d_debug_fp = fopen (file, "wb")) == 0){
      perror (file);
      exit (1);
    }
  }

}

atsci_interpolator::~atsci_interpolator ()
{
  if (d_debug_fp)
    fclose (d_debug_fp);
}

void
atsci_interpolator::reset ()
{
  d_w = d_nominal_ratio_of_rx_clock_to_symbol_freq;
  d_mu = 0.5;
  d_incr = 0;
  d_loop.reset ();
}

bool
atsci_interpolator::update (
	   const sssr::sample_t input_samples[],	// I: vector of samples
	   int	     		nsamples,		// I: total number of samples
	   int	    	       *index,			// I/O: current input index
	   double     		timing_adjustment,	// I: how much to bump timing
	   sssr::sample_t      *output_sample)
{
  if (*index + (int) ntaps () > nsamples)
    return false;

  // FIXME Confirm that this is right.  I think it is.  It was (1-d_mu)
  *output_sample = d_interp.interpolate (&input_samples[*index], d_mu);

  double filter_out = 0;

#if 0

  filter_out = d_loop.filter (timing_adjustment);
  d_w = d_w + ADJUSTMENT_GAIN * filter_out * 1e-3;

#elif 1

  filter_out = d_loop.filter (timing_adjustment);
  d_mu = d_mu + ADJUSTMENT_GAIN * 10e3 * filter_out;

#else

  static const double alpha = 0.01;
  static const double beta = alpha * alpha / 16;

  double x = ADJUSTMENT_GAIN * 10e3 * timing_adjustment;

  d_mu = d_mu + alpha * x;	// conceptually "phase"
  d_w  = d_w  + beta * x;	// conceptually "frequency"

#endif

  double s = d_mu + d_w;
  double float_incr = floor (s);
  d_mu = s - float_incr;
  d_incr = (int) float_incr;

  assert (d_incr >= 1 && d_incr <= 3);
  *index += d_incr;

  if (_SSSR_DIAG_OUTPUT_){
    float	iodata[6];
    iodata[0] = timing_adjustment;
    iodata[1] = filter_out;
    iodata[2] = d_w;
    iodata[3] = d_mu;
    iodata[4] = d_incr;
    iodata[5] = *output_sample;
    if (fwrite (iodata, sizeof (iodata), 1, d_debug_fp) != 1){
      perror ("fwrite: interpolate");
      exit (1);
    }
  }

  return true;
}


//State Vector- variable acceleration model
	// X(0) = quad_position.point.x;
    // X(1) = quad_velocity.point.x;
    // X(2) = quad_acc.point.x;
    // X(3) = quad_position.point.y;
    // X(4) = quad_velocity.point.y;
    // X(5) = quad_acc.point.y;
    // X(6) = quad_position.point.z;
    // X(7) = quad_velocity.point.z;
    // X(8) = quad_acc.point.z;
	// X(9)  = leader.point.x;
    // X(10) = leader.point.y;
    // X(11) = leader.point.z;
	// X(12) = neighbor.point.x;
	// X(13) = neighbor.point.y;
	// X(14) = neighbor.point.z;
///measeurement Vector
    //1.
	// Y(0) = quad.imu.acc.x
	// Y(1) = 0
	// Y(2) = 0
	// Y(3) = quad.imu.acc.y
	// Y(4) = 0
	// Y(5) = 0
	// Y(6) = quad.imu.acc.z
	// Y(7) = 0
	// Y(8) = 0
    //2.
    // Y(0) = quad.imu.vel.x
    // Y(1) = 0
    // Y(2) = 0
    // Y(3) = quad.imu.vel.y
    // Y(4) = 0
    // Y(5) = 0
    // Y(6) = quad.imu.vel.z
    // Y(7) = 0
    // Y(8) = 0
    //3.
	// Y() = leader_quad.point.x
	// Y() = leader_quad.point.y
	// Y() = leader_quad.point.z
    //4.
	// Y() = neighbor_quad.point.x
	// Y() = neighbor_quad.point.y
	// Y() = neighbor_quad.point.z

#include <Eigen/Dense>
#include <Eigen/Core>
#include <Eigen/Geometry>
#include <bits/stdc++.h>
#include <ros/ros.h>
#include <math.h>
#include <stdio.h>
#include <dwm1001/anchor.h>
#include <sensor_msgs/Imu.h>
#include "std_msgs/MultiArrayLayout.h"
#include "std_msgs/MultiArrayDimension.h"
#include "std_msgs/Float32MultiArray.h"
#include "geometry_msgs/PoseStamped.h"
#include "geometry_msgs/Vector3Stamped.h"
#include "geometry_msgs/TwistStamped.h"

// Declarations

double r_pred_l,r_meas_l;
double R_l;
double m_R_scale_l,sigma_r_l;
//double error_l;
double precisionRangeErrEst_l;
Eigen::MatrixXd H_l(1,15);
Eigen::VectorXd K_l(15);
Eigen::VectorXd Ze_r_l(15);
Eigen::MatrixXd Om_r_l(15,15);
Eigen::VectorXd X_r_l(15);
//
double r_pred_n,r_meas_n;
double R_n;
double m_R_scale_n,sigma_r_n;
//double error_n;
double precisionRangeErrEst_n;
Eigen::MatrixXd H_n(1,15);
Eigen::VectorXd K_n(15);
Eigen::VectorXd Ze_r_n(15);
Eigen::MatrixXd Om_r_n(15,15);
Eigen::VectorXd X_r_n(15);
//
Eigen::Quaternionf q;
Eigen::MatrixXf R_mat(3,3);
Eigen::MatrixXf imuacc(3,1);
Eigen::MatrixXf acc(3,1);
double r_acc;
Eigen::MatrixXd r_meas_acc(3,1);
Eigen::MatrixXd H_acc(3,15);
Eigen::MatrixXd R_acc(3,3);
Eigen::MatrixXd K_acc(15,3);
Eigen::MatrixXd Om_acc(15,15);
Eigen::VectorXd Ze_acc(15);
Eigen::VectorXd X_acc(15);
//
double r_vel;
Eigen::MatrixXd r_meas_vel(3,1);
Eigen::MatrixXd H_vel(3,15);
Eigen::MatrixXd R_vel(3,3);
Eigen::MatrixXd K_vel(15,3);
Eigen::MatrixXd Om_vel(15,15);
Eigen::VectorXd Ze_vel(15);
Eigen::VectorXd X_vel(15);
//
double r_pos_l;
Eigen::MatrixXd pos_meas_l(3,1);
Eigen::MatrixXd H_pos_l(3,15);
Eigen::MatrixXd R_pos_l(3,3);
Eigen::MatrixXd K_pos_l(15,3);
Eigen::MatrixXd Om_pos_l(15,15);
Eigen::VectorXd Ze_pos_l(15);
Eigen::VectorXd X_pos_l(15);
//
double r_pos_n;
Eigen::MatrixXd pos_meas_n(3,1);
Eigen::MatrixXd H_pos_n(3,15);
Eigen::MatrixXd R_pos_n(3,3);
Eigen::MatrixXd K_pos_n(15,3);
Eigen::MatrixXd Om_pos_n(15,15);
Eigen::VectorXd Ze_pos_n(15);
Eigen::VectorXd X_pos_n(15);
//
double R_h,r_h,h_meas;
Eigen::MatrixXd H_h(1,15);
Eigen::MatrixXd Om_h(15,15);
Eigen::VectorXd Ze_h(15);
Eigen::VectorXd X_h(15);
//
Eigen::VectorXd X(15);
Eigen::VectorXd u(15);
Eigen::VectorXd X_e(15);
Eigen::MatrixXd F(15,15);
Eigen::MatrixXd block_F(3,3);
Eigen::MatrixXd B(15,15);
Eigen::MatrixXd block_B(3,3);
Eigen::MatrixXd Q(15,15);
Eigen::MatrixXd block_Q(3,3);
Eigen::MatrixXd Q_lead(3,3);
Eigen::MatrixXd Q_neigh(3,3);
Eigen::MatrixXd Om_p(15,15);
Eigen::VectorXd Ze_p(15);
Eigen::MatrixXd Om(15,15);
Eigen::VectorXd Ze(15);
geometry_msgs::Vector3Stamped pose;
geometry_msgs::Vector3Stamped vel;
geometry_msgs::Vector3Stamped acc_;
std_msgs::Float32MultiArray output;
double tao_bias,m_y_damping_factor;
double T,error;
double m_last_range_time;
double error_threshold;
double T_sq,m_tao_bias_sqrt,T_cub,m_z_damping_factor,m_Q_scale;
double x,y,z,x_l,y_l,z_l,x_n,y_n,z_n;
double q_l,q_n;

void setVector(Eigen::VectorXd Y)
{
    Ze = Y;

}
void setMatrix(Eigen::MatrixXd S)
{
    Om = S;
}

//void prediction_step(const sensor_msgs::Imu::ConstPtr& msg)
void prediction_step()
{
    T = ros::Time::now().toSec() - m_last_range_time;
    int static count =0 ;
    if(count < 5)
       {ROS_WARN("prediction_step");
        count++;
        T = 0.0301735 ;}

    T_sq = std::pow(T,2);
    T_cub = std::pow(T,3);
    // F is a 9x9 State Transition Matrix
    F = Eigen::MatrixXd::Zero(15,15);
    block_F << 1, T, -T_sq/2.0,
               0, 1, -T,
               0, 0, 1;
    F.block<3,3>(0,0) = block_F;
    F.block<3,3>(3,3) = block_F;
    F.block<3,3>(6,6) = block_F;
    F.block<6,6>(9,9) = Eigen::MatrixXd::Identity(6,6);

    //std::cout << F << '\n'<<'\n';

    // Q is the acceleration model
    tao_bias = m_tao_bias_sqrt * m_tao_bias_sqrt;
    Q = Eigen::MatrixXd::Zero(15,15);

    block_Q << (T_cub*T_sq)/20.0,   (T_sq*T_sq)/8.0 ,   -T_cub/6,
        	   (T_sq*T_sq)/8.0 ,    (T_cub)/3,          -T_sq/2,
         	    T_cub/6.0,		     -T_sq/2 ,               T  ;
    block_Q *= tao_bias;

    Q.block<3,3>(0,0) = block_Q;
    Q.block<3,3>(3,3) = block_Q * m_y_damping_factor ;
    Q.block<3,3>(6,6) = block_Q * m_z_damping_factor;
    Q.block<3,3>(9,9) = Q_lead;
    Q.block<3,3>(12,12) = Q_neigh;
    Q *= m_Q_scale;
    //ROS_WARN("covariance:## %f, ## %f, ## %f",(T_cub*tao_acc/3.0)+(T_cub*T_sq)*tao_bias/20.0 , T*tao_acc+(T_cub*tao_bias/3), T*tao_bias);

    // if(count%50 == 1)
    // std::cout << X << '\n'<<'\n';
    // M is the predicted covariance matrix
    // X is the predicted state vector
    X = Om.inverse() * Ze;
    Om_p = (F *Om.inverse() * F.transpose()) + Q ;
    Om_p = Om_p.inverse();
    X = F * X ;
    Ze_p = Om_p * X;
   //  q = Eigen::Quaternionf(msg->orientation.w, msg->orientation.x, msg->orientation.y, msg->orientation.z);
   //  R_mat= q.toRotationMatrix();
   //  imuacc << msg->linear_acceleration.x,msg->linear_acceleration.y,msg->linear_acceleration.z;
   //  acc= R_mat*imuacc;
   //  B = Eigen::MatrixXd::Zero(15,15);
   //  block_B << T_sq/2.0,  0,  0,
   //             T      ,  0,  0,
   //             1      ,  0,  0;
   //  B.block<3,3>(0,0) = block_B;
   //  B.block<3,3>(3,3) = block_B;
   //  B.block<3,3>(6,6) = block_B;
   //  B.block<6,6>(9,9) = Eigen::MatrixXd::Identity(6,6);;
   //  u<<acc(0,0),0,0,
   //    -acc(1,0),0,0,
   // acc(2,0)-9.8,0,0,
   // 0,0,0,0,0,0;

   //  X = F * X + B * u;
   //  Ze_p = Om_p * X;
    // time update

    //count++;
    // setState(X_e);
    // setCovariance(M);

        pose.vector.x=X(0);
        pose.vector.y=X(3);
        pose.vector.z=X(6);

        vel.vector.x=X(1);
        vel.vector.y=X(4);
        vel.vector.z=X(7);

        acc_.vector.x=X(2);
        acc_.vector.y=X(5);
        acc_.vector.z=X(8);
    //    fused.publish(output);

    // if(error < error_threshold){
    //     //ROS_WARN("\n sucess too large: %f", error);
        setVector(Ze_p);
        setMatrix(Om_p);
        //std::cout << X << '\n'<<'\n';
    //     return ;
    // } else {

    //     ROS_WARN("\n Estimate too large: %f", error);
    //     return ;
    // }
    // count++;
    m_last_range_time = ros::Time::now().toSec();
}

void correction_step_imu(const sensor_msgs::Imu::ConstPtr& msg)
{
    int static count =0 ;
    if(count == 0)
       {ROS_WARN("imu_step");
        count++; }

    //conerts the acceleration from body frame to earth frame(NED)
    q = Eigen::Quaternionf(msg->orientation.w, msg->orientation.x, msg->orientation.y, msg->orientation.z);
    R_mat= q.toRotationMatrix();

    imuacc << msg->linear_acceleration.x,msg->linear_acceleration.y,msg->linear_acceleration.z;
    acc= R_mat*imuacc;
    acc(0,0)=acc(0,0);
    acc(1,0)=acc(1,0);
    acc(2,0)=(acc(2,0)-9.8);

    r_meas_acc << acc(0,0),acc(1,0),acc(2,0);
    // K is the Kalman Gain
    R_acc = Eigen::MatrixXd::Identity(3,3) * (r_acc*r_acc);
    R_acc(1,1) = R_acc(1,1) ;
    R_acc(2,2) = R_acc(2,2) * m_z_damping_factor;
    //ROS_WARN("R_matrix:## %f",R);
    // Update the state
    R_acc = R_acc.inverse();
    Om_acc = Om + (H_acc.transpose() * R_acc * H_acc) ;
    Ze_acc = Ze + (H_acc.transpose() * R_acc * r_meas_acc) ;

    X_acc = Om_acc.inverse() * Ze_acc;

if(abs(X_acc(0)-X(0)) < error_threshold && abs(X_acc(3)-X(3)) < error_threshold && abs(X_acc(6)-X(6)) < error_threshold )
    {   setVector(Ze_acc);
        setMatrix(Om_acc);
    }

}

void correction_step_vel(const geometry_msgs::TwistStamped::ConstPtr& msg)
{
    int static count =0 ;
    if(count == 0)
       {ROS_WARN("vel_step");
        count++; }

    r_meas_vel << msg->twist.linear.x,msg->twist.linear.y,msg->twist.linear.z;
    // K is the Kalman Gain
    R_vel = Eigen::MatrixXd::Identity(3,3) * (r_vel*r_vel);
    R_vel(1,1) = R_vel(1,1) ;
    R_vel(2,2) = R_vel(2,2) * m_z_damping_factor;
    //ROS_WARN("R_matrix:## %f",R);
    R_vel = R_vel.inverse();
    Om_vel = Om + (H_vel.transpose() * R_vel * H_vel) ;
    Ze_vel = Ze + (H_vel.transpose() * R_vel * r_meas_vel) ;

        setVector(Ze_vel);
        setMatrix(Om_vel);
    X_vel = Om_vel.inverse() * Ze_vel;

if(abs(X_vel(0)-X(0)) < error_threshold && abs(X_vel(3)-X(3)) < error_threshold && abs(X_vel(6)-X(6)) < error_threshold )
    {   setVector(Ze_vel);
        setMatrix(Om_vel);
    }

}

void correction_step_leader(const dwm1001::anchor::ConstPtr& msg)
{
    int static count =0 ;
    if(count == 0)
       {ROS_WARN("leader_r_step");
        count++; }

	r_pred_l = std::sqrt( std::pow(X(0) -  X(9), 2) +
                          std::pow(X(3) - X(10), 2) +
                          std::pow(X(6) - X(11), 2) ) + 1e-5;

	r_meas_l = msg->range;
    // H is the linearized measurement matrix
    H_l << (X(0) -  X(9))/r_pred_l, 0, 0,
           (X(3) - X(10))/r_pred_l, 0, 0,
           (X(6) - X(11))/r_pred_l, 0, 0,
           (X(9) - X(0))/r_pred_l, (X(10) - X(3))/r_pred_l, (X(11) - X(6))/r_pred_l,
         	0,	0,	0;

    // K is the Kalman Gain
    sigma_r_l = double(precisionRangeErrEst_l) / 1000.0;
    R_l = std::pow(sigma_r_l,2) * m_R_scale_l;
    //ROS_WARN("R_matrix:## %f",R);
    R_l = 1/R_l;
    Om_r_l = Om + (H_l.transpose() * R_l * H_l) ;
    Ze_r_l = Ze + (H_l.transpose() * R_l * (r_meas_l - r_pred_l + (H_l*X)(0,0) )) ;

    X_r_l = Om_r_l.inverse() * Ze_r_l;

if(abs(X_r_l(0)-X(0)) < error_threshold && abs(X_r_l(3)-X(3)) < error_threshold && abs(X_r_l(6)-X(6)) < error_threshold )
    {   setVector(Ze_r_l);
        setMatrix(Om_r_l);
    }
}

void correction_step_neigh(const dwm1001::anchor::ConstPtr& msg)
{
    int static count =0 ;
    if(count == 0)
       {ROS_WARN("neghbor_r_step");
        count++; }

	r_pred_n = std::sqrt( std::pow(X(0) - X(12), 2) +
                          std::pow(X(3) - X(13), 2) +
                          std::pow(X(6) - X(14), 2) ) + 1e-5;

	r_meas_n = msg->range;
    // H is the linearized measurement matrix
    H_n << (X(0) - X(12))/r_pred_n, 0, 0,
           (X(3) - X(13))/r_pred_n, 0, 0,
           (X(6) - X(14))/r_pred_n, 0, 0,
         	0,	0,	0,
           (X(12) - X(0))/r_pred_n, (X(13) - X(3))/r_pred_n, (X(14) - X(6))/r_pred_n ;

    // K is the Kalman Gain
    sigma_r_n = double(precisionRangeErrEst_n) / 1000.0;
    R_n = std::pow(sigma_r_n,2) * m_R_scale_n;
    //ROS_WARN("R_matrix:## %f",R);
    R_n = 1/R_n;
    Om_r_n = Om + (H_n.transpose() * R_n * H_n) ;
    Ze_r_n = Ze + (H_n.transpose() * R_n * (r_meas_n - r_pred_n + (H_n*X)(0,0) )) ;

    X_r_n = Om_r_n.inverse() * Ze_r_n;

if(abs(X_r_n(0)-X(0)) < error_threshold && abs(X_r_n(3)-X(3)) < error_threshold && abs(X_r_n(6)-X(6)) < error_threshold )
    {   setVector(Ze_r_n);
        setMatrix(Om_r_n);
    }

}

void position_lead_cb(const geometry_msgs::PoseStamped::ConstPtr& msg)
{
    int static count =0 ;
    if(count == 0)
       {ROS_WARN("leader_pos_step");
        count++; }
    pos_meas_l << msg->pose.position.x,msg->pose.position.y,msg->pose.position.z;
    // K is the Kalman Gain
    R_pos_l = Eigen::MatrixXd::Identity(3,3) * (r_pos_l*r_pos_l);
    //ROS_WARN("R_matrix:## %f",R);

    R_pos_l = R_pos_l.inverse();
    Om_pos_l = Om + (H_pos_l.transpose() * R_pos_l * H_pos_l) ;
    Ze_pos_l = Ze + (H_pos_l.transpose() * R_pos_l * pos_meas_l) ;

    X_pos_l = Om_pos_l.inverse() * Ze_pos_l;

if(abs(X_pos_l(0)-X(0)) < error_threshold && abs(X_pos_l(3)-X(3)) < error_threshold && abs(X_pos_l(6)-X(6)) < error_threshold )
    {   setVector(Ze_pos_l);
        setMatrix(Om_pos_l);
    }
}

void position_neigh_cb(const geometry_msgs::PoseStamped::ConstPtr& msg)
{
    int static count =0 ;
    if(count == 0)
       {ROS_WARN("neighbor_pos_step");
        count++; }
    pos_meas_n << msg->pose.position.x,msg->pose.position.y,msg->pose.position.z;
    // K is the Kalman Gain
    R_pos_n = Eigen::MatrixXd::Identity(3,3) * (r_pos_n*r_pos_n);
    //ROS_WARN("R_matrix:## %f",R);

    R_pos_n = R_pos_n.inverse();
    Om_pos_n = Om + (H_pos_n.transpose() * R_pos_n * H_pos_n) ;
    Ze_pos_n = Ze + (H_pos_n.transpose() * R_pos_n * pos_meas_n) ;

    X_pos_n = Om_pos_n.inverse() * Ze_pos_n;

if(abs(X_pos_n(0)-X(0)) < error_threshold && abs(X_pos_n(3)-X(3)) < error_threshold && abs(X_pos_n(6)-X(6)) < error_threshold )
    {   setVector(Ze_pos_n);
        setMatrix(Om_pos_n);
    }

}

void height_cb(const geometry_msgs::PoseStamped::ConstPtr& msg)
{
    int static count =0 ;
    if(count == 0)
       {ROS_WARN("height_cb");
        count++; }
    h_meas = msg->pose.position.z;
    // K is the Kalman Gain
    R_h = r_h*r_h;
    //ROS_WARN("R_matrix:## %f",R);

    R_h = 1/R_h;
    Om_h = Om + (H_h.transpose() * R_h * H_h) ;
    Ze_h = Ze + (H_h.transpose() * R_h * h_meas) ;

        setVector(Ze_h);
        setMatrix(Om_h);
    X_h = Om_h.inverse() * Ze_h;

if(abs(X_h(0)-X(0)) < error_threshold && abs(X_h(3)-X(3)) < error_threshold && abs(X_h(6)-X(6)) < error_threshold )
    {   setVector(Ze_h);
        setMatrix(Om_h);
    }

}

void param(ros::NodeHandle& nh)
{
    nh.getParam("IKalmanFilter_swarm/start_x", x);
    nh.getParam("IKalmanFilter_swarm/start_y", y);
    nh.getParam("IKalmanFilter_swarm/start_z", z);
    nh.getParam("IKalmanFilter_swarm/lead_x", x_l);
    nh.getParam("IKalmanFilter_swarm/lead_y", y_l);
    nh.getParam("IKalmanFilter_swarm/lead_z", z_l);
    nh.getParam("IKalmanFilter_swarm/neighbor_x", x_n);
    nh.getParam("IKalmanFilter_swarm/neighbor_y", y_n);
    nh.getParam("IKalmanFilter_swarm/neighbor_z", z_n);
    nh.getParam("IKalmanFilter_swarm/leader_covariance", q_l);
    nh.getParam("IKalmanFilter_swarm/neighbor_covariance", q_n);
    nh.getParam("IKalmanFilter_swarm/m_tao_bias_sqrt", m_tao_bias_sqrt );
    nh.getParam("IKalmanFilter_swarm/m_y_damping_factor", m_y_damping_factor);
    nh.getParam("IKalmanFilter_swarm/m_z_damping_factor", m_z_damping_factor);
    nh.getParam("IKalmanFilter_swarm/m_Q_scale", m_Q_scale);
    nh.getParam("IKalmanFilter_swarm/error_threshold", error_threshold);
    nh.getParam("IKalmanFilter_swarm/m_R_scale_leader", m_R_scale_l);
    nh.getParam("IKalmanFilter_swarm/precisionRangeErrEst_leader", precisionRangeErrEst_l);
    nh.getParam("IKalmanFilter_swarm/m_R_scale_neighbor", m_R_scale_n);
    nh.getParam("IKalmanFilter_swarm/precisionRangeErrEst_neighbor", precisionRangeErrEst_n);
    nh.getParam("IKalmanFilter_swarm/r_acc", r_acc);
    nh.getParam("IKalmanFilter_swarm/r_vel", r_vel);
    nh.getParam("IKalmanFilter_swarm/r_pos_l", r_pos_l);
    nh.getParam("IKalmanFilter_swarm/r_pos_n", r_pos_n);
    nh.getParam("IKalmanFilter_swarm/r_h", r_h);

    ROS_WARN("%f, %f, %f ",x,y,z);
}

void Initialize(ros::NodeHandle& nh)
{

    param(nh);

    X<<x,0,0,y,0,0,z,0,0,x_l,y_l,z_l,x_n,y_n,z_n;
    Eigen::MatrixXd nine_cov = Eigen::MatrixXd::Identity(15,15);
    nine_cov(0,0) = 0.001;
    nine_cov(3,3) = 0.001;
    nine_cov(6,6) = 0.001;
    // set cov of vel
    nine_cov(1,1) = 0.01;
    nine_cov(4,4) = 0.01;
    nine_cov(7,7) = 0.01;
    // set cov of acc_bias
    nine_cov(2,2) = 0.01;
    nine_cov(5,5) = 0.01;
    nine_cov(8,8) = 0.01;

    nine_cov(9,9)   = 0.001;
    nine_cov(10,10) = 0.001;
    nine_cov(11,11) = 0.001;

    nine_cov(12,12) = 0.001;
    nine_cov(13,13) = 0.001;
    nine_cov(14,14) = 0.001;
    nine_cov = nine_cov*100;
    nine_cov = nine_cov.inverse();
    setMatrix(nine_cov);
    setVector(nine_cov*X);

    Q_lead <<  q_l,  0,  0,
                0, q_l, 0,
                0,  0,  q_l ;
    Q_neigh << q_n,  0,  0,
                0, q_n, 0,
                0,  0,  q_n ;
        // H is the linearized measurement matrix
    H_acc << 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0;
        // H is the linearized measurement matrix
    H_vel << 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
             0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0;
        // H is the linearized measurement matrix
    H_pos_l << 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0,
               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0,
               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0;
        // H is the linearized measurement matrix
    H_pos_n << 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0,
               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0,
               0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1;

    H_h << 0,0,0,0,0,0,1,0,0,0,0,0,0,0,0;

    m_last_range_time = ros::Time::now().toSec();

}

int main(int argc, char** argv){

    ros::init(argc,argv,"IKalmanFilter_swarm");
    ros::NodeHandle nh;
    //Initializinging the parameters
    Initialize(nh);
    //param(nh);
    //Subscriber and Publisher for the data. remapped in the launch file to the topic required
    ros::Subscriber quad_imu = nh.subscribe("imu",20,correction_step_imu);
    ros::Subscriber quad_vel = nh.subscribe("velocity",20,correction_step_vel);
    ros::Subscriber anchor_1 = nh.subscribe("anchor_lead", 20,correction_step_leader);
    ros::Subscriber anchor_2 = nh.subscribe("anchor_neigh",20,correction_step_neigh);
    ros::Subscriber anchor_3 = nh.subscribe("position_leader",20,position_lead_cb);
    ros::Subscriber anchor_4 = nh.subscribe("position_neighour",20,position_neigh_cb);
    ros::Subscriber height = nh.subscribe("height",20,height_cb);

    //ros::Publisher fused = nh.advertise<std_msgs::Float32MultiArray>("Filtered_data", 100);
    ros::Publisher fused_pose = nh.advertise<geometry_msgs::Vector3Stamped>("Filtered_pose", 10);
    ros::Publisher fused_vel = nh.advertise<geometry_msgs::Vector3Stamped>("Filtered_velocity", 10);
    ros::Publisher fused_acc = nh.advertise<geometry_msgs::Vector3Stamped>("Filtered_acc", 10);

    ros::Rate loop_rate(50);
    while(ros::ok()){
    //	output.data.clear();
	//	for (int i = 0; i < 15; i++)
	//		output.data.push_back(X(i));
        prediction_step();
		pose.header.stamp = ros::Time::now();
        vel.header.stamp = ros::Time::now();
        acc_.header.stamp = ros::Time::now();
        fused_pose.publish(pose);
        fused_vel.publish(vel);
        fused_acc.publish(acc_);

        ros::spinOnce();
        loop_rate.sleep();
    }
    return 0;
}


#ifndef BOOST_GEOMETRY_PROJECTIONS_TMERC_HPP
#define BOOST_GEOMETRY_PROJECTIONS_TMERC_HPP

// Boost.Geometry - extensions-gis-projections (based on PROJ4)
// This file is automatically generated. DO NOT EDIT.

// Copyright (c) 2008-2012 Barend Gehrels, Amsterdam, the Netherlands.

// Use, modification and distribution is 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)

// This file is converted from PROJ4, http://trac.osgeo.org/proj
// PROJ4 is originally written by Gerald Evenden (then of the USGS)
// PROJ4 is maintained by Frank Warmerdam
// PROJ4 is converted to Boost.Geometry by Barend Gehrels

// Original copyright notice:

// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:

// The above copyright notice and this permission notice shall be included
// in all copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

#include <boost/concept_check.hpp>
#include <boost/math/special_functions/hypot.hpp>

#include <boost/geometry/extensions/gis/projections/impl/base_static.hpp>
#include <boost/geometry/extensions/gis/projections/impl/base_dynamic.hpp>
#include <boost/geometry/extensions/gis/projections/impl/projects.hpp>
#include <boost/geometry/extensions/gis/projections/impl/factory_entry.hpp>
#include <boost/geometry/extensions/gis/projections/impl/function_overloads.hpp>
#include <boost/geometry/extensions/gis/projections/impl/pj_mlfn.hpp>

#include <boost/geometry/extensions/gis/projections/epsg_traits.hpp>

namespace boost { namespace geometry { namespace projections
{
    #ifndef DOXYGEN_NO_DETAIL
    namespace detail { namespace tmerc{
            static const double EPS10 = 1.e-10;
            static const double FC1 = 1.;
            static const double FC2 = .5;
            static const double FC3 = .16666666666666666666;
            static const double FC4 = .08333333333333333333;
            static const double FC5 = .05;
            static const double FC6 = .03333333333333333333;
            static const double FC7 = .02380952380952380952;
            static const double FC8 = .01785714285714285714;

            struct par_tmerc
            {
                double    esp;
                double    ml0;
                double    en[EN_SIZE];
            };

            // template class, using CRTP to implement forward/inverse
            template <typename Geographic, typename Cartesian, typename Parameters>
            struct base_tmerc_ellipsoid : public base_t_fi<base_tmerc_ellipsoid<Geographic, Cartesian, Parameters>,
                     Geographic, Cartesian, Parameters>
            {

                 typedef double geographic_type;
                 typedef double cartesian_type;

                par_tmerc m_proj_parm;

                inline base_tmerc_ellipsoid(const Parameters& par)
                    : base_t_fi<base_tmerc_ellipsoid<Geographic, Cartesian, Parameters>,
                     Geographic, Cartesian, Parameters>(*this, par) {}

                inline void fwd(geographic_type& lp_lon, geographic_type& lp_lat, cartesian_type& xy_x, cartesian_type& xy_y) const
                {
                    double al, als, n, cosphi, sinphi, t;

                        /*
                         * Fail if our longitude is more than 90 degrees from the
                         * central meridian since the results are essentially garbage.
                         * Is error -20 really an appropriate return value?
                         *
                         *  http://trac.osgeo.org/proj/ticket/5
                         */
                        if( lp_lon < -HALFPI || lp_lon > HALFPI )
                        {
                            xy_x = HUGE_VAL;
                            xy_y = HUGE_VAL;
                            throw proj_exception( -14 );
                            return;
                        }

                    sinphi = sin(lp_lat); cosphi = cos(lp_lat);
                    t = fabs(cosphi) > 1e-10 ? sinphi/cosphi : 0.;
                    t *= t;
                    al = cosphi * lp_lon;
                    als = al * al;
                    al /= sqrt(1. - this->m_par.es * sinphi * sinphi);
                    n = this->m_proj_parm.esp * cosphi * cosphi;
                    xy_x = this->m_par.k0 * al * (FC1 +
                        FC3 * als * (1. - t + n +
                        FC5 * als * (5. + t * (t - 18.) + n * (14. - 58. * t)
                        + FC7 * als * (61. + t * ( t * (179. - t) - 479. ) )
                        )));
                    xy_y = this->m_par.k0 * (pj_mlfn(lp_lat, sinphi, cosphi, this->m_proj_parm.en) - this->m_proj_parm.ml0 +
                        sinphi * al * lp_lon * FC2 * ( 1. +
                        FC4 * als * (5. - t + n * (9. + 4. * n) +
                        FC6 * als * (61. + t * (t - 58.) + n * (270. - 330 * t)
                        + FC8 * als * (1385. + t * ( t * (543. - t) - 3111.) )
                        ))));
                }

                inline void inv(cartesian_type& xy_x, cartesian_type& xy_y, geographic_type& lp_lon, geographic_type& lp_lat) const
                {
                    double n, con, cosphi, d, ds, sinphi, t;

                    lp_lat = pj_inv_mlfn(this->m_proj_parm.ml0 + xy_y / this->m_par.k0, this->m_par.es, this->m_proj_parm.en);
                    if (fabs(lp_lat) >= HALFPI) {
                        lp_lat = xy_y < 0. ? -HALFPI : HALFPI;
                        lp_lon = 0.;
                    } else {
                        sinphi = sin(lp_lat);
                        cosphi = cos(lp_lat);
                        t = fabs(cosphi) > 1e-10 ? sinphi/cosphi : 0.;
                        n = this->m_proj_parm.esp * cosphi * cosphi;
                        d = xy_x * sqrt(con = 1. - this->m_par.es * sinphi * sinphi) / this->m_par.k0;
                        con *= t;
                        t *= t;
                        ds = d * d;
                        lp_lat -= (con * ds / (1.-this->m_par.es)) * FC2 * (1. -
                            ds * FC4 * (5. + t * (3. - 9. *  n) + n * (1. - 4 * n) -
                            ds * FC6 * (61. + t * (90. - 252. * n +
                                45. * t) + 46. * n
                           - ds * FC8 * (1385. + t * (3633. + t * (4095. + 1574. * t)) )
                            )));
                        lp_lon = d*(FC1 -
                            ds*FC3*( 1. + 2.*t + n -
                            ds*FC5*(5. + t*(28. + 24.*t + 8.*n) + 6.*n
                           - ds * FC7 * (61. + t * (662. + t * (1320. + 720. * t)) )
                        ))) / cosphi;
                    }
                }
            };

            // template class, using CRTP to implement forward/inverse
            template <typename Geographic, typename Cartesian, typename Parameters>
            struct base_tmerc_spheroid : public base_t_fi<base_tmerc_spheroid<Geographic, Cartesian, Parameters>,
                     Geographic, Cartesian, Parameters>
            {

                 typedef double geographic_type;
                 typedef double cartesian_type;

                par_tmerc m_proj_parm;

                inline base_tmerc_spheroid(const Parameters& par)
                    : base_t_fi<base_tmerc_spheroid<Geographic, Cartesian, Parameters>,
                     Geographic, Cartesian, Parameters>(*this, par) {}

                inline void fwd(geographic_type& lp_lon, geographic_type& lp_lat, cartesian_type& xy_x, cartesian_type& xy_y) const
                {
                    double b, cosphi;

                        /*
                         * Fail if our longitude is more than 90 degrees from the
                         * central meridian since the results are essentially garbage.
                         * Is error -20 really an appropriate return value?
                         *
                         *  http://trac.osgeo.org/proj/ticket/5
                         */
                        if( lp_lon < -HALFPI || lp_lon > HALFPI )
                        {
                            xy_x = HUGE_VAL;
                            xy_y = HUGE_VAL;
                            throw proj_exception( -14 );
                            return;
                        }

                    b = (cosphi = cos(lp_lat)) * sin(lp_lon);
                    if (fabs(fabs(b) - 1.) <= EPS10) throw proj_exception();;
                    xy_x = this->m_proj_parm.ml0 * log((1. + b) / (1. - b));
                    if ((b = fabs( xy_y = cosphi * cos(lp_lon) / sqrt(1. - b * b) )) >= 1.) {
                        if ((b - 1.) > EPS10) throw proj_exception();
                        else xy_y = 0.;
                    } else
                        xy_y = acos(xy_y);
                    if (lp_lat < 0.) xy_y = -xy_y;
                    xy_y = this->m_proj_parm.esp * (xy_y - this->m_par.phi0);
                }

                inline void inv(cartesian_type& xy_x, cartesian_type& xy_y, geographic_type& lp_lon, geographic_type& lp_lat) const
                {
                    double h, g;

                    h = exp(xy_x / this->m_proj_parm.esp);
                    g = .5 * (h - 1. / h);
                    h = cos(this->m_par.phi0 + xy_y / this->m_proj_parm.esp);
                    lp_lat = asin(sqrt((1. - h * h) / (1. + g * g)));
                    if (xy_y < 0.) lp_lat = -lp_lat;
                    lp_lon = (g || h) ? atan2(g, h) : 0.;
                }
            };

            template <typename Parameters>
            void setup(Parameters& par, par_tmerc& proj_parm)  /* general initialization */
            {
                boost::ignore_unused_variable_warning(par);
                boost::ignore_unused_variable_warning(proj_parm);
                if (par.es) {
                    pj_enfn(par.es, proj_parm.en);

                    proj_parm.ml0 = pj_mlfn(par.phi0, sin(par.phi0), cos(par.phi0), proj_parm.en);
                    proj_parm.esp = par.es / (1. - par.es);
                // par.inv = e_inverse;
                // par.fwd = e_forward;
                } else {
                    proj_parm.esp = par.k0;
                    proj_parm.ml0 = .5 * proj_parm.esp;
                // par.inv = s_inverse;
                // par.fwd = s_forward;
                }
            }

            // Transverse Mercator
            template <typename Parameters>
            void setup_tmerc(Parameters& par, par_tmerc& proj_parm)
            {
                setup(par, proj_parm);
            }

            // Universal Transverse Mercator (UTM)
            template <typename Parameters>
            void setup_utm(Parameters& par, par_tmerc& proj_parm)
            {
                int zone;
                if (!par.es) throw proj_exception(-34);
                par.y0 = pj_param(par.params, "bsouth").i ? 10000000. : 0.;
                par.x0 = 500000.;
                if (pj_param(par.params, "tzone").i) /* zone input ? */
                    if ((zone = pj_param(par.params, "izone").i) > 0 && zone <= 60)
                        --zone;
                    else
                        throw proj_exception(-35);
                else /* nearest central meridian input */
                    if ((zone = int_floor((adjlon(par.lam0) + PI) * 30. / PI)) < 0)
                        zone = 0;
                    else if (zone >= 60)
                        zone = 59;
                par.lam0 = (zone + .5) * PI / 30. - PI;
                par.k0 = 0.9996;
                par.phi0 = 0.;
                setup(par, proj_parm);
            }

        }} // namespace detail::tmerc
    #endif // doxygen

    /*!
        \brief Transverse Mercator projection
        \ingroup projections
        \tparam Geographic latlong point type
        \tparam Cartesian xy point type
        \tparam Parameters parameter type
        \par Projection characteristics
         - Cylindrical
         - Spheroid
         - Ellipsoid
        \par Example
        \image html ex_tmerc.gif
    */
    template <typename Geographic, typename Cartesian, typename Parameters = parameters>
    struct tmerc_ellipsoid : public detail::tmerc::base_tmerc_ellipsoid<Geographic, Cartesian, Parameters>
    {
        inline tmerc_ellipsoid(const Parameters& par) : detail::tmerc::base_tmerc_ellipsoid<Geographic, Cartesian, Parameters>(par)
        {
            detail::tmerc::setup_tmerc(this->m_par, this->m_proj_parm);
        }
    };

    /*!
        \brief Universal Transverse Mercator (UTM) projection
        \ingroup projections
        \tparam Geographic latlong point type
        \tparam Cartesian xy point type
        \tparam Parameters parameter type
        \par Projection characteristics
         - Cylindrical
         - Spheroid
         - zone= south
        \par Example
        \image html ex_utm.gif
    */
    template <typename Geographic, typename Cartesian, typename Parameters = parameters>
    struct utm_ellipsoid : public detail::tmerc::base_tmerc_ellipsoid<Geographic, Cartesian, Parameters>
    {
        inline utm_ellipsoid(const Parameters& par) : detail::tmerc::base_tmerc_ellipsoid<Geographic, Cartesian, Parameters>(par)
        {
            detail::tmerc::setup_utm(this->m_par, this->m_proj_parm);
        }
    };

    /*!
        \brief Transverse Mercator projection
        \ingroup projections
        \tparam Geographic latlong point type
        \tparam Cartesian xy point type
        \tparam Parameters parameter type
        \par Projection characteristics
         - Cylindrical
         - Spheroid
         - Ellipsoid
        \par Example
        \image html ex_tmerc.gif
    */
    template <typename Geographic, typename Cartesian, typename Parameters = parameters>
    struct tmerc_spheroid : public detail::tmerc::base_tmerc_spheroid<Geographic, Cartesian, Parameters>
    {
        inline tmerc_spheroid(const Parameters& par) : detail::tmerc::base_tmerc_spheroid<Geographic, Cartesian, Parameters>(par)
        {
            detail::tmerc::setup_tmerc(this->m_par, this->m_proj_parm);
        }
    };

    #ifndef DOXYGEN_NO_DETAIL
    namespace detail
    {

        // Factory entry(s)
        template <typename Geographic, typename Cartesian, typename Parameters>
        class tmerc_entry : public detail::factory_entry<Geographic, Cartesian, Parameters>
        {
            public :
                virtual projection<Geographic, Cartesian>* create_new(const Parameters& par) const
                {
                    if (par.es)
                        return new base_v_fi<tmerc_ellipsoid<Geographic, Cartesian, Parameters>, Geographic, Cartesian, Parameters>(par);
                    else
                        return new base_v_fi<tmerc_spheroid<Geographic, Cartesian, Parameters>, Geographic, Cartesian, Parameters>(par);
                }
        };

        template <typename Geographic, typename Cartesian, typename Parameters>
        class utm_entry : public detail::factory_entry<Geographic, Cartesian, Parameters>
        {
            public :
                virtual projection<Geographic, Cartesian>* create_new(const Parameters& par) const
                {
                    return new base_v_fi<utm_ellipsoid<Geographic, Cartesian, Parameters>, Geographic, Cartesian, Parameters>(par);
                }
        };

        template <typename Geographic, typename Cartesian, typename Parameters>
        inline void tmerc_init(detail::base_factory<Geographic, Cartesian, Parameters>& factory)
        {
            factory.add_to_factory("tmerc", new tmerc_entry<Geographic, Cartesian, Parameters>);
            factory.add_to_factory("utm", new utm_entry<Geographic, Cartesian, Parameters>);
        }

    } // namespace detail
    // Create EPSG specializations
    // (Proof of Concept, only for some)

    template<typename LatLongRadian, typename Cartesian, typename Parameters>
    struct epsg_traits<2000, LatLongRadian, Cartesian, Parameters>
    {
        typedef tmerc_ellipsoid<LatLongRadian, Cartesian, Parameters> type;
        static inline std::string par()
        {
            return "+proj=tmerc +lat_0=0 +lon_0=-62 +k=0.9995000000000001 +x_0=400000 +y_0=0 +ellps=clrk80 +units=m";
        }
    };

    template<typename LatLongRadian, typename Cartesian, typename Parameters>
    struct epsg_traits<2001, LatLongRadian, Cartesian, Parameters>
    {
        typedef tmerc_ellipsoid<LatLongRadian, Cartesian, Parameters> type;
        static inline std::string par()
        {
            return "+proj=tmerc +lat_0=0 +lon_0=-62 +k=0.9995000000000001 +x_0=400000 +y_0=0 +ellps=clrk80 +units=m";
        }
    };

    template<typename LatLongRadian, typename Cartesian, typename Parameters>
    struct epsg_traits<2002, LatLongRadian, Cartesian, Parameters>
    {
        typedef tmerc_ellipsoid<LatLongRadian, Cartesian, Parameters> type;
        static inline std::string par()
        {
            return "+proj=tmerc +lat_0=0 +lon_0=-62 +k=0.9995000000000001 +x_0=400000 +y_0=0 +ellps=clrk80 +towgs84=725,685,536,0,0,0,0 +units=m";
        }
    };

    template<typename LatLongRadian, typename Cartesian, typename Parameters>
    struct epsg_traits<2003, LatLongRadian, Cartesian, Parameters>
    {
        typedef tmerc_ellipsoid<LatLongRadian, Cartesian, Parameters> type;
        static inline std::string par()
        {
            return "+proj=tmerc +lat_0=0 +lon_0=-62 +k=0.9995000000000001 +x_0=400000 +y_0=0 +ellps=clrk80 +towgs84=72,213.7,93,0,0,0,0 +units=m";
        }
    };

    template<typename LatLongRadian, typename Cartesian, typename Parameters>
    struct epsg_traits<2039, LatLongRadian, Cartesian, Parameters>
    {
        typedef tmerc_ellipsoid<LatLongRadian, Cartesian, Parameters> type;
        static inline std::string par()
        {
            return "+proj=tmerc +lat_0=31.73439361111111 +lon_0=35.20451694444445 +k=1.0000067 +x_0=219529.584 +y_0=626907.39 +ellps=GRS80 +towgs84=-48,55,52,0,0,0,0 +units=m";
        }
    };

    template<typename LatLongRadian, typename Cartesian, typename Parameters>
    struct epsg_traits<29118, LatLongRadian, Cartesian, Parameters>
    {
        typedef utm_ellipsoid<LatLongRadian, Cartesian, Parameters> type;
        static inline std::string par()
        {
            return "+proj=utm +zone=18 +ellps=GRS67 +units=m";
        }
    };

    template<typename LatLongRadian, typename Cartesian, typename Parameters>
    struct epsg_traits<29119, LatLongRadian, Cartesian, Parameters>
    {
        typedef utm_ellipsoid<LatLongRadian, Cartesian, Parameters> type;
        static inline std::string par()
        {
            return "+proj=utm +zone=19 +ellps=GRS67 +units=m";
        }
    };

    #endif // doxygen

}}} // namespace boost::geometry::projections

#endif // BOOST_GEOMETRY_PROJECTIONS_TMERC_HPP


// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT license.

#include "conv.h"
#include "common.h"
#include "gemm.h"
#include <Eigen/Core>
#include <Eigen/Dense>
#include <vector>

// Wasm interop method
void conv_f32(void *data) {
  uint32_t *dataIndex = static_cast<uint32_t *>(data);
  uint32_t const argc = dataIndex[0];
  // TODO: Support muti-dimensional convolution (1D and 3D atleast)
  conv2D_f32_imp(
      PARAM_FLOAT_PTR(data, dataIndex[1]), PARAM_INT32_PTR(data, dataIndex[2]),
      PARAM_FLOAT_PTR(data, dataIndex[3]), PARAM_INT32_PTR(data, dataIndex[4]),
      PARAM_FLOAT_PTR(data, dataIndex[5]), PARAM_INT32_PTR(data, dataIndex[6]),
      PARAM_FLOAT_PTR(data, dataIndex[7]), PARAM_INT32_PTR(data, dataIndex[8]),
      PARAM_INT32(data, dataIndex[9]), PARAM_INT32_PTR(data, dataIndex[10]),
      PARAM_INT32_PTR(data, dataIndex[11]));
}

// Core operator implementation
void conv2D_f32_imp(float *X, int *X_shape, float *W, int *W_shape, float *Y,
                    int *Y_shape, float *bias, int *dilations, int group,
                    int *pads, int *strides) {
  const int input_num = X_shape[0];
  const int input_channels = X_shape[1];
  const int input_height = X_shape[2];
  const int input_width = X_shape[3];
  const int input_size =
      input_num * input_channels * input_height * input_width;

  const int filter_num = W_shape[0];
  const int filter_channels = W_shape[1];
  const int filter_height = W_shape[2];
  const int filter_width = W_shape[3];
  const int filter_size =
      filter_num * filter_channels * filter_height * filter_width;
  std::vector<int> kernel_shape;
  kernel_shape.push_back(filter_height);
  kernel_shape.push_back(filter_width);

  const int output_num = Y_shape[0];
  const int output_channels = Y_shape[1];
  const int output_height = Y_shape[2];
  const int output_width = Y_shape[3];
  const int output_size =
      output_num * output_channels * output_height * output_width;

  const int input_image_size = input_height * input_width;
  const int output_image_size = output_height * output_width;
  const int kernel_size = kernel_shape[0] * kernel_shape[1];
  const int X_offset = input_channels / group * input_image_size;
  const int Y_offset = output_size / output_num / group;
  const int W_offset = filter_size / group;
  const int kernel_dim = input_channels / group * kernel_size;
  const int col_buffer_size = kernel_dim * output_image_size;

  float *col_buffer_data = new float[col_buffer_size]();

  for (int image_id = 0; image_id < input_num; ++image_id) {
    for (int group_id = 0; group_id < group; ++group_id) {
      im2col_f32(X + group_id * X_offset, input_channels / group, input_height,
                 input_width, kernel_shape[0], kernel_shape[1], dilations[0],
                 dilations[1], pads[0], pads[1], pads[2], pads[3], strides[0],
                 strides[1], col_buffer_data);

      gemm_f32_imp(false, false, filter_num / group, output_image_size,
                   kernel_dim, 1, W + group_id * W_offset, col_buffer_data, 0,
                   Y + group_id * Y_offset);
    }

    if (bias != nullptr) {
      auto Ymatrix =
          Eigen::Map<Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic>>(
              Y, output_image_size, filter_num);
      auto Bvec = Eigen::Map<const Eigen::Matrix<float, Eigen::Dynamic, 1>>(
          bias, filter_num);
      Ymatrix.rowwise() += Bvec.transpose();
    }

    X += X_offset * group;
    Y += Y_offset * group;
  }

  delete[] col_buffer_data;
}

// Some helpers specific to conv operator
void im2col_f32(const float *data_im, const int channels, const int height,
                const int width, const int kernel_h, const int kernel_w,
                const int dilation_h, const int dilation_w, const int pad_t,
                const int pad_l, const int pad_b, const int pad_r,
                const int stride_h, const int stride_w, float *data_col) {
  const int output_h =
      (height + pad_b + pad_t - (dilation_h * (kernel_h - 1) + 1)) / stride_h +
      1;
  const int output_w =
      (width + pad_l + pad_r - (dilation_w * (kernel_w - 1) + 1)) / stride_w +
      1;

  // Fast path for zero padding and no dilation
  // From Torch, THNN_(unfolded_copy)
  if (dilation_h == 1 && dilation_w == 1 && pad_l == 0 && pad_r == 0 &&
      pad_t == 0 && pad_b == 0) {
    for (auto k = 0; k < channels * kernel_h * kernel_w; k++) {
      const auto nip = k / (kernel_h * kernel_w);
      const auto rest = k % (kernel_h * kernel_w);
      const auto kh = rest / kernel_w;
      const auto kw = rest % kernel_w;
      auto *dst = data_col + nip * (kernel_h * kernel_w * output_h * output_w) +
                  kh * (kernel_w * output_h * output_w) +
                  kw * (output_h * output_w);
      const auto *src = data_im + nip * (height * width);
      for (auto y = 0; y < output_h; y++) {
        const auto iy = y * stride_h + kh;
        const auto ix = kw;
        if (stride_w == 1) {
          memcpy(dst + (y * output_w), src + (iy * width + ix),
                 sizeof(float) * output_w);
        } else {
          for (auto x = 0; x < output_w; x++) {
            memcpy(dst + (y * output_w + x),
                   src + (iy * width + ix + x * stride_w), sizeof(float));
          }
        }
      }
    }
    return;
  }

  // Fast path for equal padding
  if (pad_l == pad_r && pad_t == pad_b) {
    // From Intel, https://github.com/BVLC/caffe/pull/3536
    const int pad_h = pad_t;
    const int pad_w = pad_l;
    const int channel_size = height * width;
    for (int channel = channels; channel--; data_im += channel_size) {
      for (int kernel_row = 0; kernel_row < kernel_h; kernel_row++) {
        for (int kernel_col = 0; kernel_col < kernel_w; kernel_col++) {
          int input_row = -pad_h + kernel_row * dilation_h;
          for (int output_rows = output_h; output_rows; output_rows--) {
            if (!is_a_ge_zero_and_a_lt_b(input_row, height)) {
              for (int output_cols = output_w; output_cols; output_cols--) {
                *(data_col++) = 0;
              }
            } else {
              int input_col = -pad_w + kernel_col * dilation_w;
              for (int output_col = output_w; output_col; output_col--) {
                if (is_a_ge_zero_and_a_lt_b(input_col, width)) {
                  *(data_col++) = data_im[input_row * width + input_col];
                } else {
                  *(data_col++) = 0;
                }
                input_col += stride_w;
              }
            }
            input_row += stride_h;
          }
        }
      }
    }
    return;
  }

  // Baseline
  const int dkernel_h = dilation_h * (kernel_h - 1) + 1;
  const int dkernel_w = dilation_w * (kernel_w - 1) + 1;

  int height_col = (height + pad_t + pad_b - dkernel_h) / stride_h + 1;
  int width_col = (width + pad_l + pad_r - dkernel_w) / stride_w + 1;

  int channels_col = channels * kernel_h * kernel_w;
  for (int c = 0; c < channels_col; ++c) {
    int w_offset = c % kernel_w;
    int h_offset = (c / kernel_w) % kernel_h;
    int c_im = c / kernel_h / kernel_w;
    for (int h = 0; h < height_col; ++h) {
      for (int w = 0; w < width_col; ++w) {
        int h_pad = h * stride_h - pad_t + h_offset * dilation_h;
        int w_pad = w * stride_w - pad_l + w_offset * dilation_w;
        if (h_pad >= 0 && h_pad < height && w_pad >= 0 && w_pad < width)
          data_col[(c * height_col + h) * width_col + w] =
              data_im[(c_im * height + h_pad) * width + w_pad];
        else
          data_col[(c * height_col + h) * width_col + w] = 0;
      }
    }
  }
}


/*
 *            Copyright 2009-2020 The VOTCA Development Team
 *                       (http://www.votca.org)
 *
 *      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.
 *
 */

// Standard includes
#include <cmath>

// Third party includes
#include <boost/math/constants/constants.hpp>

// VOTCA includes
#include <votca/tools/constants.h>

// Local VOTCA includes
#include "votca/xtp/sigma_base.h"
#include "votca/xtp/threecenter.h"

namespace votca {
namespace xtp {

Eigen::MatrixXd Sigma_base::CalcExchangeMatrix() const {
  Eigen::MatrixXd result = Eigen::MatrixXd::Zero(_qptotal, _qptotal);
  Index occlevel = _opt.homo - _opt.rpamin + 1;
  Index qpmin = _opt.qpmin - _opt.rpamin;
#pragma omp parallel for schedule(dynamic)
  for (Index gw_level1 = 0; gw_level1 < _qptotal; gw_level1++) {
    const Eigen::MatrixXd& Mmn1 = _Mmn[gw_level1 + qpmin];
    for (Index gw_level2 = gw_level1; gw_level2 < _qptotal; gw_level2++) {
      const Eigen::MatrixXd& Mmn2 = _Mmn[gw_level2 + qpmin];
      double sigma_x =
          -(Mmn1.topRows(occlevel).cwiseProduct(Mmn2.topRows(occlevel))).sum();
      result(gw_level2, gw_level1) = sigma_x;
    }
  }
  result = result.selfadjointView<Eigen::Lower>();
  return result;
}

Eigen::VectorXd Sigma_base::CalcCorrelationDiag(
    const Eigen::VectorXd& frequencies) const {
  Eigen::VectorXd result = Eigen::VectorXd::Zero(_qptotal);
#pragma omp parallel for schedule(dynamic)
  for (Index gw_level = 0; gw_level < _qptotal; gw_level++) {
    result(gw_level) =
        CalcCorrelationDiagElement(gw_level, frequencies[gw_level]);
  }
  return result;
}

Eigen::MatrixXd Sigma_base::CalcCorrelationOffDiag(
    const Eigen::VectorXd& frequencies) const {
  Eigen::MatrixXd result = Eigen::MatrixXd::Zero(_qptotal, _qptotal);
#pragma omp parallel for schedule(dynamic)
  for (Index gw_level1 = 0; gw_level1 < _qptotal; gw_level1++) {
    for (Index gw_level2 = gw_level1 + 1; gw_level2 < _qptotal; gw_level2++) {
      double sigma_c = CalcCorrelationOffDiagElement(
          gw_level1, gw_level2, frequencies[gw_level1], frequencies[gw_level2]);
      result(gw_level2, gw_level1) = sigma_c;
    }
  }
  result = result.selfadjointView<Eigen::Lower>();
  return result;
}

}  // namespace xtp
}  // namespace votca


/**
 * @file
 * @copyright This code is licensed under the 3-clause BSD license.\n
 *            Copyright ETH Zurich, Laboratory for Physical Chemistry, Reiher Group.\n
 *            See LICENSE.txt for details.
 */

#include "Ecqpp.h"
#include <Eigen/QR>
#include <algorithm>
#include <cassert>
#include <iostream>
#include <limits>

namespace Scine {
namespace Utils {

Ecqpp::Ecqpp(const Eigen::MatrixXd& B, const Eigen::VectorXd& E)
  : B_(B), E_(E), dimension_(static_cast<unsigned>(E.size())) {
  assert(B.rows() == dimension_ && B.cols() == dimension_ && "B matrix and E vector do not have compatible dimensions.");
}

Eigen::VectorXd Ecqpp::calculateOptimalCoefficients() {
  solveAllConstrainedProblems();
  return solution_;
}

void Ecqpp::solveAllConstrainedProblems() {
  bestSolutionEnergy_ = std::numeric_limits<double>::max();
  for (unsigned i = 0; i < dimension_; ++i)
    solveAllConstrainedProblemsForNumberZeros(i);
}

void Ecqpp::solveAllConstrainedProblemsForNumberZeros(unsigned int numberZeros) {
  std::vector<bool> indexConsidered(dimension_, true);
  for (unsigned i = 0; i < numberZeros; ++i)
    indexConsidered[i] = false;

  do {
    generatePreviousIndexesVector(indexConsidered, numberZeros);
    generateReducedObjects();
    solveConstrainedProblem();
    if (solutionIsValid())
      addSolution();
  } while (std::next_permutation(indexConsidered.begin(), indexConsidered.end()));
}

void Ecqpp::generatePreviousIndexesVector(const std::vector<bool>& consideredIndexes, unsigned numberZeros) {
  unsigned writeIndex = 0;
  unsigned newMatrixDimension = dimension_ - numberZeros;
  previousIndexes_.resize(newMatrixDimension);
  for (unsigned readIndex = 0; readIndex < dimension_; ++readIndex) {
    if (consideredIndexes[readIndex]) {
      previousIndexes_[writeIndex] = readIndex;
      ++writeIndex;
    }
  }
}

void Ecqpp::generateReducedObjects() {
  auto reducedDimension = static_cast<unsigned>(previousIndexes_.size());

  reducedB_.resize(reducedDimension, reducedDimension);
  reducedE_.resize(reducedDimension);

  for (unsigned i = 0; i < reducedDimension; ++i) {
    reducedE_(i) = E_(previousIndexes_[i]);
    for (unsigned j = 0; j < reducedDimension; ++j) {
      reducedB_(i, j) = B_(previousIndexes_[i], previousIndexes_[j]);
    }
  }
}

void Ecqpp::solveConstrainedProblem() {
  auto rDim = reducedB_.cols();

  Eigen::MatrixXd M(rDim + 1, rDim + 1);
  Eigen::VectorXd b(rDim + 1);

  M.block(0, 0, rDim, rDim) = reducedB_;
  M.block(rDim, 0, 1, rDim).setOnes();
  M.block(0, rDim, rDim, 1).setOnes();
  M(rDim, rDim) = 0;
  b.head(rDim) = reducedE_;
  b(rDim) = 1;

  Eigen::VectorXd solution = M.colPivHouseholderQr().solve(b);
  reducedSolution_ = solution.head(rDim);
}

bool Ecqpp::solutionIsValid() const {
  for (unsigned i = 0; i < reducedSolution_.size(); ++i)
    if (reducedSolution_[i] < 0)
      return false;
  return true;
}

void Ecqpp::addSolution() {
  generateSolutionFromReducedSolution();
  setBestSolutionIfHasLowerEnergy();
}

void Ecqpp::generateSolutionFromReducedSolution() {
  fullSolution_ = Eigen::VectorXd::Zero(dimension_);
  for (unsigned i = 0; i < previousIndexes_.size(); ++i) {
    fullSolution_[previousIndexes_[i]] = reducedSolution_[i];
  }
}

void Ecqpp::setBestSolutionIfHasLowerEnergy() {
  double currentEnergy = E_.dot(fullSolution_) - 0.5 * fullSolution_.transpose() * B_ * fullSolution_;
  if (currentEnergy < bestSolutionEnergy_) {
    solution_ = fullSolution_;
    bestSolutionEnergy_ = currentEnergy;
  }
}
} // namespace Utils
} // namespace Scine


#pragma once

#include <boost/progress.hpp>
#include <boost/filesystem.hpp>
#include <unordered_map>

#include "dint_configuration.hpp"
#include "statistics_collectors.hpp"
#include "binary_blocks_collection.hpp"
#include "hash_utils.hpp"
#include "util.hpp"

namespace ds2i {

static const double codeword_bits = std::log2(constants::num_entries);
static const double initial_bpi = 3 * codeword_bits;
static const double eps = 0.0001;

double cost(uint32_t block_size, uint32_t block_frequency) {
    return block_frequency * (initial_bpi * block_size - codeword_bits);
}

double compute_saving(uint32_t block_size, uint32_t block_frequency,
                      uint64_t total_integers) {
    return cost(block_size, block_frequency) / total_integers;
};

struct cost_filter {
    cost_filter(double threshold = eps) : m_threshold(threshold) {}

    bool operator()(block_type const& block, uint64_t total_integers) const {
        return compute_saving(block.data.size(), block.freq, total_integers) >
               m_threshold;
    }

private:
    double m_threshold;
};

template <typename Dictionary, typename Statistics>
struct decreasing_static_frequencies {
    typedef Dictionary dictionary_type;
    typedef Statistics statistics_type;

    static std::string type() {
        return "DSF-" + std::to_string(dictionary_type::num_entries) + "-" +
               std::to_string(dictionary_type::max_entry_size);
    }

    static auto filter() {
        cost_filter filter(eps / 1000);
        return filter;
    }

    static void build(typename dictionary_type::builder& dict_builder,
                      statistics_type& stats) {
        logger() << "building " << type() << " dictionary for "
                 << stats.total_integers << " integers" << std::endl;

        dict_builder.init();
        for (uint64_t s = 0; s != stats.blocks.size(); ++s) {
            uint64_t n = dictionary_type::num_entries;
            if (stats.blocks[s].size() < n) {
                n = stats.blocks[s].size();
            }

            auto it = stats.blocks[s].begin();
            for (uint64_t i = 0; i != n; ++i, ++it) {
                auto const& block = *it;
                dict_builder.append(block.data.data(), block.data.size(), s);
            }
        }

        dict_builder.build();
    }
};
}  // namespace ds2i