/*
  Copyright (C) 2007 Steven L. Scott

  This library is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.

  This library 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 GNU
  Lesser General Public License for more details.

  You should have received a copy of the GNU Lesser General Public
  License along with this library; if not, write to the Free Software
  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA
*/

#ifndef BOOM_MULTINOMIAL_LOGIT_VARIABLE_SELECTION_HPP
#define BOOM_MULTINOMIAL_LOGIT_VARIABLE_SELECTION_HPP


#include <Models/PosteriorSamplers/PosteriorSampler.hpp>
#include <Models/Glm/VariableSelectionPrior.hpp>
#include <Models/Glm/PosteriorSamplers/MLVS_data_imputer.hpp>
#include <Models/PosteriorSamplers/Imputer.hpp>

namespace BOOM{

  class MultinomialLogitModel;
  class MvnBase;
  class ChoiceData;

  //------------------------------------------------------------
  class MLVS : public PosteriorSampler {
    // Draws the parameters of a multinomial logit model using the
    // approximate method from Fruewirth-Schnatter and Fruewirth,
    // Computational Statistics and Data Analysis 2007, 3508-3528.

    // This implementation only stores the complete data sufficient
    // statistics and some workspace.  It does not store the imputed
    // latent data.
  public:

    // Args:
    //   Mod:  The multinomial logit model to be sampled.
    //   Pri: The conditional prior distribution for the coefficients
    //     of *Mod, given inclusion.  The dimension of this prior must
    //     match the dimension of Mod->beta(), which is: (number of
    //     choices - 1) * (number of subject level predictors) +
    //     (number of choice level predictors).
    //   vPri: The prior distribution on the include/exclude pattern
    //     for Mod->beta().  See above for the required dimension.
    //   nthreads: The number of threads to use when imputing the
    //     latent data.
    //   check_initial_condition: If true (which is the default) an
    //     error will be reported if the initial value of Mod->beta()
    //     is inconsistent with the prior distribution.  The usual
    //     cause of something like this is if certain coefficients
    //     have been forced in by setting vPri = 1 for that
    //     coefficient, but the coefficient is excluded by Mod.  If
    //     false then no error checking will take place on the initial
    //     condition.
    MLVS(MultinomialLogitModel *Mod,
         Ptr<MvnBase> Pri,
         Ptr<VariableSelectionPrior> vPri,
         uint nthreads=1,
         bool check_initial_condition = true);

    virtual void draw();
    virtual double logpri()const;

    // The individual steps needed to implement the draw.
    void impute_latent_data();
    void draw_beta();

    // Sets the number of workers to use for multi-threaded data
    // augmentation.  Implementations are permitted to use less than
    // this number of threads.
    void set_number_of_workers(int n);

    // Functions to control Bayesian variable selection.  If
    // supress_model_selection is called then the current
    // include/exclude state of the coefficients will not be modified
    // by the sampler.  A call to allow_model_selection() turns model
    // selection back on.  It is on by default.
    void supress_model_selection();
    void allow_model_selection();

    // If the predictor space is very large, then you can save time by
    // not sampling the include/exclude decision for all possible
    // variables.  Calling limit_model_selection(nflips) will cause
    // the algorithm to only try to modify 'nflips' randomly chosen
    // locations each iteration.  By default the algorithm tries all
    // variables each time.
    void limit_model_selection(uint nflips);

    // Returns the number of variables that the model selection
    // algorithm will attempt to modify at each iteration.
    uint max_nflips()const;

  private:
    MultinomialLogitModel *mod_;
    Ptr<MvnBase> pri;
    Ptr<VariableSelectionPrior> vpri;
    typedef MultinomialLogitCompleteDataSufficientStatistics LocalSuf;
    LocalSuf suf_;
    ParallelLatentDataImputer<ChoiceData,
                              LocalSuf,
                              MultinomialLogitModel> parallel_imputer_;

    const Vec & log_sampling_probs_;
    const bool downsampling_;
    bool select_;
    uint max_nflips_;

    Spd Ominv;
    Spd iV_tilde_;
    virtual void draw_inclusion_vector();
    double log_model_prob(const Selector &inc);
  };

}
#endif//  BOOM_MULTINOMIAL_LOGIT_VARIABLE_SELECTION_HPP
