/*
 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.
*/

#include <ored/model/crcirbuilder.hpp>
#include <ored/model/crlgmbuilder.hpp>
#include <ored/model/crossassetmodelbuilder.hpp>
#include <ored/model/eqbsbuilder.hpp>
#include <ored/model/fxbsbuilder.hpp>
#include <ored/model/inflation/infdkbuilder.hpp>
#include <ored/model/inflation/infjybuilder.hpp>
#include <ored/model/inflation/infjydata.hpp>
#include <ored/model/lgmbuilder.hpp>
#include <ored/model/structuredmodelerror.hpp>
#include <ored/model/utilities.hpp>
#include <ored/utilities/correlationmatrix.hpp>
#include <ored/utilities/log.hpp>
#include <ored/utilities/parsers.hpp>

#include <qle/cashflows/jyyoyinflationcouponpricer.hpp>
#include <qle/models/cpicapfloorhelper.hpp>
#include <qle/models/fxbsconstantparametrization.hpp>
#include <qle/models/fxbspiecewiseconstantparametrization.hpp>
#include <qle/models/fxeqoptionhelper.hpp>
#include <qle/models/irlgm1fconstantparametrization.hpp>
#include <qle/models/irlgm1fpiecewiseconstanthullwhiteadaptor.hpp>
#include <qle/models/irlgm1fpiecewiseconstantparametrization.hpp>
#include <qle/models/irlgm1fpiecewiselinearparametrization.hpp>
#include <qle/models/yoycapfloorhelper.hpp>
#include <qle/models/yoyswaphelper.hpp>
#include <qle/pricingengines/analyticcclgmfxoptionengine.hpp>
#include <qle/pricingengines/analyticdkcpicapfloorengine.hpp>
#include <qle/pricingengines/analyticjycpicapfloorengine.hpp>
#include <qle/pricingengines/analyticjyyoycapfloorengine.hpp>
#include <qle/pricingengines/analyticlgmswaptionengine.hpp>
#include <qle/pricingengines/analyticxassetlgmeqoptionengine.hpp>

#include <ql/math/optimization/levenbergmarquardt.hpp>
#include <ql/models/shortrate/calibrationhelpers/swaptionhelper.hpp>
#include <ql/pricingengines/swap/discountingswapengine.hpp>
#include <ql/quotes/simplequote.hpp>
#include <ql/utilities/dataformatters.hpp>

#include <boost/algorithm/string/case_conv.hpp>
#include <boost/lexical_cast.hpp>

using QuantExt::AnalyticJyCpiCapFloorEngine;
using QuantExt::AnalyticJyYoYCapFloorEngine;
using QuantExt::CpiCapFloorHelper;
using QuantExt::InfDkParametrization;
using QuantExt::IrLgm1fParametrization;
using QuantExt::JyYoYInflationCouponPricer;
using QuantExt::YoYCapFloorHelper;
using QuantExt::YoYSwapHelper;
using QuantLib::DiscountingSwapEngine;
using std::vector;

namespace ore {
namespace data {

CrossAssetModelBuilder::CrossAssetModelBuilder(
    const boost::shared_ptr<ore::data::Market>& market, const boost::shared_ptr<CrossAssetModelData>& config,
    const std::string& configurationLgmCalibration, const std::string& configurationFxCalibration,
    const std::string& configurationEqCalibration, const std::string& configurationInfCalibration,
    const std::string& configurationCrCalibration, const std::string& configurationFinalModel,
    const DayCounter& dayCounter, const bool dontCalibrate, const bool continueOnError,
    const std::string& referenceCalibrationGrid)
    : market_(market), config_(config), configurationLgmCalibration_(configurationLgmCalibration),
      configurationFxCalibration_(configurationFxCalibration), configurationEqCalibration_(configurationEqCalibration),
      configurationInfCalibration_(configurationInfCalibration),
      configurationCrCalibration_(configurationCrCalibration), configurationFinalModel_(configurationFinalModel),
      dayCounter_(dayCounter), dontCalibrate_(dontCalibrate), continueOnError_(continueOnError),
      referenceCalibrationGrid_(referenceCalibrationGrid),
      optimizationMethod_(boost::shared_ptr<OptimizationMethod>(new LevenbergMarquardt(1E-8, 1E-8, 1E-8))),
      endCriteria_(EndCriteria(1000, 500, 1E-8, 1E-8, 1E-8)) {
    buildModel();
    registerWithSubBuilders();
    // register market observer with correlations
    marketObserver_ = boost::make_shared<MarketObserver>();
    for (auto const& c : config->correlations())
        marketObserver_->addObservable(c.second);
    // reset market observer's updated flag
    marketObserver_->hasUpdated(true);
}

Handle<QuantExt::CrossAssetModel> CrossAssetModelBuilder::model() const {
    calculate();
    return model_;
}

const std::vector<Real>& CrossAssetModelBuilder::swaptionCalibrationErrors() {
    calculate();
    return swaptionCalibrationErrors_;
}
const std::vector<Real>& CrossAssetModelBuilder::fxOptionCalibrationErrors() {
    calculate();
    return fxOptionCalibrationErrors_;
}
const std::vector<Real>& CrossAssetModelBuilder::eqOptionCalibrationErrors() {
    calculate();
    return eqOptionCalibrationErrors_;
}
const std::vector<Real>& CrossAssetModelBuilder::inflationCalibrationErrors() {
    calculate();
    return inflationCalibrationErrors_;
}

void CrossAssetModelBuilder::unregisterWithSubBuilders() {
    for (auto okv : subBuilders_)
        for (auto ikv : okv.second)
            unregisterWith(ikv.second);
}

void CrossAssetModelBuilder::registerWithSubBuilders() {
    for (auto okv : subBuilders_)
        for (auto ikv : okv.second)
            registerWith(ikv.second);
}

bool CrossAssetModelBuilder::requiresRecalibration() const {
    for (auto okv : subBuilders_)
        for (auto ikv : okv.second)
            if (ikv.second->requiresRecalibration())
                return true;

    return marketObserver_->hasUpdated(false);
}

void CrossAssetModelBuilder::performCalculations() const {
    // if any of the sub models requires a recalibration, we rebuilt the model
    // TODO we could do this more selectively
    if (!dontCalibrate_ && requiresRecalibration()) {
        // reset market observer update flag
        marketObserver_->hasUpdated(true);
        // the cast is a bit ugly, but we pretty much know what we are doing here
        const_cast<CrossAssetModelBuilder*>(this)->unregisterWithSubBuilders();
        buildModel();
        const_cast<CrossAssetModelBuilder*>(this)->registerWithSubBuilders();
    }
}

void CrossAssetModelBuilder::buildModel() const {

    QL_REQUIRE(market_ != NULL, "CrossAssetModelBuilder: no market given");
    LOG("Start building CrossAssetModel");
    DLOG("configurations: LgmCalibration "
         << configurationLgmCalibration_ << ", FxCalibration " << configurationFxCalibration_ << ", EqCalibration "
         << configurationEqCalibration_ << ", InfCalibration " << configurationInfCalibration_ << ", CrCalibration"
         << configurationCrCalibration_ << ", FinalModel " << configurationFinalModel_);
    if (dontCalibrate_) {
        DLOG("Calibration of the model is disabled.");
    }

    QL_REQUIRE(config_->irConfigs().size() > 0, "missing IR configurations");
    QL_REQUIRE(config_->irConfigs().size() == config_->fxConfigs().size() + 1,
               "FX configuration size " << config_->fxConfigs().size() << " inconsisitent with IR configuration size "
                                        << config_->irConfigs().size());

    swaptionBaskets_.resize(config_->irConfigs().size());
    optionExpiries_.resize(config_->irConfigs().size());
    swaptionMaturities_.resize(config_->irConfigs().size());
    swaptionCalibrationErrors_.resize(config_->irConfigs().size());
    fxOptionBaskets_.resize(config_->fxConfigs().size());
    fxOptionExpiries_.resize(config_->fxConfigs().size());
    fxOptionCalibrationErrors_.resize(config_->fxConfigs().size());
    eqOptionBaskets_.resize(config_->eqConfigs().size());
    eqOptionExpiries_.resize(config_->eqConfigs().size());
    eqOptionCalibrationErrors_.resize(config_->eqConfigs().size());
    inflationCalibrationErrors_.resize(config_->infConfigs().size());

    subBuilders_.clear();

    // Store information on the number of factors for each process. This is used when requesting a correlation matrix
    // from the CorrelationMatrixBuilder below.
    using ProcessInfo = CorrelationMatrixBuilder::ProcessInfo;
    namespace CT = QuantExt::CrossAssetModelTypes;
    ProcessInfo processInfo;

    /*******************************************************
     * Build the IR parametrizations and calibration baskets
     */
    std::vector<boost::shared_ptr<QuantExt::IrLgm1fParametrization>> irParametrizations;
    std::vector<RelinkableHandle<YieldTermStructure>> irDiscountCurves;
    std::vector<std::string> currencies, regions, crNames, eqNames, infIndices;
    std::vector<boost::shared_ptr<LgmBuilder>> irBuilder;

    for (Size i = 0; i < config_->irConfigs().size(); i++) {
        boost::shared_ptr<IrLgmData> ir = config_->irConfigs()[i];
        DLOG("IR Parametrization " << i << " ccy " << ir->ccy());
        boost::shared_ptr<LgmBuilder> builder =
            boost::make_shared<LgmBuilder>(market_, ir, configurationLgmCalibration_, config_->bootstrapTolerance(),
                                           continueOnError_, referenceCalibrationGrid_);
        if (dontCalibrate_)
            builder->freeze();
        irBuilder.push_back(builder);
        boost::shared_ptr<QuantExt::IrLgm1fParametrization> parametrization = builder->parametrization();
        swaptionBaskets_[i] = builder->swaptionBasket();
        currencies.push_back(ir->ccy());
        irParametrizations.push_back(parametrization);
        irDiscountCurves.push_back(builder->discountCurve());
        subBuilders_[CT::IR][i] = builder;
        processInfo[CT::IR].emplace_back(ir->ccy(), 1);
    }

    QL_REQUIRE(irParametrizations.size() > 0, "missing IR parametrizations");

    QuantLib::Currency domesticCcy = irParametrizations[0]->currency();

    /*******************************************************
     * Build the FX parametrizations and calibration baskets
     */
    std::vector<boost::shared_ptr<QuantExt::FxBsParametrization>> fxParametrizations;
    for (Size i = 0; i < config_->fxConfigs().size(); i++) {
        DLOG("FX Parametrization " << i);
        boost::shared_ptr<FxBsData> fx = config_->fxConfigs()[i];
        QuantLib::Currency ccy = ore::data::parseCurrency(fx->foreignCcy());
        QuantLib::Currency domCcy = ore::data::parseCurrency(fx->domesticCcy());

        QL_REQUIRE(ccy.code() == irParametrizations[i + 1]->currency().code(),
                   "FX parametrization currency[" << i << "]=" << ccy << " does not match IR currrency[" << i + 1
                                                  << "]=" << irParametrizations[i + 1]->currency().code());

        QL_REQUIRE(domCcy == domesticCcy, "FX parametrization [" << i << "]=" << ccy << "/" << domCcy
                                                                 << " does not match domestic ccy " << domesticCcy);

        boost::shared_ptr<FxBsBuilder> builder =
            boost::make_shared<FxBsBuilder>(market_, fx, configurationFxCalibration_, referenceCalibrationGrid_);
        boost::shared_ptr<QuantExt::FxBsParametrization> parametrization = builder->parametrization();

        fxOptionBaskets_[i] = builder->optionBasket();
        fxParametrizations.push_back(parametrization);
        subBuilders_[CT::FX][i] = builder;
        processInfo[CT::FX].emplace_back(ccy.code() + domCcy.code(), 1);
    }

    /*******************************************************
     * Build the EQ parametrizations and calibration baskets
     */
    std::vector<boost::shared_ptr<QuantExt::EqBsParametrization>> eqParametrizations;
    for (Size i = 0; i < config_->eqConfigs().size(); i++) {
        DLOG("EQ Parametrization " << i);
        boost::shared_ptr<EqBsData> eq = config_->eqConfigs()[i];
        string eqName = eq->eqName();
        QuantLib::Currency eqCcy = ore::data::parseCurrency(eq->currency());
        QL_REQUIRE(std::find(currencies.begin(), currencies.end(), eqCcy.code()) != currencies.end(),
                   "Currency (" << eqCcy << ") for equity " << eqName << " not covered by CrossAssetModelData");
        boost::shared_ptr<EqBsBuilder> builder = boost::make_shared<EqBsBuilder>(
            market_, eq, domesticCcy, configurationEqCalibration_, referenceCalibrationGrid_);
        boost::shared_ptr<QuantExt::EqBsParametrization> parametrization = builder->parametrization();
        eqOptionBaskets_[i] = builder->optionBasket();
        eqParametrizations.push_back(parametrization);
        eqNames.push_back(eqName);
        subBuilders_[CT::EQ][i] = builder;
        processInfo[CT::EQ].emplace_back(eqName, 1);
    }

    // Build the INF parametrizations and calibration baskets
    vector<boost::shared_ptr<Parametrization>> infParameterizations;
    for (Size i = 0; i < config_->infConfigs().size(); i++) {
        boost::shared_ptr<InflationModelData> imData = config_->infConfigs()[i];
        DLOG("Inflation parameterisation (" << i << ") for index " << imData->index());
        if (auto dkData = boost::dynamic_pointer_cast<InfDkData>(imData)) {
            boost::shared_ptr<InfDkBuilder> builder = boost::make_shared<InfDkBuilder>(
                market_, dkData, configurationInfCalibration_, referenceCalibrationGrid_, dontCalibrate_);
            if (dontCalibrate_)
                builder->freeze();
            infParameterizations.push_back(builder->parametrization());
            subBuilders_[CT::INF][i] = builder;
            processInfo[CT::INF].emplace_back(dkData->index(), 1);
        } else if (auto jyData = boost::dynamic_pointer_cast<InfJyData>(imData)) {
            boost::shared_ptr<InfJyBuilder> builder = boost::make_shared<InfJyBuilder>(
                market_, jyData, configurationInfCalibration_, referenceCalibrationGrid_);
            infParameterizations.push_back(builder->parameterization());
            subBuilders_[CT::INF][i] = builder;
            processInfo[CT::INF].emplace_back(jyData->index(), 2);
        } else {
            QL_FAIL("CrossAssetModelBuilder expects either DK or JY inflation model data.");
        }
        infIndices.push_back(imData->index());
    }

    /*******************************************************
     * Build the CR parametrizations and calibration baskets
     */
    // LGM (if any)
    std::vector<boost::shared_ptr<QuantExt::CrLgm1fParametrization>> crLgmParametrizations;
    for (Size i = 0; i < config_->crLgmConfigs().size(); ++i) {
        LOG("CR LGM Parametrization " << i);
        boost::shared_ptr<CrLgmData> cr = config_->crLgmConfigs()[i];
        string crName = cr->name();
        boost::shared_ptr<CrLgmBuilder> builder =
            boost::make_shared<CrLgmBuilder>(market_, cr, configurationCrCalibration_);
        boost::shared_ptr<QuantExt::CrLgm1fParametrization> parametrization = builder->parametrization();
        crLgmParametrizations.push_back(parametrization);
        crNames.push_back(crName);
        subBuilders_[CT::CR][i] = builder;
        processInfo[CT::CR].emplace_back(crName, 1);
    }

    // CIR (if any)
    std::vector<boost::shared_ptr<QuantExt::CrCirppParametrization>> crCirParametrizations;
    for (Size i = 0; i < config_->crCirConfigs().size(); ++i) {
        LOG("CR CIR Parametrization " << i);
        boost::shared_ptr<CrCirData> cr = config_->crCirConfigs()[i];
        string crName = cr->name();
        boost::shared_ptr<CrCirBuilder> builder =
            boost::make_shared<CrCirBuilder>(market_, cr, configurationCrCalibration_);
        boost::shared_ptr<QuantExt::CrCirppParametrization> parametrization = builder->parametrization();
        crCirParametrizations.push_back(parametrization);
        crNames.push_back(crName);
        subBuilders_[CT::CR][i] = builder;
        processInfo[CT::CR].emplace_back(crName, 1);
    }

    std::vector<boost::shared_ptr<QuantExt::Parametrization>> parametrizations;
    for (Size i = 0; i < irParametrizations.size(); i++)
        parametrizations.push_back(irParametrizations[i]);
    for (Size i = 0; i < fxParametrizations.size(); i++)
        parametrizations.push_back(fxParametrizations[i]);
    for (Size i = 0; i < eqParametrizations.size(); i++)
        parametrizations.push_back(eqParametrizations[i]);
    parametrizations.insert(parametrizations.end(), infParameterizations.begin(), infParameterizations.end());
    for (Size i = 0; i < crLgmParametrizations.size(); i++)
        parametrizations.push_back(crLgmParametrizations[i]);
    for (Size i = 0; i < crCirParametrizations.size(); i++)
        parametrizations.push_back(crCirParametrizations[i]);

    QL_REQUIRE(fxParametrizations.size() == irParametrizations.size() - 1, "mismatch in IR/FX parametrization sizes");

    Measure::Type measure = Measure::LGM;
    if (config_->measure() == "BA") {
        measure = Measure::BA;
        DLOG("Setting measure to BA");
    } else if (config_->measure() == "LGM") {
        measure = Measure::LGM;
        DLOG("Setting measure to BA");
    } else if (config_->measure() == "") {
        DLOG("Defaulting to LGM measure");
    } else {
        QL_FAIL("Measure " << config_->measure() << " not recognized");
    }

    // Tag on the parametrization and process info for the auxiliary state variable in the bank account measure
    if (measure == Measure::BA) {
        parametrizations.push_back(
            irParametrizations[0]); // FIXME: Is index 0 safe to reference the domestic IR parameters?
        processInfo[CT::AUX].emplace_back(config_->domesticCurrency(), 1);
    }

    /******************************
     * Build the correlation matrix
     */
    DLOG("CrossAssetModelBuilder: adding correlations.");
    CorrelationMatrixBuilder cmb;

    // Perfect instantaneous correlation of auxiliary variable and domestic LGM state variable in the bank account
    // measure
    CorrelationFactor domesticFactorIR = {CrossAssetModelTypes::IR, config_->domesticCurrency(), 0};
    CorrelationFactor domesticFactorAUX = {CrossAssetModelTypes::AUX, config_->domesticCurrency(), 0};
    if (measure == Measure::BA)
        cmb.addCorrelation(domesticFactorAUX, domesticFactorIR, 1.0);

    for (auto it = config_->correlations().begin(); it != config_->correlations().end(); it++) {
        cmb.addCorrelation(it->first.first, it->first.second, it->second);
        if (measure == Measure::BA) {
            // Copy correlation(domesticIR, other) to correlation(domesticAUX, other)
            if (it->first.first == domesticFactorIR)
                cmb.addCorrelation(domesticFactorAUX, it->first.second, it->second);
            if (it->first.second == domesticFactorIR)
                cmb.addCorrelation(it->first.first, domesticFactorAUX, it->second);
        }
    }

    Matrix corrMatrix = cmb.correlationMatrix(processInfo);

    TLOG("CAM correlation matrix:");
    TLOGGERSTREAM << corrMatrix;

    /*****************************
     * Build the cross asset model
     */

    SalvagingAlgorithm::Type salvaging = SalvagingAlgorithm::None;

    model_.linkTo(boost::make_shared<QuantExt::CrossAssetModel>(parametrizations, corrMatrix, salvaging, measure));

    /*************************
     * Calibrate IR components
     */

    for (Size i = 0; i < irBuilder.size(); i++) {
        DLOG("IR Calibration " << i);
        swaptionCalibrationErrors_[i] = irBuilder[i]->error();
    }

    /*************************
     * Relink LGM discount curves to curves used for FX calibration
     */

    for (Size i = 0; i < irParametrizations.size(); i++) {
        auto p = irParametrizations[i];
        irDiscountCurves[i].linkTo(*market_->discountCurve(p->currency().code(), configurationFxCalibration_));
        DLOG("Relinked discounting curve for " << p->currency().code() << " for FX calibration");
    }

    /*************************
     * Calibrate FX components
     */

    for (Size i = 0; i < fxParametrizations.size(); i++) {
        boost::shared_ptr<FxBsData> fx = config_->fxConfigs()[i];

        if (fx->calibrationType() == CalibrationType::None || !fx->calibrateSigma()) {
            DLOG("FX Calibration " << i << " skipped");
            continue;
        }

        DLOG("FX Calibration " << i);

        // attach pricing engines to helpers
        boost::shared_ptr<QuantExt::AnalyticCcLgmFxOptionEngine> engine =
            boost::make_shared<QuantExt::AnalyticCcLgmFxOptionEngine>(*model_, i);
        // enable caching for calibration
        // TODO: review this
        engine->cache(true);
        for (Size j = 0; j < fxOptionBaskets_[i].size(); j++)
            fxOptionBaskets_[i][j]->setPricingEngine(engine);

        if (!dontCalibrate_) {

            if (fx->calibrationType() == CalibrationType::Bootstrap && fx->sigmaParamType() == ParamType::Piecewise)
                model_->calibrateBsVolatilitiesIterative(CrossAssetModelTypes::FX, i, fxOptionBaskets_[i],
                                                         *optimizationMethod_, endCriteria_);
            else
                model_->calibrateBsVolatilitiesGlobal(CrossAssetModelTypes::FX, i, fxOptionBaskets_[i],
                                                      *optimizationMethod_, endCriteria_);

            DLOG("FX " << fx->foreignCcy() << " calibration errors:");
            fxOptionCalibrationErrors_[i] = getCalibrationError(fxOptionBaskets_[i]);
            if (fx->calibrationType() == CalibrationType::Bootstrap) {
                if (fabs(fxOptionCalibrationErrors_[i]) < config_->bootstrapTolerance()) {
                    // we check the log level here to avoid unncessary computations
                    if (Log::instance().filter(ORE_DATA)) {
                        TLOGGERSTREAM << "Calibration details:";
                        TLOGGERSTREAM << getCalibrationDetails(fxOptionBaskets_[i], fxParametrizations[i],
                                                               irParametrizations[0]);
                        TLOGGERSTREAM << "rmse = " << fxOptionCalibrationErrors_[i];
                    }
                } else {
                    std::string exceptionMessage = "FX BS " + std::to_string(i) + " calibration error " +
                                                   std::to_string(fxOptionCalibrationErrors_[i]) +
                                                   " exceeds tolerance " +
                                                   std::to_string(config_->bootstrapTolerance());
                    WLOG(StructuredModelErrorMessage("Failed to calibrate FX BS Model", exceptionMessage));
                    WLOGGERSTREAM << "Calibration details:";
                    WLOGGERSTREAM << getCalibrationDetails(fxOptionBaskets_[i], fxParametrizations[i],
                                                           irParametrizations[0]);
                    WLOGGERSTREAM << "rmse = " << fxOptionCalibrationErrors_[i];
                    if (!continueOnError_)
                        QL_FAIL(exceptionMessage);
                }
            }
        }
    }

    /*************************
     * Relink LGM discount curves to curves used for EQ calibration
     */

    for (Size i = 0; i < irParametrizations.size(); i++) {
        auto p = irParametrizations[i];
        irDiscountCurves[i].linkTo(*market_->discountCurve(p->currency().code(), configurationEqCalibration_));
        DLOG("Relinked discounting curve for " << p->currency().code() << " for EQ calibration");
    }

    /*************************
     * Calibrate EQ components
     */

    for (Size i = 0; i < eqParametrizations.size(); i++) {
        boost::shared_ptr<EqBsData> eq = config_->eqConfigs()[i];
        if (!eq->calibrateSigma()) {
            DLOG("EQ Calibration " << i << " skipped");
            continue;
        }
        DLOG("EQ Calibration " << i);
        // attach pricing engines to helpers
        Currency eqCcy = eqParametrizations[i]->currency();
        Size eqCcyIdx = model_->ccyIndex(eqCcy);
        boost::shared_ptr<QuantExt::AnalyticXAssetLgmEquityOptionEngine> engine =
            boost::make_shared<QuantExt::AnalyticXAssetLgmEquityOptionEngine>(*model_, i, eqCcyIdx);
        for (Size j = 0; j < eqOptionBaskets_[i].size(); j++)
            eqOptionBaskets_[i][j]->setPricingEngine(engine);

        if (!dontCalibrate_) {

            if (eq->calibrationType() == CalibrationType::Bootstrap && eq->sigmaParamType() == ParamType::Piecewise)
                model_->calibrateBsVolatilitiesIterative(CrossAssetModelTypes::EQ, i, eqOptionBaskets_[i],
                                                         *optimizationMethod_, endCriteria_);
            else
                model_->calibrateBsVolatilitiesGlobal(CrossAssetModelTypes::EQ, i, eqOptionBaskets_[i],
                                                      *optimizationMethod_, endCriteria_);
            DLOG("EQ " << eq->eqName() << " calibration errors:");
            eqOptionCalibrationErrors_[i] = getCalibrationError(eqOptionBaskets_[i]);
            if (eq->calibrationType() == CalibrationType::Bootstrap) {
                if (fabs(eqOptionCalibrationErrors_[i]) < config_->bootstrapTolerance()) {
                    // we check the log level here to avoid unncessary computations
                    if (Log::instance().filter(ORE_DATA)) {
                        TLOGGERSTREAM << "Calibration details:";
                        TLOGGERSTREAM << getCalibrationDetails(eqOptionBaskets_[i], eqParametrizations[i],
                                                               irParametrizations[0]);
                        TLOGGERSTREAM << "rmse = " << eqOptionCalibrationErrors_[i];
                    }
                } else {
                    std::string exceptionMessage = "EQ BS " + std::to_string(i) + " calibration error " +
                                                   std::to_string(eqOptionCalibrationErrors_[i]) +
                                                   " exceeds tolerance " +
                                                   std::to_string(config_->bootstrapTolerance());
                    WLOG(StructuredModelErrorMessage("Failed to calibrate EQ BS Model", exceptionMessage));
                    WLOGGERSTREAM << "Calibration details:";
                    WLOGGERSTREAM << getCalibrationDetails(eqOptionBaskets_[i], eqParametrizations[i],
                                                           irParametrizations[0]);
                    WLOGGERSTREAM << "rmse = " << eqOptionCalibrationErrors_[i];
                    if (!continueOnError_)
                        QL_FAIL(exceptionMessage);
                }
            }
        }
    }

    /*************************
     * Relink LGM discount curves to curves used for INF calibration
     */

    for (Size i = 0; i < irParametrizations.size(); i++) {
        auto p = irParametrizations[i];
        irDiscountCurves[i].linkTo(*market_->discountCurve(p->currency().code(), configurationInfCalibration_));
        DLOG("Relinked discounting curve for " << p->currency().code() << " for INF calibration");
    }

    // Calibrate INF components
    for (Size i = 0; i < infParameterizations.size(); i++) {
        boost::shared_ptr<InflationModelData> imData = config_->infConfigs()[i];
        if (auto dkData = boost::dynamic_pointer_cast<InfDkData>(imData)) {
            auto dkParam = boost::dynamic_pointer_cast<InfDkParametrization>(infParameterizations[i]);
            QL_REQUIRE(dkParam, "Expected DK model data to have given a DK parameterisation.");
            const auto& builder = subBuilders_.at(CT::INF).at(i);
            const auto& dkBuilder = boost::dynamic_pointer_cast<InfDkBuilder>(builder);
            calibrateInflation(*dkData, i, dkBuilder->optionBasket(), dkParam);
        } else if (auto jyData = boost::dynamic_pointer_cast<InfJyData>(imData)) {
            auto jyParam = boost::dynamic_pointer_cast<InfJyParameterization>(infParameterizations[i]);
            QL_REQUIRE(jyParam, "Expected JY model data to have given a JY parameterisation.");
            const auto& builder = subBuilders_.at(CT::INF).at(i);
            const auto& jyBuilder = boost::dynamic_pointer_cast<InfJyBuilder>(builder);
            calibrateInflation(*jyData, i, jyBuilder, jyParam);
        } else {
            QL_FAIL("CrossAssetModelBuilder expects either DK or JY inflation model data.");
        }
    }

    /*************************
     * Relink LGM discount curves to final model curves
     */

    for (Size i = 0; i < irParametrizations.size(); i++) {
        auto p = irParametrizations[i];
        irDiscountCurves[i].linkTo(*market_->discountCurve(p->currency().code(), configurationFinalModel_));
        DLOG("Relinked discounting curve for " << p->currency().code() << " as final model curves");
    }

    // play safe (although the cache of the model should be empty at
    // this point from all what we know...)
    model_->update();

    DLOG("Building CrossAssetModel done");
}

void CrossAssetModelBuilder::forceRecalculate() {
    forceCalibration_ = true;
    ModelBuilder::forceRecalculate();
    forceCalibration_ = false;
}

void CrossAssetModelBuilder::calibrateInflation(const InfDkData& data, Size modelIdx,
                                                const vector<boost::shared_ptr<BlackCalibrationHelper>>& cb,
                                                const boost::shared_ptr<InfDkParametrization>& inflationParam) const {

    LOG("Calibrate DK inflation model for inflation index " << data.index());

    if ((!data.volatility().calibrate() && !data.reversion().calibrate()) ||
        (data.calibrationType() == CalibrationType::None)) {
        LOG("Calibration of DK inflation model for inflation index " << data.index() << " not requested.");
        return;
    }

    Handle<ZeroInflationIndex> zInfIndex =
        market_->zeroInflationIndex(model_->infdk(modelIdx)->name(), configurationInfCalibration_);
    Real baseCPI = dontCalibrate_ ? 100. : zInfIndex->fixing(zInfIndex->zeroInflationTermStructure()->baseDate());
    auto engine = boost::make_shared<QuantExt::AnalyticDkCpiCapFloorEngine>(*model_, modelIdx, baseCPI);
    for (Size j = 0; j < cb.size(); j++)
        cb[j]->setPricingEngine(engine);

    if (dontCalibrate_)
        return;

    if (data.volatility().calibrate() && !data.reversion().calibrate()) {
        if (data.calibrationType() == CalibrationType::Bootstrap && data.volatility().type() == ParamType::Piecewise) {
            model_->calibrateInfDkVolatilitiesIterative(modelIdx, cb, *optimizationMethod_, endCriteria_);
        } else {
            model_->calibrateInfDkVolatilitiesGlobal(modelIdx, cb, *optimizationMethod_, endCriteria_);
        }
    } else if (!data.volatility().calibrate() && data.reversion().calibrate()) {
        if (data.calibrationType() == CalibrationType::Bootstrap && data.reversion().type() == ParamType::Piecewise) {
            model_->calibrateInfDkReversionsIterative(modelIdx, cb, *optimizationMethod_, endCriteria_);
        } else {
            model_->calibrateInfDkReversionsGlobal(modelIdx, cb, *optimizationMethod_, endCriteria_);
        }
    } else {
        model_->calibrate(cb, *optimizationMethod_, endCriteria_);
    }

    DLOG("INF (DK) " << data.index() << " calibration errors:");
    inflationCalibrationErrors_[modelIdx] = getCalibrationError(cb);
    if (data.calibrationType() == CalibrationType::Bootstrap) {
        if (fabs(inflationCalibrationErrors_[modelIdx]) < config_->bootstrapTolerance()) {
            // we check the log level here to avoid unncessary computations
            if (Log::instance().filter(ORE_DATA)) {
                TLOGGERSTREAM << "Calibration details:";
                TLOGGERSTREAM << getCalibrationDetails(cb, inflationParam);
                TLOGGERSTREAM << "rmse = " << inflationCalibrationErrors_[modelIdx];
            }
        } else {
            string exceptionMessage = "INF (DK) " + std::to_string(modelIdx) + " calibration error " +
                                      std::to_string(inflationCalibrationErrors_[modelIdx]) + " exceeds tolerance " +
                                      std::to_string(config_->bootstrapTolerance());
            WLOG(StructuredModelErrorMessage("Failed to calibrate INF DK Model", exceptionMessage));
            WLOGGERSTREAM << "Calibration details:";
            WLOGGERSTREAM << getCalibrationDetails(cb, inflationParam);
            WLOGGERSTREAM << "rmse = " << inflationCalibrationErrors_[modelIdx];
            if (!continueOnError_)
                QL_FAIL(exceptionMessage);
        }
    }
}

void CrossAssetModelBuilder::calibrateInflation(const InfJyData& data, Size modelIdx,
                                                const boost::shared_ptr<InfJyBuilder>& jyBuilder,
                                                const boost::shared_ptr<InfJyParameterization>& inflationParam) const {

    LOG("Calibrate JY inflation model for inflation index " << data.index());

    const auto& rrVol = data.realRateVolatility();
    const auto& rrRev = data.realRateReversion();
    const auto& idxVol = data.indexVolatility();

    // Check if calibration is needed at all.
    if ((!rrVol.calibrate() && !rrRev.calibrate() && !idxVol.calibrate()) ||
        (data.calibrationType() == CalibrationType::None)) {
        LOG("Calibration of JY inflation model for inflation index " << data.index() << " not requested.");
        return;
    }

    // We will need the 2 baskets of helpers
    auto rrBasket = jyBuilder->realRateBasket();
    auto idxBasket = jyBuilder->indexBasket();

    // Attach engines to the helpers.
    setJyPricingEngine(modelIdx, rrBasket);
    setJyPricingEngine(modelIdx, idxBasket);

    if (dontCalibrate_)
        return;

    // Single basket of helpers is useful in various places below.
    vector<boost::shared_ptr<CalibrationHelper>> allHelpers = rrBasket;
    allHelpers.insert(allHelpers.end(), idxBasket.begin(), idxBasket.end());

    // Calibration configuration.
    const auto& cc = data.calibrationConfiguration();

    if (data.calibrationType() == CalibrationType::BestFit) {

        // If calibration type is BestFit, do a global optimisation on the parameters that need to be calibrated.
        DLOG("Calibration BestFit of JY inflation model for inflation index " << data.index() << " requested.");

        // Indicate the parameters to calibrate
        map<Size, bool> toCalibrate;
        toCalibrate[0] = rrVol.calibrate();
        toCalibrate[1] = rrRev.calibrate();
        toCalibrate[2] = idxVol.calibrate();

        // Calibrate the model.
        model_->calibrateInfJyGlobal(modelIdx, allHelpers, *optimizationMethod_, endCriteria_, toCalibrate);

    } else {

        // Calibration type is now Bootstrap, there are multiple options.
        QL_REQUIRE(data.calibrationType() == CalibrationType::Bootstrap,
                   "JY inflation calibration expected a "
                       << "calibration type of None, BestFit or Bootstrap.");
        QL_REQUIRE(!(rrRev.calibrate() && rrVol.calibrate()),
                   "Calibrating both the "
                       << "real rate reversion and real rate volatility using Bootstrap is not supported.");

        if ((!rrVol.calibrate() && !rrRev.calibrate()) && idxVol.calibrate()) {

            // Bootstrap the inflation index volatility only.
            DLOG("Bootstrap calibration of JY index volatility for index " << data.index() << ".");
            QL_REQUIRE(idxVol.type() == ParamType::Piecewise, "Index volatility parameter should be Piecewise for "
                                                                  << "a Bootstrap calibration.");
            model_->calibrateInfJyIterative(modelIdx, 2, idxBasket, *optimizationMethod_, endCriteria_);

        } else if (rrVol.calibrate() && !idxVol.calibrate()) {

            // Bootstrap the real rate volatility only
            DLOG("Bootstrap calibration of JY real rate volatility for index " << data.index() << ".");
            QL_REQUIRE(rrVol.type() == ParamType::Piecewise, "Real rate volatility parameter should be "
                                                                 << "Piecewise for a Bootstrap calibration.");
            model_->calibrateInfJyIterative(modelIdx, 0, rrBasket, *optimizationMethod_, endCriteria_);

        } else if (rrRev.calibrate() && !idxVol.calibrate()) {

            // Bootstrap the real rate reversion only
            DLOG("Bootstrap calibration of JY real rate reversion for index " << data.index() << ".");
            QL_REQUIRE(rrRev.type() == ParamType::Piecewise, "Real rate reversion parameter should be "
                                                                 << "Piecewise for a Bootstrap calibration.");
            model_->calibrateInfJyIterative(modelIdx, 1, rrBasket, *optimizationMethod_, endCriteria_);

        } else if ((rrVol.calibrate() && idxVol.calibrate()) || (rrRev.calibrate() && idxVol.calibrate())) {

            if (rrVol.calibrate()) {
                DLOG("Bootstrap calibration of JY real rate volatility and index volatility for index " << data.index()
                                                                                << ".");
            } else {
                DLOG("Bootstrap calibration of JY real rate reversion and index volatility for index " << data.index()
                                                                                << ".");
            }

            // Bootstrap the real rate volatility and the index volatility
            Size rrIdx = rrVol.calibrate() ? 0 : 1;
            Size numIts = 0;
            inflationCalibrationErrors_[modelIdx] = getCalibrationError(allHelpers);

            while (inflationCalibrationErrors_[modelIdx] > cc.rmseTolerance() && numIts < cc.maxIterations()) {
                model_->calibrateInfJyIterative(modelIdx, 2, idxBasket, *optimizationMethod_, endCriteria_);
                model_->calibrateInfJyIterative(modelIdx, rrIdx, rrBasket, *optimizationMethod_, endCriteria_);
                numIts++;
                inflationCalibrationErrors_[modelIdx] = getCalibrationError(allHelpers);
            }

            DLOG("Bootstrap calibration of JY model stopped with number of iterations "
                 << numIts << " and rmse equal to " << std::scientific << std::setprecision(6)
                 << inflationCalibrationErrors_[modelIdx] << ".");

        } else {
            QL_FAIL("JY inflation bootstrap calibration does not support the combination of real rate volatility = "
                    << std::boolalpha << rrVol.calibrate() << ", real rate reversion = " << rrRev.calibrate() << " and "
                    << "index volatility = " << idxVol.calibrate() << ".");
        }
    }

    // Log the calibration details.
    DLOG("INF (JY) " << data.index() << " calibration errors:");
    inflationCalibrationErrors_[modelIdx] = getCalibrationError(allHelpers);
    if (data.calibrationType() == CalibrationType::Bootstrap) {
        if (fabs(inflationCalibrationErrors_[modelIdx]) < cc.rmseTolerance()) {
            // we check the log level here to avoid unncessary computations
            if (Log::instance().filter(ORE_DATA)) {
                TLOGGERSTREAM << "Calibration details:";
                TLOGGERSTREAM << getCalibrationDetails(rrBasket, idxBasket, inflationParam, rrVol.calibrate());
                TLOGGERSTREAM << "rmse = " << inflationCalibrationErrors_[modelIdx];
            }
        } else {
            std::stringstream ss;
            ss << "INF (JY) " << modelIdx << " calibration error " << std::scientific
               << inflationCalibrationErrors_[modelIdx] << " exceeds tolerance " << cc.rmseTolerance();
            string exceptionMessage = ss.str();
            WLOG(StructuredModelErrorMessage("Failed to calibrate INF JY Model", exceptionMessage));
            WLOGGERSTREAM << "Calibration details:";
            WLOGGERSTREAM << getCalibrationDetails(rrBasket, idxBasket, inflationParam, rrVol.calibrate());
            WLOGGERSTREAM << "rmse = " << inflationCalibrationErrors_[modelIdx];
            if (!continueOnError_)
                QL_FAIL(exceptionMessage);
        }
    }

    LOG("Finished calibrating JY inflation model for inflation index " << data.index());
}

void CrossAssetModelBuilder::setJyPricingEngine(
    Size modelIdx, const vector<boost::shared_ptr<CalibrationHelper>>& calibrationBasket) const {

    DLOG("Start setting pricing engines on JY calibration instruments.");

    // JY supports three types of calibration helpers. Generally, all of the calibration instruments in a basket will
    // be of the same type but we support all three here.
    boost::shared_ptr<PricingEngine> cpiCapFloorEngine;
    boost::shared_ptr<PricingEngine> yoyCapFloorEngine;
    boost::shared_ptr<PricingEngine> yoySwapEngine;
    boost::shared_ptr<InflationCouponPricer> yoyCouponPricer;

    for (auto& ci : calibrationBasket) {

        if (boost::shared_ptr<CpiCapFloorHelper> h = boost::dynamic_pointer_cast<CpiCapFloorHelper>(ci)) {
            if (!cpiCapFloorEngine) {
                cpiCapFloorEngine = boost::make_shared<AnalyticJyCpiCapFloorEngine>(*model_, modelIdx);
            }
            h->setPricingEngine(cpiCapFloorEngine);
            continue;
        }

        if (boost::shared_ptr<YoYCapFloorHelper> h = boost::dynamic_pointer_cast<YoYCapFloorHelper>(ci)) {
            if (!yoyCapFloorEngine) {
                yoyCapFloorEngine = boost::make_shared<AnalyticJyYoYCapFloorEngine>(*model_, modelIdx);
            }
            h->setPricingEngine(yoyCapFloorEngine);
            continue;
        }

        if (boost::shared_ptr<YoYSwapHelper> h = boost::dynamic_pointer_cast<YoYSwapHelper>(ci)) {
            // Here we need to attach the coupon pricer to all the YoY coupons and then the generic discounting swap
            // engine to the helper.
            if (!yoyCouponPricer) {
                yoyCouponPricer = boost::make_shared<JyYoYInflationCouponPricer>(*model_, modelIdx);

                Size irIdx = model_->ccyIndex(model_->infjy(modelIdx)->currency());
                auto yts = model_->irlgm1f(irIdx)->termStructure();
                yoySwapEngine = boost::make_shared<DiscountingSwapEngine>(yts);
            }

            const auto& yoyLeg = h->yoySwap()->yoyLeg();
            for (const auto& cf : yoyLeg) {
                if (auto yoyCoupon = boost::dynamic_pointer_cast<YoYInflationCoupon>(cf))
                    yoyCoupon->setPricer(yoyCouponPricer);
            }

            h->setPricingEngine(yoySwapEngine);
            continue;
        }

        QL_FAIL("Only CPI cap floors, YoY cap floors and YoY swaps are supported for JY calibration.");
    }

    DLOG("Finished setting pricing engines on JY calibration instruments.");
}

} // namespace data
} // namespace ore


/**
* @Author: Marcel Ruhf <marcelruhf>
* @Email:  m.ruhf@protonmail.ch
* @Project: SoccerRL
* @Filename: Simulator.cpp
* @License: Licensed under the Apache 2.0 license (see LICENSE.md)
* @Copyright: Copyright (c) 2017 Marcel Ruhf
*/

#include <iostream>
#include <string>
#include <boost/random/mersenne_twister.hpp>
#include <boost/random/uniform_int_distribution.hpp>
#include <boost/random/uniform_real_distribution.hpp>
#include <Simulator.hpp>
#include <Functions.hpp>

namespace mr
{
    Simulator::Simulator()
    {
        boost::random::uniform_int_distribution<> dist(-20, 20);
        distX = dist(gen);
        while (distX == 0) distX = dist(gen);
        distY = 50;
        velocity = dist(gen);
        while (velocity == 0) velocity = dist(gen);
    }

    float Simulator::getDistX()
    {
        return distX;
    }

    float Simulator::getDistY()
    {
        return distY;
    }

    float Simulator::getVelocity()
    {
        return velocity;
    }

    void Simulator::move(int robotVelocity)
    {
        distX += (velocity - robotVelocity);
        if (distX < -20) distX  = -20;
        if (distX > 20)  distX  = 20;
        if (distY > 0)   distY -= 1;
        if (distY < 0)   distY  = 0;
    }

    void Simulator::nextEpisode()
    {
        boost::random::uniform_int_distribution<> dist(-20, 20);
        distX = dist(gen);
        while (distX == 0) distX = dist(gen);
        distY = 50;
        velocity = dist(gen);
        while (velocity == 0) velocity = dist(gen);
    }

}


#include <boost/graph/distributed/delta_stepping_shortest_paths.hpp>


#include <aslam/calibration/model/sensors/Imu.h>

#include <glog/logging.h>

#include <boost/make_shared.hpp>
#include <bsplines/NsecTimePolicy.hpp>
#include <bsplines/EuclideanBSpline.hpp>
#include <aslam/backend/QuadraticIntegralError.hpp>
#include <aslam/backend/Vector2RotationQuaternionExpressionAdapter.hpp>
#include <aslam/splines/OPTBSpline.hpp>

#include "aslam/calibration/calibrator/CalibratorI.h"
#include "aslam/calibration/data/AccelerometerMeasurement.h"
#include "aslam/calibration/data/GyroscopeMeasurement.h"
#include <aslam/calibration/DesignVariableReceiver.h>
#include <aslam/calibration/algo/splinesToFile.h>
#include <aslam/calibration/model/Model.h>
#include <aslam/calibration/model/fragments/TrajectoryCarrier.h>
#include <aslam/calibration/model/ModuleTools.h>
#include <aslam/calibration/calibrator/StateCarrier.h>
#include "aslam/calibration/error-terms/ErrorTermAccelerometer.h"
#include "aslam/calibration/error-terms/ErrorTermGyroscope.h"
#include "aslam/calibration/tools/ErrorTermStatisticsWithProblemAndPredictor.h"
#include "aslam/calibration/tools/SplineWriter.h"

namespace aslam {
namespace calibration {

Imu::Imu(Model& model, const std::string& name, sm::value_store::ValueStoreRef config) :
    Sensor(model, name, config),
    measurements_(std::make_shared<Measurements>()),
    useAcc_(getMyConfig().getBool("acc/used", true)),
    useGyro_(getMyConfig().getBool("gyro/used", true)),
    covAcc_{getMyConfig().getChild("acc/noise/cov"), 3, useAcc_},
    covGyro_{getMyConfig().getChild("gyro/noise/cov"), 3, useGyro_},
    enforceAccCovariance_(getMyConfig().getBool("acc/enforceCovariance", false)),
    enforceGyroCovariance_(getMyConfig().getBool("gyro/enforceCovariance", false)),
    accBias(*this, "accBias", getMyConfig().getChild("acc")),
    gyroBias(*this, "gyroBias", getMyConfig().getChild("gyro")),
    minimalMeasurementsPerBatch(getMyConfig().getInt("minimalMeasurementsPerBatch", 100)),
    inertiaFrame(getModel().getFrame(getMyConfig().getString("inertiaFrame")))
{
  if(isUsed()){
    SM_ASSERT_GE(std::runtime_error, minimalMeasurementsPerBatch, 0, "");
    if(useAcc_){
      accRandomWalk = getMyConfig().getDouble("acc/noise/biasRandomWalk");
    }
    if(useGyro_){
      gyroRandomWalk = getMyConfig().getDouble("gyro/noise/biasRandomWalk");
    }
  }

//TODO C support MESTIMATORS:  setMEstimator(boost::shared_ptr<aslam::backend::MEstimator>(new aslam::backend::CauchyMEstimator(10)));
}

void Imu::writeConfig(std::ostream& out) const {
  Sensor::writeConfig(out);

  MODULE_WRITE_PARAM(inertiaFrame);

  //TODO D write bias mode

  MODULE_WRITE_PARAM(useAcc_);
  if(useAcc_){
    MODULE_WRITE_PARAM(covAcc_);
    MODULE_WRITE_PARAM(enforceAccCovariance_);
    if(accBias.isUsed()){
      MODULE_WRITE_PARAM(accRandomWalk);
    }
  }
  MODULE_WRITE_PARAM(useGyro_);
  if(useGyro_){
    MODULE_WRITE_PARAM(covGyro_);
    MODULE_WRITE_PARAM(enforceGyroCovariance_);
    if(accBias.isUsed()){
      MODULE_WRITE_PARAM(gyroRandomWalk);
    }
  }

  MODULE_WRITE_PARAM(minimalMeasurementsPerBatch);
}

void Imu::registerWithModel() {
  Sensor::registerWithModel();
  if(useAcc_)
    accBias.registerCalibrationVariables(getModel());
  if(useGyro_)
    gyroBias.registerCalibrationVariables(getModel());
}

Imu::~Imu() = default;

typedef aslam::splines::OPTBSpline<typename bsplines::EuclideanBSpline<Eigen::Dynamic, 3, bsplines::NsecTimePolicy>::CONF>::BSpline BiasSpline;

class BiasBatchState : public BatchState {
 public:
  BiasBatchState(const TrajectoryCarrier & carrier, const std::string & name);
  void writeToFile(const CalibratorI & calib, const std::string & pathPrefix) const override;
  void addToProblem(DesignVariableReceiver & problem) const;

 private:
  std::string name_;
  BiasSpline biasSpline;
  friend Bias;
  friend Imu;
};

Bias::Bias(Module & m, const std::string & name, sm::value_store::ValueStoreRef config) :
    NamedMinimal(name)
{
  if (config.getBool("hasBias", true)) {
    biasVector = m.createCVIfUsed<EuclideanPointCv>(config.getChild("biasVector"), name);
    if (biasVector) {
      mode_ = Mode::Vector;
      biasVectorExpression = biasVector->toExpression();
    } else if (m.isUsed()) {
      mode_ = Mode::Spline;
      biasSplineCarrier.reset(new TrajectoryCarrier(config.getChild("biasSpline")));
    }
  } else {
    mode_ = Mode::None;
  }
}

void Bias::registerCalibrationVariables(Model& model) {
  if (biasVector) {
    model.addCalibrationVariables({biasVector});
  }
}

void Bias::initState(CalibratorI & calib){
  if (isUsingSpline()) {
    state_ = std::make_shared<BiasBatchState>(*biasSplineCarrier, getName());
    const auto & interval = calib.getCurrentEffectiveBatchInterval();
    const double elapsedTime = interval.getElapsedTime();
    const int numSegments = std::ceil(biasSplineCarrier->getKnotsPerSecond() * elapsedTime);
    LOG(INFO)<< "using IMU bias numSegments=" << numSegments << " for " << elapsedTime << " seconds";

    //TODO D make bias initial guess based on IMU measurements!

    state_->biasSpline.initConstantUniformSpline(interval.start, interval.end, numSegments, Eigen::Vector3d::Zero());
  }
}

void Bias::addToBatch(bool stateActive, BatchStateReceiver & batchStateReceiver, DesignVariableReceiver & problem) {
  if(state_){
    problem.addSplineDesignVariables(state_->biasSpline, stateActive);
    batchStateReceiver.addBatchState(*this, state_);
  }
}

bool Imu::initState(CalibratorI& calib) {
  if(isUsed()){
    accBias.initState(calib);
    gyroBias.initState(calib);
  }
  return true;
}
void Imu::addToBatch(const Activator & stateActivator, BatchStateReceiver & batchStateReceiver, DesignVariableReceiver & problem) {
  const bool stateActive = stateActivator.isActive(*this);
  if(useAcc_){
    accBias.setActive(stateActive);
    accBias.addToBatch(stateActive, batchStateReceiver, problem);
  }
  if(useGyro_){
    gyroBias.setActive(stateActive);
    gyroBias.addToBatch(stateActive, batchStateReceiver, problem);
  }
}

void Imu::setActive(bool spatial, bool temporal) {
  Sensor::setActive(spatial, temporal);
}

BiasBatchState::BiasBatchState(const TrajectoryCarrier & carrier, const std::string & name)
    : name_(name), biasSpline(carrier.getSplineOrder())
{
}

void BiasBatchState::writeToFile(const CalibratorI & calib, const std::string& pathPrefix) const {
  writeSpline(biasSpline, calib.getOptions().getSplineOutputSamplePeriod(), pathPrefix + name_);
}

void Imu::addAccelerometerMeasurement(CalibratorI & calib, const AccelerometerMeasurement& data, Timestamp timestamp) const {
  calib.addMeasurementTimestamp(timestamp, *this);
  measurements_->accelerometer.emplace_back(timestamp, data);
}

void Imu::addGyroscopeMeasurement(CalibratorI & calib, const GyroscopeMeasurement& data, Timestamp timestamp) const {
  calib.addMeasurementTimestamp(timestamp, *this);
  measurements_->gyroscope.emplace_back(timestamp, data);
}

void Imu::addInputTo(Timestamp t, const AccelerometerMeasurement& input, ModuleStorage&) const {
  measurements_->accelerometer.emplace_back(t, input);
}

void Imu::addInputTo(Timestamp t, const GyroscopeMeasurement& input, ModuleStorage&) const {
  measurements_->gyroscope.emplace_back(t, input);
}

void Imu::clearMeasurements() {
  measurements_->accelerometer.clear();
  measurements_->gyroscope.clear();
}

using namespace aslam::backend;

template<typename BSpline_>
struct BiasModelIntegrationErrorExpressionFactory {
  typedef BSpline_ BSplineT;
  typedef typename BSplineT::time_t TimeT;
  enum { PointSize = 3 };
  typedef VectorExpression<PointSize> value_expression_t;
  typedef Eigen::Matrix<double, 1, 1>  value_t;

  const BSplineT & _bspline;
  const Eigen::Matrix<double, PointSize, PointSize> _sqrtInvR;

  BiasModelIntegrationErrorExpressionFactory(const BSplineT & bspline, Eigen::Matrix<double, PointSize, PointSize> sqrtInvR) : _bspline(bspline), _sqrtInvR(sqrtInvR) {}

  inline value_expression_t operator()(TimeT time) const {
    return _bspline.template getExpressionFactoryAt<1>(time).getValueExpression(1);
  }
};

template <typename SplineT>
void addBiasModelErrorTerms(CalibratorI & calib, std::string name, ErrorTermReceiver & errorTermReceiver, const SplineT & spline, const Eigen::MatrixXd & sqrtInvR, bool observeOnly, int numberOfIntegrationPoints = -1){
  BiasModelIntegrationErrorExpressionFactory<SplineT> integrationFunctor(spline, sqrtInvR);
  if(numberOfIntegrationPoints < 0){
    numberOfIntegrationPoints = (spline.getAbsoluteNumberOfSegments() + spline.getSplineOrder()) * 2;
  }

  ErrorTermStatisticsWithProblemAndPredictor stat(calib, name + "Bias", errorTermReceiver, observeOnly);

  aslam::backend::integration::addQuadraticIntegralExpressionErrorTerms<aslam::backend::integration::DefaultAlgorithm>(stat, spline.getMinTime(), spline.getMaxTime(), numberOfIntegrationPoints, integrationFunctor, sqrtInvR);

  stat.printInto(LOG(INFO));
}

template <typename T, typename ErrorTermFactory>
void addImuErrorTerms(CalibratorI & calib, const Imu & imu, std::string name, const MeasurementsContainer<T>& measurements, ErrorTermFactory etFactory, ErrorTermReceiver & errorTermReceiver, bool observeOnly) {
  LOG(INFO) << "Adding " << measurements.size() << " " << name << " error terms";

  ErrorTermStatisticsWithProblemAndPredictor statWPAP(calib, name, errorTermReceiver, observeOnly);
  const Interval & interval = calib.getCurrentEffectiveBatchInterval();

  Timestamp minTime = interval.end, maxTime = interval.start;

  for (auto & m : measurements) {
    Timestamp timestamp = m.first;
    if (!interval.contains(timestamp, imu)){
      LOG(INFO) << name << " measurement out of spline range at " << calib.secsSinceStart(timestamp) << "s.";
      continue;
    }

    if(minTime > timestamp) minTime = timestamp;
    if(maxTime < timestamp) maxTime = timestamp;

    auto e = etFactory(timestamp, m.second);

    if(imu.getMEstimator()){
      e->setMEstimatorPolicy(imu.getMEstimator());
    }

    VLOG(2) << "Cost function " << name << " : " << e->evaluateError() << " count: " << statWPAP.getCounter() << " timestamp: " << calib.secsSinceStart(timestamp) << "s,";
    statWPAP.add(timestamp, e);
  }
  statWPAP.printInto(LOG(INFO)) << " Between " << calib.secsSinceStart(minTime) << "s and " << calib.secsSinceStart(maxTime) << "s.";
}

template <typename T>
bool isCovarianceAvailable(T&t){
  return t.cov(0,0) > 0.0;
}

void Imu::addMeasurementErrorTerms(CalibratorI & calib, const CalibrationConfI & /*ec*/, ErrorTermReceiver & errorTermReceiver, bool observeOnly) const {
  if(useAcc_){
    const std::string accelerometerName = getName() + "Accelerometer";
    if(accBias.isUsingSpline()){
      addBiasModelErrorTerms(calib, accelerometerName, errorTermReceiver, accBias.state_->biasSpline, Eigen::Matrix3d::Identity() / accRandomWalk, observeOnly);
    }

    //TODO C solve gravity vector problem. Each model has a gravity vector?
    auto g_m = calib.getModel().getGravity().getVectorExpression();
    Eigen::Matrix3d covarianceMatrix = covAcc_.getValue();
    ErrorTermGroupReference etgr(getName() + "Acc");
    addImuErrorTerms(
        calib, *this, accelerometerName, measurements_->accelerometer,
        [&, this](const Timestamp timestamp, const AccelerometerMeasurement & m){
          auto modelAt = calib.getModelAt(*this, timestamp, 2, {false});
          return boost::make_shared<ErrorTermAccelerometer>(
              modelAt.getAcceleration(getFrame(), inertiaFrame),
              getTransformationExpressionTo(modelAt, inertiaFrame).toRotationExpression().inverse(),
              g_m,
              accBias.getBiasExpression(timestamp),
              m.a,
              isCovarianceAvailable(m) && !enforceAccCovariance_ ? m.cov : covarianceMatrix,
              etgr);
        },
        errorTermReceiver,
        observeOnly
      );
  }
  if(useGyro_){
    const std::string gyroscopeName = getName() + "Gyroscope";
    if(gyroBias.isUsingSpline()){
      addBiasModelErrorTerms(calib, gyroscopeName, errorTermReceiver, gyroBias.state_->biasSpline, Eigen::Matrix3d::Identity() / gyroRandomWalk, observeOnly);
    }
    Eigen::Matrix3d covarianceMatrix = covGyro_.getValue();
    ErrorTermGroupReference etgr(getName() + "Gyro");
    addImuErrorTerms(
        calib, *this, gyroscopeName, measurements_->gyroscope,
        [&, this](const Timestamp timestamp, const GyroscopeMeasurement & m){
          auto modelAt = calib.getModelAt(*this,  timestamp, 1, {false});
          return boost::make_shared<ErrorTermGyroscope>(
              getTransformationExpressionTo(modelAt, inertiaFrame).toRotationExpression().inverse() * modelAt.getAngularVelocity(getReferenceFrame(), inertiaFrame),
              gyroBias.getBiasExpression(timestamp),
              m.w,
              isCovarianceAvailable(m) && !enforceGyroCovariance_ ? m.cov : covarianceMatrix,
              etgr
            );
        },
        errorTermReceiver,
        observeOnly
      );
  }
}

double Imu::getMaximalTimeGap() const {
  if(measurements_){
    Timestamp gGap(measurements_->gyroscope.getMaximalTimeGap());
    Timestamp aGap(measurements_->accelerometer.getMaximalTimeGap());
    return std::max(double(gGap), double(aGap));
  } else {
    return 0.0;
  }
}

bool Imu::hasTooFewMeasurements() const {
  return getGyroscopeMeasurements().size() < size_t(minimalMeasurementsPerBatch);
}

aslam::backend::EuclideanExpression Bias::getBiasExpression(Timestamp t) const {
  if(isUsingSpline()){
    SM_ASSERT_NOTNULL(std::runtime_error, state_, "");
    return state_->biasSpline.getExpressionFactoryAt<0>(t).getValueExpression();
  } else {
    return biasVectorExpression;
  }
}

void Imu::addPriorFactors(CalibratorI & calib, ErrorTermReceiver & errorTermReceiver, double priorFactor) const {
  const double invSigma = 1e-2 * priorFactor;
  Timestamp startTime = calib.getCurrentEffectiveBatchInterval().start;
  if(accBias.isUsingSpline()) errorTermReceiver.addErrorTerm(backend::toErrorTerm(accBias.getBiasExpression(startTime), Eigen::Matrix3d::Identity() * invSigma));
  if(gyroBias.isUsingSpline()) errorTermReceiver.addErrorTerm(backend::toErrorTerm(gyroBias.getBiasExpression(startTime), Eigen::Matrix3d::Identity() * invSigma));
}

} /* namespace calibration */
} /* namespace aslam */


// Boost.Geometry (aka GGL, Generic Geometry Library)

// Copyright (c) 2007-2014 Barend Gehrels, Amsterdam, the Netherlands.
// Copyright (c) 2008-2014 Bruno Lalande, Paris, France.
// Copyright (c) 2009-2014 Mateusz Loskot, London, UK.
// Copyright (c) 2013-2014 Adam Wulkiewicz, Lodz, Poland

// This file was modified by Oracle on 2013-2014.
// Modifications copyright (c) 2013-2014, Oracle and/or its affiliates.

// Contributed and/or modified by Adam Wulkiewicz, on behalf of Oracle
// Contributed and/or modified by Menelaos Karavelas, on behalf of Oracle

// Parts of Boost.Geometry are redesigned from Geodan's Geographic Library
// (geolib/GGL), copyright (c) 1995-2010 Geodan, 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)

#ifndef BOOST_GEOMETRY_ALGORITHMS_DETAIL_DISJOINT_POINT_POINT_HPP
#define BOOST_GEOMETRY_ALGORITHMS_DETAIL_DISJOINT_POINT_POINT_HPP

#include <cstddef>

#include <boost/geometry/core/access.hpp>
#include <boost/geometry/core/coordinate_dimension.hpp>
#include <boost/geometry/core/tags.hpp>

#include <boost/geometry/util/math.hpp>

#include <boost/geometry/algorithms/dispatch/disjoint.hpp>

namespace boost { namespace geometry
{

#ifndef DOXYGEN_NO_DETAIL
namespace detail { namespace disjoint
{

template
<
    typename Point1, typename Point2,
    std::size_t Dimension, std::size_t DimensionCount
>
struct point_point
{
    static inline bool apply(Point1 const& p1, Point2 const& p2)
    {
        if (! geometry::math::equals(get<Dimension>(p1), get<Dimension>(p2)))
        {
            return true;
        }
        return point_point
            <
                Point1, Point2,
                Dimension + 1, DimensionCount
            >::apply(p1, p2);
    }
};

template <typename Point1, typename Point2, std::size_t DimensionCount>
struct point_point<Point1, Point2, DimensionCount, DimensionCount>
{
    static inline bool apply(Point1 const& , Point2 const& )
    {
        return false;
    }
};

/*!
    \brief Internal utility function to detect of points are disjoint
    \note To avoid circular references
 */
template <typename Point1, typename Point2>
inline bool disjoint_point_point(Point1 const& point1, Point2 const& point2)
{
    return point_point
        <
            Point1, Point2,
            0, dimension<Point1>::type::value
        >::apply(point1, point2);
}

}} // namespace detail::disjoint
#endif // DOXYGEN_NO_DETAIL

#ifndef DOXYGEN_NO_DISPATCH
namespace dispatch
{

template <typename Point1, typename Point2, std::size_t DimensionCount>
struct disjoint<Point1, Point2, DimensionCount, point_tag, point_tag, false>
    : detail::disjoint::point_point<Point1, Point2, 0, DimensionCount>
{};

} // namespace dispatch
#endif // DOXYGEN_NO_DISPATCH

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_ALGORITHMS_DETAIL_DISJOINT_POINT_POINT_HPP


#pragma once
#include <boost/optional.hpp>
#include <map>
#include <memory>
#include <regex>
#include "ear/common/point_source_panner.hpp"
#include "ear/common/screen_edge_lock.hpp"
#include "ear/helpers/output_gains.hpp"
#include "ear/layout.hpp"
#include "ear/metadata.hpp"
#include "ear/warnings.hpp"

namespace ear {

  class GainCalculatorDirectSpeakersImpl {
   public:
    GainCalculatorDirectSpeakersImpl(
        const Layout& layout,
        std::map<std::string, std::string> additionalSubstitutions = {});

    void calculate(const DirectSpeakersTypeMetadata& metadata,
                   OutputGains& direct,
                   const WarningCB& warning_cb = default_warning_cb);

    template <typename T>
    void calculate(const DirectSpeakersTypeMetadata& metadata,
                   std::vector<T>& direct,
                   const WarningCB& warning_cb = default_warning_cb) {
      OutputGainsT<T> direct_wrapped(direct);
      calculate(metadata, direct_wrapped, warning_cb);
    }

   private:
    /** @brief Get the bs.2051 speaker label from an ADM speaker label.
     *
     * This parses URNs, and deals with alternative notations for LFE channels.
     */
    std::string _nominalSpeakerLabel(const std::string& label);
    /** @brief Get the index of the candidate speaker closest to a given
    position.
     *
     * If there are multiple speakers that are considered equally close with
     * respect to a given tolerance, no decision on the closest speaker can be
     * made.
     *
     * @param position  Target position
     * @param candidates Subset of speakers to be considered
     * @param tol  tolerance for defintion of "closest".
     *
     * @returns  Index to the speaker in speakers that is closest to the
     * target position or `None` if no such speaker can be uniquely defined.
     */
    boost::optional<int> _closestChannelIndex(const SpeakerPosition& position,
                                              std::vector<bool> candidates,
                                              double tol);

    /** @brief Determine if type_metadata is an LFE channel, issuing a warning
     * is there's a discrepancy between the speakerLabel and the frequency
     * element.
     */
    bool _isLfeChannel(const DirectSpeakersTypeMetadata& metadata,
                       const WarningCB& warning_cb);

    SpeakerPosition _applyScreenEdgeLock(SpeakerPosition position);
    boost::optional<int> _findChannelWithinBounds(
        const SpeakerPosition& position, bool isLfe, double tol);
    std::vector<std::pair<int, double>> _findCandidates(
        const PolarSpeakerPosition& position, bool isLfe, double tol);
    std::vector<std::pair<int, double>> _findCandidates(
        const CartesianSpeakerPosition& position, bool isLfe, double tol);

    Layout _layout;
    std::shared_ptr<PointSourcePanner> _pointSourcePanner;
    ScreenEdgeLockHandler _screenEdgeLockHandler;
    int _nChannels;
    std::vector<std::string> _channelNames;
    Eigen::VectorXd _azimuths;
    Eigen::VectorXd _elevations;
    Eigen::VectorXd _distances;
    Eigen::MatrixXd _positions;
    Eigen::Array<bool, Eigen::Dynamic, 1> _isLfe;
    std::map<std::string, std::string> _substitutions;
    const std::regex SPEAKER_URN_REGEX =
        std::regex("^urn:itu:bs:2051:[0-9]+:speaker:(.*)$");
  };

}  // namespace ear


// Copyright (c) 2020 Bitcoin Association.
// Distributed under the Open BSV software license, see the accompanying file LICENSE.

#include <boost/test/unit_test.hpp>

#include "validation.h"
#include "config.h"
#include "script/instruction.h"
#include "script/instruction_iterator.h"
#include "mining/journal.h"
#include "mining/journal_builder.h"
#include "mining/journal_change_set.h"
#include "mining/assembler.h"

#include "test/test_bitcoin.h"
#include "test/mempool_test_access.h"

namespace{

CTxMemPoolEntry MakeEntry(
    CFeeRate feerate,
    std::vector<std::tuple<TxId, size_t, Amount>> inChainInputs,
    std::vector<std::tuple<CTransactionRef, int>> inMempoolInputs,
    size_t nOutputs, size_t additionalSize=0, Amount feeAlreadyPaid=Amount{1},
    size_t opReturnSize=0)
{
    CMutableTransaction tx;
    Amount totalInput;
    for(const auto& input: inChainInputs)
    {
        auto[id, ndx, amount] = input;
        tx.vin.push_back(CTxIn(id, ndx, CScript()));
        totalInput += amount;
    }

    for(const auto& input: inMempoolInputs)
    {
        auto[txInput, ndx] = input;
        tx.vin.push_back(CTxIn(txInput->GetId(), ndx, CScript()));
        totalInput += txInput->vout[ndx].nValue;
    }

    for(size_t i = 0; i < nOutputs; i++)
    {
        CScript script;
        script << OP_TRUE;
        tx.vout.push_back(CTxOut(Amount{1}, script));
    }

    if(opReturnSize != 0)
    {
        CScript script;
        script << OP_FALSE << OP_RETURN;
        script << std::vector<uint8_t>(opReturnSize);
        tx.vout.push_back(CTxOut(Amount(), script));
    }

    auto txSize = ::GetSerializeSize(tx, SER_NETWORK, PROTOCOL_VERSION) + additionalSize;
    auto totalFee = feerate.GetFee(txSize);
    auto perOputput = (totalInput - totalFee) / int64_t(nOutputs) + Amount(1);

    for(auto& output: tx.vout)
    {
        if(output.scriptPubKey.begin_instructions()->opcode() != OP_FALSE)
        {
            output.nValue = perOputput;
        }
    }

    auto txRef = MakeTransactionRef(tx);
    CTxMemPoolEntry entry(txRef, totalFee, int64_t(0), false, false, LockPoints());
    return entry;
}

TxId MakeId(uint16_t n)
{
    TxId id;
    *id.begin() = n >> 1;
    *(id.begin() + 1) = n | 0x00ff;
    return id;
}

std::vector<std::tuple<TxId, size_t, Amount>> MakeConfirmedInputs(size_t count, Amount value)
{
    static uint16_t nextTxid = 1;
    std::vector<std::tuple<TxId, size_t, Amount>> inputs;
    for (size_t i = 0; i < count; i++)
    {
        inputs.push_back(std::make_tuple(MakeId(nextTxid++), i, value));
    }
    return inputs;
}

bool checkGroupContinuity(const mining::CJournalChangeSetPtr& changeSet)
{
    std::unordered_set<mining::GroupID> seenGroups;
    const auto& changes = changeSet->getChangeSet();
    const auto disjoint = std::adjacent_find(changes.cbegin(), changes.cend(),
                                             [&seenGroups](const auto& aChange, const auto&bChange) {
        // adjacent_find returns iterator where condition is true
        auto a = aChange.second.getGroupId();
        auto b = bChange.second.getGroupId();
        return (a != b                                 // we are at a boundary
                && a                                   // ending sequence was a group
                && !seenGroups.insert(a).second);      // that we have already seen
    });
    return disjoint == changes.cend();
}

CFeeRate DefaultFeeRate()
{
    return CFeeRate(DEFAULT_BLOCK_MIN_TX_FEE);
}

CTxMemPoolTestAccess::txiter AddToMempool(CTxMemPoolEntry& entry)
{
    mempool.AddUnchecked(entry.GetSharedTx()->GetId(), entry, TxStorage::memory, {});
    CTxMemPoolTestAccess testAccess(mempool);
    auto it = testAccess.mapTx().find(entry.GetSharedTx()->GetId());
    BOOST_ASSERT(it != testAccess.mapTx().end());
    return it;
}

std::unique_ptr<mining::CBlockTemplate> CreateBlock()
{
    CBlockIndex* pindexPrev {nullptr};
    std::unique_ptr<mining::CBlockTemplate> pblocktemplate;
    CScript scriptPubKey =
    CScript() << ParseHex("04678afdb0fe5548271967f1a67130b7105cd6a828e03909"
                            "a67962e0ea1f61deb649f6bc3f4cef38c4f35504e51ec112"
                            "de5c384df7ba0b8d578a4c702b6bf11d5f")
                << OP_CHECKSIG;
    BOOST_CHECK(pblocktemplate = mining::g_miningFactory->GetAssembler()->CreateNewBlock(scriptPubKey, pindexPrev));
    return pblocktemplate;
}

using JournalEntry = mining::CJournalEntry;
using JournalTester = mining::CJournalTester;
}

BOOST_FIXTURE_TEST_SUITE(cpfp_tests, TestingSetup)

mining::CJournalPtr CheckMempoolRebuild(CTxMemPoolTestAccess& testAccess)
{
    auto oldJournal = testAccess.getJournalBuilder().getCurrentJournal();
    auto contentsBefore = JournalTester(oldJournal).getContents();
    auto oldMapTx = testAccess.mapTx();

    auto changeSet = mempool.RebuildMempool();

    BOOST_CHECK(checkGroupContinuity(changeSet));
    BOOST_CHECK(changeSet->CheckTopoSort());
    changeSet->apply();

    auto newJournal = testAccess.getJournalBuilder().getCurrentJournal();
    auto contentsAfter = JournalTester(newJournal).getContents();
    BOOST_CHECK(contentsBefore == contentsAfter);

    auto newMapTx = testAccess.mapTx();

    BOOST_TEST(oldMapTx.size() == newMapTx.size());
    const auto& oldAccess = oldMapTx.get<insertion_order>();
    const auto& newAccess = newMapTx.get<insertion_order>();
    for(auto it1 = oldAccess.begin(),it2 =  newAccess.begin(); it1 != oldAccess.end(); it1++, it2++)
    {
        BOOST_CHECK(it1->GetTxId() == it2->GetTxId());
        BOOST_CHECK(it1->IsCPFPGroupMember() == it2->IsCPFPGroupMember());
        BOOST_CHECK(it1->IsInPrimaryMempool() == it2->IsInPrimaryMempool());
    }

    return newJournal;
}

BOOST_AUTO_TEST_CASE(group_forming_and_disbanding)
{
    //           |                  |
    //           |            entryNotPaying
    //           |                  |
    //   entryNotPaying3    entryPayForItself
    //           |            |     |    |
    //           +------------+     |    +-------------+
    //           |                  |                  |
    //    entryNotPaying4    entryPayForGroup   entryNotPaying2
    //           |
    //  entryPayingFor3And4
    //
    //  entries in group1 (entering primary mempool): entryNotPaying, entryPayForItself and entryPayForGroup
    //  entries in group2 (entering primary mempool): entryNotPaying3, entryNotPaying4 and entryPayingFor3And4
    //  entry still in secondary: entryNotPaying2

    mempool.SetSanityCheck(0);

    auto entryNotPaying = MakeEntry(CFeeRate{}, MakeConfirmedInputs(1, Amount{1000000}), {}, 1);

    auto entryPayForItself = MakeEntry(DefaultFeeRate(), {}, {std::make_tuple(entryNotPaying.GetSharedTx(), 0)}, 3);

    auto entryNotPaying2 = MakeEntry(CFeeRate{}, {}, {std::make_tuple(entryPayForItself.GetSharedTx(), 1)}, 1);

    auto entryNotPaying3 = MakeEntry(CFeeRate{}, MakeConfirmedInputs(1, Amount{1000000}), {}, 1);

    auto entryNotPaying4 = MakeEntry(CFeeRate{}, {}, {std::make_tuple(entryPayForItself.GetSharedTx(), 2), std::make_tuple(entryNotPaying3.GetSharedTx(), 0)}, 1);

    auto sizeOfNotPaying3and4 = entryNotPaying3.GetSharedTx()->GetTotalSize() + entryNotPaying4.GetSharedTx()->GetTotalSize();
    auto feeOfNotPaying3and4 = entryNotPaying3.GetModifiedFee() + entryNotPaying4.GetModifiedFee();
    auto entryPayingFor3And4 = MakeEntry(DefaultFeeRate(), {}, {std::make_tuple(entryNotPaying4.GetSharedTx(), 0)}, 1, sizeOfNotPaying3and4, feeOfNotPaying3and4);

    auto sizeSoFar = entryNotPaying.GetSharedTx()->GetTotalSize() + entryPayForItself.GetSharedTx()->GetTotalSize();
    auto feeSoFar = entryNotPaying.GetModifiedFee() + entryPayForItself.GetModifiedFee();
    auto entryPayForGroup = MakeEntry(DefaultFeeRate(), {}, {std::make_tuple(entryPayForItself.GetSharedTx(), 0)}, 1, sizeSoFar, feeSoFar);

    CTxMemPoolTestAccess testAccess(mempool);
    auto journal = testAccess.getJournalBuilder().getCurrentJournal();

    auto notPayingIt = AddToMempool(entryNotPaying);
    BOOST_ASSERT(!notPayingIt->IsInPrimaryMempool());

    auto payForItselfIt = AddToMempool(entryPayForItself);
    BOOST_ASSERT(!payForItselfIt->IsInPrimaryMempool());

    auto notPaying2It = AddToMempool(entryNotPaying2);
    BOOST_ASSERT(!notPaying2It->IsInPrimaryMempool());

    auto notPaying3It = AddToMempool(entryNotPaying3);
    BOOST_ASSERT(!notPaying3It->IsInPrimaryMempool());

    auto notPaying4It = AddToMempool(entryNotPaying4);
    BOOST_ASSERT(!notPaying4It->IsInPrimaryMempool());

    // entryPayingFor3And4 pays for entryNotPaying4 and entryNotPaying3
    // but not enough for entryPayForItself and entryNotPaying
    // so it will not be able to form a group yet
    auto payFor3And4It = AddToMempool(entryPayingFor3And4);
    BOOST_ASSERT(!payFor3And4It->IsInPrimaryMempool());

    // still nothing is accepted to primary mempool
    BOOST_ASSERT(JournalTester(journal).journalSize() == 0);

    // now we will add payForGroupIt which pays enough for entryPayForItself and entryNotPaying
    // this will cause forming a group
    auto payForGroupIt = AddToMempool(entryPayForGroup);
    BOOST_ASSERT(payForGroupIt->IsInPrimaryMempool() && payForGroupIt->IsCPFPGroupMember());
    BOOST_ASSERT(payForItselfIt->IsInPrimaryMempool() && payForItselfIt->IsCPFPGroupMember());
    BOOST_ASSERT(notPayingIt->IsInPrimaryMempool() && notPayingIt->IsCPFPGroupMember());

    BOOST_ASSERT(payForGroupIt->GetCPFPGroup() == payForItselfIt->GetCPFPGroup());
    BOOST_ASSERT(payForGroupIt->GetCPFPGroup() == notPayingIt->GetCPFPGroup());

    // as the entryNotPaying4 (and consequently entryPayingFor3And4) is no longer obliged to pay
    // for entryPayForItself and entryNotPaying, a new group can be formed
    // (entryPayingFor3And4 pays for entryNotPaying4 and entryNotPaying3)
    BOOST_ASSERT(payFor3And4It->IsInPrimaryMempool() && payFor3And4It->IsCPFPGroupMember());
    BOOST_ASSERT(notPaying4It->IsInPrimaryMempool() && notPaying4It->IsCPFPGroupMember());
    BOOST_ASSERT(notPaying3It->IsInPrimaryMempool() && notPaying3It->IsCPFPGroupMember());

    BOOST_ASSERT(payFor3And4It->GetCPFPGroup() == notPaying4It->GetCPFPGroup());
    BOOST_ASSERT(payFor3And4It->GetCPFPGroup() == notPaying3It->GetCPFPGroup());

    // check that they are not part of the same group
    BOOST_ASSERT(payForGroupIt->GetCPFPGroup() != payFor3And4It->GetCPFPGroup());

    // nobody payed for notPaying2It, still in secondary mempool
    BOOST_ASSERT(!notPaying2It->IsInPrimaryMempool());

    // journal is no longer empty
    BOOST_ASSERT(JournalTester(journal).journalSize() != 0);

    // check content of the journal
    BOOST_ASSERT(JournalTester(journal).checkTxnExists(JournalEntry{*notPayingIt}));
    BOOST_ASSERT(JournalTester(journal).checkTxnExists(JournalEntry{*payForItselfIt}));
    BOOST_ASSERT(JournalTester(journal).checkTxnExists(JournalEntry{*payForGroupIt}));

    BOOST_ASSERT(JournalTester(journal).checkTxnExists(JournalEntry{*payFor3And4It}));
    BOOST_ASSERT(JournalTester(journal).checkTxnExists(JournalEntry{*notPaying4It}));
    BOOST_ASSERT(JournalTester(journal).checkTxnExists(JournalEntry{*notPaying3It}));

    BOOST_ASSERT(! JournalTester(journal).checkTxnExists(JournalEntry{*notPaying2It}));

    // remove payFor3And4It, notPaying4It from mempool
    CTxMemPoolTestAccess::setEntries entriesToRemove = {payFor3And4It, notPaying4It};
    auto changeSet = testAccess.getJournalBuilder().getNewChangeSet(mining::JournalUpdateReason::UNKNOWN);
    testAccess.removeStagedNL(entriesToRemove, *changeSet, CTransactionConflict{}, MemPoolRemovalReason::UNKNOWN);

    changeSet->apply();
    changeSet->clear();

    // entries which we have removed, they should removed from mempool and also from the journal
    for(auto entry: {entryNotPaying4, entryPayingFor3And4})
    {
        BOOST_ASSERT(testAccess.mapTx().find(entry.GetTxId()) == testAccess.mapTx().end());
        BOOST_ASSERT(!JournalTester(journal).checkTxnExists(JournalEntry{entry}));
    }

    // unaffected entries, they should stay in the mempool and journal
    for(auto entry: {entryNotPaying, entryPayForItself, entryPayForGroup})
    {
        BOOST_ASSERT(testAccess.mapTx().find(entry.GetTxId()) != testAccess.mapTx().end());
        BOOST_ASSERT(JournalTester(journal).checkTxnExists(JournalEntry{entry}));
    }

    // notPaying3It is still in the mempool
    BOOST_ASSERT(testAccess.mapTx().find(entryNotPaying3.GetTxId()) != testAccess.mapTx().end());
    // but not in the journal
    BOOST_ASSERT(!JournalTester(journal).checkTxnExists(JournalEntry{entryNotPaying3}));

    // return removed transactions back to mempool

    notPaying4It = AddToMempool(entryNotPaying4);
    payFor3And4It = AddToMempool(entryPayingFor3And4);
    BOOST_ASSERT(notPaying4It->IsInPrimaryMempool());
    BOOST_ASSERT(payFor3And4It->IsInPrimaryMempool());

    // things should be as before removal
    for(auto entry: {entryNotPaying, entryPayForItself, entryPayForGroup, entryNotPaying3, entryNotPaying4, entryPayingFor3And4})
    {
        BOOST_ASSERT(testAccess.mapTx().find(entry.GetTxId()) != testAccess.mapTx().end());
        BOOST_ASSERT(JournalTester(journal).checkTxnExists(JournalEntry{entry}));
    }

    // now remove entryPayForGroup
    entriesToRemove.clear();
    entriesToRemove.insert(payForGroupIt);
    testAccess.removeStagedNL(entriesToRemove, *changeSet, CTransactionConflict{}, MemPoolRemovalReason::UNKNOWN);
    changeSet->apply();

    // everything should be removed from journal
    BOOST_ASSERT(JournalTester(journal).journalSize() == 0);

    // and nothing should stay in the primary mempool
    for(const auto& entryIt: testAccess.mapTx())
    {
        BOOST_ASSERT(!entryIt.IsInPrimaryMempool());
    }

    // now raise modified fee for the entryPayingFor3And4 so that it can pay for all ancestors (entryNotPaying4, entryNotPaying3, entryPaysForItself, entryNotPaying)
    mempool.PrioritiseTransaction(entryPayingFor3And4.GetTxId(), entryPayingFor3And4.GetTxId().GetHex(), Amount{10000});
    for(const auto& entryIt: {notPayingIt, notPaying3It, notPaying4It, payForItselfIt, payFor3And4It})
    {
        BOOST_ASSERT(entryIt->IsInPrimaryMempool());
        BOOST_ASSERT(JournalTester(journal).checkTxnExists(JournalEntry{*entryIt}));
    }

    CheckMempoolRebuild(testAccess);
};

BOOST_AUTO_TEST_CASE(group_recalculation_when_removing_for_block)
{
    //
    //  entryNotPaying1            entryNotPaying3            RemoveForBlock
    //        |                           |
    // ----------------------------------------------------------
    //        |                           |
    //  entryNotPaying2            entryPaysForItself
    //        |                           |
    //   entryPaysFor2               entryPaysFor3
    //

    // before: 1. entryPaysFor2 can not form a group
    //         2. entryPaysFor3 forms a group
    //
    // after: 1. entryPaysFor2 forms a group (got rid of the entryNotPaying1 debt)
    //        2  entryPaysFor3 group is disbanded, and entryPaysForItself and entryPaysFor3 are accepted as standalone

    mempool.SetSanityCheck(0);

    auto entryNotPaying1 = MakeEntry(CFeeRate{}, MakeConfirmedInputs(1, Amount{1000000}), {}, 1);
    auto entryNotPaying2 = MakeEntry(CFeeRate{}, {}, {std::make_tuple(entryNotPaying1.GetSharedTx(), 0)}, 1);
    auto entryPaysFor2 = MakeEntry(DefaultFeeRate(), {}, {std::make_tuple(entryNotPaying2.GetSharedTx(), 0)}, 1, entryNotPaying2.GetSharedTx()->GetTotalSize());
    auto entryNotPaying3 = MakeEntry(CFeeRate{}, MakeConfirmedInputs(1, Amount{1000000}), {}, 1);
    auto entryPayForItself = MakeEntry(DefaultFeeRate(), {}, {std::make_tuple(entryNotPaying3.GetSharedTx(), 0)}, 3);
    auto entryPaysFor3 = MakeEntry(CFeeRate{Amount{10000}}, {}, {std::make_tuple(entryPayForItself.GetSharedTx(), 0)}, 3);

    auto notPaying1   = AddToMempool(entryNotPaying1);
    auto notPaying2   = AddToMempool(entryNotPaying2);
    auto paysFor2     = AddToMempool(entryPaysFor2);
    auto notPaying3   = AddToMempool(entryNotPaying3);
    auto payForItself = AddToMempool(entryPayForItself);
    auto paysFor3     = AddToMempool(entryPaysFor3);

    CTxMemPoolTestAccess testAccess(mempool);
    auto journal = testAccess.getJournalBuilder().getCurrentJournal();

    for(auto entryIt: {notPaying1, notPaying2, paysFor2})
    {
        BOOST_ASSERT(!entryIt->IsInPrimaryMempool());
        BOOST_ASSERT(!JournalTester(journal).checkTxnExists(JournalEntry{*entryIt}));
    }

    for(auto entryIt: {notPaying3, payForItself, paysFor3})
    {
        BOOST_ASSERT(entryIt->IsCPFPGroupMember());
        BOOST_ASSERT(JournalTester(journal).checkTxnExists(JournalEntry{*entryIt}));
    }

    mempool.RemoveForBlock({entryNotPaying1.GetSharedTx(), entryNotPaying3.GetSharedTx()}, mining::CJournalChangeSetPtr{}, uint256{});

    for(auto entryIt: {notPaying2, paysFor2})
    {
        BOOST_ASSERT(entryIt->IsInPrimaryMempool());
        BOOST_ASSERT(JournalTester(journal).checkTxnExists(JournalEntry{*entryIt}));
    }

    for(auto entryIt: {payForItself, paysFor3})
    {
        BOOST_ASSERT(!entryIt->IsCPFPGroupMember());
        BOOST_ASSERT(entryIt->IsInPrimaryMempool());
        BOOST_ASSERT(JournalTester(journal).checkTxnExists(JournalEntry{*entryIt}));
    }

    CheckMempoolRebuild(testAccess);
};

BOOST_AUTO_TEST_CASE(mempool_rebuild)
{
    //
    //  entry1    + -------------- entryGroup1Tx1
    //       |    |                       |
    //  entryGroup2Tx1 ----- + --- entryGroup1Tx2
    //       |               |            |
    //  entryGroup2Tx2    entry2   entryNonPaying1
    //                                    |
    //                             entryNonPaying2

    mempool.SetSanityCheck(0);

    auto entry1 = MakeEntry(DefaultFeeRate(), MakeConfirmedInputs(1, Amount{1000000}), {}, 1);
    auto entryGroup1Tx1 = MakeEntry(CFeeRate{}, MakeConfirmedInputs(1, Amount{1000000}), {}, 2);
    auto entryGroup1Tx2 = MakeEntry(DefaultFeeRate(), {}, {std::make_tuple(entryGroup1Tx1.GetSharedTx(), 1)}, 2, entryGroup1Tx1.GetSharedTx()->GetTotalSize());
    auto entryGroup2Tx1 = MakeEntry(CFeeRate{}, {}, {std::make_tuple(entry1.GetSharedTx(), 0), std::make_tuple(entryGroup1Tx1.GetSharedTx(), 0)}, 2);
    auto entryGroup2Tx2 = MakeEntry(DefaultFeeRate(), {}, {std::make_tuple(entryGroup2Tx1.GetSharedTx(), 0)}, 5, entryGroup2Tx1.GetSharedTx()->GetTotalSize());
    auto entryNonPaying1 = MakeEntry(CFeeRate{}, {}, {std::make_tuple(entryGroup1Tx2.GetSharedTx(), 1)}, 1);
    auto entryNonPaying2 = MakeEntry(CFeeRate{}, {}, {std::make_tuple(entryNonPaying1.GetSharedTx(), 0)}, 1);
    auto entry2 = MakeEntry(DefaultFeeRate(), {}, {std::make_tuple(entryGroup2Tx1.GetSharedTx(), 1), std::make_tuple(entryGroup1Tx2.GetSharedTx(), 0)}, 3);

    auto tx1 = AddToMempool(entry1);
    auto tx2 = AddToMempool(entryGroup1Tx1);
    auto tx3 = AddToMempool(entryGroup1Tx2);
    auto tx4 = AddToMempool(entryGroup2Tx1);
    auto tx5 = AddToMempool(entryGroup2Tx2);
    auto tx6 = AddToMempool(entryNonPaying1);
    auto tx7 = AddToMempool(entryNonPaying2);
    auto tx8 = AddToMempool(entry2);

    CTxMemPoolTestAccess testAccess(mempool);
    mempool.SetSanityCheck(0);// CheckMempool checks coins also. we do not have coins in this tests
    auto journal = testAccess.getJournalBuilder().getCurrentJournal();

    for(auto entryIt: {tx1, tx2, tx3, tx4, tx5, tx8})
    {
        BOOST_ASSERT(entryIt->IsInPrimaryMempool());
        BOOST_ASSERT(JournalTester(journal).checkTxnExists(JournalEntry{*entryIt}));
    }

    for(auto entryIt: {tx6, tx7})
    {
        BOOST_ASSERT(!entryIt->IsInPrimaryMempool());
        BOOST_ASSERT(!JournalTester(journal).checkTxnExists(JournalEntry{*entryIt}));
    }

    for(auto entryIt: {tx2, tx3, tx4, tx5})
    {
        BOOST_ASSERT(entryIt->IsCPFPGroupMember());
    }

    CheckMempoolRebuild(testAccess);

};

BOOST_AUTO_TEST_CASE(journal_groups)
{
    //
    //  entry1     entryGroup1Tx1
    //                   |
    //             entryGroup1Tx1
    //
    //

    CTxMemPoolTestAccess testAccess(mempool);

    mempool.SetSanityCheck(0);
    testConfig.SetMaxGeneratedBlockSize(250000);
    CheckMempoolRebuild(testAccess);

    auto journal = testAccess.getJournalBuilder().getCurrentJournal();

    auto entry1 = MakeEntry(DefaultFeeRate(), MakeConfirmedInputs(1, Amount(1000000)), {}, 1, 0, Amount{0}, 100000);
    auto entryGroup1Tx1 = MakeEntry(CFeeRate(), MakeConfirmedInputs(1, Amount(1000000)), {}, 2, 0, Amount{0}, 100000);
    auto entryGroup1Tx2 = MakeEntry(DefaultFeeRate(), {}, {std::make_tuple(entryGroup1Tx1.GetSharedTx(), 1)}, 2, entryGroup1Tx1.GetSharedTx()->GetTotalSize(), Amount{0}, 100000);

    auto tx1 = AddToMempool(entry1);
    auto tx2 = AddToMempool(entryGroup1Tx1);
    auto tx3 = AddToMempool(entryGroup1Tx2);

    BOOST_CHECK(tx2->GetCPFPGroupId() == tx3->GetCPFPGroupId());

    std::this_thread::sleep_for(std::chrono::seconds(1));

    for(auto entryIt: {tx1, tx2, tx3})
    {
        BOOST_CHECK(entryIt->IsInPrimaryMempool());
        BOOST_CHECK(JournalTester(journal).checkTxnExists(JournalEntry{*entryIt}));
    }

    for(auto entryIt: {tx2, tx3})
    {
        BOOST_CHECK(entryIt->IsCPFPGroupMember());
    }

    auto blockTemplatePtr = CreateBlock();
    auto vtx = blockTemplatePtr->GetBlockRef()->vtx;
    BOOST_CHECK(vtx.size() == 2);
    BOOST_CHECK(entry1.GetTxId() == vtx[1]->GetId());

    testAccess.RemoveRecursive(*entry1.GetSharedTx(), {nullptr});

    blockTemplatePtr = CreateBlock();
    vtx = blockTemplatePtr->GetBlockRef()->vtx;
    BOOST_CHECK(vtx.size() == 3);
    BOOST_CHECK(entryGroup1Tx1.GetTxId() == vtx[1]->GetId());
    BOOST_CHECK(entryGroup1Tx2.GetTxId() == vtx[2]->GetId());

};

BOOST_AUTO_TEST_CASE(conflicts)
{
    mempool.SetSanityCheck(0);
    CTxMemPoolTestAccess testAccess(mempool);
    auto journal = testAccess.getJournalBuilder().getCurrentJournal();

    //               |                            |
    //     entryDoubleSpendMempool         entryToBeMined
    //               |                       |        |
    //     ------------------------------------------------------ CONTENT OF THE BLOCK(entryDoubleSpendBlock, entryToBeMined)
    //               |                       |        |
    //               +-------+        +------+        +---+
    //                       |        |                   |
    //                   entryDoubleSpendChild        entryStayInMempool
    //
    // entryDoubleSpendMempool is in conflict with entryDoubleSpendBlock causing that entryDoubleSpendChild will be removed from the mempool
    //
    //

    auto inputForDoubleSpend = MakeConfirmedInputs(1, Amount{1000000});

    auto entryDoubleSpendMempool = MakeEntry(CFeeRate(CFeeRate()), inputForDoubleSpend, {}, 1);
    auto entryDoubleSpendBlock = MakeEntry(CFeeRate(CFeeRate()), inputForDoubleSpend, {}, 2);
    auto entryToBeMined = MakeEntry(CFeeRate(), MakeConfirmedInputs(1, Amount{1000000}), {}, 2);
    auto entryDoubleSpendChild = MakeEntry(CFeeRate(Amount(20000)), {}, {std::make_tuple(entryDoubleSpendMempool.GetSharedTx(), 0), std::make_tuple(entryToBeMined.GetSharedTx(), 0)}, 1);
    auto entryStayInMempool = MakeEntry(CFeeRate(Amount(20000)), {}, {std::make_tuple(entryToBeMined.GetSharedTx(), 1)}, 1);

    auto tx1 = AddToMempool(entryDoubleSpendMempool);
    auto tx2 = AddToMempool(entryToBeMined);
    auto tx3 = AddToMempool(entryDoubleSpendChild);
    auto tx4 = AddToMempool(entryStayInMempool);

    for(auto entryIt: {tx1, tx2, tx3, tx4})
    {
        BOOST_CHECK(entryIt->IsInPrimaryMempool());
        BOOST_CHECK(JournalTester(journal).checkTxnExists(JournalEntry{*entryIt}));
    }

    mempool.RemoveForBlock({entryDoubleSpendBlock.GetSharedTx(), entryToBeMined.GetSharedTx()}, mining::CJournalChangeSetPtr{}, uint256{});

    BOOST_CHECK(tx4->IsInPrimaryMempool());
    BOOST_CHECK(JournalTester(journal).checkTxnExists(JournalEntry{*tx4}));
    BOOST_CHECK(mempool.GetTransactions().size() == 1);
}

BOOST_AUTO_TEST_SUITE_END()


#include <boost/graph/distributed/boman_et_al_graph_coloring.hpp>


#include "keystore.h"
#include "core/proto.h"
#include <boost/filesystem.hpp>
#include <stdexcept>
#include <stdio.h>

namespace beam {

namespace {

using PrivKey = ECC::Scalar::Native;
using KeyPairs = std::map<PubKey, ECC::NoLeak<PrivKey>>;
using Nonce = ECC::Scalar::Native;
using PasswordHash = ECC::NoLeak<ECC::Hash::Value>;

void gen_nonce(Nonce& nonce) {
    ECC::Scalar sc;
    uint64_t seed;

    // here we want to read as little as possible from slow sources, TODO: review this
    ECC::GenRandom(&seed, 8);
    ECC::Hash::Processor() << seed >> sc.m_Value;

    nonce.Import(sc);
}

void hash_from_password(PasswordHash& out, const void* password, size_t passwordLen) {
    ECC::Hash::Processor()
		<< Blob(password, static_cast<uint32_t>(passwordLen))
		>> out.V;
}

void aes_decrypt(void* buffer, size_t bufferLen, const PasswordHash& key) {
    AES::Encoder enc;
    enc.Init(key.V.m_pData);
    AES::Decoder dec;
    dec.Init(enc);
    uint8_t* p = (uint8_t*)buffer;
    uint8_t* end = p + bufferLen;
    for (; p<end; p+=AES::s_BlockSize) {
        dec.Proceed(p, p);
    }
}

void aes_encrypt(void* buffer, size_t bufferLen, const PasswordHash& key) {
    AES::Encoder enc;
    enc.Init(key.V.m_pData);
    uint8_t* p = (uint8_t*)buffer;
    uint8_t* end = p + bufferLen;
    for (; p<end; p+=AES::s_BlockSize) {
        enc.Proceed(p, p);
    }
}

struct AutoClose {
    FILE* f=0;
    ~AutoClose() { if(f) fclose(f); }
};

void xor_32_bytes(uint8_t* out, const uint8_t* mask) {
    for (int i=0; i<32; ++i) {
        *(out++) ^= *(mask++);
    }
}

FILE* open_file(const std::string& fileName, bool newFile)
{
    FILE* f = nullptr;
#ifdef WIN32
    if (_wfopen_s(&f, Utf8toUtf16(fileName.c_str()).c_str(), newFile ? L"w+b" : L"a+b") != 0)
    {
        f = nullptr;
    }
#else
    f = fopen(fileName.c_str(), newFile ? "w+b" : "a+b");
#endif

    if (!f) {
        throw KeyStoreException(std::string("keystore: cannot open file ") + fileName);
    }
    return f;
}

void read_keystore_file(KeyPairs& out, const std::string& fileName, const PasswordHash& key) {
    AutoClose a;
    a.f = open_file(fileName, false);

    fseek(a.f, 0, SEEK_END);
    size_t size = ftell(a.f);

    if (size == 0) {
        return;
    }

    static const size_t MAX_FILE_SIZE = 64*2000;

    if ((size % 64) != 0 || size > MAX_FILE_SIZE) {
        fclose(a.f);
        throw KeyStoreException(std::string("keystore: invalid file size: ") + fileName);
    }

    void* buffer = alloca(size);

    rewind(a.f);
    auto bytesRead = fread(buffer, 1, size, a.f);
    if (bytesRead != size) {
        throw KeyStoreException(std::string("keystore: file read error: ") + fileName);
    }

    aes_decrypt(buffer, size, key);

    uint8_t* p = (uint8_t*)buffer;
    const uint8_t* end = p + size;
    PubKey pubKey;
    for (; p<end; p += 64) {
        xor_32_bytes(p, p + 32);
        memcpy(pubKey.m_pData, p, 32);
        memcpy(&(out[pubKey].V), p+32, 32);
    }
}

void write_keystore_file(const KeyPairs& in, const std::string& fileName, const PasswordHash& key) {
    std::string newFileName = fileName + ".new";

    size_t size = in.size() * 64;

    {
        AutoClose a;
        a.f = open_file(newFileName, true);

        if (size == 0)
            return;

        void* buffer = alloca(size);

        uint8_t* p = (uint8_t*)buffer;
        for (const auto& kp : in) {
            memcpy(p, kp.first.m_pData, 32);
            memcpy(p + 32, &(kp.second.V), 32);
            xor_32_bytes(p, p + 32);
            p += 64;
        }

        aes_encrypt(buffer, size, key);

        if (size != fwrite(buffer, 1, size, a.f)) {
            throw KeyStoreException(std::string("keystore: cannot write file ") + newFileName);
        }
    }

#ifdef WIN32
    boost::filesystem::rename(Utf8toUtf16(newFileName.c_str()), Utf8toUtf16(fileName.c_str()));
#else
    boost::filesystem::rename(newFileName, fileName);
#endif
}

} //namespace

class LocalFileKeystore : public IKeyStore {
public:
    LocalFileKeystore(const IKeyStore::Options& options, const void* password, size_t passwordLen) :
        _fileName(options.fileName)
    {
        hash_from_password(_pass, password, passwordLen);
        bool allEnabled = (options.flags & Options::enable_all_keys) != 0;
        if (allEnabled) {
            read_keystore_file_and_check();
        } else {
            enable_keys(options.enableKeys);
        }
    }

private:
    void read_keystore_file_and_check() {
        read_keystore_file(_keyPairs, _fileName, _pass);
        if (_keyPairs.empty()) return;
        static const char data[] = "xoxoxoxo";
        ByteBuffer buf;
        std::string errorMsg(std::string("keystore: file corrupted: ") + _fileName);
        if (!encrypt(buf, data, sizeof(data), _keyPairs.begin()->first)) {
            throw KeyStoreException(errorMsg);
        }
        uint8_t* out=0;
        uint32_t size=0;
        if (
            !decrypt(out, size, buf, _keyPairs.begin()->first) ||
            size != sizeof(data) ||
            memcmp(data, out, size) != 0
        ) {
            throw KeyStoreException(errorMsg);
        }
    }

    void gen_keypair(PubKey& pubKey) override {
        ECC::NoLeak<PrivKey> privKey;
        gen_nonce(privKey.V);
        proto::Sk2Pk(pubKey, privKey.V);
        memcpy(&(_unsaved[pubKey].V), &privKey.V, 32);
    }

    void save_keypair(const PubKey& pubKey, bool enable) override {
        auto it = _unsaved.find(pubKey);
        if (it == _unsaved.end()) {
            return;
        }
        memcpy(&(_keyPairs[pubKey].V), &(it->second.V), 32);
        write_keystore_file(_keyPairs, _fileName, _pass);
        if (!enable) {
            _keyPairs.erase(pubKey);
        }
        _unsaved.erase(it);
    }

    size_t size() override {
        return _keyPairs.size();
    }

    void change_password(const void* password, size_t passwordLen) override {
        KeyPairs savedKeys;
        read_keystore_file(savedKeys, _fileName, _pass);
        hash_from_password(_pass, password, passwordLen);
        write_keystore_file(savedKeys, _fileName, _pass);
    }

    void get_enabled_keys(std::set<PubKey>& enabledKeys) override {
        enabledKeys.clear();
        for (const auto& p : _keyPairs) {
            enabledKeys.insert(p.first);
        }
    }

    void enable_keys(const std::set<PubKey>& enableKeys) override {
        _keyPairs.clear();
        if (enableKeys.empty())
            return;
        read_keystore_file_and_check();
        if (_keyPairs.empty())
            return;
        std::set<PubKey> toBeErased;
        for (const auto& p : _keyPairs) {
            if (enableKeys.count(p.first) == 0) {
                toBeErased.insert(p.first);
            }
        }
        for (const auto& k : toBeErased) {
            _keyPairs.erase(k);
        }
    }

    void disable_key(const PubKey& pubKey) override {
        _keyPairs.erase(pubKey);
    }

    void erase_key(const PubKey& pubKey) override {
        size_t s = _keyPairs.size();
        _keyPairs.erase(pubKey);
        if (s != _keyPairs.size()) {
            write_keystore_file(_keyPairs, _fileName, _pass);
        }
    }

    bool encrypt(ByteBuffer& out, const void* data, size_t size, const PubKey& pubKey) override {
        Nonce nonce;
        gen_nonce(nonce);
        return proto::BbsEncrypt(out, pubKey, nonce, data, static_cast<uint32_t>(size));
    }

    bool encrypt(ByteBuffer& out, const io::SerializedMsg& in, const PubKey& pubKey) override {
        io::SharedBuffer msg = io::normalize(in, false);
        return encrypt(out, msg.data, msg.size, pubKey);
    }

    bool decrypt(uint8_t*& out, uint32_t& size, ByteBuffer& buffer, const PubKey& pubKey) override {
        auto it = _keyPairs.find(pubKey);
        if (it == _keyPairs.end()) {
            return false;
        }
        out = &buffer.at(0);
        size = static_cast<uint32_t>(buffer.size());
        return proto::BbsDecrypt(out, size, (PrivKey&)it->second.V);
    }

    std::string _fileName;
    KeyPairs _keyPairs;
    KeyPairs _unsaved;

    // TODO: use locked in memory secure buffer
    PasswordHash _pass;
};

IKeyStore::Ptr IKeyStore::create(const IKeyStore::Options& options, const void* password, size_t passwordLen) {
    static const std::string errMsgPrefix("keystore create: ");

    Ptr ptr;

    if (options.flags & Options::local_file) {
        if (options.fileName.empty() || passwordLen == 0) {
            throw KeyStoreException(errMsgPrefix + "empty file name or key");
        }
        ptr = std::make_shared<LocalFileKeystore>(options, password, passwordLen);

    } else {
        throw KeyStoreException(errMsgPrefix + "invalid options");
    }

    return ptr;
}

} //namespace