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

/*
 Copyright (C) 2008 Yee Man Chan

 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/math/distributions/chisquaredistribution.hpp>
#include <ql/math/distributions/normaldistribution.hpp>
#include <ql/processes/eulerdiscretization.hpp>
#include <ql/processes/gjrgarchprocess.hpp>
#include <ql/quotes/simplequote.hpp>
#include <utility>

namespace QuantLib {

    GJRGARCHProcess::GJRGARCHProcess(Handle<YieldTermStructure> riskFreeRate,
                                     Handle<YieldTermStructure> dividendYield,
                                     Handle<Quote> s0,
                                     Real v0,
                                     Real omega,
                                     Real alpha,
                                     Real beta,
                                     Real gamma,
                                     Real lambda,
                                     Real daysPerYear,
                                     Discretization d)
    : StochasticProcess(ext::shared_ptr<discretization>(new EulerDiscretization)),
      riskFreeRate_(std::move(riskFreeRate)), dividendYield_(std::move(dividendYield)),
      s0_(std::move(s0)), v0_(v0), omega_(omega), alpha_(alpha), beta_(beta), gamma_(gamma),
      lambda_(lambda), daysPerYear_(daysPerYear), discretization_(d) {
        registerWith(riskFreeRate_);
        registerWith(dividendYield_);
        registerWith(s0_);
    }

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

    Array GJRGARCHProcess::initialValues() const {
        return { s0_->value(), daysPerYear_*v0_ };
    }

    Array GJRGARCHProcess::drift(Time t, const Array& x) const {
        const Real N = CumulativeNormalDistribution()(lambda_);
        const Real n = std::exp(-lambda_*lambda_/2.0)/std::sqrt(2*M_PI);
        const Real q2 = 1.0 + lambda_*lambda_;
        const Real q3 = lambda_*n + N + lambda_*lambda_*N;
        const Real vol = (x[1] > 0.0) ? std::sqrt(x[1])
                         : (discretization_ == Reflection) ? Real(-std::sqrt(-x[1]))
                         : 0.0;

        return {
            riskFreeRate_->forwardRate(t, t, Continuous).rate()
               - dividendYield_->forwardRate(t, t, Continuous).rate()
               - 0.5 * vol * vol,
            daysPerYear_*daysPerYear_*omega_ + daysPerYear_*(beta_ 
                                             + alpha_*q2 + gamma_*q3 - 1.0) *
               ((discretization_==PartialTruncation) ? x[1] : vol*vol)
        };
    }

    Matrix GJRGARCHProcess::diffusion(Time, const Array& x) const {
        /* the correlation matrix is
           |  1   rho |
           | rho   1  |
           whose square root (which is used here) is
           |  1          0       |
           | rho   std::sqrt(1-rho^2) |
        */
        Matrix tmp(2,2);
        const Real N = CumulativeNormalDistribution()(lambda_);
        const Real n = std::exp(-lambda_*lambda_/2.0)/std::sqrt(2*M_PI);
        const Real sigma2 = 2.0 + 4.0*lambda_*lambda_;
        const Real q3 = lambda_*n + N + lambda_*lambda_*N;
        const Real Eml_e4 = lambda_*lambda_*lambda_*n + 5.0*lambda_*n 
            + 3.0*N + lambda_*lambda_*lambda_*lambda_*N 
            + 6.0*lambda_*lambda_*N;
        const Real sigma3 = Eml_e4 - q3*q3;
        const Real sigma12 = -2.0*lambda_;
        const Real sigma13 = -2.0*n - 2*lambda_*N;
        const Real sigma23 = 2.0*N + sigma12*sigma13;
        const Real vol = (x[1] > 0.0) ? std::sqrt(x[1])
                         : (discretization_ == Reflection) ? Real(- std::sqrt(-x[1]))
                         : 1e-8; // set vol to (almost) zero but still
                                 // expose some correlation information
        const Real rho1 = std::sqrt(daysPerYear_)*(alpha_*sigma12 
                                            + gamma_*sigma13) * vol * vol;
        const Real rho2 = vol*vol*std::sqrt(daysPerYear_)
            *std::sqrt(alpha_*alpha_*(sigma2 - sigma12*sigma12) 
                       + gamma_*gamma_*(sigma3 - sigma13*sigma13) 
                       + 2.0*alpha_*gamma_*(sigma23 - sigma12*sigma13)); 

            // tmp[0][0], tmp[0][1] are the coefficients of dW_1 and dW_2 
            // in asset return stochastic process
        tmp[0][0] = vol;  tmp[0][1] = 0.0;
        tmp[1][0] = rho1; tmp[1][1] = rho2;
        return tmp;
    }

    Array GJRGARCHProcess::apply(const Array& x0,
                                 const Array& dx) const {
        return { x0[0] * std::exp(dx[0]), x0[1] + dx[1] };
    }

    Array GJRGARCHProcess::evolve(Time t0, const Array& x0,
                                  Time dt, const Array& dw) const {
        Array retVal(2);
        Real vol, mu, nu;

        const Real sdt = std::sqrt(dt);
        const Real N = CumulativeNormalDistribution()(lambda_);
        const Real n = std::exp(-lambda_*lambda_/2.0)/std::sqrt(2*M_PI);
        const Real sigma2 = 2.0 + 4.0*lambda_*lambda_;
        const Real q2 = 1.0 + lambda_*lambda_;
        const Real q3 = lambda_*n + N + lambda_*lambda_*N;
        const Real Eml_e4 = lambda_*lambda_*lambda_*n + 5.0*lambda_*n 
            + 3.0*N + lambda_*lambda_*lambda_*lambda_*N 
            + 6.0*lambda_*lambda_*N;
        const Real sigma3 = Eml_e4 - q3*q3;
        const Real sigma12 = -2.0*lambda_;
        const Real sigma13 = -2.0*n - 2*lambda_*N;
        const Real sigma23 = 2.0*N + sigma12*sigma13;
        const Real rho1 = std::sqrt(daysPerYear_)*(alpha_*sigma12 + gamma_*sigma13);
        const Real rho2 = std::sqrt(daysPerYear_)
            *std::sqrt(alpha_*alpha_*(sigma2 - sigma12*sigma12) 
                       + gamma_*gamma_*(sigma3 - sigma13*sigma13) 
                       + 2.0*alpha_*gamma_*(sigma23 - sigma12*sigma13));

        switch (discretization_) {
          // For the definition of PartialTruncation, FullTruncation
          // and Reflection  see Lord, R., R. Koekkoek and D. van Dijk (2006),
          // "A Comparison of biased simulation schemes for
          //  stochastic volatility models",
          // Working Paper, Tinbergen Institute
          case PartialTruncation:
            vol = (x0[1] > 0.0) ? Real(std::sqrt(x0[1])) : 0.0;
            mu =    riskFreeRate_->forwardRate(t0, t0+dt, Continuous).rate()
                  - dividendYield_->forwardRate(t0, t0+dt, Continuous).rate()
                    - 0.5 * vol * vol;
            nu = daysPerYear_*daysPerYear_*omega_ 
                + daysPerYear_*(beta_ + alpha_*q2 + gamma_*q3 - 1.0) * x0[1];

            retVal[0] = x0[0] * std::exp(mu*dt+vol*dw[0]*sdt);
            retVal[1] = x0[1] + nu*dt + sdt*vol*vol*(rho1*dw[0] + rho2*dw[1]);
            break;
          case FullTruncation:
            vol = (x0[1] > 0.0) ? Real(std::sqrt(x0[1])) : 0.0;
            mu =    riskFreeRate_->forwardRate(t0, t0+dt, Continuous).rate()
                  - dividendYield_->forwardRate(t0, t0+dt, Continuous).rate()
                    - 0.5 * vol * vol;
            nu = daysPerYear_*daysPerYear_*omega_ 
                + daysPerYear_*(beta_ + alpha_*q2 + gamma_*q3 - 1.0) * vol *vol;

            retVal[0] = x0[0] * std::exp(mu*dt+vol*dw[0]*sdt);
            retVal[1] = x0[1] + nu*dt + sdt*vol*vol*(rho1*dw[0] + rho2*dw[1]);
            break;
          case Reflection:
            vol = std::sqrt(std::fabs(x0[1]));
            mu =    riskFreeRate_->forwardRate(t0, t0+dt, Continuous).rate()
                  - dividendYield_->forwardRate(t0, t0+dt, Continuous).rate()
                    - 0.5 * vol*vol;
            nu = daysPerYear_*daysPerYear_*omega_ 
                + daysPerYear_*(beta_ + alpha_*q2 + gamma_*q3 - 1.0) * vol * vol;

            retVal[0] = x0[0]*std::exp(mu*dt+vol*dw[0]*sdt);
            retVal[1] = vol*vol
                        +nu*dt + sdt*vol*vol*(rho1*dw[0] + rho2*dw[1]);
            break;
          default:
            QL_FAIL("unknown discretization schema");
        }

        return retVal;
    }

    const Handle<Quote>& GJRGARCHProcess::s0() const {
        return s0_;
    }

    const Handle<YieldTermStructure>& GJRGARCHProcess::dividendYield() const {
        return dividendYield_;
    }

    const Handle<YieldTermStructure>& GJRGARCHProcess::riskFreeRate() const {
        return riskFreeRate_;
    }

    Time GJRGARCHProcess::time(const Date& d) const {
        return riskFreeRate_->dayCounter().yearFraction(
                                           riskFreeRate_->referenceDate(), d);
    }

}
