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

/*
 Copyright (C) 2012, 2013 Klaus Spanderen
 Copyright (C) 2014 Johannes Göttker-Schnetmann

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

#include <ql/methods/finitedifferences/operators/fdmhestonfwdop.hpp>
#include <ql/methods/finitedifferences/operators/modtriplebandlinearop.hpp>
#include <ql/methods/finitedifferences/meshers/fdmmesher.hpp>
#include <ql/methods/finitedifferences/operators/fdmlinearoplayout.hpp>
#include <ql/methods/finitedifferences/operators/firstderivativeop.hpp>
#include <ql/methods/finitedifferences/operators/secondderivativeop.hpp>
#include <ql/methods/finitedifferences/operators/secondordermixedderivativeop.hpp>
#include <ql/processes/hestonprocess.hpp>
#include <cmath>
#include <utility>

using std::exp;

namespace QuantLib {

    FdmHestonFwdOp::FdmHestonFwdOp(const ext::shared_ptr<FdmMesher>& mesher,
                                   const ext::shared_ptr<HestonProcess>& process,
                                   FdmSquareRootFwdOp::TransformationType type,
                                   ext::shared_ptr<LocalVolTermStructure> leverageFct,
                                   const Real mixingFactor)
    : type_(type), kappa_(process->kappa()), theta_(process->theta()), sigma_(process->sigma()),
      rho_(process->rho()), v0_(process->v0()), mixedSigma_(mixingFactor * sigma_),
      rTS_(process->riskFreeRate().currentLink()), qTS_(process->dividendYield().currentLink()),
      varianceValues_(0.5 * mesher->locations(1)),
      dxMap_(ext::make_shared<FirstDerivativeOp>(0, mesher)),
      dxxMap_(ext::make_shared<ModTripleBandLinearOp>(TripleBandLinearOp(
          type == FdmSquareRootFwdOp::Log ?
              SecondDerivativeOp(0, mesher).mult(0.5 * Exp(mesher->locations(1))) :
              SecondDerivativeOp(0, mesher).mult(0.5 * mesher->locations(1))))),
      boundary_(ext::make_shared<ModTripleBandLinearOp>(TripleBandLinearOp(
          SecondDerivativeOp(0, mesher).mult(Array(mesher->locations(0).size(), 0.0))))),
      mapX_(ext::make_shared<TripleBandLinearOp>(0, mesher)),
      mapY_(ext::make_shared<FdmSquareRootFwdOp>(mesher, kappa_, theta_, mixedSigma_, 1, type)),
      correlation_(ext::make_shared<NinePointLinearOp>(
          type == FdmSquareRootFwdOp::Log ?
              SecondOrderMixedDerivativeOp(0, 1, mesher)
                  .mult(Array(mesher->layout()->size(), rho_ * mixedSigma_)) :
              SecondOrderMixedDerivativeOp(0, 1, mesher)
                  .mult(rho_ * mixedSigma_ * mesher->locations(1)))),
      leverageFct_(std::move(leverageFct)), mesher_(mesher) {
        // zero flux boundary condition
        const Size n = mesher->layout()->dim()[1];
        const Real lowerBoundaryFactor = mapY_->lowerBoundaryFactor(type);
        const Real upperBoundaryFactor = mapY_->upperBoundaryFactor(type);

        const Real logFacLow = type == FdmSquareRootFwdOp::Log ? Real(exp(mapY_->v(0))) : 1.0;
        const Real logFacUpp = type == FdmSquareRootFwdOp::Log ? Real(exp(mapY_->v(n+1))) : 1.0;

        const Real alpha = -2*rho_/mixedSigma_*lowerBoundaryFactor*logFacLow;
        const Real beta  = -2*rho_/mixedSigma_*upperBoundaryFactor*logFacUpp;

        ModTripleBandLinearOp fDx(FirstDerivativeOp(0, mesher));

        for (const auto& iter : *mesher->layout()) {
            if (iter.coordinates()[1] == 0) {
                const Size idx = iter.index();
                if (!leverageFct_) {
                    dxxMap_->upper(idx) += alpha*fDx.upper(idx);
                    dxxMap_->diag(idx) += alpha*fDx.diag(idx);
                    dxxMap_->lower(idx) += alpha*fDx.lower(idx);
                }
                boundary_->upper(idx)= alpha*fDx.upper(idx);
                boundary_->diag(idx) = alpha*fDx.diag(idx);
                boundary_->lower(idx) = alpha*fDx.lower(idx);
            }
            else if (iter.coordinates()[1] == n-1) {
                const Size idx = iter.index();

                if (!leverageFct_) {
                    dxxMap_->upper(idx)+= beta*fDx.upper(idx);
                    dxxMap_->diag(idx) += beta*fDx.diag(idx);
                    dxxMap_->lower(idx) += beta*fDx.lower(idx);
                }
                boundary_->upper(idx)= beta*fDx.upper(idx);
                boundary_->diag(idx) = beta*fDx.diag(idx);
                boundary_->lower(idx) = beta*fDx.lower(idx);
            }
        }
    }

    Size FdmHestonFwdOp::size() const {
        return 2;
    }

    void FdmHestonFwdOp::setTime(Time t1, Time t2){
        const Rate r = rTS_->forwardRate(t1, t2, Continuous).rate();
        const Rate q = qTS_->forwardRate(t1, t2, Continuous).rate();
        if (leverageFct_ != nullptr) {
            L_ = getLeverageFctSlice(t1, t2);
            Array Lsquare = L_*L_;
            if (type_ == FdmSquareRootFwdOp::Plain) {
                mapX_->axpyb( Array(1, -r + q), *dxMap_,
                    dxxMap_->multR(Lsquare).add(boundary_->multR(L_))
                    .add(dxMap_->multR(rho_*mixedSigma_*L_))
                    .add(dxMap_->mult(varianceValues_).multR(Lsquare)),
                              Array());
            } else if (type_ == FdmSquareRootFwdOp::Power) {
                mapX_->axpyb( Array(1, -r + q), *dxMap_,
                    dxxMap_->multR(Lsquare).add(boundary_->multR(L_))
                    .add(dxMap_->multR(rho_*2.0*kappa_*theta_/(mixedSigma_)*L_))
                    .add(dxMap_->mult(varianceValues_).multR(Lsquare)), Array());
            } else if (type_ == FdmSquareRootFwdOp::Log) {
                mapX_->axpyb( Array(1, -r + q), *dxMap_,
                    dxxMap_->multR(Lsquare).add(boundary_->multR(L_))
                    .add(dxMap_->mult(0.5*Exp(2.0*varianceValues_)).multR(Lsquare)),
                              Array());
            }
        } else {
            if (type_ == FdmSquareRootFwdOp::Plain) {
                mapX_->axpyb( - r + q + rho_*mixedSigma_ + varianceValues_, *dxMap_,
                        *dxxMap_, Array());
            } else if (type_ == FdmSquareRootFwdOp::Power) {
                mapX_->axpyb( - r + q + rho_*2.0*kappa_*theta_/(mixedSigma_) + varianceValues_,
                              *dxMap_, *dxxMap_, Array());
            } else if (type_ == FdmSquareRootFwdOp::Log) {
                mapX_->axpyb( - r + q + 0.5*Exp(2.0*varianceValues_), *dxMap_,
                        *dxxMap_, Array());
            }
        }
    }

    Array FdmHestonFwdOp::apply(const Array& u) const {
        if (leverageFct_ != nullptr) {
            return mapX_->apply(u)
                    + mapY_->apply(u)
                    + correlation_->apply(L_*u);
        } else {
            return mapX_->apply(u)
                    + mapY_->apply(u)
                    + correlation_->apply(u);
        }
    }

    Array FdmHestonFwdOp::apply_mixed(const Array& u) const{
        if (leverageFct_ != nullptr) {
            return correlation_->apply(L_*u);
        } else {
            return correlation_->apply(u);
        }
    }

    Array FdmHestonFwdOp::apply_direction(
        Size direction, const Array& u) const {

        if (direction == 0)
            return mapX_->apply(u) ;
        else if (direction == 1)
            return mapY_->apply(u) ;
        else
            QL_FAIL("direction too large");
    }

    Array FdmHestonFwdOp::solve_splitting(
        Size direction, const Array& u, Real s) const{
        if (direction == 0) {
            return mapX_->solve_splitting(u, s, 1.0);
        }
        else if (direction == 1) {
            return mapY_->solve_splitting(1, u, s);
        }
        else
            QL_FAIL("direction too large");
    }

    Array FdmHestonFwdOp::preconditioner(
        const Array& u, Real dt) const{
        return solve_splitting(1, u, dt);
    }

    Array FdmHestonFwdOp::getLeverageFctSlice(Time t1, Time t2) const {
        Array v(mesher_->layout()->size(), 1.0);

        if (!leverageFct_)
            return v;

        const Real t = 0.5*(t1+t2);
        const Time time = std::min(leverageFct_->maxTime(), t);
                                   //std::max(leverageFct_->minTime(), t));

        for (const auto& iter : *mesher_->layout()) {
            const Size nx = iter.coordinates()[0];

            if (iter.coordinates()[1] == 0) {
                const Real x = std::exp(mesher_->location(iter, 0));
                const Real spot = std::min(leverageFct_->maxStrike(),
                                           std::max(leverageFct_->minStrike(), x));
                v[nx] = std::max(0.01, leverageFct_->localVol(time, spot, true));
            }
            else {
                v[iter.index()] = v[nx];
            }
        }
        return v;
    }

    std::vector<SparseMatrix> FdmHestonFwdOp::toMatrixDecomp() const {

        std::vector<SparseMatrix> retVal(3);

        retVal[0] = mapX_->toMatrix();
        retVal[1] = mapY_->toMatrix();
        retVal[2] = correlation_->toMatrix();

        return retVal;
    }

}
