/*
  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_MLVS_DATA_IMPUTER_HPP
#define BOOM_MLVS_DATA_IMPUTER_HPP

#include <Models/Glm/ChoiceData.hpp>
#include <Models/Glm/MultinomialLogitModel.hpp>
#include <Models/Glm/PosteriorSamplers/MultinomialLogitCompleteDataSuf.hpp>
#include <Models/PosteriorSamplers/Imputer.hpp>

namespace BOOM{

  class MlvsDataImputer
      : public LatentDataImputer<ChoiceData,
                                 MultinomialLogitCompleteDataSufficientStatistics> {
  public:
    typedef MultinomialLogitCompleteDataSufficientStatistics LocalSuf;

    // Args:
    //   Mod:  A pointer to the model being data-agumented.
    MlvsDataImputer(MultinomialLogitModel *Mod);

    // Impute latent data for a single observation, and add the
    // results to the complete data sufficient statistics.
    virtual void impute_latent_data(const ChoiceData &observed_data,
                                    LocalSuf *suf,
                                    RNG &rng) const;

    // Used to decompose latent utilities into a mixture of Gaussians.
    // Args:
    //   rng: Random number genrerator.  Must not be shared by any
    //     other threads.
    //   u:  Latent utility.
    // Returns:
    //   The index of the imputed Gaussian mixture component
    //   responsible for u.
    uint unmix(RNG &rng, double u) const;

  private:
    MultinomialLogitModel * model_;

    const Vec mu_;                   // mean for EV approx
    const Vec sigsq_inv_;            // inverse variance for EV approx
    const Vec sd_;                   // standard deviations for EV approx
    const Vec log_mixing_weights_;   // log of mixing weights for EV approx
    const Vec & log_sampling_probs_;
    const bool downsampling_;

    // Workspace for impute_latent_data, to avoid reallocating space
    // each time.
    mutable Vec post_prob_;
    mutable Vec u;
    mutable Vec eta;
    mutable Vec wgts;
  };

}  // namespace BOOM

#endif// BOOM_MLVS_DATA_IMPUTER_HPP
