//   -*- C++ -*-
//
//======================================================================
//  File:         bayes.h
//  Author:       Thomas G. Dietterich
//  Description:  Declares classes for manipulating Bayesian networks.
//                This code assumes all variables are binary.
//                Variables are numbered starting at zero.  No more
//                than 32 variables are supported.
//
//  Copyright (c) 1995 by Thomas G. Dietterich.  All Rights Reserved.
//       may be reproduced by any non-commercial purpose.
//======================================================================

#ifndef BAYES_H
#define BAYES_H

#include "vector.h"
#include "clstring.h"
#include "ptable.h"
#include <iostream.h>
#include "jt.h"
#include "clrandom.h"

class node;
class factor;
class bayesNet;


//----------------------------------------------------------------------
//
// class node
//
// A node in a Bayesian Network
//
// The node contains an initial probability table and a current
// probability table.  The current table reflects any observations
// that have been made.
//
// The node also contains links in the network to support identifying
// relevant nodes.
//

istream & operator >> (istream & str, node & n);
ostream & operator << (ostream & str, node & n);

class node
{
public:
  // constructors
  node(int var): flag(0), index(var), references(1 << var), jtcontainer(0) {};
  node(const node & other);
  virtual ~node();
  node & operator = (const node & other);

  // data 
  probabilityTable * initialTable;
  probabilityTable * currentTable;
  list<int>  parents;
  list<int>  children;
  variableSet references;    // variables whose ptable mentions this variable.
			     // all such ptables must be part of a
			     // factor before we can sum over this variable.
  string name;               // name of node
  list<string> valueNames;   // names of values of this node
  jtnode *   jtcontainer;    // junction tree node whose table
                             // contains our table

  // member functions
  void mark() { flag = 1;};
  void clear() { flag = 0; };
  int test() {return flag;};
  int variable() {return index;};

  void observe(int var, int value);
  void save();
  void reset();
  void fill(float value) {
    if (currentTable) currentTable->fill(value);
  }
  void randomize(clrandom & rng);

  // friends
  friend istream & operator >> (istream & str, node & n);
  friend ostream & operator << (ostream & str, node & n);

protected:
  int flag;   // used for marking relevant variables
  int index;  // index number of this variable
};
  

//----------------------------------------------------------------------
//
// class factor: support for the set factoring algorithm
//
// A factor is a term in a factoring of the joint distribution (via
// the chain rule).  Each factor involves taking a conformal product
// of two other factors possibly followed by taking some sums (to
// marginalize away some of the variables).
// 

ostream & operator << (ostream & str, factor & f);

class factor
{
public:
  // constructors
  factor(node * nodeptr);
  factor(factor * leftFactor, factor * rightFactor);
  virtual ~factor();
  
  // data items
  node * leaf;
  factor * left;
  factor * right;
  variableSet  vars;  // union of variables from left and right children
  variableSet  references;  // output variables of all ptables in this subtree
  variableSet  sums;  // sums that can be performed on resulting ptable
  variableSet  sharedVariables;  // intersection of left and right variables
  probabilityTable * ptable;  // result of conformal product of left and right
  list<jtnode *> separators;  // for building JT, separators
                              // multiplied to create this factor.
  variableSet  jtReferences;  // original tables that have been
                              // multiplied together to create this node.

  // member functions
  void computeSums(bayesNet & net);
};



//----------------------------------------------------------------------
//
// class bayesNet
//
//

istream & operator >> (istream & str, bayesNet & net);
ostream & operator << (ostream & str, bayesNet & net);

class bayesNet
{
public:
  // constructors
  bayesNet(char filename[]);  // read network from a file
  bayesNet(int n = 0);
  bayesNet(const bayesNet & other);
  virtual ~bayesNet();
  bayesNet & operator = (const bayesNet & right);
  
  // actions
  void observe(int var, int value);     // update to reflect observed value
  float query(int var, int value);      // compute a probability
  float query(list<int> & queryVars, list<int> & queryValues);
  float queryInternal(variableSet queryVars, variableSet queryValues);
  float fullQuery(int var, int value);  // query via full joint distr
  void reset();
  void save();
  void fill(float value){               // set all cells in all
                                        // ptables to this value
    for (int i = 0; i < nodes.length(); i++) {
      nodes[i]->fill(value);
    }
  };
  void randomize(clrandom & rng);
  int size() {return nvars;};

  // methods for creating junction trees
  void createClusterNode(junctionTree & jt,
			 factor * fact,
			 variableSet sums);
  void mergeSeparators(junctionTree & jt,
		       factor * fact,
		       variableSet sums);
  void buildJT(junctionTree & jt);
  void updateJT();
  void normalizeAllNodes();
  void accumulateEStep();
  void accumulateHardEStep();

  // friends
  friend istream & operator >> (istream & str, bayesNet & net);
  friend ostream & operator << (ostream & str, bayesNet & net);
  friend class factor;

  // public data
  vector<node *> nodes;      // the nodes, one for each variable

protected:
  int  nvars;		     // number of variables
  variableSet vars;	     // a set containing all variables in this net

  // members for recording observations
  vector<int> isObserved;    // indicates whether variable has been observed
  variableSet observedVars;  // also indicates whether variable has
			     // been observed.
  vector<int> observedValue; // value observed for variable

  // members used during query processing
  variableSet desiredSums;   // one bit for each variable that we
			     // would like to sum over.  Set during a query.
  list<factor *> factors;    // factors computed so far
  list<factor *> candidates; // candidates for the next factor to compute
  void addFactor(factor *);  // add a new element to factors, update candidates

  void printFactors(list <factor *> f);  // print a factors list
};



#endif
