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

/*

 Copyright (C) 2006 Banca Profilo S.p.A.

 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.

*/

#include <ql/processes/g2process.hpp>

#include <ql/processes/eulerdiscretization.hpp>

namespace QuantLib {

    G2Process::G2Process(Real a, Real sigma, Real b, Real eta, Real rho,

                         const Handle<YieldTermStructure>& termStructure)

    : a_(a), sigma_(sigma), b_(b), eta_(eta), rho_(rho),

      xProcess_(new QuantLib::OrnsteinUhlenbeckProcess(a, sigma, 0.0)),

      yProcess_(new QuantLib::OrnsteinUhlenbeckProcess(b, eta, 0.0)),

      termStructure_(termStructure) {

        registerWith(termStructure_);

    }

    Size G2Process::size() const {

        return 2;

    }

    Array G2Process::initialValues() const {

        Real z1_0 = termStructure_.empty() ? x0_ : phi(0.0);

        return { z1_0, y0_ };

    }

    Array G2Process::drift(Time t, const Array& z) const {

        // Drift in shifted coordinates z1 = x + phi(t), z2 = y:

        //   dz1 = (-a*z1 + a*phi(t) + phi'(t)) dt + sigma dW1

        //   dz2 = -b*z2 dt + eta dW2

        Real shiftDrift = 0.0;

        if (!termStructure_.empty()) {

            const Real h = 1.0e-4;

            Real phi_t  = phi(t);

            Real phi_th = phi(t + h);

            Real phiPrime = (phi_th - phi_t) / h;

            shiftDrift = a_ * phi_t + phiPrime;

        }

        return {

            xProcess_->drift(t, z[0]) + shiftDrift,

            yProcess_->drift(t, z[1])

        };

    }

    Matrix G2Process::diffusion(Time, const Array&) const {

        /* the correlation matrix is

           |  1   rho |

           | rho   1  |

           whose square root (which is used here) is

           |  1          0       |

           | rho   sqrt(1-rho^2) |

        */

        Matrix tmp(2,2);

        Real sigma1 = sigma_;

        Real sigma2 = eta_;

        tmp[0][0] = sigma1;       tmp[0][1] = 0.0;

        tmp[1][0] = rho_*sigma1;  tmp[1][1] = std::sqrt(1.0-rho_*rho_)*sigma2;

        return tmp;

    }

    Array G2Process::expectation(Time t0, const Array& z0,

                                 Time dt) const {

        // E[z1(t0+dt) | z1(t0)] = z1(t0)*exp(-a*dt)

        //                       + phi(t0+dt) - phi(t0)*exp(-a*dt)

        // E[z2(t0+dt) | z2(t0)] = z2(t0)*exp(-b*dt)

        Real shiftExp = 0.0;

        if (!termStructure_.empty()) {

            shiftExp = phi(t0 + dt) - phi(t0) * std::exp(-a_ * dt);

        }

        return {

            xProcess_->expectation(t0, z0[0], dt) + shiftExp,

            yProcess_->expectation(t0, z0[1], dt)

        };

    }

    Matrix G2Process::stdDeviation(Time t0, const Array& x0, Time dt) const {

        /* the correlation matrix is

           |  1   rho |

           | rho   1  |

           whose square root (which is used here) is

           |  1          0       |

           | rho   sqrt(1-rho^2) |

        */

        Matrix tmp(2,2);

        Real sigma1 = xProcess_->stdDeviation(t0, x0[0], dt);

        Real sigma2 = yProcess_->stdDeviation(t0, x0[1], dt);

        Real expa = std::exp(-a_*dt), expb = std::exp(-b_*dt);

        Real H = (rho_*sigma_*eta_)/(a_+b_)*(1-expa*expb);

        Real den =

            (0.5*sigma_*eta_)*std::sqrt((1-expa*expa)*(1-expb*expb)/(a_*b_));

        Real newRho = H/den;

        tmp[0][0] = sigma1;

        tmp[0][1] = 0.0;

        tmp[1][0] = newRho*sigma2;

        tmp[1][1] = std::sqrt(1.0-newRho*newRho)*sigma2;

        return tmp;

    }

    Matrix G2Process::covariance(Time t0, const Array& x0, Time dt) const {

        Matrix sigma = stdDeviation(t0, x0, dt);

        Matrix result = sigma*transpose(sigma);

        return result;

    }

    Real G2Process::x0() const {

        return termStructure_.empty() ? x0_ : phi(0.0);

    }

    Real G2Process::y0() const {

        return y0_;

    }

    Real G2Process::a() const {

        return a_;

    }

    Real G2Process::sigma() const {

        return sigma_;

    }

    Real G2Process::b() const {

        return b_;

    }

    Real G2Process::eta() const {

        return eta_;

    }

    Real G2Process::rho() const {

        return rho_;

    }

    const Handle<YieldTermStructure>& G2Process::termStructure() const {

        return termStructure_;

    }

    Real G2Process::phi(Time t) const {

        QL_REQUIRE(!termStructure_.empty(),

                   "no term structure given to G2Process");

        Rate forward = termStructure_->forwardRate(t, t, Continuous, NoFrequency);

        Real temp1 = sigma_*(1.0-std::exp(-a_*t))/a_;

        Real temp2 = eta_ *(1.0-std::exp(-b_*t))/b_;

        return 0.5*temp1*temp1 + 0.5*temp2*temp2 + rho_*temp1*temp2 + forward;

    }

    Rate G2Process::shortRate(Time, Real z1, Real z2) const {

        // The simulated state already includes phi(t) in z1, so r = z1 + z2.

        return z1 + z2;

    }

    G2ForwardProcess::G2ForwardProcess(Real a, Real sigma, Real b, Real eta, Real rho,

                                       const Handle<YieldTermStructure>& termStructure)

    : a_(a), sigma_(sigma), b_(b), eta_(eta), rho_(rho),

      xProcess_(new QuantLib::OrnsteinUhlenbeckProcess(a, sigma, 0.0)),

      yProcess_(new QuantLib::OrnsteinUhlenbeckProcess(b, eta, 0.0)),

      termStructure_(termStructure) {

        registerWith(termStructure_);

    }

    Size G2ForwardProcess::size() const {

        return 2;

    }

    Array G2ForwardProcess::initialValues() const {

        Real z1_0 = termStructure_.empty() ? x0_ : phi(0.0);

        return { z1_0, y0_ };

    }

    Array G2ForwardProcess::drift(Time t, const Array& z) const {

        Real shiftDrift = 0.0;

        if (!termStructure_.empty()) {

            const Real h = 1.0e-4;

            Real phi_t  = phi(t);

            Real phi_th = phi(t + h);

            Real phiPrime = (phi_th - phi_t) / h;

            shiftDrift = a_ * phi_t + phiPrime;

        }

        return {

            xProcess_->drift(t, z[0]) + xForwardDrift(t, T_) + shiftDrift,

            yProcess_->drift(t, z[1]) + yForwardDrift(t, T_)

        };

    }

    Matrix G2ForwardProcess::diffusion(Time, const Array&) const {

        Matrix tmp(2,2);

        Real sigma1 = sigma_;

        Real sigma2 = eta_;

        tmp[0][0] = sigma1;       tmp[0][1] = 0.0;

        tmp[1][0] = rho_*sigma1;  tmp[1][1] = std::sqrt(1.0-rho_*rho_)*sigma2;

        return tmp;

    }

    Array G2ForwardProcess::expectation(Time t0, const Array& z0,

                                        Time dt) const {

        Real shiftExp = 0.0;

        if (!termStructure_.empty()) {

            shiftExp = phi(t0 + dt) - phi(t0) * std::exp(-a_ * dt);

        }

        return {

            xProcess_->expectation(t0, z0[0], dt) - Mx_T(t0, t0+dt, T_) + shiftExp,

            yProcess_->expectation(t0, z0[1], dt) - My_T(t0, t0+dt, T_)

        };

    }

    Matrix G2ForwardProcess::stdDeviation(Time t0, const Array& x0, Time dt) const {

        Matrix tmp(2,2);

        Real sigma1 = xProcess_->stdDeviation(t0, x0[0], dt);

        Real sigma2 = yProcess_->stdDeviation(t0, x0[1], dt);

        Real expa = std::exp(-a_*dt), expb = std::exp(-b_*dt);

        Real H = (rho_*sigma_*eta_)/(a_+b_)*(1-expa*expb);

        Real den =

            (0.5*sigma_*eta_)*std::sqrt((1-expa*expa)*(1-expb*expb)/(a_*b_));

        Real newRho = H/den;

        tmp[0][0] = sigma1;

        tmp[0][1] = 0.0;

        tmp[1][0] = newRho*sigma2;

        tmp[1][1] = std::sqrt(1.0-newRho*newRho)*sigma2;

        return tmp;

    }

    Matrix G2ForwardProcess::covariance(Time t0, const Array& x0, Time dt) const {

        Matrix sigma = stdDeviation(t0, x0, dt);

        Matrix result = sigma*transpose(sigma);

        return result;

    }

    Real G2ForwardProcess::xForwardDrift(Time t, Time T) const {

        Real expatT = std::exp(-a_*(T-t));

        Real expbtT = std::exp(-b_*(T-t));

        return -(sigma_*sigma_/a_) * (1-expatT)

              - (rho_*sigma_*eta_/b_) * (1-expbtT);

    }

    Real G2ForwardProcess::yForwardDrift(Time t, Time T) const {

        Real expatT = std::exp(-a_*(T-t));

        Real expbtT = std::exp(-b_*(T-t));

        return -(eta_*eta_/b_) * (1-expbtT)

              - (rho_*sigma_*eta_/a_) * (1-expatT);

    }

    Real G2ForwardProcess::Mx_T(Real s, Real t, Real T) const {

        Real M;

        M = ( (sigma_*sigma_)/(a_*a_) + (rho_*sigma_*eta_)/(a_*b_) )

          * (1-std::exp(-a_*(t-s)));

        M += -(sigma_*sigma_)/(2*a_*a_) *

              (std::exp(-a_*(T-t))-std::exp(-a_*(T+t-2*s)));

        M += -(rho_*sigma_*eta_)/(b_*(a_+b_))

            * (std::exp(-b_*(T-t)) -std::exp(-b_*T-a_*t+(a_+b_)*s));

        return M;

    }

    Real G2ForwardProcess::My_T(Real s, Real t, Real T) const {

        Real M;

        M = ( (eta_*eta_)/(b_*b_) + (rho_*sigma_*eta_)/(a_*b_) )

          * (1-std::exp(-b_*(t-s)));

        M += -(eta_*eta_)/(2*b_*b_) *

              (std::exp(-b_*(T-t))-std::exp(-b_*(T+t-2*s)));

        M += -(rho_*sigma_*eta_)/(a_*(a_+b_))

            * (std::exp(-a_*(T-t))-std::exp(-a_*T-b_*t+(a_+b_)*s));

        return M;

    }

    const Handle<YieldTermStructure>& G2ForwardProcess::termStructure() const {

        return termStructure_;

    }

    Real G2ForwardProcess::phi(Time t) const {

        QL_REQUIRE(!termStructure_.empty(),

                   "no term structure given to G2ForwardProcess");

        Rate forward = termStructure_->forwardRate(t, t, Continuous, NoFrequency);

        Real temp1 = sigma_*(1.0-std::exp(-a_*t))/a_;

        Real temp2 = eta_ *(1.0-std::exp(-b_*t))/b_;

        return 0.5*temp1*temp1 + 0.5*temp2*temp2 + rho_*temp1*temp2 + forward;

    }

    Rate G2ForwardProcess::shortRate(Time, Real z1, Real z2) const {

        return z1 + z2;

    }

}