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

/*
 Copyright (C) 2010 Klaus Spanderen

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


#include <ql/processes/blackscholesprocess.hpp>
#include <ql/methods/finitedifferences/meshers/fdmmesher.hpp>
#include <ql/methods/finitedifferences/operators/fdmlinearoplayout.hpp>
#include <ql/methods/finitedifferences/operators/fdm2dblackscholesop.hpp>
#include <ql/methods/finitedifferences/operators/secondordermixedderivativeop.hpp>
#include <boost/numeric/ublas/matrix.hpp>

namespace QuantLib {

    Fdm2dBlackScholesOp::Fdm2dBlackScholesOp(
            const ext::shared_ptr<FdmMesher>& mesher,
            const ext::shared_ptr<GeneralizedBlackScholesProcess>& p1,
            const ext::shared_ptr<GeneralizedBlackScholesProcess>& p2,
            Real correlation,
            Time /*maturity*/,
            bool localVol,
            Real illegalLocalVolOverwrite)
    : mesher_(mesher),
      p1_(p1),
      p2_(p2),
      
      localVol1_((localVol) ? p1->localVolatility().currentLink()
                            : ext::shared_ptr<LocalVolTermStructure>()),
      localVol2_((localVol) ? p2->localVolatility().currentLink()
                            : ext::shared_ptr<LocalVolTermStructure>()),
                            
      x_((localVol) ? Array(Exp(mesher->locations(0))) : Array()),
      y_((localVol) ? Array(Exp(mesher->locations(1))) : Array()),

      opX_(mesher, p1, p1->x0(), localVol, illegalLocalVolOverwrite, 0),
      opY_(mesher, p2, p2->x0(), localVol, illegalLocalVolOverwrite, 1),
      
      corrMapT_(0, 1, mesher),
      corrMapTemplate_(SecondOrderMixedDerivativeOp(0, 1, mesher)
                      .mult(Array(mesher->layout()->size(), correlation))),
                      
      illegalLocalVolOverwrite_(illegalLocalVolOverwrite) { 
    }
           
    Size Fdm2dBlackScholesOp::size() const {
        return 2;
    }
    
    void Fdm2dBlackScholesOp::setTime(Time t1, Time t2) {
        opX_.setTime(t1, t2);
        opY_.setTime(t1, t2);

        if (localVol1_ != nullptr) {
            Array vol1(mesher_->layout()->size()), vol2(mesher_->layout()->size());
            for (const auto& iter : *mesher_->layout()) {
                const Size i = iter.index();

                if (illegalLocalVolOverwrite_ < 0.0) {
                    vol1[i] = localVol1_->localVol(0.5*(t1+t2), x_[i], true);
                    vol2[i] = localVol2_->localVol(0.5*(t1+t2), y_[i], true);
                }
                else {
                    try {
                        vol1[i] = localVol1_->localVol(0.5*(t1+t2), x_[i],true);
                    } catch (Error&) {
                        vol1[i] = illegalLocalVolOverwrite_;
                    }
                    try {
                        vol2[i] = localVol2_->localVol(0.5*(t1+t2), y_[i],true);
                    } catch (Error&) {
                        vol2[i] = illegalLocalVolOverwrite_;
                    }

                }
            }
            corrMapT_ = corrMapTemplate_.mult(vol1*vol2);
        } else {
            const Real vol1 = p1_
                    ->blackVolatility()->blackForwardVol(t1, t2, p1_->x0());
    
            const Real vol2 = p2_
                    ->blackVolatility()->blackForwardVol(t1, t2, p2_->x0());
            
            corrMapT_ = corrMapTemplate_
                      .mult(Array(mesher_->layout()->size(), vol1*vol2));
        }

        currentForwardRate_ = p1_->riskFreeRate()
                                 ->forwardRate(t1, t2, Continuous).rate();
    }

    Array Fdm2dBlackScholesOp::apply(const Array& x) const {
        return opX_.apply(x) + opY_.apply(x) + apply_mixed(x);
    }
    
    Array Fdm2dBlackScholesOp::apply_mixed(const Array& x) const {
        return corrMapT_.apply(x) + currentForwardRate_*x;
    }
    
    Array Fdm2dBlackScholesOp::apply_direction(
                                       Size direction, const Array& x) const {
        if (direction == 0) {
            return opX_.apply(x);
        }
        else if (direction == 1) {
            return opY_.apply(x);
        }
        else {
            QL_FAIL("direction is too large");
        }
    }
    
    Array Fdm2dBlackScholesOp::solve_splitting(Size direction,
                                               const Array& x, Real s) const {
        if (direction == 0) {
            return opX_.solve_splitting(direction, x, s);
        }
        else if (direction == 1) {
            return opY_.solve_splitting(direction, x, s);
        }
        else
            QL_FAIL("direction is too large");
    }
    
    Array Fdm2dBlackScholesOp::preconditioner(const Array& r, 
                                              Real dt) const {
        return solve_splitting(0, r, dt);
    }

    std::vector<SparseMatrix> Fdm2dBlackScholesOp::toMatrixDecomp() const {
        return {
            opX_.toMatrix(),
            opY_.toMatrix(),
            corrMapT_.toMatrix() +
            currentForwardRate_*boost::numeric::ublas::identity_matrix<Real>(
                    mesher_->layout()->size())
        };
    }

}
