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

/*

 Copyright (C) 2008 Andreas Gaida

 Copyright (C) 2008 Ralph Schreyer

 Copyright (C) 2008, 2019 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.

*/

#include <ql/math/functional.hpp>

#include <ql/math/distributions/normaldistribution.hpp>

#include <ql/math/distributions/chisquaredistribution.hpp>

#include <ql/math/interpolations/linearinterpolation.hpp>

#include <ql/math/integrals/gausslobattointegral.hpp>

#include <ql/termstructures/volatility/equityfx/localvoltermstructure.hpp>

#include <ql/methods/finitedifferences/meshers/fdmhestonvariancemesher.hpp>

#include <boost/accumulators/accumulators.hpp>

#include <boost/accumulators/statistics/mean.hpp>

#include <boost/accumulators/statistics/stats.hpp>

#include <set>

#include <algorithm>

namespace QuantLib {

    namespace {

        struct interpolated_volatility {

            interpolated_volatility(const std::vector<Real>& pGrid,

                                    const std::vector<Real>& vGrid)

            : variance(pGrid.begin(), pGrid.end(), vGrid.begin()) {}

            Real operator()(Real x) const {

                return std::sqrt(variance(x, true));

            }

            LinearInterpolation variance;

        };

    }

    FdmHestonVarianceMesher::FdmHestonVarianceMesher(

        Size size,

        const ext::shared_ptr<HestonProcess> & process,

        Time maturity, Size tAvgSteps, Real epsilon,

        Real mixingFactor)

        : Fdm1dMesher(size) {

        std::vector<Real> vGrid(size, 0.0), pGrid(size, 0.0);

        const Real mixedSigma = process->sigma()*mixingFactor;

        const Real df  = 4*process->theta()*process->kappa()/squared(mixedSigma);

        try {

            std::multiset<std::pair<Real, Real> > grid;

            for (Size l=1; l<=tAvgSteps; ++l) {

                const Real t = (maturity*l)/tAvgSteps;

                const Real ncp = 4*process->kappa()*std::exp(-process->kappa()*t)/(squared(mixedSigma)

                    *(1-std::exp(-process->kappa()*t)))*process->v0();

                const Real k = squared(mixedSigma)

                    *(1-std::exp(-process->kappa()*t))/(4*process->kappa());

                const Real qMin = 0.0; // v_min = 0.0;

                const Real qMax = std::max(process->v0(),

                    k*InverseNonCentralCumulativeChiSquareDistribution(

                                            df, ncp, 100, 1e-8)(1-epsilon));

                const Real minVStep=(qMax-qMin)/(50*size);

                Real ps,p = 0.0;

                Real vTmp = qMin;

                grid.insert(std::pair<Real, Real>(qMin, epsilon));

                for (Size i=1; i < size; ++i) {

                    ps = (1 - epsilon - p)/(size-i);

                    p += ps;

                    const Real tmp = k*InverseNonCentralCumulativeChiSquareDistribution(

                        df, ncp, 100, 1e-8)(p);

                    const Real vx = std::max(vTmp+minVStep, tmp);

                    p = NonCentralCumulativeChiSquareDistribution(df, ncp)(vx/k);

                    vTmp=vx;

                    grid.insert(std::pair<Real, Real>(vx, p));

                }

            }

            QL_REQUIRE(grid.size() == size*tAvgSteps, 

                       "something wrong with the grid size");

            const std::vector<std::pair<Real, Real> > tp(grid.begin(), grid.end());

            for (Size i=0; i < size; ++i) {

                const Size b = (i*tp.size())/size;

                const Size e = ((i+1)*tp.size())/size;

                for (Size j=b; j < e; ++j) {

                    vGrid[i]+=tp[j].first/(e-b);

                    pGrid[i]+=tp[j].second/(e-b);

                }

            }

        } 

        catch (const Error&) {

            // use default mesh

            const Real vol = mixedSigma*

                std::sqrt(process->theta()/(2*process->kappa()));

            const Real mean = process->theta();

            const Real upperBound = std::max(process->v0()+4*vol, mean+4*vol);

            const Real lowerBound

                = std::max(0.0, std::min(process->v0()-4*vol, mean-4*vol));

            for (Size i=0; i < size; ++i) {

                pGrid[i] = i/(size-1.0);

                vGrid[i] = lowerBound + i*(upperBound-lowerBound)/(size-1.0);

            }

        }

        Real skewHint = ((process->kappa() != 0.0) 

                ? Real(std::max(1.0, mixedSigma/process->kappa())) : 1.0);

        std::sort(pGrid.begin(), pGrid.end());

        volaEstimate_ = GaussLobattoIntegral(100000, 1e-4)(

            interpolated_volatility(pGrid, vGrid),

                pGrid.front(), pGrid.back())*std::pow(skewHint, 1.5);

        const Real v0 = process->v0();

        for (Size i=1; i<vGrid.size(); ++i) {

            if (vGrid[i-1] <= v0 && vGrid[i] >= v0) {

                if (std::fabs(vGrid[i-1] - v0) < std::fabs(vGrid[i] - v0))

                    vGrid[i-1] = v0;

                else

                    vGrid[i] = v0;

            }

        }

        std::copy(vGrid.begin(), vGrid.end(), locations_.begin());

        for (Size i=0; i < size-1; ++i) {

            dminus_[i+1] = dplus_[i] = vGrid[i+1] - vGrid[i];

        }

        dplus_.back() = dminus_.front() = Null<Real>();

    }

    FdmHestonLocalVolatilityVarianceMesher::FdmHestonLocalVolatilityVarianceMesher(

        Size size,

        const ext::shared_ptr<HestonProcess>& process,

        const ext::shared_ptr<LocalVolTermStructure>& leverageFct,

        Time maturity, Size tAvgSteps, Real epsilon,

        Real mixingFactor)

     : Fdm1dMesher(size) {

        const FdmHestonVarianceMesher mesher(

            size, process, maturity, tAvgSteps, epsilon, mixingFactor);

        for (Size i=0; i < size; ++i) {

            dplus_[i] = mesher.dplus(i);

            dminus_[i] = mesher.dminus(i);

            locations_[i] = mesher.location(i);

        }

        volaEstimate_ = mesher.volaEstimate();

        if (leverageFct != nullptr) {

            typedef boost::accumulators::accumulator_set<

                Real, boost::accumulators::stats<

                    boost::accumulators::tag::mean> >

                accumulator_set;

            accumulator_set acc;

            const Real s0 = process->s0()->value();

            acc(leverageFct->localVol(0, s0, true));

            const Handle<YieldTermStructure> rTS = process->riskFreeRate();

            const Handle<YieldTermStructure> qTS = process->dividendYield();

            for (Size l=1; l <= tAvgSteps; ++l) {

                const Real t = (maturity*l)/tAvgSteps;

                const Real vol = volaEstimate_ * boost::accumulators::mean(acc);

                const Real fwd = s0*qTS->discount(t)/rTS->discount(t);

                const Size sAvgSteps = 50;

                std::vector<Real> u(sAvgSteps), sig(sAvgSteps);

                for (Size i=0; i < sAvgSteps; ++i) {

                    u[i] = epsilon + ((1-2*epsilon)/(sAvgSteps-1))*i;

                    const Real x = InverseCumulativeNormal()(u[i]);

                    const Real gf = x*vol*std::sqrt(t);

                    const Real f = fwd*std::exp(gf);

                    sig[i] = squared(leverageFct->localVol(t, f, true));

                }

                const Real leverageAvg =

                    GaussLobattoIntegral(10000, 1e-4)(

                        interpolated_volatility(u, sig), u.front(), u.back())

                    / (1-2*epsilon);

                acc(leverageAvg);

            }

            volaEstimate_ *= boost::accumulators::mean(acc);

        }

    }

}