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

/*
 Copyright (C) 2001, 2002, 2003 Sadruddin Rejeb
 Copyright (C) 2005 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 trinomialtree.hpp
    \brief Trinomial tree class
*/

#ifndef quantlib_trinomial_tree_hpp
#define quantlib_trinomial_tree_hpp

#include <ql/stochasticprocess.hpp>
#include <ql/methods/lattices/tree.hpp>
#include <ql/timegrid.hpp>

namespace QuantLib {

    //! Recombining trinomial tree class
    /*! This class defines a recombining trinomial tree approximating a
        1-D stochastic process.
        \warning The diffusion term of the SDE must be independent of the
                 underlying process.

        \ingroup lattices
    */
    class TrinomialTree : public Tree<TrinomialTree> {
        class Branching;
      public:
        enum Branches { branches = 3 };
        TrinomialTree(const boost::shared_ptr<StochasticProcess1D>& process,
                      const TimeGrid& timeGrid,
                      bool isPositive = false);
        Real dx(Size i) const { return dx_[i]; }
        const TimeGrid& timeGrid() const { return timeGrid_; }

        Size size(Size i) const;
        Real underlying(Size i, Size index) const;
        Size descendant(Size i, Size index, Size branch) const;
        Real probability(Size i, Size index, Size branch) const;

      protected:
        std::vector<Branching> branchings_;
        Real x0_;
        std::vector<Real> dx_;
        TimeGrid timeGrid_;

      private:
        /* Branching scheme for a trinomial node.  Each node has three
           descendants, with the middle branch linked to the node
           which is closest to the expectation of the variable. */
        class Branching {
          public:
            Branching();
            Size descendant(Size index, Size branch) const;
            Real probability(Size index, Size branch) const;
            Size size() const;
            Integer jMin() const;
            Integer jMax() const;
            void add(Integer k, Real p1, Real p2, Real p3);
          private:
            std::vector<Integer> k_;
            std::vector<std::vector<Real> > probs_;
            Integer kMin_, jMin_, kMax_, jMax_;
        };
    };

    // inline definitions

    inline Size TrinomialTree::size(Size i) const {
        return i==0 ? 1 : branchings_[i-1].size();
    }

    inline Real TrinomialTree::underlying(Size i, Size index) const {
        if (i==0)
            return x0_;
        else
            return x0_ + (branchings_[i-1].jMin() +
                          static_cast<Real>(index))*dx(i);
    }

    inline Size TrinomialTree::descendant(Size i, Size index,
                                          Size branch) const {
        return branchings_[i].descendant(index, branch);
    }

    inline Real TrinomialTree::probability(Size i, Size j, Size b) const {
        return branchings_[i].probability(j, b);
    }

    inline TrinomialTree::Branching::Branching()
    : probs_(3), kMin_(QL_MAX_INTEGER), jMin_(QL_MAX_INTEGER),
                 kMax_(QL_MIN_INTEGER), jMax_(QL_MIN_INTEGER) {}

    inline Size TrinomialTree::Branching::descendant(Size index,
                                                     Size branch) const {
        return k_[index] - jMin_ - 1 + branch;
    }

    inline Real TrinomialTree::Branching::probability(Size index,
                                                      Size branch) const {
        return probs_[branch][index];
    }

    inline Size TrinomialTree::Branching::size() const {
        return jMax_ - jMin_ + 1;
    }

    inline Integer TrinomialTree::Branching::jMin() const {
        return jMin_;
    }

    inline Integer TrinomialTree::Branching::jMax() const {
        return jMax_;
    }

    inline void TrinomialTree::Branching::add(Integer k,
                                              Real p1, Real p2, Real p3) {
        // store
        k_.push_back(k);
        probs_[0].push_back(p1);
        probs_[1].push_back(p2);
        probs_[2].push_back(p3);
        // maintain invariants
        kMin_ = std::min(kMin_, k);
        jMin_ = kMin_ - 1;
        kMax_ = std::max(kMax_, k);
        jMax_ = kMax_ + 1;
    }

}


#endif