/*
 Copyright (C) 2003 Niels Elken Sønderby

 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 ferdinando@ametrano.net
 The license is also available online at http://quantlib.org/html/license.html

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

/*
 NesQuant is an extension to QuantLib
 http://www.nielses.dk/quantlib/nesquant
*/

#include <nq/PricingEngines/lsmvanillaengine.hpp>
#include <nq/Math/polynomialfit.hpp>
#include <iostream>

using namespace std;

namespace NesQuant {

    void LSMVanillaEngine::calculate() const {

        double s0 = arguments_.underlying;
        double maturity = arguments_.maturity;
        Handle<Payoff> &payoff = arguments_.payoff;

        Handle<DiffusionProcess> bs(new BlackScholesProcess(
                                arguments_.riskFreeTS,
                                arguments_.dividendTS,
                                arguments_.volTS,
                                s0));

        GaussianRandomSequenceGenerator gen =
            PseudoRandom::make_sequence_generator(timeSteps_, seed_);

        Handle<GaussianPathGenerator> pathGenerator =
            Handle<GaussianPathGenerator> (new GaussianPathGenerator(
                  bs, maturity, timeSteps_, gen));

        // if antithetic paths are used, twice as many paths are used
        Size modSamples; // modified samples
        if (antitheticVariate_)
            modSamples = 2 * samples_;
        else
            modSamples = samples_;

        // Initialize control variate
        double controlVariateValue;
        std::vector<double> controlVariatePathValue(modSamples);
        arguments_.exerciseType = Exercise::European;
        Handle<PlainVanillaPayoff> payoff2 = arguments_.payoff;
        EuropeanPathPricer controlPP(payoff2->optionType(), s0,
                                     payoff2->strike(), arguments_.riskFreeTS);
        if (controlVariate_) {
            AnalyticEuropeanEngine controlPE;

            VanillaOptionArguments* controlArguments =
                dynamic_cast<VanillaOptionArguments*>(
                    controlPE.arguments());
            *controlArguments = arguments_;
            controlPE.calculate();

            const VanillaOptionResults* controlResults =
                dynamic_cast<const VanillaOptionResults*>(controlPE.results());
            controlVariateValue = controlResults->value;
        }

        // generate paths
        std::vector<std::vector<double> >
            spaths(modSamples, std::vector<double>(timeSteps_));

        GaussianPathGenerator::sample_type pathSample = pathGenerator->next();

        for (Size i = 0; i < modSamples; ++i) {
            if (antitheticVariate_ && i%2 == 1)
                pathSample = pathGenerator->antithetic();
            else
                pathSample = pathGenerator->next();

            double s = s0;
            for (Size j = 0; j < timeSteps_; ++j) {
                s *= QL_EXP(pathSample.value[j]);
                spaths[i][j] = s;
            }

            if (controlVariate_) {
                controlVariatePathValue[i] = controlPP(pathSample.value);
            }
        }

        // set values at maturity
        std::vector<double> U(modSamples);
        std::vector<Size> stoppingTime(modSamples, -1); // never exercised = -1

        Size k = timeSteps_;
        DiscountFactor discount = arguments_.riskFreeTS->discount(maturity);

        for (i = 0; i < modSamples; ++i) {
            U[i] = discount * (*payoff)(spaths[i][k-1]);
        }

        // move backwards
        double dt = maturity / timeSteps_;

        for (k = timeSteps_ - 1; k > 0; --k) {
            cout << "k : " << k << endl;

            // find in-the-money paths and (x, y) regression data
            std::vector<Size> itmpaths;
            std::vector<double> xdata;
            std::vector<double> ydata;

            for (i = 0; i < modSamples; ++i) {
                double s = spaths[i][k-1];
                if ((*payoff)(s) > 0) {
                    itmpaths.push_back(i);
                    xdata.push_back(s);
                    ydata.push_back(U[i]);
                }
            }

            if (itmpaths.size() > 0) {
                // do regression
                PolynomialFit pf(xdata, ydata, degree_);
                const Array contval = pf.value();

                // compare continuation value with exercise value
                discount = arguments_.riskFreeTS->discount(k * dt);
                for (Size l = 0; l < itmpaths.size(); ++l) {
                    i = itmpaths[l];
                    double exerval = discount * (*payoff)(xdata[l]);

                    if (exerval > contval[l]) {
                        U[i] = exerval;
                        stoppingTime[i] = k;
                    }
                }
            }
        }

        // adjust for control variate
        if (controlVariate_)
            for (i = 0; i < modSamples; i++) {
                U[i] += controlVariateValue - controlVariatePathValue[i];
            }

        // calculate average
        if (antitheticVariate_)
            for (i = 0; i < modSamples; i+=2) {
                sampleAccumulator_.add((U[i] + U[i+1]) / 2);
            }
        else
            for (i = 0; i < samples_; i++) {
                sampleAccumulator_.add(U[i]);
            }
        double value = sampleAccumulator_.mean();

        // compare with exercise at time 0
        if ((*payoff)(s0) > value) {
            value = (*payoff)(s0);
            sampleAccumulator_.reset();
            for (i = 0; i < modSamples; i++) {
                if ((!antitheticVariate_) || i%2==0)
                    sampleAccumulator_.add(value);
                stoppingTime[i] = 0;
            }
        }

        // return value
        results_.value = value;
        results_.errorEstimate = sampleAccumulator_.errorEstimate();
    }

}