/* 
 * Scythe Statistical Library Copyright (C) 2000-2002 Andrew D. Martin
 * and Kevin M. Quinn; 2002-present Andrew D. Martin, Kevin M. Quinn,
 * and Daniel Pemstein.  All Rights Reserved.
 *
 * This program is free software; you can redistribute it and/or
 * modify under the terms of the GNU General Public License as
 * published by Free Software Foundation; either version 2 of the
 * License, or (at your option) any later version.  See the text files
 * COPYING and LICENSE, distributed with this source code, for further
 * information.
 * --------------------------------------------------------------------
 *  scythestat/rng/rtmvnorm.h
 *
 */

/*!
 * \file rng/rtmvnorm.h
 *
 * \brief A truncated multivariate normal random number generator.
 * 
 * This file provides the class definition for the rtmvnorm class, a
 * functor that generates random variates from truncated multivariate
 * normal distributions.
 *
 */

#ifndef SCYTHE_RTMVNORM_H
#define SCYTHE_RTMVNORM_H

#include <iostream>
#include <cmath>

#ifdef SCYTHE_COMPILE_DIRECT
#include "matrix.h"
#include "rng.h"
#include "error.h"
#include "algorithm.h"
#include "ide.h"
#else
#include "scythestat/matrix.h"
#include "scythestat/rng.h"
#include "scythestat/error.h"
#include "scythestat/algorithm.h"
#include "scythestat/ide.h"
#endif
namespace scythe
{
  /* Truncated Multivariate Normal Distribution by Gibbs sampling
   * (Geweke 1991).  This is a functor that allows one to
   * initialize---and optionally burn in---a sampler for a given
   * truncated multivariate normal distribution on construction
   * and then make (optionally thinned) draws with calls to the ()
   * operator.
   *
   */
   /*! \brief Truncated multivariate normal distribution random number
    * generator.
    *
    * This class is a functor that allows one to initialize, and
    * optionally burn in, a Gibbs sampler (Geweke 1991) for a given
    * truncated multivariate normal distribution on construction and
    * then make optionally thinned draws from the distribution with
    * calls to the () operator.
    */

  template <class RNGTYPE>
  class rtmvnorm {
    public:

       /*! \brief Standard constructor.
        *
        * This method constructs a functor capable of generating
        * linearly constrained variates of the form: \f$x \sim
        * N_n(\mu, \Sigma), a \le Dx \le b\f$.  That is, it generates
        * an object capable of simulating random variables from an
        * n-variate normal distribution defined by \a mu
        * (\f$\mu\f$) and \a sigma (\f$\Sigma\f$) subject to fewer
        * than \f$n\f$ linear constraints, defined by the Matrix \a D
        * and the bounds vectors \a a and \a b.
        *
        * The user may pass optional burn in and thinning
        * parameters to the constructor.  The \a burnin parameter
        * indicates the number of draws that the sampler should
        * initially make and throw out on construction.  The \a thin
        * parameter controls the behavior of the functor's ()
        * operator.  A thinning parameter of 1 indicates that each
        * call to operator()() should return the random variate
        * generated by one iteration of the Gibbs sampler, while a
        * value of 2 indicates that the sampler should throw every
        * other variate out, a value of 3 causes operator()() to
        * iterate the sampler three times before returning, and so on.
        *
        * Finally, this constructor inverts \a D before proceeding.
        * If you have pre-inverted \a D, you can set the \a
        * preinvertedD flag to true and the functor will not redo the
        * operation.  This helps optimize common cases; for example,
        * when \a D is simply the identity matrix (and thus equal to
        * its own inverse), there is no need to compute the inverse.
        *
        * \param mu An n x 1 vector of means. \param sigma An n x n
        * variance-covariance matrix. \param D An n x n linear
        * constraint definition matrix; should be of rank n. \param a
        * An n x 1 lower bound vector (may contain infinity or
        * negative infinity). \param b An n x 1 upper bound vector (may
        * contain infinity or negative infinity). \param generator
        * Reference to an rng object \param burnin Optional burnin
        * parameter; default value is 0. \param thin Optional thinning
        * parameter; default value is 1. \param preinvertedD Optional
        * flag with default value of false; if set to true, functor
        * will not invert \a D.
        *
        * \throw scythe_dimension_error (Level 1)
        * \throw scythe_conformation_error (Level 1)
        * \throw scythe_invalid_arg (Level 1)
        *
        * \see operator()()
        * \see rng
        */
      template <matrix_order PO1, matrix_style PS1, matrix_order PO2,
                matrix_style PS2, matrix_order PO3, matrix_style PS3,
                matrix_order PO4, matrix_style PS5, matrix_order PO5,
                matrix_style PS4>
      rtmvnorm (const Matrix<double, PO1,PS1>& mu,
                const Matrix<double, PO2, PS2>& sigma,
                const Matrix<double, PO3, PS3>& D,
                const Matrix<double, PO4, PS4>& a,
                const Matrix<double, PO5, PS5>& b, rng<RNGTYPE>& generator, 
                unsigned int burnin = 0, unsigned int thin = 1,
                bool preinvertedD = false) 
      : mu_ (mu), C_ (mu.rows(), mu.rows(), false), 
        h_ (mu.rows(), 1, false), z_ (mu.rows(), 1, true, 0), 
        generator_ (generator), n_ (mu.rows()), thin_ (thin), iter_ (0)
      {
        SCYTHE_CHECK_10(thin == 0, scythe_invalid_arg,
            "thin must be >= 1");
        SCYTHE_CHECK_10(! mu.isColVector(), scythe_dimension_error,
            "mu not column vector");
        SCYTHE_CHECK_10(!  sigma.isSquare(), scythe_dimension_error,
            "sigma not square");
        SCYTHE_CHECK_10(!  D.isSquare(), scythe_dimension_error,
            "D not square");
        SCYTHE_CHECK_10(!  a.isColVector(), scythe_dimension_error,
            "a not column vector");
        SCYTHE_CHECK_10(!  b.isColVector(), scythe_dimension_error,
            "b not column vector");
        SCYTHE_CHECK_10(sigma.rows() != n_ || D.rows() != n_ || 
            a.rows() != n_ || b.rows() != n_, scythe_conformation_error,
            "mu, sigma, D, a, and b not conformable");

        // TODO will D * sigma * t(D) always be positive definite,
        // allowing us to use the faster invpd?
        if (preinvertedD)
          Dinv_ = D;
        else
          Dinv_ = inv(D);
        Matrix<> Tinv = inv(D * sigma * t(D));
        alpha_ = a - D * mu;
        beta_ =  b - D * mu;

        // Check truncation bounds
        if (SCYTHE_DEBUG > 0) {
          for (unsigned int i = 0; i < n_; ++i) {
            SCYTHE_CHECK(alpha_(i) >= beta_(i), scythe_invalid_arg,
                "Truncation bound " << i 
                << " not logically consistent");
          }
        }

        // Precompute some stuff (see Geweke 1991 pg 7).
        for (unsigned int i = 0; i < n_; ++i) {
          C_(i, _) = -(1 / Tinv(i, i)) % Tinv(i, _);
          C_(i, i) = 0; // not really clever but probably too clever
          h_(i) = std::sqrt(1 / Tinv(i, i));
          SCYTHE_CHECK_30(std::isnan(h_(i)), scythe_invalid_arg,
              "sigma is not positive definite");
        }

        // Do burnin
        for (unsigned int i = 0; i < burnin; ++i)
          sample ();
      }

       /*! \brief Generate random variates.
        *
        * Iterates the Gibbs sampler and returns a Matrix containing a
        * single draw from the truncated multivariate random number
        * generator encapsulated by the instantiated object.  Thinning
        * of sampler draws is specified at construction.
        *
        * \see rtmvnorm()
        */
      template <matrix_order O, matrix_style S>
      Matrix<double, O, S> operator() ()
      {
        do { sample (); } while (iter_ % thin_ != 0);

        return (mu_ + Dinv_ * z_);
      }

       /*! \brief Generate random variates.
        *
        * Default template. See general template for details.
        *
        * \see operator()().
        */
      Matrix<double,Col,Concrete> operator() ()
      {
        return operator()<Col, Concrete>();
      }

    protected:
      /* Does one step of the Gibbs sampler (see Geweke 1991 p 6) */
      void sample ()
      {
        double czsum;
        double above;
        double below;
        for (unsigned int i = 0; i < n_; ++i) {

          // Calculate sum_{j \ne i} c_{ij} z_{j}
          czsum = 0;
          for (unsigned int j = 0; j < n_; ++j) {
            if (i == j) continue;
            czsum += C_(i, j) * z_(j);
          }

          // Calc truncation of conditional univariate std normal
          below = (alpha_(i) - czsum) / h_(i);
          above = (beta_(i) - czsum) / h_(i);
          
          // Draw random variate z_i
          z_(i) = h_(i);
          if (above == std::numeric_limits<double>::infinity()){
            if (below == -std::numeric_limits<double>::infinity())
              z_(i) *= generator_.rnorm(0, 1); // untruncated
            else
              z_(i) *= generator_.rtbnorm_combo(0, 1, below);
          } else if (below == 
              -std::numeric_limits<double>::infinity())
            z_(i) *= generator_.rtanorm_combo(0, 1, above);
          else
            z_(i) *= generator_.rtnorm_combo(0, 1, below, above);

          z_(i) += czsum;
        }

        ++iter_;
      }

      /* Instance variables */
      // Various reused computation matrices with names from
      // Geweke 1991.
      Matrix<> mu_;   Matrix<> Dinv_; 
      Matrix<> C_; Matrix<> alpha_; Matrix<> beta_; Matrix<> h_; 

      Matrix<> z_; // The current draw of the posterior

      rng<RNGTYPE>& generator_; // Refernce to random number generator

      unsigned int n_;  // The dimension of the distribution
      unsigned int thin_; // thinning parameter
      unsigned int iter_; // The current post-burnin iteration
  };
} // end namespace scythe
#endif
