//  -*- C++ -*-
//
//
//======================================================================
//  File:         jt.h
//  Author:       Thomas G. Dietterich
//  Description:  Declares classes for reasoning with junction trees
//                using the Hugin propagation algorithm.
//
//  Copyright (c) 2000 by Thomas G. Dietterich.  All Rights Reserved.
//       May be reproduced by any non-commercial purpose.
//======================================================================
//
// 
#ifndef JT_H
#define JT_H

#include "ptable.h"
#include <iostream.h>

// forward references
class bayesNet;

class jtnode
{
public:
  // constructor
  jtnode(probabilityTable * t, int isCluster, variableSet & refs):
    table(t),
    initialTable(0),
    isClusterNode(isCluster),
    references(refs) {};
  // destructor
  ~jtnode() {
    if (table) delete table;
    if (initialTable) delete initialTable;
  }
  
  // primary methods
  void huginPropagation();
  probabilityTable * collectEvidence(jtnode * caller);
  void distributeEvidence(jtnode * caller, probabilityTable * message);
  void maxPropagation();
  probabilityTable * collectMax(jtnode * caller);
  void distributeMax(jtnode * caller, probabilityTable * message);

  // utility methods
  void print(ostream & str, int indent, jtnode * caller);
  void printSelf(ostream &);
  void save() {
    if (initialTable) delete initialTable;
    initialTable = table->duplicate();
  }
  void reset() {
    if (table) delete table;
    table = initialTable->duplicate();
  }
  void fill(float val) {
    if (table) table->fill(val);
  }

  // data fields
  list<jtnode *> neighbors;
  int isClusterNode;   // 1 => cluster node; 0 => separator node
  probabilityTable * table;
  probabilityTable * initialTable;  
  variableSet references;         // variables whose ptables were
                                  // multiplied to create this table.
};

ostream & operator << (ostream &, jtnode &);

class junctionTree
{
public:
  // constructor
  junctionTree(bayesNet * net);    // convert belief network into JT

  // primary methods
  void observe(int var, int value); // update to reflect observed value
  float query(int var, int value);  // compute a probability
  jtnode * propagate();             // Perform huginPropagation
  jtnode * maxPropagate();	    // Perform maxPropagation
  void EM(vector< vector<int> > & trainingData,
	  vector< vector<int> > & testData); // fit the JT to data
  void HardEM(vector< vector<int> > & trainingData,
	      vector< vector<int> > & testData); // fit the JT to data


  // utility methods
  void applyExampleAsEvidence(vector<int> & example);
  void EMIteration(vector< vector<int> > & data);
  void HardEMIteration(vector< vector<int> > & data);
  double logLikelihood(vector< vector<int> > & data);
  double likelihood(vector<int> & example);

  void save();                      // save current tables
  void reset();                     // restore saved tables
  void fill(float val);		    // fill all tables

  void print(ostream & str);       // display the entire tree
  void printTables(ostream & str);  // display probability tables
  void deleteUselessNodes();


  // data fields
  list<jtnode *> nodes;            // nodes in the tree:  includes
                                   // separator nodes and cluster nodes.
  int consistent;                  // true if tree is consistent.
  bayesNet * network;		   // network used to build the JT
};

ostream & operator << (ostream & str, junctionTree & jt);

#endif
