#include <iostream>

#include <opengm/opengm.hxx>
#include <opengm/graphicalmodel/graphicalmodel.hxx>
#include <opengm/operations/adder.hxx>
#include <opengm/functions/potts.hxx>

#include <opengm/inference/icm.hxx>
#include <opengm/inference/messagepassing/messagepassing.hxx>
#include <opengm/inference/trws/trws_trws.hxx>
#include <opengm/inference/trws/smooth_nesterov.hxx>

//*******************
//** Typedefs
//*******************
typedef double                                                                 ValueType;          // type used for values
typedef size_t                                                                 IndexType;          // type used for indexing nodes and factors (default : size_t)
typedef size_t                                                                 LabelType;          // type used for labels (default : size_t)
typedef opengm::Multiplier                                                     OpType;             // operation used to combine terms
typedef opengm::ExplicitFunction<ValueType,IndexType,LabelType>                ExplicitFunction;   // shortcut for explicite function
typedef opengm::PottsFunction<ValueType,IndexType,LabelType>                   PottsFunction;      // shortcut for Potts function
typedef opengm::meta::TypeListGenerator<ExplicitFunction,PottsFunction>::type  FunctionTypeList;   // list of all function the model cal use (this trick avoids virtual methods) - here only one
typedef opengm::DiscreteSpace<IndexType, LabelType>                            SpaceType;          // type used to define the feasible statespace
typedef opengm::GraphicalModel<ValueType,OpType,FunctionTypeList,SpaceType>    Model;              // type of the model
typedef Model::FunctionIdentifier                                              FunctionIdentifier; // type of the function identifier

//Build a model that corresponds to x_1 ... x_N being drawn IID from the uniform distribution on M labels
Model buildIID(size_t N, size_t M)
{
   LabelType numLabel = M;
   std::vector<LabelType> numbersOfLabels(N,numLabel);
   Model gm(SpaceType(numbersOfLabels.begin(), numbersOfLabels.end()));
   // Add 1st order functions and factors to the model
   // Univariate factors are constant
   for(IndexType variable = 0; variable < gm.numberOfVariables(); ++variable) {
      // construct 1st order function
      const LabelType shape[] = {gm.numberOfLabels(variable)};
      ExplicitFunction f(shape, shape + 1);
      for (size_t i=0; i<numLabel; ++i)
	f(i) = 2;
	//   f(i) = static_cast<ValueType>(rand()) / (RAND_MAX) * 1000.0;
      FunctionIdentifier id = gm.addFunction(f);
      IndexType variableIndex[] = {variable};
      gm.addFactor(id, variableIndex, variableIndex + 1);
   }
   // two-variable factors are constant
   for(IndexType variable = 0; variable < gm.numberOfVariables()-1; ++variable) {
      // construct 1st order function
      const IndexType vars[]  = {variable,variable+1};
      const LabelType shape[] = {gm.numberOfLabels(variable),gm.numberOfLabels(variable+1)};
      ExplicitFunction f(shape, shape + 2);
      for (size_t i=0; i<numLabel; ++i)
         for (size_t j=0; j<numLabel; ++j)
	   f(i,j) = 1;
      FunctionIdentifier id = gm.addFunction(f);
      gm.addFactor(id, vars, vars + 2);
   }

   return gm;
}


void inferBP(const Model& gm, bool normalization = true){
   typedef opengm::BeliefPropagationUpdateRules<Model, opengm::Integrator> UpdateRules;
   typedef opengm::MessagePassing<Model, opengm::Integrator, UpdateRules, opengm::MaxDistance>  LBP;

   LBP::Parameter parameter(size_t(100)); //maximal number of iterations=0
   parameter.useNormalization_ = normalization;
   LBP lbp(gm, parameter);

   lbp.infer();


   Model::IndependentFactorType marg;
   for(IndexType var=0; var<gm.numberOfVariables(); ++var)
   {
      std::cout<< "Sum of marginal distribution P(x_"<<var<<") :";
      lbp.marginal(var,marg);
      double s = 0;
      for(LabelType i=0; i<gm.numberOfLabels(var); ++i)
	s+=marg(&i);
      std::cout<< s << std::endl;
   }
}


int main(int argc, char** argv) {
   // Infer with LBP
   Model gm_small = buildIID(3,15);
   Model gm_large = buildIID(3,350);
   std::cout << "Sum of marginals with 15 labels" <<std::endl;
   inferBP(gm_small);
   std::cout << "Sum of marginals with 350 labels" <<std::endl;
   inferBP(gm_large);;
   return 0;
};