/* -*- mode: c++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */

/*
 Copyright (C) 2019 SoftSolutions! S.r.l.
 Copyright (C) 2025 Peter Caspers

 This file is part of QuantLib, a free-software/open-source library
 for financial quantitative analysts and developers - http://quantlib.org/

 QuantLib is free software: you can redistribute it and/or modify it
 under the terms of the QuantLib license.  You should have received a
 copy of the license along with this program; if not, please email
 <quantlib-dev@lists.sf.net>. The license is also available online at
 <https://www.quantlib.org/license.shtml>.

 This program is distributed in the hope that it will be useful, but WITHOUT
 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 FOR A PARTICULAR PURPOSE.  See the license for more details.
*/

/*! \file globalbootstrap.hpp
    \brief global bootstrap, with additional restrictions
*/

#ifndef quantlib_global_bootstrap_hpp
#define quantlib_global_bootstrap_hpp

#include <ql/math/interpolations/linearinterpolation.hpp>
#include <ql/math/optimization/levenbergmarquardt.hpp>
#include <ql/termstructures/bootstraphelper.hpp>
#include <ql/utilities/dataformatters.hpp>
#include <algorithm>
#include <functional>
#include <utility>

namespace QuantLib {

class MultiCurveBootstrap;

class MultiCurveBootstrapContributor {
public:
    virtual ~MultiCurveBootstrapContributor() = default;
    virtual void
    setParentBootstrapper(const ext::shared_ptr<MultiCurveBootstrap>& b) const = 0;
    virtual Array setupCostFunction() const = 0;
    virtual void setCostFunctionArgument(const Array& v) const = 0;
    virtual Array evaluateCostFunction() const = 0;
    virtual void setToValid() const = 0;
};

class MultiCurveBootstrap : public ext::enable_shared_from_this<MultiCurveBootstrap> {
  public:
    explicit MultiCurveBootstrap(Real accuracy);
    explicit MultiCurveBootstrap(ext::shared_ptr<OptimizationMethod> optimizer = nullptr,
                        ext::shared_ptr<EndCriteria> endCriteria = nullptr);
    void add(const MultiCurveBootstrapContributor* c);
    void addObserver(Observer* o);
    void runMultiCurveBootstrap();
    void setOtherContributorsToValid() const;
    void finalizeCalculation();

  private:
    ext::shared_ptr<OptimizationMethod> optimizer_;
    ext::shared_ptr<EndCriteria> endCriteria_;
    std::vector<const MultiCurveBootstrapContributor*> contributors_;
    std::vector<Observer*> observers_;
};

class AdditionalBootstrapVariables {
  public:
    virtual ~AdditionalBootstrapVariables() = default;
    // Initialize variables to initial guesses and return them.
    virtual Array initialize(bool validData) = 0;
    // Update variables to given values.
    virtual void update(const Array& x) = 0;
};

/*! Global boostrapper, with additional restrictions

  The additionalDates functor must return a set of additional dates to add to the
  interpolation grid. These dates must only depend on the global evaluation date.

  The additionalPenalties functor must yield at least as many values such that

  number of (usual, alive) rate helpers + number of additional values >= number of data points - 1

  (note that the data points contain t=0). These values are treated as additional
  error terms in the optimization. The usual rate helpers return quoteError here.
  All error terms are equally weighted.

  The additionalHelpers are registered with the curve like the usual rate helpers,
  but no pillar dates or error terms are added for them. Pillars and error terms
  have to be added by additionalDates and additionalPenalties.

  The additionalVariables interface manages a set of additional variables to add
  to the optimization. This is useful to optimize model parameters used by rate
  helpers, for example, convexity adjustments for futures. See SimpleQuoteVariables
  for a concrete implementation of this interface.

  WARNING: This class is known to work with Traits Discount, ZeroYield, Forward,
  i.e. the usual IR curves traits in QL. It requires Traits::transformDirect()
  and Traits::transformInverse() to be implemented. Also, check the usage of
  Traits::updateGuess(), Traits::guess() in this class.
*/
template <class Curve> class GlobalBootstrap final : public MultiCurveBootstrapContributor {
    typedef typename Curve::traits_type Traits;             // ZeroYield, Discount, ForwardRate
    typedef typename Curve::interpolator_type Interpolator; // Linear, LogLinear, ...
    typedef std::function<Array(const std::vector<Time>&, const std::vector<Real>&)>
        AdditionalPenalties;

  public:
    GlobalBootstrap(Real accuracy = Null<Real>(),
                    ext::shared_ptr<OptimizationMethod> optimizer = nullptr,
                    ext::shared_ptr<EndCriteria> endCriteria = nullptr,
                    std::vector<Real> instrumentWeights = {});
    GlobalBootstrap(std::vector<ext::shared_ptr<typename Traits::helper>> additionalHelpers,
                    std::function<std::vector<Date>()> additionalDates,
                    AdditionalPenalties additionalPenalties,
                    Real accuracy = Null<Real>(),
                    ext::shared_ptr<OptimizationMethod> optimizer = nullptr,
                    ext::shared_ptr<EndCriteria> endCriteria = nullptr,
                    ext::shared_ptr<AdditionalBootstrapVariables> additionalVariables = nullptr,
                    std::vector<Real> instrumentWeights = {});
    GlobalBootstrap(std::vector<ext::shared_ptr<typename Traits::helper>> additionalHelpers,
                    std::function<std::vector<Date>()> additionalDates,
                    std::function<Array()> additionalPenalties,
                    Real accuracy = Null<Real>(),
                    ext::shared_ptr<OptimizationMethod> optimizer = nullptr,
                    ext::shared_ptr<EndCriteria> endCriteria = nullptr,
                    ext::shared_ptr<AdditionalBootstrapVariables> additionalVariables = nullptr,
                    std::vector<Real> instrumentWeights = {});
    void setup(Curve *ts);
    void calculate() const;

  private:
    void initialize() const;
    void
    setParentBootstrapper(const ext::shared_ptr<MultiCurveBootstrap>& b) const override;
    Array setupCostFunction() const override;
    void setCostFunctionArgument(const Array& v) const override;
    Array evaluateCostFunction() const override;
    void setToValid() const override;
    Curve* ts_;
    Real accuracy_;
    ext::shared_ptr<OptimizationMethod> optimizer_;
    ext::shared_ptr<EndCriteria> endCriteria_;
    std::vector<ext::shared_ptr<typename Traits::helper>> additionalHelpers_;
    mutable std::vector<ext::shared_ptr<typename Traits::helper>> aliveInstruments_;
    mutable std::vector<ext::shared_ptr<typename Traits::helper>> aliveAdditionalHelpers_;
    std::function<std::vector<Date>()> additionalDates_;
    AdditionalPenalties additionalPenalties_;
    ext::shared_ptr<AdditionalBootstrapVariables> additionalVariables_;
    mutable std::vector<Real> instrumentWeights_;
    mutable std::vector<Real> aliveInstrumentWeights_;
    mutable bool initialized_ = false, validCurve_ = false;
    mutable ext::shared_ptr<MultiCurveBootstrap> parentBootstrapper_ = nullptr;
};

// template definitions

template <class Curve>
GlobalBootstrap<Curve>::GlobalBootstrap(Real accuracy,
                                        ext::shared_ptr<OptimizationMethod> optimizer,
                                        ext::shared_ptr<EndCriteria> endCriteria,
                                        std::vector<Real> instrumentWeights)
: ts_(nullptr), accuracy_(accuracy), optimizer_(std::move(optimizer)),
  endCriteria_(std::move(endCriteria)), instrumentWeights_(std::move(instrumentWeights)) {}

template <class Curve>
GlobalBootstrap<Curve>::GlobalBootstrap(
    std::vector<ext::shared_ptr<typename Traits::helper>> additionalHelpers,
    std::function<std::vector<Date>()> additionalDates,
    AdditionalPenalties additionalPenalties,
    Real accuracy,
    ext::shared_ptr<OptimizationMethod> optimizer,
    ext::shared_ptr<EndCriteria> endCriteria,
    ext::shared_ptr<AdditionalBootstrapVariables> additionalVariables,
    std::vector<Real> instrumentWeights)
: ts_(nullptr), accuracy_(accuracy), optimizer_(std::move(optimizer)),
  endCriteria_(std::move(endCriteria)), additionalHelpers_(std::move(additionalHelpers)),
  additionalDates_(std::move(additionalDates)),
  additionalPenalties_(std::move(additionalPenalties)),
  additionalVariables_(std::move(additionalVariables)),
  instrumentWeights_(std::move(instrumentWeights)) {}

template <class Curve>
GlobalBootstrap<Curve>::GlobalBootstrap(
    std::vector<ext::shared_ptr<typename Traits::helper>> additionalHelpers,
    std::function<std::vector<Date>()> additionalDates,
    std::function<Array()> additionalPenalties,
    Real accuracy,
    ext::shared_ptr<OptimizationMethod> optimizer,
    ext::shared_ptr<EndCriteria> endCriteria,
    ext::shared_ptr<AdditionalBootstrapVariables> additionalVariables,
    std::vector<Real> instrumentWeights)
: GlobalBootstrap(std::move(additionalHelpers),
                  std::move(additionalDates),
                  additionalPenalties ?
                      [f = std::move(additionalPenalties)](
                          const std::vector<Time>&, const std::vector<Real>&) { return f(); } :
                      AdditionalPenalties(),
                  accuracy,
                  std::move(optimizer),
                  std::move(endCriteria),
                  std::move(additionalVariables),
                  std::move(instrumentWeights)) {}

template <class Curve>
void GlobalBootstrap<Curve>::setParentBootstrapper(const ext::shared_ptr<MultiCurveBootstrap>& b) const {
    parentBootstrapper_ = b;
}

template <class Curve> void GlobalBootstrap<Curve>::setToValid() const { validCurve_ = true; }

template <class Curve> void GlobalBootstrap<Curve>::setup(Curve* ts) {
    ts_ = ts;
    for (Size j = 0; j < ts_->instruments_.size(); ++j)
        ts_->registerWithObservables(ts_->instruments_[j]);
    for (Size j = 0; j < additionalHelpers_.size(); ++j)
        ts_->registerWithObservables(additionalHelpers_[j]);

    // setup optimizer and EndCriteria
    Real accuracy = accuracy_ != Null<Real>() ? accuracy_ : ts_->accuracy_;
    if (!optimizer_) {
        optimizer_ = ext::make_shared<LevenbergMarquardt>(accuracy, accuracy, accuracy);
    }
    if (!endCriteria_) {
        endCriteria_ = ext::make_shared<EndCriteria>(1000, 10, accuracy, accuracy, accuracy);
    }

    // check number of instrument weights
    QL_REQUIRE(instrumentWeights_.empty() || instrumentWeights_.size() == ts_->instruments_.size(),
               "GlobalBootstrap: number of instrument weights ("
                   << instrumentWeights_.size() << ") must match number of instruments ("
                   << ts_->instruments_.size() << ")");
    instrumentWeights_.resize(ts_->instruments_.size(), 1.0);

    // do not initialize yet: instruments could be invalid here
    // but valid later when bootstrapping is actually required
}

template <class Curve> void GlobalBootstrap<Curve>::initialize() const {

    const Date firstDate = Traits::initialDate(ts_);

    // set up alive instruments and weights
    aliveInstruments_.clear();
    aliveInstrumentWeights_.clear();
    for(Size i=0;i<ts_->instruments_.size();++i) {
        if(ts_->instruments_[i]->pillarDate() > firstDate) {
            aliveInstruments_.push_back(ts_->instruments_[i]);
            aliveInstrumentWeights_.push_back(instrumentWeights_[i]);
        }
    }

    // set up alive additional helpers
    aliveAdditionalHelpers_.clear();
    std::copy_if(additionalHelpers_.begin(), additionalHelpers_.end(),
                 std::back_inserter(aliveAdditionalHelpers_),
                 [&firstDate](const ext::shared_ptr<typename Traits::helper>& h) {
                     return h->pillarDate() > firstDate;
                 });

    // skip expired additional dates
    std::vector<Date> additionalDates;
    if (additionalDates_)
        additionalDates = additionalDates_();
    if (!additionalDates.empty()) {
        additionalDates.erase(
            std::remove_if(additionalDates.begin(), additionalDates.end(),
                           [=](const Date& date) { return date <= firstDate; }),
            additionalDates.end()
        );
    }

    // calculate dates and times
    std::vector<Date> &dates = ts_->dates_;
    std::vector<Time> &times = ts_->times_;

    // first populate the dates vector and make sure they are sorted and unique
    dates.clear();
    dates.push_back(firstDate);
    std::transform(
        aliveInstruments_.begin(), aliveInstruments_.end(), std::back_inserter(dates),
        [](const ext::shared_ptr<typename Traits::helper>& h) { return h->pillarDate(); });
    dates.insert(dates.end(), additionalDates.begin(), additionalDates.end());
    std::sort(dates.begin(), dates.end());
    dates.erase(std::unique(dates.begin(), dates.end()), dates.end());

    // check if there are enough interpolation points
    QL_REQUIRE(dates.size() >= Interpolator::requiredPoints,
               "GlobalBootstrap: not enough curve points ("
                   << dates.size() << ") for interpolation requiring at least "
                   << Interpolator::requiredPoints);

    // build times vector
    times.clear();
    std::transform(dates.begin(), dates.end(), std::back_inserter(times),
                   [this](const Date& d) { return ts_->timeFromReference(d); });

    // determine maxDate ensuring all instruments and additional helpers are covered
    ts_->maxDate_ = dates.back();
    for (auto const& h : aliveInstruments_)
        ts_->maxDate_ = std::max(ts_->maxDate_, h->latestRelevantDate());
    for (auto const& h : aliveAdditionalHelpers_)
        ts_->maxDate_ = std::max(ts_->maxDate_, h->latestRelevantDate());

    // set initial guess only if the current curve cannot be used as guess
    if (!validCurve_ || ts_->data_.size() != dates.size()) {
        // ts_->data_[0] is the only relevant item,
        // but reasonable numbers might be needed for the whole data vector
        // because, e.g., of interpolation's early checks
        ts_->data_ = std::vector<Real>(dates.size(), Traits::initialValue(ts_));
        validCurve_ = false;
    }
    initialized_ = true;
}

template <class Curve> Array GlobalBootstrap<Curve>::setupCostFunction() const {

    // for single-curve boostrap, this was done in LazyObject::calculate() already, but for
    // multi-curve boostrap we have to do this manually for all contributing curves except
    // the main one, because calculate() is never triggered for them
    ts_->setCalculated(true);

    // we might have to call initialize even if the curve is initialized
    // and not moving, just because helpers might be date relative and change
    // with evaluation date change.
    // anyway it makes little sense to use date relative helpers with a
    // non-moving curve if the evaluation date changes
    if (!initialized_ || ts_->moving_)
        initialize();

    // setup helpers
    for (auto const& helper : aliveInstruments_) {
        // check for valid quote
        QL_REQUIRE(helper->quote()->isValid(),
                   "instrument (maturity: " << helper->maturityDate() << ", pillar: "
                                            << helper->pillarDate() << ") has an invalid quote");
        // don't try this at home!
        // This call creates helpers, and removes "const".
        // There is a significant interaction with observability.
        helper->setTermStructure(const_cast<Curve*>(ts_));
    }

    // setup additional helpers
    for (auto const& helper : aliveAdditionalHelpers_) {
        QL_REQUIRE(helper->quote()->isValid(),
                   "additional instrument (maturity: " << helper->maturityDate()
                                                       << ") has an invalid quote");
        helper->setTermStructure(const_cast<Curve*>(ts_));
    }

    // setup interpolation
    if (!validCurve_) {
        ts_->interpolation_ = ts_->interpolator_.interpolate(ts_->times_.begin(), ts_->times_.end(),
                                                             ts_->data_.begin());
    }

    // Initial guess. We have guesses for the curve values first (numberPillars),
    // followed by guesses for the additional variables.
    Array additionalGuesses;
    if (additionalVariables_) {
        additionalGuesses = additionalVariables_->initialize(validCurve_);
    }
    Array guess(ts_->times_.size() - 1 + additionalGuesses.size());
    for (Size i = 0; i < ts_->times_.size() - 1; ++i) {
        // just pass zero as the first alive helper, it's not used in the standard QL traits anyway
        // update ts_->data_ since Traits::guess() usually depends on previous values
        Traits::updateGuess(ts_->data_, Traits::guess(i + 1, ts_, validCurve_, 0), i + 1);
        guess[i] = Traits::transformInverse(ts_->data_[i + 1], i + 1, ts_);
    }
    std::copy(additionalGuesses.begin(), additionalGuesses.end(),
              guess.begin() + ts_->times_.size() - 1);
    return guess;
}

template <class Curve>
void GlobalBootstrap<Curve>::setCostFunctionArgument(const Array& x) const {
    // x has the same layout as guess above: the first numberPillars values go into
    // the curve, while the rest are new values for the additional variables.
    for (Size i = 0; i < ts_->times_.size() - 1; ++i) {
        Traits::updateGuess(ts_->data_, Traits::transformDirect(x[i], i + 1, ts_), i + 1);
    }
    ts_->interpolation_.update();
    if (additionalVariables_) {
        additionalVariables_->update(Array(x.begin() + ts_->times_.size() - 1, x.end()));
    }
}

template <class Curve>
Array GlobalBootstrap<Curve>::evaluateCostFunction() const {
    Array additionalErrors;
    if (additionalPenalties_) {
        additionalErrors = additionalPenalties_(ts_->times_, ts_->data_);
    }
    Array result(aliveInstruments_.size() + additionalErrors.size());
    for (Size i = 0; i < aliveInstruments_.size(); ++i)
        result[i] = aliveInstruments_[i]->quoteError() * aliveInstrumentWeights_[i];
    std::copy(additionalErrors.begin(), additionalErrors.end(),
              result.begin() + aliveInstruments_.size());
    return result;
}

template <class Curve>
void GlobalBootstrap<Curve>::calculate() const {

    if (parentBootstrapper_) {
        parentBootstrapper_->runMultiCurveBootstrap();
        return;
    }

    // single curve boostrap

    Array guess = setupCostFunction();

    NoConstraint noConstraint;

    SimpleCostFunction costFunction([this](const Array& x) {
        this->setCostFunctionArgument(x);
        return this->evaluateCostFunction();
    });

    Problem problem(costFunction, noConstraint, guess);
    EndCriteria::Type endType = optimizer_->minimize(problem, *endCriteria_);
    QL_REQUIRE(EndCriteria::succeeded(endType),
               "global bootstrap failed to minimize to required accuracy: " << endType);
    validCurve_ = true;
}

} // namespace QuantLib

#endif