/*!
  \file
  \brief Bayesian Filtering for mixtures of exponential family (EF) members
  \author Vaclav Smidl.

  -----------------------------------
  BDM++ - C++ library for Bayesian Decision Making under Uncertainty

  Using IT++ for numerical operations
  -----------------------------------
*/

#ifndef MIXTURES_H
#define MIXTURES_H


#include "../math/functions.h"
#include "../stat/exp_family.h"
#include "../stat/emix.h"
#include "arx.h"

namespace bdm {

//! enum switch for internal approximation used in MixEF
enum MixEF_METHOD { EM = 0, QB = 1};

/*!
* \brief Mixture of Exponential Family Densities

An approximate estimation method for models with latent discrete variable, such as
mixture models of the following kind:
\f[
f(y_t|\psi_t,     heta) = \sum_{i=1}^{n} w_i f(y_t|\psi_t,     heta_i)
\f]
where  \f$\psi\f$ is a known function of past outputs, \f$w=[w_1,\ldots,w_n]\f$ are component weights, and component parameters \f$    heta_i\f$ are assumed to be mutually independent. \f$    heta\f$ is an aggregation af all component parameters and weights, i.e. \f$    heta = [    heta_1,\ldots,    heta_n,w]\f$.

The characteristic feature of this model is that if the exact values of the latent variable were known, estimation of the parameters can be handled by a single model. For example, for the case of mixture models, posterior density for each component parameters would be a BayesianModel from Exponential Family.

This class uses EM-style type algorithms for estimation of its parameters. Under this simplification, the posterior density is a product of exponential family members, hence under EM-style approximate estimation this class itself belongs to the exponential family.

Two methods are provided:
 - QB where the probability of being from one component is computed using predictors,
 - EM where the data is assigned to component with the highest likelihood (winner takes all)

These methods are stored in attribute options:
 - method: ["EM","QB"] estimation method as mentioned above, QB is default
 - max_niter: maximum of iterations in bayes_batch()

*/
class MixEF: public BMEF {
protected:
    //! Models for Components of \f$    heta_i\f$
    Array<BMEF*> Coms;
    //! Statistics for weights
    multiBM weights;
    //aux
    friend class eprod_mix;

    //! \brief Posterior on component parameters
    class eprod_mix: public eprod_base {
    protected:
        const MixEF &mix; // pointer to parent.n
    public:
        eprod_mix(const MixEF &m):mix(m) {}
        const epdf* factor(int i) const {
            return (i==(mix.Coms.length()-1)) ? &mix.weights.posterior() : &mix.Coms(i)->posterior();
        }
        const int no_factors()const {
            return mix.Coms.length()+1;
        }
    } est;
    ////!indices of component rvc in common rvc

    class Options: public root {
    public:
        //! Flag for a method that is used in the inference
        MixEF_METHOD method;

        //! maximum number of iterations
        int max_niter;

        Options():method(QB),max_niter(10) {};

         /*! Create object from the following structure

        \code
        --- optional fields ---        
        method = '...';            % 'EM' or 'QB', methods of computing bayes
        max_iter = [];             % maximum number of iterations in bayes_batch
        --- inherited fields ---
        bdm::root::from_setting
        \endcode
        If the optional fields are not given, they will be filled as follows:
        \code
        max_niter = 10;            
        method = 'QB';
        \endcode
        */
        void from_setting(const Setting &set) {
            string meth;
            UI::get(meth,set,"method",UI::optional);
            if (meth=="EM") {
                method=EM;
            }
            max_niter =10;
            UI::get(max_niter,set,"max_niter",UI::optional);
        };

        void to_setting(Setting &set)const {
            string meth=(method==EM ? "EM" : "QB");
            UI::save(meth,set,"method");
            UI::save(max_niter,set,"max_niter");
        };
    };

    Options options;
public:
    //! Full constructor
    MixEF ( const Array<BMEF*> &Coms0, const vec &alpha0 ) :
        BMEF ( ), Coms ( Coms0.length() ),
        weights (), est(*this), options() {
        for ( int i = 0; i < Coms0.length(); i++ ) {
            Coms ( i ) = ( BMEF* ) Coms0 ( i )->_copy();
        }
        weights.set_parameters(alpha0);
        weights.validate();
    }

    //! Constructor of empty mixture
    MixEF () :
        BMEF ( ), Coms ( 0 ),
        weights (), est(*this), options() {
    }
    //! Copy constructor
    MixEF ( const MixEF &M2 ) : BMEF ( ),  Coms ( M2.Coms.length() ),
        weights ( M2.weights ), est(*this), options(M2.options) {
        for ( int i = 0; i < M2.Coms.length(); i++ ) {
            Coms ( i ) = (BMEF*) M2.Coms ( i )->_copy();
        }
    }

    //! Initializing the mixture by a random pick of centroids from data
    //! \param Com0 Initial component - necessary to determine its type.
    //! \param Data Data on which the initialization will be done
    //! \param c Initial number of components, default=5
    //! when the number of data records (columns of Data) is equal to the number of requested components, each data is used, otherwise, they are picked randomly.
    void init ( BMEF* Com0, const mat &Data, const int c = 5 );
    //Destructor
    //! Recursive EM-like algorithm (QB-variant), see Karny et. al, 2006
    void bayes ( const vec &yt, const vec &cond );
    //! batch weighted Bayes rule
    double bayes_batch_weighted ( const mat &yt, const mat &cond, const vec &wData );
    double bayes_batch ( const mat &yt, const mat &cond) {
        return bayes_batch_weighted(yt,cond,ones(yt.cols()));
    };
    double logpred ( const vec &yt, const vec &cond ) const;
    //! return correctly typed posterior (covariant return)
    const eprod_mix& posterior() const {
        return est;
    }

    emix* epredictor(const vec &cond=vec()) const;
    //! Flatten the density as if it was not estimated from the data
    void flatten ( const BMEF* M2, double weight );
    //! Access function
    BMEF* _Coms ( int i ) {
        return Coms ( i );
    }

    //!Set which method is to be used
    void set_method ( MixEF_METHOD M ) {
        options.method = M;
    }

    void to_setting ( Setting &set ) const {
        BMEF::to_setting( set );
        UI::save ( Coms, set, "Coms" );
        UI::save ( &weights, set, "weights" );
        UI::save (options, set, "options");
    }

    /*! Create object from the following structure

    \code
    class = 'MixEF';
    Coms = { list of bdm::BMEF };                    % list of components containing any offsprings of Bayesian models BMEF, bdm::BMEF::from_setting
    weights = [...];                                 % vector containing weights of components
    --- optional fields ---
    options = configuration of bdm::MixEF::Options;  % see MixEF::Options, bdm::MixEF::Options::from_setting
    --- inherited fields ---
    bdm::BMEF::from_setting
    \endcode
    */
    void from_setting (const  Setting &set ) {
        BMEF::from_setting( set );
        UI::get ( Coms, set, "Coms" );
        UI::get ( weights, set, "weights" );
        UI::get (options, set, "options",UI::optional);
    }
};
UIREGISTER ( MixEF );

//! \brief Product of ARX models forming an ARX model with potentially reduced structure
class ARXprod: public ProdBMBase {
    Array<shared_ptr<ARX> > arxs;
public:
    ARX* bm(int i) {
        return arxs(i).get();
    }
    int no_bms() {
        return arxs.length();
    }
};
UIREGISTER(ARXprod);

}
#endif // MIXTURES_H