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

/*

 Copyright (C) 2009 Ralph Schreyer

 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 fdmarithmeticaveragecondition.cpp

    \brief step condition to handle arithmetic average

*/

#include <ql/math/interpolations/cubicinterpolation.hpp>

#include <ql/methods/finitedifferences/operators/fdmlinearoplayout.hpp>

#include <ql/methods/finitedifferences/stepconditions/fdmarithmeticaveragecondition.hpp>

#include <utility>

namespace QuantLib {

    FdmArithmeticAverageCondition::FdmArithmeticAverageCondition(

        std::vector<Time> averageTimes,

        Real,

        Size pastFixings,

        const ext::shared_ptr<FdmMesher>& mesher,

        Size equityDirection)

    : x_(mesher->layout()->dim()[equityDirection]),

      a_(mesher->layout()->dim()[equityDirection == 0 ? 1 : 0]),

      averageTimes_(std::move(averageTimes)), pastFixings_(pastFixings), mesher_(mesher),

      equityDirection_(equityDirection) {

        QL_REQUIRE(mesher->layout()->dim().size()==2, "2D allowed only");

        QL_REQUIRE(equityDirection == 0 || equityDirection == 1,

                   "equityDirection has to be 0 or 1");

        const Size xSpacing = mesher_->layout()->spacing()[equityDirection];

        Array tmp = mesher_->locations(equityDirection);

        for (Size i = 0; i < x_.size(); ++i) {

            x_[i] = std::exp(tmp[i*xSpacing]);

        }

        const Size averageDirection = equityDirection == 0 ? 1 : 0;

        const Size aSpacing = mesher_->layout()->spacing()[averageDirection];

        tmp = mesher_->locations(averageDirection);

        for (Size i = 0; i < a_.size(); ++i) {

            a_[i] = std::exp(tmp[i*aSpacing]);

        }

    }

    void FdmArithmeticAverageCondition::applyTo(Array& a, Time t) const {

        QL_REQUIRE(mesher_->layout()->size() == a.size(),

                   "inconsistent array dimensions");

        const auto iter

            = std::find(averageTimes_.begin(), averageTimes_.end(), t);

        const Size nTimes

            = std::count(averageTimes_.begin(), averageTimes_.end(), t);

        if (nTimes > 0) {

            Array aCopy(a);

            const Size iT = iter - averageTimes_.begin() + 1 + pastFixings_;

            const Size averageDirection = equityDirection_ == 0 ? 1 : 0;

            const Size xSpacing = mesher_->layout()->spacing()[equityDirection_];

            const Size aSpacing = mesher_->layout()->spacing()[averageDirection];

            Array tmp(a_.size());

            for (Size i=0; i<x_.size(); ++i) {

                for (Size j=0; j<a_.size(); ++j) {

                    Size index = i*xSpacing + j*aSpacing;

                    tmp[j] = aCopy[index];

                }

                MonotonicCubicNaturalSpline interp(a_.begin(), a_.end(),

                                           tmp.begin());

                for (Size j=0; j<a_.size(); ++j) {

                    Size index = i*xSpacing + j*aSpacing;

                    a[index] = interp((iT-nTimes)/(double)(iT)*a_[j] +

                                      nTimes/(double)(iT)*x_[i], true);

                }

            }

        }

    }

}