/*
  BeliefPropagationHyperEA.cpp
  
  Belief propagation for the Edwards-Anderson model on a hypergraph.
  
  The energy of a given configuration is       E = sum_a E_a
  with the energy of clause a being            E_a = -J_a prod_{i in a} s_i
  J_a is either +J or -J, and s_i is either +1 or -1.
  Energy unit is set to be J.
 
  Some of the results obtained by this program are reported in Section 3.6
  of the book:
  Hai-Jun Zhou, "Spin Glass and Message Passing" 
  (Science Press, Beijing, 2016).

  LOG:
  08.05.2013--10.05.2013
  
  PROGRAMMER:
  Haijun Zhou
  Institute of Theoretical Physics, Chinese Academy of Sciences,
  Zhong-Guan-Cun East Road 55, Beijing 100190, China
  http://www.itp.ac.cn/~zhouhj
  email: zhouhj@itp.ac.cn

*/


#include <cmath>
#include <exception>
#include <fstream>
#include <iostream>
#include <string>
#include <valarray>

#include "zhjrandom.h"
//#include "/home/Haijun/bin/Programs/zhjrandom.h"

using namespace std;



/*                                        --- interaction (clause) structure */
struct clausestruct
{
  int Jval;                                           //coupling constant +/- 1
  int degree;                                        //number of attached edges
  int begin;                  //begin address of input and output message array
  clausestruct(void);
};

clausestruct::clausestruct(void)
{
  Jval=1;
  degree=0;
  begin=0;
  return;
}


/*                                                    --- variable structure */
struct variablestruct
{
  int degree;                                        //number of attached edges
  int begin;                               //begin address of VimPosition array
  variablestruct(void);
};

variablestruct::variablestruct(void)
{
  degree =0;
  begin  =0;
}



class BPforHEA
{
public:
  BPforHEA(ZHJRANDOMv3*);

  bool ReadGraph( string&);
  void InitPopulation(void);
  void SetBeta(double);
  void SetDamping(double);

  bool PopulationConverge(int, double);
  void Thermodynamics( string&);

private:

  int ClauseNumber;
  int VariableNumber;
  int EdgeNumber;                             //number of clause-variable pairs
  int MaxVariableDegree;                              //maximal variable degree

  double Damping;                                              //damping factor

  double Beta;                                            //inverse temperature
  double TanhBetaJ;                                              //tanh(beta J)
  double LogCoshBetaJ;                                       //ln[cosh(beta J)]

  valarray<clausestruct> Clause;
  valarray<variablestruct> Variable;

  valarray<double> Pai; /* =tanh(beta J_a) prod_{j in a\i}(2 q_{ja}^+ -1)
			   The array records all the edge outputs of the
			   first clause, then those of the second clause,
			   ..., then those of the the last clause. */
  valarray<double> Qia; /*variable cavity probability of spin +1 (expt a)
			  The array records all the edge outputs of the
			  variables. The edges are ordered in following
			  way: first all edges of the first clause, then
			  all edges of the second clause, ..., and last
			  all edges of the last clause. */
  valarray<int> VimPosition; /*index of the edges attached to the variables.
			       variable input message position */


  valarray<double> pProduct;                                // auxiliary array
  valarray<double> mProduct;                                // auxiliary array

  ZHJRANDOMv3            *rd_ptr;

  void UpdateClauses(void); //update the message array Pai (clause to variable)
  void UpdateVariables(double&, double&);  //update message Qia (var to clause)

};




int main(int argc, char ** argv)
{
  int rdseed  =323456789;
  int rdseed2 =197654321;

  cout
    <<"rdseed= ";
  cout.flush();
  cin
    >>rdseed;

  cout
    <<"rdvalue equilibriation =";
  cout.flush();
  cin
    >>rdseed2;

  ZHJRANDOMv3 rdgenerator(rdseed);
  for(int i=0; i<rdseed2; ++i)
    rdgenerator.rdflt();

  BPforHEA bp_zhj( &rdgenerator );

  cout
    <<"instance file ";
  cout.flush();
  string inputfilename;
  cin
    >>inputfilename;

  bp_zhj.ReadGraph(inputfilename);

  double Beta;

  do
    {

      bool tocontinue=true;
      double epsilon=1.0e-6;
      int iterations=100;

      cout<<"Beta = ";
      cin >>Beta;

      if(Beta<=0) continue;

      bp_zhj.SetBeta( Beta );

      bp_zhj.InitPopulation();

      while(tocontinue)
	{
	  cout
	    <<"convergence criterion epsilon= ";
	  cout.flush();
	  cin
	    >>epsilon;

	  cout
	    <<"iterations=";
	  cout.flush();
	  cin
	    >>iterations;

	  double dpfactor;                                     //damping factor
	  cout<<"Damping [0,1) =";
	  cin>>dpfactor;
	  bp_zhj.SetDamping( dpfactor );

	  if(bp_zhj.PopulationConverge(iterations, epsilon) )
	    tocontinue=false;
	}

      string outputfilename;
      cout
	<<"results output file ";
      cout.flush();
      cin
	>>outputfilename;

      bp_zhj.Thermodynamics(outputfilename);
    }
  while( Beta > 0);

  return 1;
}



BPforHEA::BPforHEA(ZHJRANDOMv3* rd)
{
  rd_ptr   = rd;

  return;
}



void BPforHEA::SetBeta(double b)
{
  Beta = b;

  TanhBetaJ = tanh(Beta);                                        //tanh(beta J)
  LogCoshBetaJ = Beta+log( 0.5e0*(1+exp(-2.0e0*Beta) ) );    //ln[cosh(beta J)]

  return ;
}



void BPforHEA::SetDamping(double d)
{
  Damping = d;
  if(Damping<0)       Damping=0;
  else if(Damping>=1) Damping=0.9999e0;

  return ;
}



/*                                                      --- read hyper-graph */
bool BPforHEA::ReadGraph( string& filename)
{
  ifstream source( filename.c_str() );
  if( !source.good() )
    {
      cerr<<"Error in cnf formula input (file possibly non-existent).\n";
      cerr.flush();
      return false;
    }

  source
    >> VariableNumber
    >> ClauseNumber;

  try { Clause.resize( ClauseNumber+1 ); }
  catch(bad_alloc) { return false; }

  try { Variable.resize(VariableNumber+1); }
  catch(bad_alloc) { return false; }

  //                                                --- first pass for counting
  EdgeNumber=0;

  int allesgut=true;
  for(int m=1; m<=ClauseNumber && allesgut; ++m)
    {
      int var;
      int Ja;
      int size=0;
      source>>Ja;                                           //coupling constant

      if(Ja != +1 && Ja != -1)
	{
	  cerr<<"Coupling constant different from expection.\n";
	  cerr.flush();
	  allesgut=false;
	}

      source>>var; 
      while(var!=0 && allesgut)
	{
	  if(var>VariableNumber || var<0)
	    {
	      cerr<<"Too many variables, or variable index negative.\n";
	      cerr.flush();
	      allesgut=false;
	    }
	  if(allesgut)
	    {
	      ++size;
	      ++(Variable[var].degree);
	      ++EdgeNumber;
	      source>>var;
	    }
	}
      Clause[m].degree=size;
    }
  source.close();

  try { Pai.resize( EdgeNumber ); }
  catch(bad_alloc) { return false; }

  try { Qia.resize( EdgeNumber ); }
  catch(bad_alloc) { return false;}

  try { VimPosition.resize( EdgeNumber); }
  catch(bad_alloc) { return false; }

  int position=0;
  for(int m=1; m<=ClauseNumber; ++m)
    {
      Clause[m].begin=position;
      position += Clause[m].degree;
    }

  MaxVariableDegree=0;

  position=0;
  for(int var=1; var<=VariableNumber; ++var)
    {
      Variable[var].begin =position;
      position += Variable[var].degree;
      if(MaxVariableDegree<Variable[var].degree)
	MaxVariableDegree=Variable[var].degree;
      Variable[var].degree=0;
    }


  //                                   --- second pass to do the actual reading
  source.open(filename.c_str());
  source
    >> VariableNumber
    >> ClauseNumber;

  //                                                     --- read clauses again
  EdgeNumber=0;
  for(int m=1; m<=ClauseNumber; ++m)
    {
      int var;
      int size=0;
      int Ja;
      source>>Ja;
      Clause[m].Jval=Ja;

      source
	>>var;
      while(var!=0)
	{
	  VimPosition[Variable[ var ].begin+Variable[ var ].degree]=EdgeNumber;
	  /* y=VimPosition[x] specifies the input message, Pai[y], from
	     clause to variable, it also specifies the output message,
	     Qia[y], from variable to clause */
	  ++(Variable[ var ].degree);
	  ++EdgeNumber;
	  source>>var;
	}
    }
  source.close();

  try { pProduct.resize(MaxVariableDegree); }
  catch(bad_alloc) { return false; }
  try { mProduct.resize(MaxVariableDegree); }
  catch(bad_alloc) { return false; }

  cout    <<"Clause number "
	  <<ClauseNumber
	  <<"  Variable number "
	  <<VariableNumber
	  <<"  Edge number "
	  <<EdgeNumber
	  <<endl;
  cout.flush();

  return true;
}



/*                                                 --- initialize population */
void BPforHEA::InitPopulation(void)
{
  for(int edge=0; edge<EdgeNumber; ++edge)
    Qia[edge] = rd_ptr->rdflt();

  UpdateClauses();

  return ;
}



/*                                                --- update all the clauses */
void BPforHEA::UpdateClauses(void)
{
  for(int cl=1; cl<=ClauseNumber; ++cl)
    {
      int degree=Clause[cl].degree;                          //degree of clause
      int position_begin=Clause[cl].begin;  /*begin-address of output array Pai
					      and that of input array Qia */
      if(Clause[cl].Jval== +1)
	{
	  for(int v_j=0; v_j<degree; ++v_j)
	    Pai[position_begin+v_j]= TanhBetaJ;
	}
      else
	{
	  for(int v_j=0; v_j<degree; ++v_j)
	    Pai[position_begin+v_j]= -TanhBetaJ;
	}

      for(int v_j=0; v_j<degree; ++v_j)
	{
	  double q_ja=Qia[position_begin+v_j];       //probability of spin j +1
	  for(int v_k=0; v_k<v_j; ++v_k)
	    Pai[position_begin+v_k] *=(2.0e0*q_ja-1.0e0);
	  for(int v_k=v_j+1; v_k<degree; ++v_k)
	    Pai[position_begin+v_k] *=(2.0e0*q_ja-1.0e0);
	}
    }
  return ;
}





/*                                                --- updating all variables */
void BPforHEA::UpdateVariables
(
 double& max_distance,                //maximal distance between two iterations
 double& mean_distance                   //mean distance between two iterations
 )
{
  max_distance =0.0e0;
  mean_distance=0.0e0;

  for(int vr=1; vr<=VariableNumber; ++vr)
    {
      int degree=Variable[vr].degree;

      if(degree==0)                                                 //no output
	continue;

      for(int c_a=0; c_a<degree; ++c_a)
	{
	  pProduct[c_a]=1.0e0;
	  mProduct[c_a]=1.0e0;
	}

      int position_begin=Variable[vr].begin;

      for(int c_a=0; c_a<degree; ++c_a)
	{
	  double p_ai=Pai[ VimPosition[position_begin+c_a] ];
	  double valp = 1.0e0+p_ai;
	  double valm = 1.0e0-p_ai;

	  for(int c_b=0; c_b<c_a; ++c_b)
	    {
	      pProduct[c_b] *= valp;
	      mProduct[c_b] *= valm;
	    }
	  for(int c_b=c_a+1; c_b<degree; ++c_b)
	    {
	      pProduct[c_b] *= valp;
	      mProduct[c_b] *= valm;
	    }
	}

      for(int c_a=0; c_a<degree; ++c_a)
	{
	  int position=VimPosition[position_begin+c_a];

	  double difference =
	    pProduct[c_a]/(pProduct[c_a]+mProduct[c_a]) - Qia[ position ];

	  Qia[position] += (1.0e0-Damping)*difference;

	  if(difference <0.0e0 ) difference  *= -1;
	  if(difference>max_distance) max_distance=difference;
	  mean_distance += difference;
	}
    }
  mean_distance /= EdgeNumber;

  UpdateClauses();

  return;
}






/*                                                   --- population dynamics */
bool BPforHEA::PopulationConverge(int repeat_number, double eps)
{
  double max_distance =2.0e0*eps+1.0e0;
  double mean_distance=0.0e0;

  int repeat=0;
  while(repeat<repeat_number && max_distance>eps)
    {
      UpdateVariables(max_distance, mean_distance);
      cout << repeat <<'\t'
	   << max_distance <<'\t'
	   << mean_distance
	   << endl;
      ++repeat;
    }

  if(max_distance<=eps)
    return true;
  else
    return false;
}



/*                                                        --- thermodynamics */
void BPforHEA::Thermodynamics
( string& filename)
{

  double total_bfe_variable=0;
  double total_bfe_clause=0;
  double total_e_clause=0;

  //                                  --- free energy contribution of variables
  for(int vr=1; vr<=VariableNumber; ++vr)
    {
      int degree=Variable[vr].degree;
      if(degree==0)
	{
	  total_bfe_variable += -log(2.0e0);
	  continue;
	}
      double p_mean=1.0e0;
      double m_mean=1.0e0;

      int position_begin=Variable[vr].begin;

      for(int c_a=0; c_a<degree; ++c_a)
	{
	  double p_ai=Pai[ VimPosition[position_begin+c_a ] ];
	  p_mean *= 1.0e0+p_ai;
	  m_mean *= 1.0e0-p_ai;
	}
      total_bfe_variable += -degree*LogCoshBetaJ-log( p_mean+m_mean );
    }

  //                         --- free energy and energy contribution of clauses
  for(int cl=1; cl<=ClauseNumber; ++cl)
    {
      int degree=Clause[cl].degree;
      int position_begin=Clause[cl].begin;

      double value=1.0e0;
      for(int v_j=0; v_j<degree; ++v_j)
	value *= 2.0e0*Qia[position_begin+v_j]-1.0e0;

      if(Clause[cl].Jval==+1)
	{
	  total_bfe_clause +=
	    (degree-1)*(-LogCoshBetaJ-log(1.0e0+TanhBetaJ*value));
	  total_e_clause += -(TanhBetaJ+value)/(1.0e0+TanhBetaJ*value);
	}
      else
	{
	  total_bfe_clause +=
	    (degree-1)*(-LogCoshBetaJ-log(1.0e0-TanhBetaJ*value));
	  total_e_clause +=(-TanhBetaJ+value)/(1.0e0-TanhBetaJ*value); 
	}
    }

  total_bfe_variable /= VariableNumber;
  total_bfe_clause /= VariableNumber;

  double bf=(total_bfe_variable-total_bfe_clause);
  double ed=total_e_clause /= VariableNumber;
  double sd=Beta*ed-bf;

  ofstream output(filename.c_str() );
  
  output<<Beta<<'\t'<<bf/Beta<<'\t'<<ed<<'\t'<<sd<<endl;
  output.close();

  return ;
}