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

/*
 Copyright (C) 2005, 2006 Theo Boafo
 Copyright (C) 2006 StatPro Italia srl

 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
 <http://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 tflattice.hpp
    \brief Trinomial Tsiveriotis-Fernandes tree model
*/

#ifndef tf_lattice_hpp
#define tf_lattice_hpp

#include "bsmlattice.hpp"
#include <ql/pricingengines/hybrid/discretizedconvertible.hpp>

namespace QuantLib {

    //! Trinomial lattice approximating the Tsiveriotis-Fernandes model
    /*! \ingroup lattices */
	template <class T>
    class TsiveriotisFernandesTrinomialLattice : public BlackScholesTrinomialLattice<T> {
      public:
        TsiveriotisFernandesTrinomialLattice(const boost::shared_ptr<T>& tree,
                                    Rate riskFreeRate,
                                    Time end,
                                    Size steps,
                                    Real creditSpread,
                                    Volatility volatility,
                                    Spread divYield);

        Rate riskFreeRate() const { return riskFreeRate_; };
        Real creditSpread() const { return creditSpread_; };
        Real dt() const {return dt_;};

      protected:
        void stepback(Size i, const Array& values,
                      const Array& conversionProbability,
                      const Array& spreadAdjustedRate,
                      Array& newValues, Array& newConversionProbability,
                      Array& newSpreadAdjustedRate) const;
        void rollback(DiscretizedAsset&, Time to) const;
        void partialRollback(DiscretizedAsset&, Time to) const;

      private:
		//new private member pm_
        Real pd_, pu_,pm_,creditSpread_,dt_;
        Rate riskFreeRate_;
    };


	// template definitions

    template <class T>
    TsiveriotisFernandesTrinomialLattice<T>::TsiveriotisFernandesTrinomialLattice(
                                             const boost::shared_ptr<T>& tree,
                                             Rate riskFreeRate,
                                             Time end, Size steps,
                                             Real creditSpread,
                                             Volatility sigma, Spread divYield)
    : BlackScholesTrinomialLattice<T>(tree,riskFreeRate,end,steps) {

		/*initializing probabilities using tree values. Previous code used fixed values for pu_ and pd_,

		Real a = std::exp((riskFreeRate-divYield)*dt_);

		Real u = std::exp(sigma*(std::sqrt(dt_)));

		Real d = 1/u;

		pu_ = (a-d)/(u-d);
		pd_ = 1 - pu_;

		which is the up and down values for Cox-Ross-Rubinstein. Any insight on why this is fixed (and not
		determined by the tree used) would be appreciated.
		*/
        dt_ = end/steps;

		pu_ = tree->probability(0,0,0);
		pm_ = tree->probability(0,0,1);
		pd_ = tree->probability(0,0,2);

        riskFreeRate_ = riskFreeRate;
        creditSpread_ = creditSpread;

        QL_REQUIRE(pu_<=1.0, "negative probability");
        QL_REQUIRE(pu_>=0.0, "negative probability");
    }

    template <class T>
    void TsiveriotisFernandesTrinomialLattice<T>::stepback(
              Size i, const Array& values, const Array& conversionProbability,
              const Array& spreadAdjustedRate, Array& newValues,
              Array& newConversionProbability,
              Array& newSpreadAdjustedRate) const {

        for (Size j=1; j<this->size(i); j++) {

            // new conversion probability is calculated via backward
            // induction using up and down probabilities on tree on
            // previous conversion probabilities, ie weighted average
            // of previous probabilities.

			//changed newConversionProbability and newValues to take values from three lattice points
            newConversionProbability[j] =
                pd_*conversionProbability[j-1] + pm_*conversionProbability[j] + pu_*conversionProbability[j+1];

            // Use blended discounting rate
            newSpreadAdjustedRate[j] =
                newConversionProbability[j] * riskFreeRate_ +
                (1-newConversionProbability[j])*(riskFreeRate_+creditSpread_);

            newValues[j] =
                (pd_*values[j-1]/(1+(spreadAdjustedRate[j-1]*dt_)))
              + (pm_*values[j]/(1+(spreadAdjustedRate[j]*dt_)))
              + (pu_*values[j+1]/(1+(spreadAdjustedRate[j+1]*dt_)));

        }
    }

    template <class T>
	void TsiveriotisFernandesTrinomialLattice<T>::rollback(
                                     DiscretizedAsset& asset, Time to) const {
        partialRollback(asset,to);
        asset.adjustValues();
    }


	template <class T>
	void TsiveriotisFernandesTrinomialLattice<T>::partialRollback(
                                     DiscretizedAsset& asset, Time to) const {

        Time from = asset.time();

        if (close(from,to))
            return;

        QL_REQUIRE(from > to,
                   "cannot roll the asset back to" << to
                   << " (it is already at t = " << from << ")");

        DiscretizedConvertible& convertible =
            dynamic_cast<DiscretizedConvertible&>(asset);

        Integer iFrom = Integer(this->t_.index(from));
        Integer iTo = Integer(this->t_.index(to));

        for (Integer i=iFrom-1; i>=iTo; --i) {

            Array newValues(this->size(i));
            Array newSpreadAdjustedRate(this->size(i));
            Array newConversionProbability(this->size(i));

            stepback(i, convertible.values(),
                     convertible.conversionProbability(),
                     convertible.spreadAdjustedRate(), newValues,
                     newConversionProbability,newSpreadAdjustedRate);

            convertible.time() = this->t_[i];
            convertible.values() = newValues;
            convertible.spreadAdjustedRate() = newSpreadAdjustedRate;
            convertible.conversionProbability() = newConversionProbability;

            // skip the very last adjustment
            if (i != iTo)
                convertible.adjustValues();
        }
    }

}

#endif
