/*
 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 <boost/timer.hpp>
#include <orea/cube/cubewriter.hpp>
#include <orea/cube/sensicube.hpp>
#include <orea/engine/sensitivityanalysis.hpp>
#include <orea/engine/valuationengine.hpp>
#include <orea/scenario/clonescenariofactory.hpp>
#include <ored/utilities/log.hpp>
#include <ored/utilities/osutils.hpp>
#include <ored/utilities/to_string.hpp>
#include <ql/errors.hpp>
#include <ql/instruments/forwardrateagreement.hpp>
#include <ql/instruments/makeois.hpp>
#include <ql/instruments/makevanillaswap.hpp>
#include <ql/math/solvers1d/newtonsafe.hpp>
#include <ql/math/comparison.hpp>
#include <ql/pricingengines/capfloor/bacheliercapfloorengine.hpp>
#include <ql/pricingengines/capfloor/blackcapfloorengine.hpp>
#include <ql/pricingengines/swap/discountingswapengine.hpp>
#include <ql/termstructures/yield/oisratehelper.hpp>
#include <qle/instruments/crossccybasisswap.hpp>
#include <qle/instruments/deposit.hpp>
#include <qle/instruments/fxforward.hpp>
#include <qle/pricingengines/crossccyswapengine.hpp>
#include <qle/pricingengines/depositengine.hpp>
#include <qle/pricingengines/discountingfxforwardengine.hpp>

using namespace QuantLib;
using namespace QuantExt;
using namespace std;
using namespace ore::data;

namespace ore {
namespace analytics {

SensitivityAnalysis::SensitivityAnalysis(
    const boost::shared_ptr<ore::data::Portfolio>& portfolio, const boost::shared_ptr<ore::data::Market>& market,
    const string& marketConfiguration, const boost::shared_ptr<ore::data::EngineData>& engineData,
    const boost::shared_ptr<ScenarioSimMarketParameters>& simMarketData,
    const boost::shared_ptr<SensitivityScenarioData>& sensitivityData, const Conventions& conventions,
    const bool recalibrateModels, const CurveConfigurations& curveConfigs,
    const TodaysMarketParameters& todaysMarketParams, const bool nonShiftedBaseCurrencyConversion,
    std::vector<boost::shared_ptr<ore::data::EngineBuilder>> extraEngineBuilders,
    std::vector<boost::shared_ptr<ore::data::LegBuilder>> extraLegBuilders, const bool continueOnError)
    : market_(market), marketConfiguration_(marketConfiguration), asof_(market->asofDate()),
      simMarketData_(simMarketData), sensitivityData_(sensitivityData), conventions_(conventions),
      recalibrateModels_(recalibrateModels), curveConfigs_(curveConfigs), todaysMarketParams_(todaysMarketParams),
      overrideTenors_(false), nonShiftedBaseCurrencyConversion_(nonShiftedBaseCurrencyConversion),
      extraEngineBuilders_(extraEngineBuilders), extraLegBuilders_(extraLegBuilders), continueOnError_(continueOnError),
      engineData_(engineData), portfolio_(portfolio), initialized_(false), computed_(false) {}

std::vector<boost::shared_ptr<ValuationCalculator>> SensitivityAnalysis::buildValuationCalculators() const {
    vector<boost::shared_ptr<ValuationCalculator>> calculators;
    if (nonShiftedBaseCurrencyConversion_) // use "original" FX rates to convert sensi to base currency
        calculators.push_back(boost::make_shared<NPVCalculatorFXT0>(simMarketData_->baseCcy(), market_));
    else // use the scenario FX rate when converting sensi to base currency
        calculators.push_back(boost::make_shared<NPVCalculator>(simMarketData_->baseCcy()));
    return calculators;
}

void SensitivityAnalysis::initialize(boost::shared_ptr<NPVSensiCube>& cube) {
    LOG("Build Sensitivity Scenario Generator and Simulation Market");
    initializeSimMarket();

    LOG("Build Engine Factory and rebuild portfolio");
    boost::shared_ptr<EngineFactory> factory = buildFactory(extraEngineBuilders_, extraLegBuilders_);
    resetPortfolio(factory);
    if (recalibrateModels_)
        modelBuilders_ = factory->modelBuilders();
    else
        modelBuilders_.clear();

    if (!cube) {
        LOG("Build the cube object to store sensitivities");
        initializeCube(cube);
    }

    sensiCube_ = boost::make_shared<SensitivityCube>(cube, scenarioGenerator_->scenarioDescriptions(),
                                                     scenarioGenerator_->shiftSizes());
    initialized_ = true;
}

void SensitivityAnalysis::generateSensitivities(boost::shared_ptr<NPVSensiCube> cube) {

    QL_REQUIRE(!initialized_, "unexpected state of SensitivitiesAnalysis object");

    // initialize the helper member objects
    initialize(cube);
    QL_REQUIRE(initialized_, "SensitivitiesAnalysis member objects not correctly initialized");
    boost::shared_ptr<DateGrid> dg = boost::make_shared<DateGrid>("1,0W", NullCalendar());
    vector<boost::shared_ptr<ValuationCalculator>> calculators = buildValuationCalculators();
    ValuationEngine engine(asof_, dg, simMarket_, modelBuilders_);
    for (auto const& i : this->progressIndicators())
        engine.registerProgressIndicator(i);
    LOG("Run Sensitivity Scenarios");
    engine.buildCube(portfolio_, cube, calculators);

    computed_ = true;
    LOG("Sensitivity analysis completed");
}

void SensitivityAnalysis::initializeSimMarket(boost::shared_ptr<ScenarioFactory> scenFact) {

    LOG("Initialise sim market for sensitivity analysis (continueOnError=" << std::boolalpha << continueOnError_
        << ")");
    simMarket_ = boost::make_shared<ScenarioSimMarket>(market_, simMarketData_, conventions_, marketConfiguration_, curveConfigs_, todaysMarketParams_, continueOnError_);

    LOG("Sim market initialised for sensitivity analysis");

    LOG("Create scenario factory for sensitivity analysis");
    boost::shared_ptr<Scenario> baseScenario = simMarket_->baseScenario();
    boost::shared_ptr<ScenarioFactory> scenarioFactory =
        scenFact ? scenFact : boost::make_shared<CloneScenarioFactory>(baseScenario);
    LOG("Scenario factory created for sensitivity analysis");

    LOG("Create scenario generator for sensitivity analysis (continueOnError=" << std::boolalpha << continueOnError_
                                                                               << ")");
    scenarioGenerator_ = boost::make_shared<SensitivityScenarioGenerator>(
        sensitivityData_, baseScenario, simMarketData_, scenarioFactory, overrideTenors_, continueOnError_);
    LOG("Scenario generator created for sensitivity analysis");

    // Set simulation market's scenario generator
    simMarket_->scenarioGenerator() = scenarioGenerator_;
}

boost::shared_ptr<EngineFactory>
SensitivityAnalysis::buildFactory(const std::vector<boost::shared_ptr<EngineBuilder>> extraBuilders,
                                  const std::vector<boost::shared_ptr<LegBuilder>> extraLegBuilders) const {
    map<MarketContext, string> configurations;
    configurations[MarketContext::pricing] = marketConfiguration_;
    boost::shared_ptr<EngineFactory> factory =
        boost::make_shared<EngineFactory>(engineData_, simMarket_, configurations, extraBuilders, extraLegBuilders);
    return factory;
}

void SensitivityAnalysis::resetPortfolio(const boost::shared_ptr<EngineFactory>& factory) {
    portfolio_->reset();
    portfolio_->build(factory);
}

void SensitivityAnalysis::initializeCube(boost::shared_ptr<NPVSensiCube>& cube) const {
    cube = boost::make_shared<DoublePrecisionSensiCube>(portfolio_->ids(), asof_, scenarioGenerator_->samples());
}

Real getShiftSize(const RiskFactorKey& key, const SensitivityScenarioData& sensiParams,
                  const boost::shared_ptr<ScenarioSimMarket>& simMarket, const string& marketConfiguration) {

    Date asof = simMarket->asofDate();
    RiskFactorKey::KeyType keytype = key.keytype;
    string keylabel = key.name;
    Real shiftSize = 0.0;
    Real shiftMult = 1.0;
    switch (keytype) {
    case RiskFactorKey::KeyType::FXSpot: {
        auto itr = sensiParams.fxShiftData().find(keylabel);
        QL_REQUIRE(itr != sensiParams.fxShiftData().end(), "shiftData not found for " << keylabel);
        shiftSize = itr->second.shiftSize;
        if (parseShiftType(itr->second.shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            shiftMult = simMarket->fxSpot(keylabel, marketConfiguration)->value();
        }
    } break;
    case RiskFactorKey::KeyType::EquitySpot: {
        auto itr = sensiParams.equityShiftData().find(keylabel);
        QL_REQUIRE(itr != sensiParams.equityShiftData().end(), "shiftData not found for " << keylabel);
        shiftSize = itr->second.shiftSize;
        if (parseShiftType(itr->second.shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            shiftMult = simMarket->equitySpot(keylabel, marketConfiguration)->value();
        }
    } break;
    case RiskFactorKey::KeyType::DiscountCurve: {
        string ccy = keylabel;
        auto itr = sensiParams.discountCurveShiftData().find(ccy);
        QL_REQUIRE(itr != sensiParams.discountCurveShiftData().end(), "shiftData not found for " << ccy);
        shiftSize = itr->second->shiftSize;
        if (parseShiftType(itr->second->shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            Size keyIdx = key.index;
            Period p = itr->second->shiftTenors[keyIdx];
            Handle<YieldTermStructure> yts = simMarket->discountCurve(ccy, marketConfiguration);
            Time t = yts->dayCounter().yearFraction(asof, asof + p);
            Real zeroRate = yts->zeroRate(t, Continuous);
            shiftMult = zeroRate;
        }
    } break;
    case RiskFactorKey::KeyType::IndexCurve: {
        string idx = keylabel;
        auto itr = sensiParams.indexCurveShiftData().find(idx);
        QL_REQUIRE(itr != sensiParams.indexCurveShiftData().end(), "shiftData not found for " << idx);
        shiftSize = itr->second->shiftSize;
        if (parseShiftType(itr->second->shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            Size keyIdx = key.index;
            Period p = itr->second->shiftTenors[keyIdx];
            Handle<YieldTermStructure> yts = simMarket->iborIndex(idx, marketConfiguration)->forwardingTermStructure();
            Time t = yts->dayCounter().yearFraction(asof, asof + p);
            Real zeroRate = yts->zeroRate(t, Continuous);
            shiftMult = zeroRate;
        }
    } break;
    case RiskFactorKey::KeyType::YieldCurve: {
        string yc = keylabel;
        auto itr = sensiParams.yieldCurveShiftData().find(yc);
        QL_REQUIRE(itr != sensiParams.yieldCurveShiftData().end(), "shiftData not found for " << yc);
        shiftSize = itr->second->shiftSize;
        if (parseShiftType(itr->second->shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            Size keyIdx = key.index;
            Period p = itr->second->shiftTenors[keyIdx];
            Handle<YieldTermStructure> yts = simMarket->yieldCurve(yc, marketConfiguration);
            Time t = yts->dayCounter().yearFraction(asof, asof + p);
            Real zeroRate = yts->zeroRate(t, Continuous);
            shiftMult = zeroRate;
        }
    } break;
    case RiskFactorKey::KeyType::DividendYield: {
        string eq = keylabel;
        auto itr = sensiParams.dividendYieldShiftData().find(eq);
        QL_REQUIRE(itr != sensiParams.dividendYieldShiftData().end(), "shiftData not found for " << eq);
        shiftSize = itr->second->shiftSize;

        if (parseShiftType(itr->second->shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            Size keyIdx = key.index;
            Period p = itr->second->shiftTenors[keyIdx];
            Handle<YieldTermStructure> ts = simMarket->equityDividendCurve(eq, marketConfiguration);
            Time t = ts->dayCounter().yearFraction(asof, asof + p);
            Real zeroRate = ts->zeroRate(t, Continuous);
            shiftMult = zeroRate;
        }
    } break;
    case RiskFactorKey::KeyType::FXVolatility: {
        string pair = keylabel;
        auto itr = sensiParams.fxVolShiftData().find(pair);
        QL_REQUIRE(itr != sensiParams.fxVolShiftData().end(), "shiftData not found for " << pair);
        shiftSize = itr->second.shiftSize;
        if (parseShiftType(itr->second.shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            vector<Real> strikes = sensiParams.fxVolShiftData().at(pair).shiftStrikes;
            QL_REQUIRE(strikes.size() == 1 && close_enough(strikes[0], 0), "Only ATM FX vols supported");
            Real atmFwd = 0.0; // hardcoded, since only ATM supported
            Size keyIdx = key.index;
            Period p = itr->second.shiftExpiries[keyIdx];
            Handle<BlackVolTermStructure> vts = simMarket->fxVol(pair, marketConfiguration);
            Time t = vts->dayCounter().yearFraction(asof, asof + p);
            Real atmVol = vts->blackVol(t, atmFwd);
            shiftMult = atmVol;
        }
    } break;
    case RiskFactorKey::KeyType::EquityVolatility: {
        string pair = keylabel;
        auto itr = sensiParams.equityVolShiftData().find(pair);
        QL_REQUIRE(itr != sensiParams.equityVolShiftData().end(), "shiftData not found for " << pair);
        shiftSize = itr->second.shiftSize;
        if (parseShiftType(itr->second.shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            Size keyIdx = key.index;
            Period p = itr->second.shiftExpiries[keyIdx];
            Handle<BlackVolTermStructure> vts = simMarket->equityVol(pair, marketConfiguration);
            Time t = vts->dayCounter().yearFraction(asof, asof + p);
            Real atmVol = vts->blackVol(t, Null<Real>());
            shiftMult = atmVol;
        }
    } break;
    case RiskFactorKey::KeyType::SwaptionVolatility: {
        string ccy = keylabel;
        auto itr = sensiParams.swaptionVolShiftData().find(ccy);
        QL_REQUIRE(itr != sensiParams.swaptionVolShiftData().end(), "shiftData not found for " << ccy);
        shiftSize = itr->second.shiftSize;
        if (parseShiftType(itr->second.shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            vector<Real> strikes = itr->second.shiftStrikes;
            vector<Period> tenors = itr->second.shiftTerms;
            vector<Period> expiries = itr->second.shiftExpiries;
            Size keyIdx = key.index;
            Size expIdx = keyIdx / (tenors.size() * strikes.size());
            Period p_exp = expiries[expIdx];
            Size tenIdx = keyIdx % tenors.size();
            Period p_ten = tenors[tenIdx];
            Real strike = Null<Real>(); // for cubes this will be ATM
            Handle<SwaptionVolatilityStructure> vts = simMarket->swaptionVol(ccy, marketConfiguration);
            // Time t_exp = vts->dayCounter().yearFraction(asof, asof + p_exp);
            // Time t_ten = vts->dayCounter().yearFraction(asof, asof + p_ten);
            // Real atmVol = vts->volatility(t_exp, t_ten, Null<Real>());
            Real vol = vts->volatility(p_exp, p_ten, strike);
            shiftMult = vol;
        }
    } break;
    case RiskFactorKey::KeyType::YieldVolatility: {
        string securityId = keylabel;
        auto itr = sensiParams.yieldVolShiftData().find(securityId);
        QL_REQUIRE(itr != sensiParams.yieldVolShiftData().end(), "shiftData not found for " << securityId);
        shiftSize = itr->second.shiftSize;
        if (parseShiftType(itr->second.shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            vector<Real> strikes = itr->second.shiftStrikes;
            QL_REQUIRE(strikes.size() == 1 && close_enough(strikes[0], 0.0),
                       "shift strikes should be {0.0} for yield volatilities");
            vector<Period> tenors = itr->second.shiftTerms;
            vector<Period> expiries = itr->second.shiftExpiries;
            Size keyIdx = key.index;
            Size expIdx = keyIdx / (tenors.size() * strikes.size());
            Period p_exp = expiries[expIdx];
            Size tenIdx = keyIdx % tenors.size();
            Period p_ten = tenors[tenIdx];
            Real strike = Null<Real>(); // for cubes this will be ATM
            Handle<SwaptionVolatilityStructure> vts = simMarket->yieldVol(securityId, marketConfiguration);
            Real vol = vts->volatility(p_exp, p_ten, strike);
            shiftMult = vol;
        }
    } break;
    case RiskFactorKey::KeyType::OptionletVolatility: {
        string ccy = keylabel;
        auto itr = sensiParams.capFloorVolShiftData().find(ccy);
        QL_REQUIRE(itr != sensiParams.capFloorVolShiftData().end(), "shiftData not found for " << ccy);
        shiftSize = itr->second.shiftSize;
        if (parseShiftType(itr->second.shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            vector<Real> strikes = itr->second.shiftStrikes;
            vector<Period> expiries = itr->second.shiftExpiries;
            QL_REQUIRE(strikes.size() > 0, "Only strike capfloor vols supported");
            Size keyIdx = key.index;
            Size expIdx = keyIdx / strikes.size();
            Period p_exp = expiries[expIdx];
            Size strIdx = keyIdx % strikes.size();
            Real strike = strikes[strIdx];
            Handle<OptionletVolatilityStructure> vts = simMarket->capFloorVol(ccy, marketConfiguration);
            Time t_exp = vts->dayCounter().yearFraction(asof, asof + p_exp);
            Real vol = vts->volatility(t_exp, strike);
            shiftMult = vol;
        }
    } break;
    case RiskFactorKey::KeyType::CDSVolatility: {
        string name = keylabel;
        auto itr = sensiParams.cdsVolShiftData().find(name);
        QL_REQUIRE(itr != sensiParams.cdsVolShiftData().end(), "shiftData not found for " << name);
        shiftSize = itr->second.shiftSize;
        if (parseShiftType(itr->second.shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            vector<Period> expiries = itr->second.shiftExpiries;
            Size keyIdx = key.index;
            Size expIdx = keyIdx;
            Period p_exp = expiries[expIdx];
            Handle<BlackVolTermStructure> vts = simMarket->cdsVol(name, marketConfiguration);
            Time t_exp = vts->dayCounter().yearFraction(asof, asof + p_exp);
            Real strike = 0.0; // FIXME
            Real atmVol = vts->blackVol(t_exp, strike);
            shiftMult = atmVol;
        }
    } break;
    case RiskFactorKey::KeyType::SurvivalProbability: {
        string name = keylabel;
        auto itr = sensiParams.creditCurveShiftData().find(name);
        QL_REQUIRE(itr != sensiParams.creditCurveShiftData().end(), "shiftData not found for " << name);
        shiftSize = itr->second->shiftSize;
        if (parseShiftType(itr->second->shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            Size keyIdx = key.index;
            Period p = itr->second->shiftTenors[keyIdx];
            Handle<DefaultProbabilityTermStructure> ts = simMarket->defaultCurve(name, marketConfiguration);
            Time t = ts->dayCounter().yearFraction(asof, asof + p);
            Real prob = ts->survivalProbability(t);
            shiftMult = -std::log(prob) / t;
        }
    } break;
    case RiskFactorKey::KeyType::BaseCorrelation: {
        string name = keylabel;
        auto itr = sensiParams.baseCorrelationShiftData().find(name);
        QL_REQUIRE(itr != sensiParams.baseCorrelationShiftData().end(), "shiftData not found for " << name);
        shiftSize = itr->second.shiftSize;
        if (parseShiftType(itr->second.shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            vector<Real> lossLevels = itr->second.shiftLossLevels;
            vector<Period> terms = itr->second.shiftTerms;
            Size keyIdx = key.index;
            Size lossLevelIdx = keyIdx / terms.size();
            Real lossLevel = lossLevels[lossLevelIdx];
            Size termIdx = keyIdx % terms.size();
            Period term = terms[termIdx];
            Handle<BilinearBaseCorrelationTermStructure> ts = simMarket->baseCorrelation(name, marketConfiguration);
            Real bc = ts->correlation(asof + term, lossLevel, true); // extrapolate
            shiftMult = bc;
        }
    } break;
    case RiskFactorKey::KeyType::ZeroInflationCurve: {
        string idx = keylabel;
        auto itr = sensiParams.zeroInflationCurveShiftData().find(idx);
        QL_REQUIRE(itr != sensiParams.zeroInflationCurveShiftData().end(), "shiftData not found for " << idx);
        shiftSize = itr->second->shiftSize;
        if (parseShiftType(itr->second->shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            Size keyIdx = key.index;
            Period p = itr->second->shiftTenors[keyIdx];
            Handle<ZeroInflationTermStructure> yts =
                simMarket->zeroInflationIndex(idx, marketConfiguration)->zeroInflationTermStructure();
            Time t = yts->dayCounter().yearFraction(asof, asof + p);
            Real zeroRate = yts->zeroRate(t);
            shiftMult = zeroRate;
        }
    } break;
    case RiskFactorKey::KeyType::YoYInflationCurve: {
        string idx = keylabel;
        auto itr = sensiParams.yoyInflationCurveShiftData().find(idx);
        QL_REQUIRE(itr != sensiParams.yoyInflationCurveShiftData().end(), "shiftData not found for " << idx);
        shiftSize = itr->second->shiftSize;
        if (parseShiftType(itr->second->shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            Size keyIdx = key.index;
            Period p = sensiParams.yoyInflationCurveShiftData().at(idx)->shiftTenors[keyIdx];
            Handle<YoYInflationTermStructure> yts =
                simMarket->yoyInflationIndex(idx, marketConfiguration)->yoyInflationTermStructure();
            Time t = yts->dayCounter().yearFraction(asof, asof + p);
            Real yoyRate = yts->yoyRate(t);
            shiftMult = yoyRate;
        }
    } break;
    case RiskFactorKey::KeyType::YoYInflationCapFloorVolatility: {
        string name = keylabel;
        auto itr = sensiParams.yoyInflationCapFloorVolShiftData().find(name);
        QL_REQUIRE(itr != sensiParams.yoyInflationCapFloorVolShiftData().end(), "shiftData not found for " << name);
        shiftSize = itr->second.shiftSize;
        if (parseShiftType(itr->second.shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            vector<Real> strikes = itr->second.shiftStrikes;
            vector<Period> expiries = itr->second.shiftExpiries;
            QL_REQUIRE(strikes.size() > 0, "Only strike yoy inflation capfloor vols supported");
            Size keyIdx = key.index;
            Size expIdx = keyIdx / strikes.size();
            Period p_exp = expiries[expIdx];
            Size strIdx = keyIdx % strikes.size();
            Real strike = strikes[strIdx];
            Handle<QuantExt::YoYOptionletVolatilitySurface> vts = simMarket->yoyCapFloorVol(name, marketConfiguration);
            Real vol = vts->volatility(p_exp, strike, vts->observationLag());
            shiftMult = vol;
        }
    } break;
    case RiskFactorKey::KeyType::CommodityCurve: {
        auto it = sensiParams.commodityCurveShiftData().find(keylabel);
        QL_REQUIRE(it != sensiParams.commodityCurveShiftData().end(), "shiftData not found for " << keylabel);
        shiftSize = it->second->shiftSize;
        if (parseShiftType(it->second->shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            Period p = it->second->shiftTenors[key.index];
            Handle<PriceTermStructure> priceCurve = simMarket->commodityPriceCurve(keylabel, marketConfiguration);
            Time t = priceCurve->dayCounter().yearFraction(asof, asof + p);
            shiftMult = priceCurve->price(t);
        }
    } break;
    case RiskFactorKey::KeyType::CommodityVolatility: {
        auto it = sensiParams.commodityVolShiftData().find(keylabel);
        QL_REQUIRE(it != sensiParams.commodityVolShiftData().end(), "shiftData not found for " << keylabel);

        shiftSize = it->second.shiftSize;
        if (parseShiftType(it->second.shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            Size moneynessIndex = key.index / it->second.shiftExpiries.size();
            Size expiryIndex = key.index % it->second.shiftExpiries.size();
            Real moneyness = it->second.shiftStrikes[moneynessIndex];
            Period expiry = it->second.shiftExpiries[expiryIndex];
            Real spotValue = simMarket->commoditySpot(keylabel, marketConfiguration)->value();
            Handle<BlackVolTermStructure> vts = simMarket->commodityVolatility(keylabel, marketConfiguration);
            Time t = vts->dayCounter().yearFraction(asof, asof + expiry);
            Real vol = vts->blackVol(t, moneyness * spotValue);
            shiftMult = vol;
        }
    } break;
    case RiskFactorKey::KeyType::SecuritySpread: {
        auto itr = sensiParams.securityShiftData().find(keylabel);
        QL_REQUIRE(itr != sensiParams.securityShiftData().end(), "shiftData not found for " << keylabel);
        shiftSize = itr->second.shiftSize;
        if (parseShiftType(itr->second.shiftType) == SensitivityScenarioGenerator::ShiftType::Relative) {
            shiftMult = simMarket->securitySpread(keylabel, marketConfiguration)->value();
        }
    } break;
    default:
        // KeyType::CPIIndex does not get shifted
        QL_FAIL("KeyType not supported yet - " << keytype);
    }
    Real realShift = shiftSize * shiftMult;
    return realShift;
}

} // namespace analytics
} // namespace ore


/* A simple simulator for a single quadcopter.
 *
 */

#include <memory>
#include <Eigen/Dense>
#include <stdio.h>
#include <unistd.h>
#include <boost/program_options.hpp>
#include <boost/property_tree/json_parser.hpp>
#include <boost/property_tree/ptree.hpp>
#include <ctime>
#include <opencv2/opencv.hpp>
#include <random>
#include <string>

#include "Common/Math/Vec3.hpp"
#include "Common/Math/Trajectory.hpp"
#include "Components/TrajectoryGenerator/RapidTrajectoryGenerator.hpp"

#include "Common/Time/ManualTimer.hpp"
#include "Common/Time/HardwareTimer.hpp"
#include "Components/Simulation/Quadcopter_T.hpp"
#include "Components/Simulation/UWBNetwork.hpp"
#include "Components/Simulation/CommunicationsDelay.hpp"
#include "Components/Simulation/ArucoCamera.hpp"

#include "Components/Offboard/MocapStateEstimator.hpp"
#include "Components/Offboard/QuadcopterController.hpp"
#include "Components/Offboard/SafetyNet.hpp"

#include "Components/Logic/QuadcopterLogic.hpp"

#include <fstream>

// AirSim includes
//#include "api/RpcLibClientBase.hpp"
#include "vehicles/multirotor/api/MultirotorRpcLibClient.hpp"
#include "common/common_utils/FileSystem.hpp"
#include <iostream> // for image saving

#include "Components/DepthImagePlanner/DepthImagePlanner.hpp"

using namespace std;
using namespace Offboard;
using namespace CommonMath;
using namespace RapidQuadrocopterTrajectoryGenerator;
using namespace RectangularPyramidPlanner;

enum {
  CTRL_ONBOARD_UWB,
  CTRL_OFFBOARD_RATES,
} controllerType;

template<typename Real>
string toCSV(const Vec3<Real> v) {
  stringstream ss;
  ss << v.x << "," << v.y << "," << v.z << ",";
  return ss.str();
}
string toCSV(const Rotationf r) {
  stringstream ss;
  ss << r[0] << "," << r[1] << "," << r[2] << "," << r[3] << ",";
  return ss.str();
}
string toCSV(const std::vector<float> v) {
  stringstream ss;
  for (int i = 0; i < v.size(); i++) {
    ss << v[i] << ",";
  }
  return ss.str();
}

class ExplorationCost {
 public:
  ExplorationCost(Vec3d estPos, Rotationd estAtt, Vec3d goalWorld,
                  Rotationd const depthCamAtt) {
    _estPos = estPos;
    _estAtt = estAtt;
    _goalWorld = goalWorld;
    _depthCamAtt = depthCamAtt;
  }
  void updateWorld(Vec3d estPos, Rotationd estAtt, Vec3d goalWorld) {
    _estPos = estPos;
    _estAtt = estAtt;
    _goalWorld = goalWorld;
  }
  static double GetTrajCostWrapper(
      void *ptr2obj,
      RapidQuadrocopterTrajectoryGenerator::RapidTrajectoryGenerator &traj) {
    ExplorationCost *explorationCost = (ExplorationCost*) ptr2obj;
    return explorationCost->GetTrajCost(traj);
  }

  double GetTrajCost(
      RapidQuadrocopterTrajectoryGenerator::RapidTrajectoryGenerator &traj) {
    double duration = traj.GetFinalTime();
    Vec3d Pi_C = traj.GetPosition(duration);  // position of camera at end of trajectory relative to start written in camera frame

    Vec3d G_W = _goalWorld;                               // goal in world frame
    Vec3d G_C = _depthCamAtt.Inverse() * _estAtt.Inverse()
        * (_goalWorld - _estPos);  // goal in camera frame
    Vec3d S_C = Vec3d(0, 0, 0);                // starting point in camera frame

    double SG = (G_C - S_C).GetNorm2();
    double PiG = (G_C - Pi_C).GetNorm2();
    return -(SG - PiG) / duration;
  }
 private:
  Vec3d _goalWorld;
  Rotationd _depthCamAtt;
  Vec3d _estPos;
  Rotationd _estAtt;
};

int main(void) {
////////////////////////////////////////////

//depthImageSetting
  // Far and near clipping plane distance in Unity
  double far = 10.0;
  double depthScale = far / (256.0);  //output depth image in uint_8
  Rotationd const depthCamAtt = Rotationd::FromEulerYPR(-90.0 * M_PI / 180.0,
                                                        0 * M_PI / 180.0,
                                                        -90 * M_PI / 180.0);

  cv::Mat depthImage;
  cv::Mat depthImage8;
  double focalLength;
  unsigned depthImageHeight;
  unsigned depthImageWidth;
  Rotationd trajAtt;
  Vec3d trajOffset;
  bool imageReady;
  unsigned plannedTrajCount = 0;

  std::shared_ptr<RapidQuadrocopterTrajectoryGenerator::RapidTrajectoryGenerator> _traj;

//Basic timing:
  const double dt = 1.0 / 500.0;  //run the simulation at this rate
  const double startFlightTime = 5.0f;  // time for rappids to begin
  const double endTime = 8.0f;  //[s]  total simulation time
  ManualTimer simTimer;  // Tracks time in the simulation
  HardwareTimer globalTimer;  // Tracks the actual wall clock time (in the real world)

//Create the quadcopter:
  uint8_t vehicleId = 1;  //For UWB network, commands, etc.
  Onboard::QuadcopterConstants::QuadcopterType quadcopterType =
      Onboard::QuadcopterConstants::GetVehicleTypeFromID(vehicleId);
  Onboard::QuadcopterConstants vehConsts(quadcopterType);
//Create the quadcopter:
  double const mass = vehConsts.mass;  //[kg]
  double const inertia_xx = vehConsts.inertia_xx;  //[kg.m**2]
  double const inertia_yy = inertia_xx;  //[kg.m**2]
  double const inertia_zz = vehConsts.inertia_zz;  //[kg.m**2]
  double armLength = vehConsts.armLength;  //[m]
  double propThrustFromSpeedSqr = vehConsts.propellerThrustFromSpeedSqr;  //[N/(rad/s)**2]
  double propTorqueFromSpeedSqr = vehConsts.propellerTorqueFromThrust
      * vehConsts.propellerThrustFromSpeedSqr;

  double motorTimeConst = vehConsts.motorTimeConst;  // [s]
  double motorInertia = vehConsts.motorInertia;  // [kg.m**2]
  double motorMinSpeed = vehConsts.motorMinSpeed;  // [rad/s]
  double motorMaxSpeed = vehConsts.motorMaxSpeed;  // [rad/s]

//Estimate radius for planning
  double const physicalVehicleRadius = armLength * 2;
  double const vehicleRadiusPlanning = armLength * 2 * 1.5;
  double const minCollisionDist = 0.5;  //[m]

  Vec3d centreOfMassError = Vec3d(0, 0, 0);

// Mocap system:
  double const periodMocapSystem = 1.0 / 200.0;  //[s] time between mocap measurments
  double const periodOffboardMainLoop = 1.0 / 100.0;  //[s] time between offboard main loop runs
  double const periodTelemetryLoop = 1.0 / 100.0;  //[s] time between telemetry update
  double const periodOnboardLogic = 1.0 / 500.0;  //[s] time between onboard logic runs
  double const timeDelayOffboardControlLoopTrue = 0.03;  //[s]
  double const timeDelayOffboardControlLoopEstimate = 0.03;  //[s]

// AirSim client setup
  using namespace msr::airlib;
  MultirotorRpcLibClient client("", 41451, 1);
  client.confirmConnection();  // Need to have env open and running (press play)

// AirSim image capture setup
  typedef ImageCaptureBase::ImageRequest ImageRequest;
  typedef ImageCaptureBase::ImageResponse ImageResponse;
  typedef ImageCaptureBase::ImageType ImageType;
  typedef common_utils::FileSystem FileSystem;
  bool recordImage = false;  // Toggle this to record the images generated by Unity.

  /*
   * Essentially, once the simulation timer reaches a time beyond the periodBetweenImages variable
   * it will wait for Unity to generate an image. It is Unity's job
   * to let this simulator know that it has generated an image, so that the simulator may proceed.
   */

  bool requestNewImage = true;
  double const desiredFrameRate = 30.0;
  double const periodBetweenImages = 1 / desiredFrameRate;

  Eigen::Matrix<double, 3, 3> inertiaMatrix;
  inertiaMatrix << inertia_xx, 0, 0, 0, inertia_yy, 0, 0, 0, inertia_zz;

//drag coefficients
  Vec3d linDragCoeffB = Vec3d(vehConsts.linDragCoeffBx,
                              vehConsts.linDragCoeffBy,
                              vehConsts.linDragCoeffBz);

  std::shared_ptr<Simulation::Quadcopter> quad;
  quad.reset(
      new Simulation::Quadcopter(&simTimer, mass, inertiaMatrix, armLength,
                                 centreOfMassError, motorMinSpeed,
                                 motorMaxSpeed, propThrustFromSpeedSqr,
                                 propTorqueFromSpeedSqr, motorTimeConst,
                                 motorInertia, linDragCoeffB, vehicleId,
                                 quadcopterType, periodOnboardLogic));

//Offboard estimation / control code:
  std::shared_ptr<MocapStateEstimator> est;
  est.reset(
      new MocapStateEstimator(&simTimer, vehicleId,
                              timeDelayOffboardControlLoopEstimate));
  QuadcopterController ctrl;
  SafetyNet safetyNet;
  ctrl.SetParameters(vehConsts.posControl_natFreq, vehConsts.posControl_damping,
                     vehConsts.attControl_timeConst_xy,
                     vehConsts.attControl_timeConst_z);

  controllerType = CTRL_OFFBOARD_RATES;

//create an initial error:
  Vec3d const initErrPos = Vec3d(0, 0, 0);
  Rotationd const initErrAtt = Rotationd::FromEulerYPR(0 * M_PI / 180.0,
                                                       0 * M_PI / 180.0, 0);

//where we want the quadcopter to fly to:
  Vec3d desiredPosition(0, 0, 3.5);
  Vec3d desiredVelocity(0, 0, 0);
  Vec3d desiredAcc(0, 0, 0);
  Vec3d desiredAngVel(0, 0, 0);
  double desiredThrust;
  double desYawAngleDeg = 0;  //[deg]
  Vec3d const goalWorld(120.0, 0.0, 3.5);

//keep track of last set of inputs
  double previousThrust;
  Vec3d previousAngVel;

//file for logging:
  ofstream logfile;
  ofstream logfile2;
  ofstream trajLog;
  ofstream airsimSetPoseTimeLog;
  ofstream airsimGetPoseTimeLog;

  logfile.open("Logs/rappids_simulator/simulation.csv");
  logfile2.open("Logs/rappids_simulator/PictureFrequencyLog.csv");
  trajLog.open("Logs/rappids_simulator/PlannedTrajectory.csv");

  airsimSetPoseTimeLog.open("Logs/rappids_simulator/airsimSetPoseTimeLog.csv");
  airsimGetPoseTimeLog.open("Logs/rappids_simulator/airsimGetPoseTimeLog.csv");

// create header, a lot of columns
  logfile
      << "t,posx,posy,posz,velx,vely,velz,attY,attP,attR,angvelx,angvely,angvelz,m1,m2,m3,m4,"
//          "tx,ty,tz,vx,vy,vz,qx,qy,qz,qw,Pr11,Pr12,Pr13,Pr22,Pr23,Pr33,Pt11,Pt12,Pt13,Pt22,Pt23,Pt33,"
          "estposx,estposy,estposz,estvelx,estvely,estvelz,esty,estp,estr,estangx,estangy,estangz,"
          "desposx,desposy,desposz,desvelx,desvely,desvelz,panic,r1,r2,r3,r4\n";
  logfile2 << "realTime,SimulationTime,picTime,picNum\n";
  printf("Starting simulation\n");

  Timer t(&simTimer);
  Timer tReal(&globalTimer);
  Timer tPic(&globalTimer);  //Time to track the seconds between pics.
  Timer tAirSimLoop(&globalTimer);  //Measures the time between set airsim pose to pose being updated;
  Timer tGetPose(&globalTimer);  //Measures the time between two instincts between getting airsim getpose;
  quad->SetPosition(initErrPos);
  quad->SetAttitude(initErrAtt);

  Simulation::CommunicationsDelay<RadioTypes::RadioMessageDecoded::RawMessage> cmdRadioChannel(
      &simTimer, timeDelayOffboardControlLoopTrue);

  Timer timerPrint(&simTimer);
  Timer timerMocap(&simTimer);
  Timer timerOffboardMainLoop(&simTimer);
  Timer timerRequestNewImage(&simTimer);
  Timer timerTelemetryLoop(&simTimer);

  Timer trackTrajTime(&simTimer);
  bool trajPlanned = false;
  double trajEndTime;

  float lastRadioCommand[4];
  for (int i = 0; i < 4; i++) {
    lastRadioCommand[i] = 0;
  }
// Hack for now: Set the start x-y position of the airsim Position as 0,0

  auto initPosition = client.getMultirotorState().getPosition();
  double const initX = initPosition.x();
  double const initY = initPosition.y();

  printf("initx=%.3fs\n", initX);
  printf("inity=%.3fs\n", initY);

  unsigned imageCount = 0;
  tPic.Reset();
// aligns airsim cameras using openVINS json camera extrinsics
// set TCtoI in airsim
  msr::airlib::Pose airsim_pose;
  airsim_pose.position[0] = initX;
  airsim_pose.position[1] = initY;
  airsim_pose.position[2] = 0.0;
  airsim_pose.orientation.w() = 1.0;
  airsim_pose.orientation.x() = 0.0;
  airsimGetPoseTimeLog << tGetPose.GetSeconds<double>() << "\n";
  airsim_pose.orientation.y() = 0.0;
  airsim_pose.orientation.z() = 0.0;
  client.simSetVehiclePose(airsim_pose, true);

  bool isStereo = false;
  bool measCSVNew = false;
  string measCSV;

  ExplorationCost expCost(Vec3d(0.0, 0.0, 0.0), Rotationd::Identity(),
                          Vec3d(0.0, 0.0, 0.0), depthCamAtt);

  while (t.GetSeconds<double>() < endTime) {
    if (requestNewImage) {
      msr::airlib::vector<ImageRequest> request = { ImageRequest(
          "0", ImageType::DepthVis, false, true) };
      const msr::airlib::vector<ImageResponse> &response = client.simGetImages(
          request);

      // We do not store the first three images, which often has erroneous setup due to the latency of initialization of Unity simulator.
      if (imageCount > 3) {

        for (const ImageResponse &image_info : response) {
          std::string path;
          path = "/home/clark/Documents/AirSim";
          char buffer[256];
          sprintf(buffer, "%04d", imageCount - 3);
          std::string str(buffer);
          std::string file_path = FileSystem::combine(path, "img" + str);
          std::ofstream file(file_path + ".png", std::ios::binary);
          file.write(
              reinterpret_cast<const char*>(image_info.image_data_uint8.data()),
              image_info.image_data_uint8.size());
          file.close();
          depthImage8 = cv::imdecode(response.at(0).image_data_uint8,
                                     cv::IMREAD_UNCHANGED);
          depthImage8.convertTo(depthImage, CV_16U);

        }

        depthImageHeight = response.at(0).height;
        depthImageWidth = response.at(0).width;
        focalLength = depthImageWidth / 2.0;

        // Plot out the smallest depth in picture
        uint16_t *depthData;
        depthData = reinterpret_cast<uint16_t*>(depthImage.data);

        unsigned smallestDepth(1000);

        for (unsigned i = 0; i < depthImageWidth; i++) {
          for (unsigned j = 0; j < depthImageHeight; j++)
            if ((depthData[i * depthImageWidth + j] < smallestDepth)
                && (depthData[i * depthImageWidth + j] > 0)) {
              smallestDepth = depthData[i * depthImageWidth + j];
            }
        }

        double closestObstacle = smallestDepth * depthScale;
        printf("closest=%.8fs\n", closestObstacle);

        requestNewImage = false;  // Continue running the simulation loop.
        imageReady = true;  // We have an image for planning for this round.
        logfile2 << tReal.GetSeconds<double>() << ",";
        logfile2 << t.GetSeconds<double>() << ",";
        logfile2 << tPic.GetSeconds<double>() << ",";
        tPic.Reset();
        logfile2 << imageCount << ",";
        logfile2 << "\n";
      }
      imageCount++;
    }

    quad->Run();
    simTimer.AdvanceMicroSeconds(uint64_t(dt * 1e6));

    Vec3d simTruthPos(quad->GetPosition());
    Rotationd simTruthAtt(quad->GetAttitude());

    msr::airlib::Pose airsim_pose;

    airsim_pose.position[0] = simTruthPos.x + initX;
    airsim_pose.position[1] = -simTruthPos.y + initY;
    airsim_pose.position[2] = -simTruthPos.z;

    airsim_pose.orientation.w() = simTruthAtt[0];
    airsim_pose.orientation.x() = simTruthAtt[1];
    airsim_pose.orientation.y() = -simTruthAtt[2];
    airsim_pose.orientation.z() = -simTruthAtt[3];

    client.simSetVehiclePose(airsim_pose, true);
    tAirSimLoop.Reset();  //Measures the time between set airsim pose to pose being updated;

    if (t.GetSeconds<double>() > startFlightTime) {
      bool wait_flag = true;
      msr::airlib::Pose receivedAirSimPose;
      Vec3d receivedPos;
      double PosDiffThreshold = 0.001;
      while (wait_flag) {
        // wait for the AirSim communication, until unity reach same state as simulation
        tGetPose.Reset();
        receivedAirSimPose = client.simGetVehiclePose();
        airsimGetPoseTimeLog << t.GetSeconds<double>() << ",";
        airsimGetPoseTimeLog << tGetPose.GetSeconds<double>() << ",";

        receivedPos[0] = receivedAirSimPose.position[0] - initX;
        receivedPos[1] = -(receivedAirSimPose.position[1] - initY);
        receivedPos[2] = -receivedAirSimPose.position[2];
        Vec3d PosDiff = receivedPos - simTruthPos;
        double PosDiff_d = PosDiff.GetNorm2();

        airsimGetPoseTimeLog << receivedPos[0] << ",";
        airsimGetPoseTimeLog << receivedPos[1] << ",";
        airsimGetPoseTimeLog << receivedPos[2] << ",";
        airsimGetPoseTimeLog << "\n";
        if (PosDiff_d < PosDiffThreshold) {
          wait_flag = false;
        }
      }
      airsimSetPoseTimeLog << tAirSimLoop.GetSeconds<double>() << "\n";
    }

    if (timerRequestNewImage.GetSeconds<double>() > periodBetweenImages) {
      requestNewImage = true;  // We wait until we get an image.
      // If we don't this part of loop won't run anymore.
      timerRequestNewImage.AdjustTimeBySeconds(-periodBetweenImages);
    }

    if (timerPrint.GetSeconds<double>() >= 1) {
      timerPrint.AdjustTimeBySeconds(-1);
      printf("Current sim time = %.1fs\n", t.GetSeconds<double>());
    }

    if (timerMocap.GetSeconds<double>() > periodMocapSystem) {
      //simulate mocap packets over the network
      timerMocap.AdjustTimeBySeconds(-periodMocapSystem);
      Vec3d measPos(quad->GetPosition());
      Rotationd measAtt(quad->GetAttitude());
      est->UpdateWithMeasurement(measPos, measAtt);
    }

    if (timerTelemetryLoop.GetSeconds<double>() > periodTelemetryLoop) {
      timerTelemetryLoop.AdjustTimeBySeconds(-periodTelemetryLoop);
      TelemetryPacket::data_packet_t dataPacketRaw1, dataPacketRaw2;
      quad->GetTelemetryDataPackets(dataPacketRaw1, dataPacketRaw2);
      TelemetryPacket::TelemetryPacket dataPacket;
      TelemetryPacket::DecodeTelemetryPacket(dataPacketRaw1, dataPacket);
      TelemetryPacket::DecodeTelemetryPacket(dataPacketRaw2, dataPacket);
    }

    MocapStateEstimator::MocapEstimatedState estState;
    estState = est->GetPrediction(timeDelayOffboardControlLoopEstimate);

    if (timerOffboardMainLoop.GetSeconds<double>() > periodOffboardMainLoop) {

      expCost.updateWorld(estState.pos, estState.att, goalWorld);

      //Simulate the offboard main loop
      timerOffboardMainLoop.AdjustTimeBySeconds(-periodOffboardMainLoop);

      if (t.GetSeconds<double>() > startFlightTime) {
        if (imageReady) {
//          cv::imwrite(
//              "/home/clark/Documents/Repos/LabCode/GeneralCode/Logs/rappids_simulator/depthImage.bmp",
//              depthImage);

          DepthImagePlanner planner(depthImage, depthScale, focalLength,
                                    depthImageWidth / 2.0,
                                    depthImageHeight / 2.0,
                                    physicalVehicleRadius,
                                    vehicleRadiusPlanning, minCollisionDist);

          Vec3d vel = depthCamAtt.Inverse() * estState.att.Inverse()
              * estState.vel;
          Vec3d acc = depthCamAtt.Inverse() * estState.att.Inverse()
              * (Vec3d(0, 0, 1) * previousThrust - Vec3d(0, 0, 9.81));
          Vec3d g = depthCamAtt.Inverse() * estState.att.Inverse()
              * Vec3d(0, 0, -9.81);
          RapidQuadrocopterTrajectoryGenerator::RapidTrajectoryGenerator candidateTraj(
              Vec3d(0, 0, 0), vel, acc, g);
          double compTime = 0.05;

          std::vector<TrajectoryTest> trajectories;
          bool planResult = planner.FindLowestCostTrajectoryRandomCandidates(
              candidateTraj, trajectories, compTime, (void*) &expCost,
              &ExplorationCost::GetTrajCostWrapper);

//          bool planResult = planner.FindLowestCostTrajectoryRandomCandidates(
//              candidateTraj, compTime, (void*) &expCost,
//              &ExplorationCost::GetTrajCostWrapper);

          Vec3d exploration_direction(0, 0, 1);

//          double compTime = 0.02;
//          bool planResult = planner.FindFastestTrajRandomCandidates(
//              candidateTraj, compTime, exploration_direction);

          // printf("Plan result = %d", planResult);
          // Record the rotation between the trajectory and the world frame

          if (planResult) {
//            printf("traj found \n");
            trajAtt = estState.att * depthCamAtt;
            double traj_y, traj_p, traj_r;
            trajAtt.ToEulerYPR(traj_y, traj_p, traj_r);
            trajOffset = estState.pos;
            trackTrajTime.Reset();
            trajPlanned = true;
            _traj = std::make_shared
                < RapidQuadrocopterTrajectoryGenerator::RapidTrajectoryGenerator
                > (candidateTraj);

            Trajectory trajToRecord = candidateTraj.GetTrajectory();
            std::vector<Vec3d> trajCoeffs = trajToRecord.GetCoeffs();
            plannedTrajCount++;

            // Record the trajectory
            trajLog << plannedTrajCount << ",";
            trajLog << trajCoeffs[0].x << "," << trajCoeffs[0].y << ","
                << trajCoeffs[0].z << ",";
            trajLog << trajCoeffs[1].x << "," << trajCoeffs[1].y << ","
                << trajCoeffs[1].z << ",";
            trajLog << trajCoeffs[2].x << "," << trajCoeffs[2].y << ","
                << trajCoeffs[2].z << ",";
            trajLog << trajCoeffs[3].x << "," << trajCoeffs[3].y << ","
                << trajCoeffs[3].z << ",";
            trajLog << trajCoeffs[4].x << "," << trajCoeffs[4].y << ","
                << trajCoeffs[4].z << ",";
            trajLog << trajCoeffs[5].x << "," << trajCoeffs[5].y << ","
                << trajCoeffs[5].z << ",";
            trajLog << traj_y << "," << traj_p << "," << traj_r << ",";
            trajLog << trajOffset.x << "," << trajOffset.y << ","
                << trajOffset.z << ",";
            trajEndTime = trajToRecord.GetEndTime();
            trajLog << trajToRecord.GetStartTime() << ","
                << trajToRecord.GetEndTime() << "\n";
            imageReady = false;
          }
        }

        double traj_t = trackTrajTime.GetSeconds<double>();
        Vec3d trajPos, trajVel, trajAcc;

        if (not trajPlanned) {
          desiredPosition = Vec3d(0, 0, 2);
          desiredVelocity = Vec3d(0, 0, 0);
          desiredAcc = Vec3d(0, 0, 0);
        } else {
          if (traj_t < trajEndTime) {
            traj_t += 0.04;
            trajPos = _traj->GetPosition(traj_t);
            trajVel = _traj->GetVelocity(traj_t);
            trajAcc = _traj->GetAcceleration(traj_t);
          } else {
            trajPos = _traj->GetPosition(trajEndTime);
            trajVel = Vec3d(0, 0, 0);
            trajAcc = Vec3d(0, 0, 0);
          }

          // Disallow the vehicle from going backwards
          // This can occur if the velocity estimate is noisy, which will result in trajectories
          // that move backwards briefly before moving into the camera FOV
          // However, applying these trajectories in a receding horizon fashion will result
          // in only the portion of the trajectory that moves backwards being executed, resulting
          // in the vehicle slowly drifting backwards
          if (trajPos.z < 0) {
            trajPos.z = 0;
            // Only saturate velocity and acceleration if the trajectory is actually behind the camera
            // (We don't want to prevent the vehicle from slowing down at the end of a trajectory)
            if (trajVel.z < 0) {
              trajVel.z = 0;
            }
            if (trajAcc.z < 0) {
              trajAcc.z = 0;
            }
          }

          Vec3d refPos = trajAtt * trajPos + trajOffset;
          Vec3d refVel = trajAtt * trajVel;
          Vec3d refAcc = trajAtt * trajAcc;
          double refThrust = _traj->GetThrust(traj_t);
          // Rotate angular velocity into body frame
          Vec3d refAngVel = estState.att.Inverse() * trajAtt
              * _traj->GetOmega(traj_t, 0.02);

          desiredPosition = refPos;
          desiredVelocity = refVel;
          desiredAcc = refAcc;
          desiredAngVel = refAngVel;
          desiredThrust = refThrust;
        }
      }

      //Publish the commands:
      RadioTypes::RadioMessageDecoded::RawMessage rawMsg;

      if (controllerType == CTRL_OFFBOARD_RATES) {
        safetyNet.UpdateWithEstimator(estState,
                                      est->GetTimeSinceLastGoodMeasurement());

        Vec3d cmdAngVel;
        double cmdThrust;
        Rotationf cmdAtt;
        uint8_t flags = 0;

        if (t.GetSeconds<double>() < startFlightTime) {

          ctrl.Run(estState.pos, estState.vel, estState.att, desiredPosition,
                   Vec3d(0, 0, 0), Vec3d(0, 0, 0),
                   desYawAngleDeg * M_PI / 180.0, cmdAngVel, cmdThrust);
        } else {

          ctrl.RunTracking(estState.pos, estState.vel, estState.att,
                           desiredPosition, desiredVelocity, desiredAcc,
                           desYawAngleDeg * M_PI / 180.0, desiredThrust,
                           desiredAngVel, cmdAngVel, cmdThrust, cmdAtt);
        }
        RadioTypes::RadioMessageDecoded::CreateRatesCommand(flags,
                                                            float(cmdThrust),
                                                            Vec3f(cmdAngVel),
                                                            rawMsg.raw);

        previousThrust = cmdThrust;
        previousAngVel = cmdAngVel;

        lastRadioCommand[0] = cmdThrust;
        lastRadioCommand[1] = cmdAngVel.x;
        lastRadioCommand[2] = cmdAngVel.y;
        lastRadioCommand[3] = cmdAngVel.z;
        est->SetPredictedValues(
            cmdAngVel,
            (estState.att * Vec3d(0, 0, 1) * cmdThrust - Vec3d(0, 0, 9.81)));

        //We disable this for now, since it triggers on the first run.
        //        if (!safetyNet.GetIsSafe()) {
        //          RadioTypes::RadioMessageDecoded::CreateKillOneCommand(vehicleId, 0,
        //                                                                rawMsg.raw);
        //        }
      } else if (controllerType == CTRL_ONBOARD_UWB) {
        uint8_t flags = 0;
        RadioTypes::RadioMessageDecoded::CreatePositionCommand(flags,
                                                               desiredPosition,
                                                               desiredVelocity,
                                                               Vec3d(0, 0, 0),  //zero acceleration
                                                               rawMsg.raw);
      }

      TelemetryPacket::data_packet_t dataPacketRaw1, dataPacketRaw2;
      quad->GetTelemetryDataPackets(dataPacketRaw1, dataPacketRaw2);

      TelemetryPacket::TelemetryPacket dataPacket;
      TelemetryPacket::DecodeTelemetryPacket(dataPacketRaw1, dataPacket);
      TelemetryPacket::DecodeTelemetryPacket(dataPacketRaw2, dataPacket);

      //add message to command radio queue
      cmdRadioChannel.AddMessage(rawMsg);

      //Write the simulation state to a file:
      logfile << t.GetSeconds<double>() << ",";
      logfile << toCSV(quad->GetPosition());
      logfile << toCSV(quad->GetVelocity());
      logfile << toCSV(quad->GetAttitude().ToEulerYPR());
//    logfile << toCSV(quad->GetAttitude());
      logfile << toCSV(quad->GetAngularVelocity());
      logfile << dataPacket.motorForces[0] << ",";
      logfile << dataPacket.motorForces[1] << ",";
      logfile << dataPacket.motorForces[2] << ",";
      logfile << dataPacket.motorForces[3] << ",";

      // get openVINS state
//        ov_est_state = ov_est->get_state(t.GetSeconds<float>());
//        logfile << toCSV(ov_est_state.pos);
//        logfile << toCSV(ov_est_state.vel);
//        logfile << toCSV(ov_est_state.q);
//        logfile << toCSV(ov_est_state.cov_upper);

      //estimator state:
      Vec3f pos;
      Vec3f vel;
      Rotationf att;
      Vec3f angVel;
      if (controllerType != CTRL_ONBOARD_UWB) {
        MocapStateEstimator::MocapEstimatedState estState = est->GetPrediction(
            0);
        pos = estState.pos;
        vel = estState.vel;
        att = estState.att;
        angVel = estState.angVel;
      } else {
        pos = Vec3d(dataPacket.position);
        vel = Vec3d(dataPacket.velocity);
        att = Rotationd::FromVectorPartOfQuaternion(Vec3d(dataPacket.attitude));
        angVel = Vec3d(dataPacket.gyro);
      }

      logfile << toCSV(pos);
      logfile << toCSV(vel);
      logfile << toCSV(att.ToEulerYPR());
      logfile << toCSV(angVel);

      //desired position:
      logfile << toCSV(desiredPosition);
      logfile << toCSV(desiredVelocity);

      logfile << int(dataPacket.panicReason) << ",";

      for (int i = 0; i < 4; i++) {
        logfile << double(lastRadioCommand[i]) << ",";
      }

      // Adding measured values to the logfile
      if (measCSVNew) {
        logfile << measCSV;
        measCSVNew = false;
      }
      logfile << "\n";
    }  //main loop

    //Check if we have a new message to transmit:
    if (cmdRadioChannel.HaveNewMessage()) {
      quad->SetCommandRadioMsg(cmdRadioChannel.GetMessage());
    }

  }  //sim loop
  printf("Done.\n");
//  ov_est->end_sim();
  logfile.close();
  logfile2.close();
  trajLog.close();
  airsimGetPoseTimeLog.close();
  airsimSetPoseTimeLog.close();
}


/*

Copyright (c) 2005-2021, University of Oxford.
All rights reserved.

University of Oxford means the Chancellor, Masters and Scholars of the
University of Oxford, having an administrative office at Wellington
Square, Oxford OX1 2JD, UK.

This file is part of Chaste.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
 * Redistributions of source code must retain the above copyright notice,
   this list of conditions and the following disclaimer.
 * Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution.
 * Neither the name of the University of Oxford nor the names of its
   contributors may be used to endorse or promote products derived from this
   software without specific prior written permission.

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

*/

#ifndef CorriasBuistICCModified_HPP_
#define CorriasBuistICCModified_HPP_

#include "ChasteSerialization.hpp"
#include <boost/serialization/base_object.hpp>
#include "AbstractCardiacCell.hpp"
#include "AbstractStimulusFunction.hpp"

/**
 * This class is a modified version of the  model of a gastric
 * Interstitial Cell of Cajal.
 *
 * Reference publication is:
 *
 * Corrias A, Buist ML.
 * "Quantitative cellular description of gastric slow wave activity."
 * Am J Physiol Gastrointest Liver Physiol. 2008 Apr;294(4):G989-95. Epub 2008 Feb 14.
 *
 * Modifications include:
 * - simplified mitochondria dynamics (assumed mitochondrial potential is almost constant
 * - ability to set K+ channels-affecting CO concentrations
 * - ability to deflect a fraction of VDDR channels into the pacemaker unit.
 */
class CorriasBuistICCModified : public AbstractCardiacCell
{
    friend class boost::serialization::access;
    /**
     * Boost Serialization method for archiving/checkpointing.
     * Archives the object and its member variables.
     *
     * @param archive  The boost archive.
     * @param version  The current version of this class.
     */
    template<class Archive>
    void serialize(Archive & archive, const unsigned int version)
    {
        archive & boost::serialization::base_object<AbstractCardiacCell >(*this);
        archive &  mFractionOfVDDRInPU;
        archive &  mIP3Concentration;
        archive &  mScaleFactorSerca;
        archive &  mScaleFactorCarbonMonoxide;
    }

private:

    /**fraction of VDDR channel in the PU, initialised to zero*/
    double mFractionOfVDDRInPU;
    /**the IP3 concentration, defaults to 0.0006 mM*/
    double mIP3Concentration;
    /**scales the flux through the SERCA pump.1.0-> control. 0.0-> blocked*/
    double mScaleFactorSerca;
    /**
     * Scale factor for CO-affected currents
     * Note that this the number that multiply the currents, hence it is not [CO],
     * but a function of [CO] (for example, 2.8*[CO] - 0.1)
     */
    double mScaleFactorCarbonMonoxide;

    /* Concentrations */
    double Ca_o;/**<  mM */
    double Cl_o;/**<  mM */
    double K_o;/**<  mM */
    double Na_o;/**<  mM */

    /* Nernst parameters */
    double R;/**<  pJ/nmol/K */
    double T;/**<  degK */
    double F;/**<  nC/nmol */
    double FoRT;/**<  1/mV */
    double RToF;/**<  mV */

    double  Cm ;/**<  pF */
    double Asurf_in_cm_square;/**< cm2 */
    double  Asurf ;/**<  mm2 */
    double  Cl_i  ;/**<  mM */
    double  K_i   ;/**<  mM */
    double  Na_i    ;/**<  mM */
    double  P_cyto;/**<  dim */
    double  Vol  ;/**<  mm3 */
    double  fc ;/**<  dim */
    double  fe  ;/**<  dim */
    double  fm ;/**<  dim */
    double  Q10Ca ;/**<  dim */
    double  Q10K ;/**<  dim */
    double  Q10Na  ;/**<  dim */
    double  T_exp  ;/**<  degK */

    double  G_max_BK  ;/**<  nS */
    double  G_max_CaCl  ;/**<  nS */
    double  G_max_ERG  ;/**<  nS */
    double  G_max_Ltype ;/**<  nS */
    double  G_max_NSCC  ;/**<  nS */
    double  G_max_Na ;/**<  nS */
    double  G_max_VDDR ;/**<  nS */
    double  G_max_bk    ;/**<  nS */
    double  G_max_kv11 ;/**<  nS */

    double  J_max_PMCA  ;/**<  mM/ms   (mM/s) * 1/1000 (s/ms) = mM/ms */
    double  J_max_PMCA_PU ;/**<  mM/ms   (mM/s) * 1/1000 (s/ms) = mM/ms */
    double  J_ERleak   ;/**<  1/ms    (1/s) * 1/1000 (ms/s) = 1/ms */
    double  J_max_leak ;/**<  1/ms    (1/s) * 1/1000 (ms/s) = 1/ms */
    double  Jmax_IP3  ;/**<  1/ms    (1/s) * 1/1000 (ms/s) = 1/ms */
    double  Jmax_NaCa ;/**<  mM/ms   (mM/s) * 1/1000 (s/ms) = mM/ms */
    double  Jmax_serca   ;/**<  mM/ms   (mM/s) * 1/1000 (s/ms) = mM/ms */
    double  Jmax_uni    ;/**<  1/ms    (1/s) * 1/1000 (ms/s) = 1/ms */

    double  NaPerm_o_Kperm   ;/**<  dim */
    double  L    ;/**<  dim */
    double  P_ER  ;/**<  dim */
    double  P_PU   ;/**<  dim */
    double  P_mito ;/**<  dim */
    double  b ;/**<  dim */
    double  na ;/**<  dim */

    double  K_Ca   ;/**<  mM */
    double  K_Na ;/**<  mM */
    double  K_act ;/**<  mM */
    double  K_trans   ;/**<  mM */
    double  k_serca ;/**<  mM */
    double  conc   ;/**<  mM */
    double  d_ACT   ;/**<  mM */
    double  d_IP3 ;/**<  mM */
    double  d_INH  ;/**<  mM */

    double  tau_d_CaCl;/**<  ms(s) * 1000 (ms/s) = ms */
    double  tau_d_NSCC ;/**<  ms(s) * 1000 (ms/s) = ms */
    double  tauh;/**<  ms(s ) * 1000 (ms/s) = ms */

    double  deltaPsi_B;/**<  mV */
    double  deltaPsi_star;/**<  mV */
    double  deltaPsi;/**<  mV */

     /////////////////////
     //Calculated constants
     ////////////////////
     /* Volumes */
     double V_cyto;              /**<  mm3 */
     double V_ER;                /**<  mm3 */
     double V_MITO;              /**<  mm3 */
     double V_PU;                /**<  mm3 */

     /* Temperature corrections */
     double T_correction_Ca;     /**<  dim */
     double T_correction_K;      /**<  dim */
     double T_correction_Na;     /**<  dim */
     double T_correction_BK;     /**<  uA/mm2 */

     /* Nernst potentials */
     double E_Na;                /**<  mV */
     double E_K;                 /**<  mV */
     double E_Cl;                /**<  mV */
     double E_NSCC;              /**<  mV */

     /* Activation gate time constants */
     double tau_d_ERG;           /**<  ms */
     double tau_d_Ltype;         /**<  ms */
     double tau_d_Na;            /**<  ms */
     double tau_d_VDDR;          /**<  ms */
     double tau_d_kv11;          /**<  ms */

     /* Inactivation gate time constants */
     double tau_f_Ltype;         /**<  ms */
     double tau_f_Na;            /**<  ms */
     double tau_f_VDDR;          /**<  ms */
     double tau_f_ca_Ltype;      /**<  ms */
     double tau_f_kv11;          /**<  ms */

     /* Speed ups */
     double e2FoRTdPsiMdPsiS;/**< speed-up constant*/
     double ebFoRTdPsiMdPsiS;/**< speed-up constant*/

public:
   /**
    * Constructor
    *
    * @param pSolver is a pointer to the ODE solver
    * @param pIntracellularStimulus is a pointer to the intracellular stimulus
    */
     CorriasBuistICCModified(boost::shared_ptr<AbstractIvpOdeSolver> pSolver, boost::shared_ptr<AbstractStimulusFunction> pIntracellularStimulus);

    /**
     * Destructor
     */
    ~CorriasBuistICCModified();

    /**
     * Now empty
     */
    void VerifyStateVariables();

    /**
     * Calculates the ionic current
     *
     * @param pStateVariables the state variables of this model
     * @return the total ionic current
     */
    double GetIIonic(const std::vector<double>* pStateVariables=NULL);

    /**
     * Compute the RHS of the FitHugh-Nagumo system of ODEs
     *
     * @param time  the current time, in milliseconds
     * @param rY  current values of the state variables
     * @param rDY  to be filled in with derivatives
     */
    void EvaluateYDerivatives(double time, const std::vector<double>& rY, std::vector<double>& rDY);

    /**
     * Sets the fraction of VDDR channels in the Pacemaker unit
     *
     * @param fraction the fraction of VDDr channels in the PU
     */
    void SetFractionOfVDDRInPU(double fraction);

    /**
     * Set the value of IP3 concentration in the cell
     *
     * @param concentration the concentration of IP3
     */
    void SetIP3Concentration(double concentration);

    /**
     * Set a multiplying factor for the influx of Ca2+ into the Er via the SERCA pump (set to zero will block the SW generation)
     *
     * @param scaleFactor the scale factor (=0 --> no Ca2+ uptake into the ER and, consequently, no SW)
     */
    void SetSercaPumpScaleFactor(double scaleFactor);

    /**
     * Set the carbon monoxide scale factor.
     * This will multiply the following currents: I_kv11, I_ERG, Ibk
     *
     * @param scaleFactor the scale factor that multiply the currents.
     */
    void SetCarbonMonoxideScaleFactor(double scaleFactor);

    /**
     * @return the Carbon Monoxide scale factor
     */
    double GetCarbonMonoxideScaleFactor();
};

// Needs to be included last
#include "SerializationExportWrapper.hpp"
CHASTE_CLASS_EXPORT(CorriasBuistICCModified)

namespace boost
{
    namespace serialization
    {
        template<class Archive>
        inline void save_construct_data(
            Archive & ar, const CorriasBuistICCModified * t, const unsigned int fileVersion)
        {
            const boost::shared_ptr<AbstractIvpOdeSolver> p_solver = t->GetSolver();
            const boost::shared_ptr<AbstractStimulusFunction> p_stimulus = t->GetStimulusFunction();
            ar << p_solver;
            ar << p_stimulus;
        }

        template<class Archive>
        inline void load_construct_data(
            Archive & ar, CorriasBuistICCModified * t, const unsigned int fileVersion)
        {
            boost::shared_ptr<AbstractIvpOdeSolver> p_solver;
            boost::shared_ptr<AbstractStimulusFunction> p_stimulus;
            ar >> p_solver;
            ar >> p_stimulus;
            ::new(t)CorriasBuistICCModified(p_solver, p_stimulus);
        }
    }
}

#endif // CorriasBuistICCModified_HPP_


//
// $Id: FitFromPepXML.cpp 1195 2009-08-14 22:12:04Z chambm $
//
// Robert Burke <robert.burke@cshs.org>
//
// Copyright 2009 Spielberg Family Center for Applied Proteomics
//   Cedars Sinai Medical Center, Los Angeles, California  90048
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#include "FitFromPepXML.hpp"
#include "math/LinearLeastSquares.hpp"
#include <boost/numeric/ublas/matrix_proxy.hpp>
#include <stdexcept>
#include <math.h>
#include <stdlib.h>
#include <iostream>
#include <iomanip>
#include <vector>

using namespace std;
using pwiz::math::LinearLeastSquares;

using namespace pwiz::pepxml;
using namespace pwiz::math::types;

namespace ublas = boost::numeric::ublas;

namespace pwiz {
namespace rtapprox {

class FitFromPepXMLImpl : public FitFromPepXML
{
public:
    FitFromPepXMLImpl();

    FitFromPepXMLImpl(const char* pepXMLFile, float tol, float rtmin = -1, float rtmax = -1, unsigned limit = 0);

    FitFromPepXMLImpl(const dmatrix& peps, const dvector& rts);

    virtual bool createSystem(vector<Peptide>& peptideList);

    virtual const dvector fit();

    virtual const dvector load();

    virtual const dvector load(const char* pepXMLFile, unsigned limit=0);

    virtual const PeptideRTSystem* getSystem()
    {
        return system.get();
    }

    virtual const dmatrix& getPeptideMatrix() const
    {
        if (system.get() == NULL)
            throw runtime_error("no system have been loaded.");

        return system->getPeptideCounts();
    }

    virtual const dvector& getRTVector() const
    {
        if (system.get() == NULL)
            throw runtime_error("no system have been loaded.");

        return system->getRetentionTimes();
    }

    virtual const vector<Peptide>& getPeptideList() const
    {
        if (extractor.get() == NULL)
            throw runtime_error("no results have been calculated.");

        return extractor->getPeptideList();
    }

    virtual bool bad() const
    {
        return bad_;
    }

    virtual float getRtMin() const
    {
        return rtmin_;
    }

    virtual void setRtMin(float rtmin)
    {
        rtmin_ = rtmin;
    }

    virtual float getRtMax() const
    {
        return rtmin_;
    }

    virtual void setRtMax(float rtmax)
    {
        rtmax_ = rtmax;
    }

    virtual unsigned getLimit() const
    {
        return limit_;
    }

    virtual void setLimit(unsigned limit)
    {
        limit_ = limit;
    }

    virtual dvector countPeptide(const char* peptide_str);

    virtual const dmatrix& getAtA()
    {
        return system->getPeptideCounts();
    }

protected:
    bool bad_;

    const char* pepxmlFile;

    float tol_;
    float rtmin_;
    float rtmax_;
    unsigned limit_;

    auto_ptr<PeptideRTSystem> system;
    LinearLeastSquares<pwiz::math::LinearLeastSquaresType_QR> lls;
    auto_ptr<DataExtractor> extractor;
};

// TODO Move the number of matrix columns into a header
const char FitFromPepXML::AMINOACIDS[] = {
    'A',       // ala
    'R',       // arg
    'N',       // asn
    'D',       // asp
    'C',       // cys
    'E',       // glu
    'Q',       // gln
    'G',       // gly
    'H',       // his
    'I',       // ile
    'L',       // leu
    'K',       // lys
    'M',       // met
    'F',       // phe
    'P',       // pro
    'S',       // ser
    'T',       // thr
    'W',       // trp
    'Y',       // tyr
    'V',       // val
//    'B',       // Asx
//    'Z',       // Glx
//    'X'        // Unk
};

const int FitFromPepXML::NUMENTRIES = sizeof(AMINOACIDS);

auto_ptr<FitFromPepXML> FitFromPepXML::create()
{
    return auto_ptr<FitFromPepXML>(new FitFromPepXMLImpl());
}

auto_ptr<FitFromPepXML> FitFromPepXML::create(const char* pepxmlFile,
                                              float tol,
                                              float rtmin, float rtmax,
                                              unsigned limit)
{
    return auto_ptr<FitFromPepXML>(new FitFromPepXMLImpl(pepxmlFile, tol, rtmin, rtmax, limit));
}

auto_ptr<FitFromPepXML> FitFromPepXML::createFromMatrix(const dmatrix& peps,
                                                    const dvector& rts)
{
    return auto_ptr<FitFromPepXML>(new FitFromPepXMLImpl(peps, rts));
}

FitFromPepXMLImpl::FitFromPepXMLImpl()
    : bad_(false), pepxmlFile(NULL),
      limit_(0),
      system(auto_ptr<PeptideRTSystem>(0)),
      extractor(auto_ptr<DataExtractor>(0))
{
}

FitFromPepXMLImpl::FitFromPepXMLImpl(const dmatrix& peps,
                                     const dvector& rts)
    : bad_(false), pepxmlFile(NULL),
      limit_(0),
      system(auto_ptr<PeptideRTSystem>(0)),
      extractor(auto_ptr<DataExtractor>(0))
{
    system = PeptideRTSystem::create(peps, rts);

    fit();
}

FitFromPepXMLImpl::FitFromPepXMLImpl(const char* pepXMLFile, float tol,
                                     float rtmin, float rtmax,
                                     unsigned limit)
    : bad_(false), pepxmlFile(pepXMLFile),
      tol_(tol), rtmin_(rtmin), rtmax_(rtmax), limit_(limit),
      system(auto_ptr<PeptideRTSystem>(0)),
      extractor(auto_ptr<DataExtractor>(0))
{
    this->pepxmlFile = pepxmlFile;

    //load(pepXMLFile);
}

const dvector FitFromPepXMLImpl::load()
{
    if (pepxmlFile)
        return load(pepxmlFile, limit_);

    return fit();
}

const dvector FitFromPepXMLImpl::load(const char* pepxmlFile, unsigned limit)
{
    bad_ = false;
    this->pepxmlFile = pepxmlFile;
    limit_ = limit;

    extractor = DataExtractor::create(pepxmlFile, tol_, rtmin_, rtmax_);

    /*
     * Ignore amino acids that we don't care about.
     */
    extractor->exclude('B');
    extractor->exclude('Z');
    extractor->exclude('X');

    extractor->parse();

    if (extractor.get() != NULL)
    {
        createSystem(extractor->getPeptideList());

        fit();
    }
    else
    {
        bad_ = true;
        throw runtime_error("Error occurred while loading pepxml");
    }

    return lls.solve(system->getPeptideCounts(),
                     system->getRetentionTimes());
}

bool FitFromPepXMLImpl::createSystem (vector<Peptide>& peptideList)
{
    // For each row, sum the occurance of different peptides.
    unsigned size = (limit_ >0 ? std::min((unsigned)peptideList.size(),
                                          limit_) :
                     peptideList.size());

    dmatrix pm(size, NUMENTRIES+1);
    dvector rts(size);

    pm.clear();
    rts.clear();

    for (size_t i= 0;i < size && !bad_; i++)
    {
        const char* rt_str = peptideList[i].getRetentionTime();
        const char* peptide_str = peptideList[i].getPeptide();

        if (rt_str == NULL)
        {
            cerr << "NULL pointer in rt_str" << endl;
            bad_ = true;
            continue;
        }

        // Add a row in the matrix and an entry in the vector
        rts(i) = atof(rt_str);
        dvector peptide = countPeptide(peptide_str);

        ublas::matrix_row<dmatrix> mr(pm, i);
        ublas::matrix_row<dmatrix>::iterator loc;
        dvector::iterator val;
        for (val=peptide.begin(),loc=mr.begin();val!=peptide.end();val++,loc++)
        {
            *loc = *val;
        }
    }

    if (!bad_)
        system = PeptideRTSystem::create(pm, rts);

    return bad_;
}

const dvector FitFromPepXMLImpl::fit()
{
    return lls.solve(system->getPeptideCounts(),
                     system->getRetentionTimes());
}

dvector FitFromPepXMLImpl::countPeptide(const char* peptide_str)
{
    dvector countv(NUMENTRIES+1);
    countv.clear();
    countv(NUMENTRIES) = 1;

    int size = strlen(peptide_str);

    for (const char* i=peptide_str; i != peptide_str+size; i++)
    {
        bool found = false;
        for (size_t j=0; j<sizeof(AMINOACIDS); j++)
        {
            if (AMINOACIDS[j] == toupper(*i))
            {
                found = true;
                countv(j)+=1;
                break;
            }
        }

        if (!found)
            cerr << "No AA for " << *i << endl;
    }

    return countv;
}

}
}


// Copyright (c) 2018 Graphcore Ltd. All rights reserved.
#define BOOST_TEST_MODULE MatMulTest

#include <boost/test/unit_test.hpp>
#include <onnx/onnx_pb.h>
#include <popart/ir.hpp>
#include <popart/op/matmul.hpp>
#include <popart/popx/op/matmulx.hpp>
#include <popart/tensor.hpp>

using namespace popart;

// Test the simple [2x2] * [2x2] matrix
BOOST_AUTO_TEST_CASE(MatMul_Case1) {

  // Setup

  popart::Ir ir;
  auto &graph = ir.getMainGraph();

  popart::MatMulOp mm(Onnx::Operators::MatMul_9,
                      {graph, ""},
                      nonstd::nullopt,
                      {},
                      OptionalDataType());

  popart::DebugContext dc;
  popart::DebugInfo di(dc, "test");

  popart::Tensor lhs("lhs", popart::TensorType::ActGrad, graph, di);
  lhs.info.set(popart::DataType::FLOAT, {2, 2});
  mm.input->insert(0, &lhs);

  popart::Tensor rhs("rhs", popart::TensorType::ActGrad, graph, di);
  rhs.info.set(popart::DataType::FLOAT, {2, 2});
  mm.input->insert(1, &rhs);

  popart::Tensor out("out", popart::TensorType::ActGrad, graph, di);
  mm.output->insert(0, &out);

  // Test the setup is correct
  mm.setup();
  BOOST_CHECK(mm.outInfo(0).dim(0) == 2);
  BOOST_CHECK(mm.outInfo(0).dim(1) == 2);
  BOOST_CHECK(mm.outInfo(0).dataType() == popart::DataType::FLOAT);
  BOOST_CHECK(mm.rhsIn() == &rhs);
  BOOST_CHECK(mm.lhsIn() == &lhs);

  // Test the clone operation
  std::unique_ptr<popart::Op> mmClone = mm.clone();
  popart::MatMulOp *mmPtr = dynamic_cast<popart::MatMulOp *>(mmClone.get());
  BOOST_CHECK(mmPtr != nullptr);
  // Clone aka copy does not copy input & output???
  // BOOST_CHECK(mmPtr->outInfo(0).dim(0) == 2);
  // BOOST_CHECK(mmPtr->outInfo(0).dim(1) == 2);
  // BOOST_CHECK(mmPtr->outInfo(0).dataType() ==
  // popart::DataType::FLOAT); BOOST_CHECK(mmPtr->rhsIn() == &rhs);
  // BOOST_CHECK(mmPtr->lhsIn() == &lhs);

  auto gradOps = mm.getGradOps();
  BOOST_CHECK(gradOps.size() == 2);

  for (auto &&gradOp : gradOps) {
    popart::Op *op = gradOp.get();
    if (typeid(*op) == typeid(popart::MatMulLhsGradOp)) {
      popart::MatMulLhsGradOp *lhsGradOp =
          dynamic_cast<popart::MatMulLhsGradOp *>(op);

      popart::Tensor lhsOut("out", popart::TensorType::ActGrad, graph, di);
      lhsGradOp->output->insert(0, &lhsOut);

      BOOST_CHECK(lhsGradOp->getGradInIndex() == 0);
      BOOST_CHECK(lhsGradOp->getRhsInIndex() == 1);

      lhsGradOp->setup();
      BOOST_CHECK(lhsGradOp->outInfo(0).dim(0) == 2);
      BOOST_CHECK(lhsGradOp->outInfo(0).dim(1) == 2);
      BOOST_CHECK(lhsGradOp->outInfo(0).dataType() == popart::DataType::FLOAT);

      auto gradInfo = lhsGradOp->gradInputInfo();
      std::vector<GradInOutMapper> expectedGradInfo = {
          {0, 0, popart::GradOpInType::GradOut},
          {1, 1, popart::GradOpInType::In}};
      BOOST_CHECK(gradInfo == expectedGradInfo);

      auto mapping                       = lhsGradOp->gradOutToNonGradIn();
      std::map<int, int> expectedMapping = {{0, 0}};
      BOOST_CHECK(mapping == expectedMapping);
    } else if (typeid(*op) == typeid(popart::MatMulRhsGradOp)) {
      popart::MatMulRhsGradOp *rhsGradOp =
          dynamic_cast<popart::MatMulRhsGradOp *>(op);

      popart::Tensor rhsOut("out", popart::TensorType::ActGrad, graph, di);
      rhsGradOp->output->insert(0, &rhsOut);

      BOOST_CHECK(rhsGradOp->getGradInIndex() == 0);
      BOOST_CHECK(rhsGradOp->getLhsInIndex() == 1);

      rhsGradOp->setup();
      BOOST_CHECK(rhsGradOp->outInfo(0).dim(0) == 2);
      BOOST_CHECK(rhsGradOp->outInfo(0).dim(1) == 2);
      BOOST_CHECK(rhsGradOp->outInfo(0).dataType() == popart::DataType::FLOAT);

      auto gradInfo = rhsGradOp->gradInputInfo();
      std::vector<GradInOutMapper> expectedGradInfo = {
          {0, 0, popart::GradOpInType::GradOut},
          {1, 0, popart::GradOpInType::In}};

      BOOST_CHECK(gradInfo == expectedGradInfo);

      auto mapping                       = rhsGradOp->gradOutToNonGradIn();
      std::map<int, int> expectedMapping = {{0, 1}};
      BOOST_CHECK(mapping == expectedMapping);

    } else {
      // Unexpected grad op
      BOOST_CHECK(false);
    }
  }
}

// Test the simple [3x2] * [2x6] matrix
BOOST_AUTO_TEST_CASE(MatMul_Case2) {

  // Setup

  popart::Ir ir;
  auto &graph = ir.getMainGraph();

  popart::MatMulOp mm(Onnx::Operators::MatMul_9,
                      {graph, ""},
                      nonstd::nullopt,
                      {},
                      OptionalDataType());

  popart::DebugContext dc;
  popart::DebugInfo di(dc, "test");

  popart::Tensor lhs("lhs", popart::TensorType::ActGrad, graph, di);
  lhs.info.set(popart::DataType::FLOAT, {3, 2});
  mm.input->insert(0, &lhs);

  popart::Tensor rhs("rhs", popart::TensorType::ActGrad, graph, di);
  rhs.info.set(popart::DataType::FLOAT, {2, 6});
  mm.input->insert(1, &rhs);

  popart::Tensor out("out", popart::TensorType::ActGrad, graph, di);
  mm.output->insert(0, &out);

  // Test the setup is correct
  mm.setup();
  BOOST_CHECK(mm.outInfo(0).dim(0) == 3);
  BOOST_CHECK(mm.outInfo(0).dim(1) == 6);
  BOOST_CHECK(mm.outInfo(0).dataType() == popart::DataType::FLOAT);
  BOOST_CHECK(mm.rhsIn() == &rhs);
  BOOST_CHECK(mm.lhsIn() == &lhs);

  // Test the clone operation
  std::unique_ptr<popart::Op> mmClone = mm.clone();
  popart::MatMulOp *mmPtr = dynamic_cast<popart::MatMulOp *>(mmClone.get());
  BOOST_CHECK(mmPtr != nullptr);
  // Clone aka copy does not copy input & output???
  // BOOST_CHECK(mmPtr->outInfo(0).dim(0) == 2);
  // BOOST_CHECK(mmPtr->outInfo(0).dim(1) == 2);
  // BOOST_CHECK(mmPtr->outInfo(0).dataType() ==
  // popart::DataType::FLOAT); BOOST_CHECK(mmPtr->rhsIn() == &rhs);
  // BOOST_CHECK(mmPtr->lhsIn() == &lhs);

  auto gradOps = mm.getGradOps();
  BOOST_CHECK(gradOps.size() == 2);

  for (auto &&gradOp : gradOps) {
    popart::Op *op = gradOp.get();
    if (typeid(*op) == typeid(popart::MatMulLhsGradOp)) {
      popart::MatMulLhsGradOp *lhsGradOp =
          dynamic_cast<popart::MatMulLhsGradOp *>(op);

      popart::Tensor lhsOut("out", popart::TensorType::ActGrad, graph, di);
      lhsGradOp->output->insert(0, &lhsOut);

      BOOST_CHECK(lhsGradOp->getGradInIndex() == 0);
      BOOST_CHECK(lhsGradOp->getRhsInIndex() == 1);

      lhsGradOp->setup();
      BOOST_CHECK(lhsGradOp->outInfo(0).dim(0) == 3);
      BOOST_CHECK(lhsGradOp->outInfo(0).dim(1) == 2);
      BOOST_CHECK(lhsGradOp->outInfo(0).dataType() == popart::DataType::FLOAT);

      auto gradInfo = lhsGradOp->gradInputInfo();
      std::vector<GradInOutMapper> expectedGradInfo = {
          {0, 0, popart::GradOpInType::GradOut},
          {1, 1, popart::GradOpInType::In}};
      BOOST_CHECK(gradInfo == expectedGradInfo);

      auto mapping                       = lhsGradOp->gradOutToNonGradIn();
      std::map<int, int> expectedMapping = {{0, 0}};
      BOOST_CHECK(mapping == expectedMapping);
    } else if (typeid(*op) == typeid(popart::MatMulRhsGradOp)) {
      popart::MatMulRhsGradOp *rhsGradOp =
          dynamic_cast<popart::MatMulRhsGradOp *>(op);

      popart::Tensor rhsOut("out", popart::TensorType::ActGrad, graph, di);
      rhsGradOp->output->insert(0, &rhsOut);

      BOOST_CHECK(rhsGradOp->getGradInIndex() == 0);
      BOOST_CHECK(rhsGradOp->getLhsInIndex() == 1);

      rhsGradOp->setup();
      BOOST_CHECK(rhsGradOp->outInfo(0).dim(0) == 2);
      BOOST_CHECK(rhsGradOp->outInfo(0).dim(1) == 6);
      BOOST_CHECK(rhsGradOp->outInfo(0).dataType() == popart::DataType::FLOAT);

      auto gradInfo = rhsGradOp->gradInputInfo();
      std::vector<GradInOutMapper> expectedGradInfo = {
          {0, 0, popart::GradOpInType::GradOut},
          {1, 0, popart::GradOpInType::In}};

      BOOST_CHECK(gradInfo == expectedGradInfo);

      auto mapping                       = rhsGradOp->gradOutToNonGradIn();
      std::map<int, int> expectedMapping = {{0, 1}};
      BOOST_CHECK(mapping == expectedMapping);

    } else {
      // Unexpected grad op
      BOOST_CHECK(false);
    }
  }
}

// Test the simple [2x1x4x3x2] * [3x1x2x6] matrix
BOOST_AUTO_TEST_CASE(MatMul_Case3) {
  // Setup

  popart::Ir ir;
  auto &graph = ir.getMainGraph();

  popart::MatMulOp mm(Onnx::Operators::MatMul_9,
                      {graph, ""},
                      nonstd::nullopt,
                      {},
                      OptionalDataType());

  popart::DebugContext dc;
  popart::DebugInfo di(dc, "test");

  popart::Tensor lhs("lhs", popart::TensorType::ActGrad, graph, di);
  lhs.info.set(popart::DataType::FLOAT, {2, 1, 4, 3, 2});
  mm.input->insert(0, &lhs);

  popart::Tensor rhs("rhs", popart::TensorType::ActGrad, graph, di);
  rhs.info.set(popart::DataType::FLOAT, {3, 1, 2, 6});
  mm.input->insert(1, &rhs);

  popart::Tensor out("out", popart::TensorType::ActGrad, graph, di);
  mm.output->insert(0, &out);

  // Test the setup is correct
  mm.setup();
  BOOST_CHECK(mm.outInfo(0).dim(0) == 2);
  BOOST_CHECK(mm.outInfo(0).dim(1) == 3);
  BOOST_CHECK(mm.outInfo(0).dim(2) == 4);
  BOOST_CHECK(mm.outInfo(0).dim(3) == 3);
  BOOST_CHECK(mm.outInfo(0).dim(4) == 6);
  BOOST_CHECK(mm.outInfo(0).dataType() == popart::DataType::FLOAT);
  BOOST_CHECK(mm.rhsIn() == &rhs);
  BOOST_CHECK(mm.lhsIn() == &lhs);

  auto gradOps = mm.getGradOps();
  BOOST_CHECK(gradOps.size() == 2);
  BOOST_CHECK(gradOps[0]->isConvertibleTo<popart::MatMulLhsGradOp>());
  BOOST_CHECK(gradOps[1]->isConvertibleTo<popart::MatMulRhsGradOp>());

  popart::Tensor t1("out", popart::TensorType::ActGrad, graph, di);
  gradOps[0]->output->reset(0, &t1);
  gradOps[0]->setup();

  popart::Tensor t2("out", popart::TensorType::ActGrad, graph, di);
  gradOps[1]->output->reset(0, &t2);
  gradOps[1]->setup();

  BOOST_CHECK(gradOps[0]->output->tensor(0)->info ==
              mm.input->tensor(popart::MatMulOp::getLhsInIndex())->info);
  BOOST_CHECK(gradOps[1]->output->tensor(0)->info ==
              mm.input->tensor(popart::MatMulOp::getRhsInIndex())->info);
}

// Test invalid rank on lhs
BOOST_AUTO_TEST_CASE(MatMul_ErrorCase1) {

  // Setup

  popart::Ir ir;
  auto &graph = ir.getMainGraph();

  popart::MatMulOp mm(Onnx::Operators::MatMul_9,
                      {graph, ""},
                      nonstd::nullopt,
                      {},
                      OptionalDataType());

  popart::DebugContext dc;
  popart::DebugInfo di(dc, "test");

  popart::Tensor lhs("lhs", popart::TensorType::ActGrad, graph, di);
  lhs.info.set(popart::DataType::FLOAT, {2, 2, 3});
  mm.input->insert(0, &lhs);

  popart::Tensor rhs("rhs", popart::TensorType::ActGrad, graph, di);
  rhs.info.set(popart::DataType::FLOAT, {2, 2});
  mm.input->insert(1, &rhs);

  popart::Tensor out("out", popart::TensorType::ActGrad, graph, di);
  mm.output->insert(0, &out);

  // Test the setup is correct
  BOOST_CHECK_THROW(mm.setup(), error);
}

// Test invalid matrix multiplication
BOOST_AUTO_TEST_CASE(MatMul_ErrorCase3) {

  // Setup

  popart::Ir ir;
  auto &graph = ir.getMainGraph();

  popart::MatMulOp mm(Onnx::Operators::MatMul_9,
                      {graph, ""},
                      nonstd::nullopt,
                      {},
                      OptionalDataType());

  popart::DebugContext dc;
  popart::DebugInfo di(dc, "test");

  popart::Tensor lhs("lhs", popart::TensorType::ActGrad, graph, di);
  lhs.info.set(popart::DataType::FLOAT, {2, 3});
  mm.input->insert(0, &lhs);

  popart::Tensor rhs("rhs", popart::TensorType::ActGrad, graph, di);
  rhs.info.set(popart::DataType::FLOAT, {10, 2});
  mm.input->insert(1, &rhs);

  popart::Tensor out("out", popart::TensorType::ActGrad, graph, di);
  mm.output->insert(0, &out);

  // Test the setup is correct
  BOOST_CHECK_THROW(mm.setup(), error);
}

// Test invalid matrix multiplication
BOOST_AUTO_TEST_CASE(MatMul_ErrorCase4) {

  // Setup

  popart::Ir ir;
  auto &graph = ir.getMainGraph();

  popart::MatMulOp mm(Onnx::Operators::MatMul_9,
                      {graph, ""},
                      nonstd::nullopt,
                      {},
                      OptionalDataType());

  popart::DebugContext dc;
  popart::DebugInfo di(dc, "test");

  popart::Tensor lhs("lhs", popart::TensorType::ActGrad, graph, di);
  lhs.info.set(popart::DataType::FLOAT, {});
  mm.input->insert(0, &lhs);

  popart::Tensor rhs("rhs", popart::TensorType::ActGrad, graph, di);
  rhs.info.set(popart::DataType::FLOAT, {10, 2});
  mm.input->insert(1, &rhs);

  popart::Tensor out("out", popart::TensorType::ActGrad, graph, di);
  mm.output->insert(0, &out);

  // Test the setup is correct
  BOOST_CHECK_THROW(mm.setup(), error);
}


#ifndef INCLUDED_kinematics_SceneBase_HH
#define INCLUDED_kinematics_SceneBase_HH

#include "scheme/types.hh"
#include "scheme/util/meta/util.hh"
#include "scheme/kinematics/ConformationBase.hh"

#include <boost/any.hpp>
#include <vector>

namespace scheme {
namespace kinematics {

template<
	class _Position,
	class _Index = uint64_t
>
struct SceneBase {
	typedef _Index Index;
	typedef _Position Position;
	typedef SceneBase<Position,Index> THIS;

	// members
		std::vector<Position> symframes_; // include identity at first position
		Index n_sym_bodies_, n_bodies_;
		std::vector<Position> positions_;

	SceneBase() : n_bodies_(0), n_sym_bodies_(0), symframes_(1,Position::Identity()) {}

	virtual ~SceneBase(){}

	virtual shared_ptr<THIS> clone_shallow() const { return clone_deep(); } // default to deep copy
	virtual shared_ptr<THIS> clone_specific_deep(std::vector<Index> deep_bodies) const { return clone_deep(); }
	virtual shared_ptr<THIS> clone_deep() const = 0;

	Position position(Index i) const
	{
		Index isym = this->sym_index_map(i);
		return this->symframes_.at(isym) * positions_.at(i);
	}
	void set_position(Index i, Position const & newp)
	{
		positions_.at(i) = newp;
	}
	template< class Actor >
	bool add_actor( Index ib, Actor const & actor )
	{
		boost::any any_actor = actor;
		return add_actor( ib, any_actor );
	}
	virtual bool add_actor( Index ib, boost::any const & a ) { return false; }

	template< class Actor >
	Actor get_actor( Index ib, Index ia ) const
	{
		boost::any any_actor = util::meta::type2type<Actor>();
		get_actor( ib, ia, any_actor );
		return boost::any_cast<Actor>( any_actor );
	}
	virtual bool get_actor( Index ib, Index ia, boost::any & a ) const { return false; }

	template< class Actor >
	Actor & get_nonconst_actor( Index ib, Index ia )
	{
		boost::any any_actor = util::meta::type2type<Actor>();
		get_nonconst_actor( ib, ia, any_actor );
		return *boost::any_cast<Actor*>( any_actor );
	}
	virtual bool get_nonconst_actor( Index ib, Index ia, boost::any & a ){ return false; }

	template< class Actor >
	int num_actors( Index ib ) const
	{
		boost::any any_actor = util::meta::type2type<Actor>();
		return num_actors( ib, any_actor );
	}
	virtual int num_actors( Index ib, boost::any const & a ) const { return -1; }

	template< class Actor >
	bool clear_actors( Index ib )
	{
		Actor a;
		boost::any any_actor = a;
		return clear_actors( ib, any_actor );
	}
	virtual bool clear_actors( Index ib, boost::any const & a ) { return false; }

	// symmetry stuff, doesn't need to be Conformation-sepcific
	void update_symmetry( Index nbodies )
	{
		n_bodies_ = nbodies;
		n_sym_bodies_ = n_bodies_*((Index)symframes_.size());
	}
	Index sym_index_map(Index & i) const
	{
		Index isym = i / n_bodies_;
		i = i % n_bodies_;
		return isym;
	}

	void set_symmetry(std::vector<Position> const & sym){ symframes_ = sym; update_symmetry(n_bodies_); }
	void add_symframe(Position const & symframe){ symframes_.push_back(symframe); update_symmetry(n_bodies_); }
	std::vector<Position> const & symframes() const { return symframes_; }
	Index nbodies() const { return n_sym_bodies_; }
	Index nbodies_asym() const { return n_bodies_; }

	// sketch courtesy of brian
	// template< class ConformationCOP >
	// void replace_body( Index ib, ConformationCOP conf) {
	// 	boost::any any_conf = conf;
	// 	replace_body( ib, any_conf );
	// }

	virtual void replace_body( Index ib, shared_ptr<ConformationBase const> cb) {}

};

}
}

#endif


#include <boost/python.hpp>

#include <gamecontroller/RoboCupGameControlData.hpp>

#include <types/JointValues.hpp>
#include <types/SensorValues.hpp>

#include "converters/array_conv.hpp"
#include "converters/EigenM33f_conv.hpp"
#include "converters/Point_conv.hpp"

#include "gamecontroller/RoboCupGameControlData.hpp"
#include "types/BroadcastData.hpp"

#include "utils/SPLDefs.hpp"

/* Used for registering c arrays of fixed size */
#define REG_ARRAY_TYPE(type, size)                                              \
   boost::python::to_python_converter<type[size],array_to_python<type,size> >();\
   array_to_python<type, size>::ready()

void register_python_converters()
{

   /* Sensor and Joint Values */
   REG_ARRAY_TYPE(float, Sensors::NUMBER_OF_SENSORS);
   REG_ARRAY_TYPE(float, Sonar::NUMBER_OF_READINGS);
   REG_ARRAY_TYPE(float, Joints::NUMBER_OF_JOINTS);

   /* Game Controller Data */
   REG_ARRAY_TYPE(PlayerInfo, MAX_NUM_PLAYERS);
   REG_ARRAY_TYPE(TeamInfo  ,               2);

   /* Receiver Blackboard Data */
   REG_ARRAY_TYPE(bool, ROBOTS_PER_TEAM);
   REG_ARRAY_TYPE(BroadcastData, ROBOTS_PER_TEAM);
   REG_ARRAY_TYPE(time_t, ROBOTS_PER_TEAM);

   boost::python::to_python_converter<
      Point,
      Point_to_python>();

   boost::python::to_python_converter<
      PointF,
      PointF_to_python>();

   boost::python::to_python_converter<
      Eigen::Matrix<float, 3, 3>,
      EigenM33f_to_python>();
}


#pragma once

#include <vector>
#include <memory>
#include <iostream>
#include <random>

#include <Eigen/Dense>

#include "Common/Math/Vec3.hpp"
#include "Common/Math/Rotation.hpp"
#include "Common/Time/BaseTimer.hpp"

#include "Components/Simulation/SimulationObject6DOF.hpp"
#include "Components/Simulation/Motor.hpp"
#include "Components/Simulation/UWBRadio.hpp"
#include "Components/Logic/QuadcopterLogic.hpp"

namespace Simulation {
template <class logicType>
class Quadcopter_T : public SimulationObject6DOF {
 public:

  Quadcopter_T(BaseTimer* const masterTimer, double mass,
             Eigen::Matrix<double, 3, 3> inertiaMatrix, double armLength,
             Vec3d centreOfMassError, double motorMinSpeed,
             double motorMaxSpeed, double propThrustFromSpeedSqr,
             double propTorqueFromSpeedSqr, double motorTimeConst,
             double motorInertia, Vec3d linDragCoeffB,
             uint8_t id,
             Onboard::QuadcopterConstants::QuadcopterType quadcopterType,
             double onboardLogicPeriod);

  virtual ~Quadcopter_T() {
  }

  virtual void Run();

  double GetMotorForce(unsigned i) const {
    //todo: this is a bit of a hack! should do something nicer than just z-component.
    return _motors[i].GetForce().z;
  }

  //! Set an external force, that remains active until replaced. Expressed in world frame
  void SetExternalForce(Vec3d in) {
    _externalForce = in;
  }

  void SetExternalTorque(Vec3d in) {
    _externalTorque = in;
  }

  void GetEstimate(Vec3f &pos, Vec3f &vel, Rotationf &att,
                   Vec3f &angVel) const {
    _logic.GetEstimate(pos, vel, att, angVel);
  }

  virtual void AddUWBRadioTarget(uint8_t id, Vec3f pos) {
    _logic.AddRangingTargetId(id, pos);
  }

  virtual void SetCommandRadioMsg(
      RadioTypes::RadioMessageDecoded::RawMessage const raw) {
    RadioTypes::RadioMessageDecoded msg = RadioTypes::RadioMessageDecoded(
        raw.raw);
    _logic.SetRadioMessage(msg);
  }

  virtual void GetTelemetryDataPackets(TelemetryPacket::data_packet_t &dataPacket1,
                               TelemetryPacket::data_packet_t &dataPacket2) {
    _logic.GetTelemetryDataPackets(dataPacket1, dataPacket2);
    return;
  }

 protected:
  //NOTE! You must add them in the order of the enum, above (FR>LE>RE>RI)
  void AddMotor(Vec3d position, Vec3d spinDirection,
                Motor::PropellerHandedness handedness, double minSpeed,
                double maxSpeed, double thrustFromSpeedSqr,
                double torqueFromSpeedSqr, double timeConst, double inertia) {
    _motors.push_back(
        Motor(_integrationTimer.GetMasterTimer(), position, spinDirection,
              handedness, minSpeed, maxSpeed, thrustFromSpeedSqr,
              torqueFromSpeedSqr, timeConst, inertia));
    _motorSpeedCommands.push_back(0);
  }

  logicType _logic;
  std::vector<float> _motorSpeedCommands;

 private:

  Eigen::Matrix<double, 3, 3> _inertiaMatrix;
  Eigen::Matrix<double, 3, 3> _inertiaMatrixInv;
  double _mass;

  Vec3d _externalForce;  //external force acting on vehicle, in [N], and expressed in world frame
  Vec3d _externalTorque; //external torque acting on vehicle, in [Nm], and expressed in world frame

  // Drag coefficients
  Vec3d _linDragCoeffB; // [N*s/m] For x, y and z components of body frame

  std::vector<Motor> _motors;

  float _battVoltage;
  float _battCurrent;

  //Noise characteristics:
  double _stdDevAccNoise, _stdDevRateGyroNoise;

  std::default_random_engine _generator;
  std::normal_distribution<double> _normalDistribution;

  Timer _timerOnboardLogic;
  double _onboardLogicPeriod;

  float _IMU_yaw; //[rad]
  float _IMU_pitch;
  float _IMU_roll;
  Matrix<float, 3, 3> _R_inverse; //rotation matrix correponds to IMU's rotation
};

typedef Quadcopter_T<Onboard::QuadcopterLogic> Quadcopter;

}  //namespace Simulation


#include <b0/utils/time_sync.h>
#include <b0/utils/env.h>

#include <boost/format.hpp>
#include <boost/date_time/gregorian/gregorian.hpp>
#include <boost/date_time/posix_time/posix_time.hpp>

namespace b0
{

TimeSync::TimeSync()
{
    target_offset_ = 0;
    max_acceptable_offset_ = b0::env::getInt("B0_TIMESYNC_MAX_OFFSET", 5 * 1000 * 1000);
    last_offset_time_ = hardwareTimeUSec();
    last_offset_value_ = 0;
    setMaxSlope(b0::env::getDouble("B0_TIMESYNC_MAX_SLOPE", 0.05));
}

TimeSync::~TimeSync()
{
}

void TimeSync::setMaxSlope(double max_slope)
{
    if(max_slope <= 0)
        throw std::runtime_error("max_slope must be strictly positive");
    if(max_slope > 1)
        throw std::runtime_error("max_slope must not be greater than one");

    max_slope_ = max_slope;
}

int64_t TimeSync::hardwareTimeUSec() const
{
    static boost::posix_time::ptime epoch(boost::gregorian::date(1970, 1, 1));
    boost::posix_time::ptime t = boost::posix_time::microsec_clock::local_time();
    return (t - epoch).total_microseconds();
}

int64_t TimeSync::timeUSec()
{
    return hardwareTimeUSec() + constantRateAdjustedOffset();
}

int64_t TimeSync::constantRateAdjustedOffset()
{
    boost::mutex::scoped_lock lock(mutex_);

    int64_t offset_delta = target_offset_ - last_offset_value_;
    int64_t slope_time = abs(offset_delta) / max_slope_;
    int64_t t = hardwareTimeUSec() - last_offset_time_;
    if(t >= slope_time)
        return target_offset_;
    else
        return last_offset_value_ + offset_delta * t / slope_time;
}

void TimeSync::updateTime(int64_t remoteTime)
{
    int64_t last_offset_value = constantRateAdjustedOffset();
    int64_t local_time = hardwareTimeUSec();

    {
        boost::mutex::scoped_lock lock(mutex_);

        last_offset_value_ = last_offset_value;
        last_offset_time_ = local_time;
        target_offset_ = remoteTime - local_time;

        if(max_acceptable_offset_ > 0 && target_offset_ > max_acceptable_offset_)
            throw std::runtime_error((boost::format("Clock offset is larger than B0_TIMESYNC_MAX_OFFSET (%ld usec)") % max_acceptable_offset_).str());
    }
}

} // namespace b0