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

/*
 Copyright (C) 2020 Lew Wei Hao

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

#include <ql/methods/finitedifferences/meshers/fdmmesher.hpp>
#include <ql/methods/finitedifferences/operators/fdmcirop.hpp>
#include <ql/methods/finitedifferences/operators/fdmlinearoplayout.hpp>
#include <ql/methods/finitedifferences/operators/secondderivativeop.hpp>
#include <ql/methods/finitedifferences/operators/secondordermixedderivativeop.hpp>
#include <ql/processes/blackscholesprocess.hpp>

namespace QuantLib {

    FdmCIREquityPart::FdmCIREquityPart(
        const ext::shared_ptr<FdmMesher>& mesher,
        const ext::shared_ptr<GeneralizedBlackScholesProcess>& bsProcess,
        Real strike)
    : dxMap_ (FirstDerivativeOp(0, mesher)),
      dxxMap_(SecondDerivativeOp(0, mesher)),
      mapT_ (0, mesher),
      mesher_(mesher),
      qTS_(bsProcess->dividendYield().currentLink()),
      strike_(strike),
      sigma1_(bsProcess->blackVolatility().currentLink()){
    }

    void FdmCIREquityPart::setTime(Time t1, Time t2) {
        const Rate q = qTS_->forwardRate(t1, t2, Continuous).rate();

        const Real v = sigma1_->blackForwardVariance(t1, t2, strike_)/(t2-t1);

        mapT_.axpyb(mesher_->locations(1) - q - 0.5*v, dxMap_,
                        dxxMap_.mult(Array(mesher_->layout()->size(), v/2)), -0.5*mesher_->locations(1));
    }

    const TripleBandLinearOp& FdmCIREquityPart::getMap() const {
        return mapT_;
    }

    FdmCIRRatesPart::FdmCIRRatesPart(
        const ext::shared_ptr<FdmMesher>& mesher,
        Real sigma, Real kappa, Real theta)
    : dyMap_(SecondDerivativeOp(1, mesher)
                   .mult(sigma*sigma*mesher->locations(1))
                   .add(FirstDerivativeOp(1, mesher)
                   .mult(kappa*(theta - mesher->locations(1))))),
      mapT_(1, mesher),
      mesher_(mesher){
    }

    void FdmCIRRatesPart::setTime(Time t1, Time t2) {
        mapT_.axpyb(Array(), dyMap_, dyMap_, -0.5*mesher_->locations(1));
    }

    const TripleBandLinearOp& FdmCIRRatesPart::getMap() const {
        return mapT_;
    }

    FdmCIRMixedPart::FdmCIRMixedPart(
        const ext::shared_ptr<FdmMesher>& mesher,
        const ext::shared_ptr<CoxIngersollRossProcess> & cirProcess,
        const ext::shared_ptr<GeneralizedBlackScholesProcess> & bsProcess,
        const Real rho,
        const Real strike)
        : dyMap_(SecondOrderMixedDerivativeOp(0, 1, mesher)
                     .mult(Array(mesher->layout()->size(), 2*rho*cirProcess->volatility()))),
          mapT_(0, 1, mesher),
          mesher_(mesher),
          sigma1_(bsProcess->blackVolatility().currentLink()),
          strike_(strike){
    }

    void FdmCIRMixedPart::setTime(Time t1, Time t2) {
        const Real v = std::sqrt(sigma1_->blackForwardVariance(t1, t2, strike_)/(t2-t1));
        NinePointLinearOp op(dyMap_.mult(Array(mesher_->layout()->size(), v)));
        mapT_.swap(op);
    }

    const NinePointLinearOp& FdmCIRMixedPart::getMap() const {
        return mapT_;
    }

    FdmCIROp::FdmCIROp(
        const ext::shared_ptr<FdmMesher>& mesher,
        const ext::shared_ptr<CoxIngersollRossProcess> & cirProcess,
        const ext::shared_ptr<GeneralizedBlackScholesProcess> & bsProcess,
        const Real rho,
        const Real strike)
    : dxMap_(mesher,
             bsProcess,
             strike),
      dyMap_(mesher,
             cirProcess->volatility(),
             cirProcess->speed(),
             cirProcess->level()),
      dzMap_(mesher,
             cirProcess,
             bsProcess,
             rho,
             strike){
    }


    void FdmCIROp::setTime(Time t1, Time t2) {
        dxMap_.setTime(t1, t2);
        dyMap_.setTime(t1, t2);
        dzMap_.setTime(t1, t2);
    }

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

    Array FdmCIROp::apply(const Array& u) const {
        Array dx = dxMap_.getMap().apply(u);
        Array dy = dyMap_.getMap().apply(u);
        Array dz = dzMap_.getMap().apply(u);

        return (dy + dx + dz);
    }

    Array FdmCIROp::apply_direction(Size direction,
                                    const Array& r) const {
        if (direction == 0)
            return dxMap_.getMap().apply(r);
        else if (direction == 1)
            return dyMap_.getMap().apply(r);
        else
            QL_FAIL("direction too large");
    }

    Array FdmCIROp::apply_mixed(const Array& r) const {
        return dzMap_.getMap().apply(r);
    }

    Array FdmCIROp::solve_splitting(Size direction,
                                    const Array& r, Real a) const {
        if (direction == 0) {
            return dxMap_.getMap().solve_splitting(r, a, 1.0);
        }
        else if (direction == 1) {
            return dyMap_.getMap().solve_splitting(r, a, 1.0);
        }
        else
            QL_FAIL("direction too large");
    }

    Array FdmCIROp::preconditioner(const Array& r, Real dt) const {
        return solve_splitting(1, solve_splitting(0, r, dt), dt) ;
    }

    std::vector<SparseMatrix> FdmCIROp::toMatrixDecomp() const {
        return {
            dxMap_.getMap().toMatrix(),
            dyMap_.getMap().toMatrix(),
            dzMap_.getMap().toMatrix()
        };
    }

}
