/**
 * @date Tue May 10 11:35:58 2011 +0200
 * @author Francois Moulin <Francois.Moulin@idiap.ch>
 * @author Laurent El Shafey <Laurent.El-Shafey@idiap.ch>
 *
 * @brief This class implements a multivariate diagonal Gaussian distribution.
 * @details See Section 2.3.9 of Bishop, "Pattern recognition and machine learning", 2006
 *
 * Copyright (C) Idiap Research Institute, Martigny, Switzerland
 */

#ifndef BOB_LEARN_EM_GMMMACHINE_H
#define BOB_LEARN_EM_GMMMACHINE_H

#include <bob.learn.em/Gaussian.h>
#include <bob.learn.em/GMMStats.h>
#include <bob.io.base/HDF5File.h>
#include <iostream>
#include <boost/shared_ptr.hpp>
#include <vector>

namespace bob { namespace learn { namespace em {

/**
 * @brief This class implements a multivariate diagonal Gaussian distribution.
 * @details See Section 2.3.9 of Bishop, "Pattern recognition and machine learning", 2006
 */
class GMMMachine
{
  public:
    /**
     * Default constructor
     */
    GMMMachine();

    /**
     * Constructor
     * @param[in] n_gaussians  The number of Gaussian components
     * @param[in] n_inputs     The feature dimensionality
     */
    GMMMachine(const size_t n_gaussians, const size_t n_inputs);

    /**
     * Copy constructor
     * (Needed because the GMM points to its constituent Gaussian members)
     */
    GMMMachine(const GMMMachine& other);

    /**
     * Constructor from a Configuration
     */
    GMMMachine(bob::io::base::HDF5File& config);

    /**
     * Assignment
     */
    GMMMachine& operator=(const GMMMachine &other);

    /**
     * Equal to
     */
    bool operator==(const GMMMachine& b) const;

    /**
     * Not equal to
     */
    bool operator!=(const GMMMachine& b) const;

    /**
     * @brief Similar to
     */
    bool is_similar_to(const GMMMachine& b, const double r_epsilon=1e-5,
      const double a_epsilon=1e-8) const;

    /**
     * Destructor
     */
    virtual ~GMMMachine();


    /**
     * Reset the input dimensionality, and the number of Gaussian components.
     * Initialises the weights to uniform distribution.
     * @param n_gaussians The number of Gaussian components
     * @param n_inputs    The feature dimensionality
     */
    void resize(const size_t n_gaussians, const size_t n_inputs);


    /////////////////////////
    // Getters
    ////////////////////////

    /**
     * Get number of inputs
     */
    size_t getNInputs() const
    { return m_n_inputs; }

    /**
     * Get the weights ("mixing coefficients") of the Gaussian components
     */
    const blitz::Array<double,1>& getWeights() const
    { return m_weights; }

    /**
     * Get the logarithm of the weights of the Gaussian components
     */
    inline const blitz::Array<double,1>& getLogWeights() const
    { return m_cache_log_weights; }


    /**
     * Get the means
     */    
    const blitz::Array<double,2> getMeans() const;
    
    /**
     * Get the mean supervector
     */
    void getMeanSupervector(blitz::Array<double,1> &mean_supervector) const;
    
     /**
     * Returns a const reference to the supervector (Put in cache)
     */
    const blitz::Array<double,1>& getMeanSupervector() const;
        
    /**
     * Get the variances
     */
    const blitz::Array<double,2> getVariances() const;
    
    /**
     * Returns a const reference to the supervector (Put in cache)
     */
    const blitz::Array<double,1>& getVarianceSupervector() const;
    

    /**
     * Get the variance flooring thresholds for each Gaussian in each dimension
     */
    const blitz::Array<double,2> getVarianceThresholds() const;



    ///////////////////////
    // Setters
    ///////////////////////

    /**
     * Set the weights
     */
    void setWeights(const blitz::Array<double,1> &weights);

    /**
     * Set the means
     */
    void setMeans(const blitz::Array<double,2> &means);
    /**
     * Set the means from a supervector
     */
    void setMeanSupervector(const blitz::Array<double,1> &mean_supervector);

    /**
     * Set the variances
     */
    void setVariances(const blitz::Array<double,2> &variances);
    /**
     * Set the variances from a supervector
     */
    void setVarianceSupervector(const blitz::Array<double,1> &variance_supervector);

    /**
     * Set the variance flooring thresholds in each dimension
     */
    void setVarianceThresholds(const double value);
    /**
     * Set the variance flooring thresholds in each dimension
     * (equal for all Gaussian components)
     */
    void setVarianceThresholds(blitz::Array<double,1> variance_thresholds);
    /**
     * Set the variance flooring thresholds for each Gaussian in each dimension
     */
    void setVarianceThresholds(const blitz::Array<double,2> &variance_thresholds);


    ////////////////
    // Methods
    /////////////////

    /**
     * Get the weights in order to be updated
     * ("mixing coefficients") of the Gaussian components
     * @warning Only trainers should use this function for efficiency reason
     */
    inline blitz::Array<double,1>& updateWeights()
    { return m_weights; }


    /**
     * Update the log of the weights in cache
     * @warning Should be used by trainer only when using updateWeights()
     */
    void recomputeLogWeights() const;



    /**
     * Output the log likelihood of the sample, x, i.e. log(p(x|GMMMachine))
     * @param[in]  x                                 The sample
     * @param[out] log_weighted_gaussian_likelihoods For each Gaussian, i: log(weight_i*p(x|Gaussian_i))
     * @return     The GMMMachine log likelihood, i.e. log(p(x|GMMMachine))
     * Dimensions of the parameters are checked
     */
    double logLikelihood(const blitz::Array<double, 1> &x, blitz::Array<double,1> &log_weighted_gaussian_likelihoods) const;

    /**
     * Output the log likelihood of the sample, x, i.e. log(p(x|GMMMachine))
     * @param[in]  x                                 The sample
     * @param[out] log_weighted_gaussian_likelihoods For each Gaussian, i: log(weight_i*p(x|Gaussian_i))
     * @return     The GMMMachine log likelihood, i.e. log(p(x|GMMMachine))
     * @warning Dimensions of the parameters are not checked
     */
    double logLikelihood_(const blitz::Array<double, 1> &x, blitz::Array<double,1> &log_weighted_gaussian_likelihoods) const;

    /**
     * Output the log likelihood of the sample, x, i.e. log(p(x|GMM))
     * @param[in]  x The sample
     * Dimension of the input is checked
     */
    double logLikelihood(const blitz::Array<double, 1> &x) const;


    /**
     * Output the averaged log likelihood of a set of samples, x, i.e. log(p(x|GMM))
     * @param[in]  x The sample
     * Dimension of the input is checked
     */
    double logLikelihood(const blitz::Array<double, 2> &x) const;


    /**
     * Output the log likelihood of the sample, x, i.e. log(p(x|GMM))
     * @param[in]  x The sample
     * @warning Dimension of the input is not checked
     */
    double logLikelihood_(const blitz::Array<double, 1> &x) const;

    /**
     * Accumulates the GMM statistics over a set of samples.
     * @see bool accStatistics(const blitz::Array<double,1> &x, GMMStats stats)
     * Dimensions of the parameters are checked
     */
    void accStatistics(const blitz::Array<double,2>& input, GMMStats &stats) const;

    /**
     * Accumulates the GMM statistics over a set of samples.
     * @see bool accStatistics(const blitz::Array<double,1> &x, GMMStats stats)
     * @warning Dimensions of the parameters are not checked
     */
    void accStatistics_(const blitz::Array<double,2>& input, GMMStats &stats) const;

    /**
     * Accumulate the GMM statistics for this sample.
     *
     * @param[in]  x     The current sample
     * @param[out] stats The accumulated statistics
     * Dimensions of the parameters are checked
     */
    void accStatistics(const blitz::Array<double,1> &x, GMMStats &stats) const;

    /**
     * Accumulate the GMM statistics for this sample.
     *
     * @param[in]  x     The current sample
     * @param[out] stats The accumulated statistics
     * @warning Dimensions of the parameters are not checked
     */
    void accStatistics_(const blitz::Array<double,1> &x, GMMStats &stats) const;


    /**
     * Get a pointer to a particular Gaussian component
     * @param[in] i The index of the Gaussian component
     * @return A smart pointer to the i'th Gaussian component
     *         if it exists, otherwise throws an exception
     */
    boost::shared_ptr<bob::learn::em::Gaussian> getGaussian(const size_t i);


    /**
     * Return the number of Gaussian components
     */
    inline size_t getNGaussians() const
    { return m_n_gaussians; }

    /**
     * Save to a Configuration
     */
    void save(bob::io::base::HDF5File& config) const;

    /**
     * Load from a Configuration
     */
    void load(bob::io::base::HDF5File& config);

    /**
     * Load/Reload mean/variance supervector in cache
     */
    void reloadCacheSupervectors() const;

    friend std::ostream& operator<<(std::ostream& os, const GMMMachine& machine);


  private:
    /**
     * Copy another GMMMachine
     */
    void copy(const GMMMachine&);

    /**
     * The number of Gaussian components
     */
    size_t m_n_gaussians;

    /**
     * The feature dimensionality
     */
    size_t m_n_inputs;

    /**
     * The Gaussian components
     */
    std::vector<boost::shared_ptr<Gaussian> > m_gaussians;

    /**
     * The weights (also known as "mixing coefficients")
     */
    blitz::Array<double,1> m_weights;

    /**
     * Update the mean and variance supervectors
     * in cache (into a 1D blitz array)
     */
    void updateCacheSupervectors() const;

    /**
     * Initialise the cache members (allocate arrays)
     */
    void initCache() const;

    /**
     * Accumulate the GMM statistics for this sample.
     * Called by accStatistics() and accStatistics_()
     *
     * @param[in]  x     The current sample
     * @param[out] stats The accumulated statistics
     * @param[in]  log_likelihood  The current log_likelihood
     * @warning Dimensions of the parameters are not checked
     */
    void accStatisticsInternal(const blitz::Array<double,1> &x,
      GMMStats &stats, const double log_likelihood) const;


    /// Some cache arrays to avoid re-allocation when computing log-likelihoods
    mutable blitz::Array<double,1> m_cache_log_weights;
    mutable blitz::Array<double,1> m_cache_log_weighted_gaussian_likelihoods;
    mutable blitz::Array<double,1> m_cache_P;
    mutable blitz::Array<double,2> m_cache_Px;

    mutable blitz::Array<double,1> m_cache_mean_supervector;
    mutable blitz::Array<double,1> m_cache_variance_supervector;
    mutable bool m_cache_supervector;

};

} } } // namespaces

#endif // BOB_LEARN_EM_GMMMACHINE_H
