//  -*- C++ -*-
//
//======================================================================
//  File:         jt.C
//  Author:       Thomas G. Dietterich
//  Description:  Implements methods 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.
//======================================================================
//
// 

#include "jt.h"
#include "bayes.h"
#include <float.h>
#include <math.h>

extern int TRACE;

//======================================================================
//
// Implementation of class jtnode
//
void jtnode::huginPropagation()
{
  // this node becomes the "root" of the Hugin tree.
  // We invoke collectEvidence on ourselves.  By passing 0 as the
  // caller, we tell collectEvidence to skip computing the return message. 
  if (TRACE) {
    cout << "huginPropagation starting at node " << *this << endl;
  }
  assert(isClusterNode);

  collectEvidence(0);
  probabilityTable * unitTable = new probabilityTable;
  // this creates a trivial probabilityTable with a single 1 element
  // By passing 0 as the caller, we cause distributeEvidence to pass
  // messages to ALL of its neighbors.
  distributeEvidence(0, unitTable);
  // The junction tree is now consistent.
}

probabilityTable * jtnode::collectEvidence(jtnode * caller)
{
  //
  if (isClusterNode) {
    listIterator<jtnode *> nitr(neighbors);
    for (nitr.init(); !nitr; ++nitr) {
      if (nitr() != caller) {
	probabilityTable * incomingMessage = nitr()->collectEvidence(this);

	if (TRACE) {
	  cout << *this << " received incoming message from "
	       << *nitr() << endl;
	  cout << *incomingMessage;
	}

	(*table) *= (*incomingMessage);
	
	if (TRACE) {
	  cout << *this << " resulting table is " << endl;
	  cout << *table;
	}
      }
    }
    // now send message to our caller, if any.
    if (caller) {
      probabilityTable * outgoingMessage = table->duplicate();
      outgoingMessage->sumOver(outgoingMessage->vars - caller->table->vars);
      if (TRACE) {
	cout << *this << " sending message to caller "
	     << *caller << " message is " << endl;
	cout << *outgoingMessage;
      }
      return outgoingMessage;
    }
    else return 0;
  }
  else {
    // we are a separator node
    listIterator<jtnode *> nitr(neighbors);
    for (nitr.init(); !nitr; ++nitr) {
      if (nitr() != caller) {
	// this is our "other" link
	probabilityTable * incomingMessage = nitr()->collectEvidence(this);
	if (TRACE) {
	  cout << *this << " received incoming message from "
	       << *nitr() << endl;
	  cout << *incomingMessage;
	}
	probabilityTable * outgoingMessage = (*incomingMessage) / (*table);
	delete table;
	table = incomingMessage;
	if (TRACE) {
	  cout << *this << " sending message to caller "
	       << *caller << " message is " << endl;
	  cout << *outgoingMessage;
	}
	return outgoingMessage;
      }
    }
  }
}
      

void jtnode::distributeEvidence(jtnode * caller,
				probabilityTable * incomingMessage)
{
  if (TRACE) {
    cout << *this << " received incoming message from ";

    if (caller) cout << *caller << endl;
    else cout << "null caller" << endl;

    if (incomingMessage) cout << *incomingMessage;
    else cout << "incoming message is null";
  }

  if (isClusterNode) {
    // conformal product
    (*table) *= (*incomingMessage);

    if (TRACE) {
      cout << *this << "resulting table is " << endl;
      cout << (*table);
    }

    listIterator<jtnode *> nitr(neighbors);
    for (nitr.init(); !nitr; ++nitr) {
      if (nitr() != caller) {
	probabilityTable * outgoingMessage = table->duplicate();
	outgoingMessage->sumOver(outgoingMessage->vars -
				  nitr()->table->vars);
	if (TRACE) {
	  cout << *this << " sending message to "
	       << *(nitr()) << " message is " << endl;
	  cout << *outgoingMessage;
	}
	
	nitr()->distributeEvidence(this, outgoingMessage);
      }
    }
  }
  else {
    // separator node
    probabilityTable * outgoingMessage = (*incomingMessage) / (*table);
    delete table;
    table = incomingMessage;
    listIterator<jtnode *> nitr(neighbors);
    for (nitr.init(); !nitr; ++nitr) {
      if (nitr() != caller) {
	if (TRACE) {
	  cout << *this << " sending outgoing message to "
	       << *(nitr()) << " message is:" << endl;
	  cout << *outgoingMessage << endl;
	}

	// this is our "other" link
	nitr()->distributeEvidence(this, outgoingMessage);
      }
    }
  }
}


void jtnode::maxPropagation()
{
  // CS539 your code here
}

probabilityTable * jtnode::collectMax(jtnode * caller)
{
  // CS539 your code here
}
      

void jtnode::distributeMax(jtnode * caller,
			   probabilityTable * incomingMessage)
{
  // CS539 your code here
}



void jtnode::print(ostream & str, int indent, jtnode * caller)
{
  if (isClusterNode) {
    int indentstart = str.tellp();
    str << "C" << table->vars << " ";
    indent += str.tellp() - indentstart;

    int nbrCount = 0;
    listIterator<jtnode *> itr(neighbors);
    for (itr.init(); !itr; ++itr) {
      if (itr() == caller) continue;
      if (nbrCount > 0) {
	for (int i = 0; i < indent; i++) str << " ";
      }
      str << "-- ";
      itr()->print(str, indent + 3, this);
      nbrCount++;
    }
    // print new line if we are a leaf
    if (nbrCount == 0) str << endl;
  }
  else { // separator node
    int indentstart = str.tellp();
    str << "S" << table->vars << " -- ";
    listIterator<jtnode *> itr(neighbors);
    for (itr.init(); !itr; ++itr) {
      if (itr() != caller) {
	itr()->print(str,
		     indent + str.tellp() - indentstart,
		     this);
      }
    }
  }
}

void jtnode::printSelf(ostream & str)
{
  str << (isClusterNode? "C" : "S") << table->vars;
  //  if (isClusterNode) str << " refs: " << references;
}

ostream & operator << (ostream & str, jtnode & node)
{
  node.printSelf(str);
  return str;
}


// ======================================================================
//
// implementation of class junctionTree
//
//

junctionTree::junctionTree(bayesNet * net)
{
  net->buildJT(*this);
  network = net;
  consistent = 0;
};

void junctionTree::observe(int var, int value)
{
  variableSet ovar((1 << var));

  listIterator<jtnode *> itr(nodes);
  for (itr.init(); !itr; ++itr) {
    if (itr()->isClusterNode &&
	(itr()->table->vars & ovar).bits) {
      // this is a cluster node that contains our variable
      if (TRACE) {
	cout << *itr() << " chosen to apply evidence" << endl;
      }
      
      // Construct a ptable for this variable
      probabilityTable evidenceTable(ovar);

      // set the unobserved value to zero.  This indexing works
      // because we know this ptable has only 2 entires in its values
      // vector. 
      evidenceTable.values[1 - value] = 0.0;
      
      // multiply the evidence table into the cluster node table
      (*itr()->table) *= evidenceTable;

      consistent = 0;
      return;
    }
  }
}

jtnode * junctionTree::propagate()
{
  // find the first cluster node in nodes, and invoke Hugin propagation
  listIterator<jtnode *> hitr(nodes);
  for (hitr.init(); !hitr; ++hitr) {
    if (hitr()->isClusterNode) {
      hitr()->huginPropagation();
      consistent = 1;
      return hitr();
    }
  }
}

jtnode * junctionTree::maxPropagate()
{
  // CS539 your code here
}

float junctionTree::query(int var, int value)
{
  // return P(var = value) given the current network (including any
  // evidence that may have been entered into it). 
  if (!consistent) propagate();

  // find a node that contains the query variable
  variableSet qvar((1 << var));  // create a varset for  query variable

  listIterator<jtnode *> itr(nodes);
  for (itr.init(); !itr; ++itr) {
    if ((itr()->table->vars & qvar).bits) {
      // this node contains our variable.
      // Copy, marginalize, and read out the answer
      if (TRACE) {
	cout << "Node " << *itr() << " chosen for answer extraction"
	     << endl; 
      }
      probabilityTable * table = itr()->table->duplicate();
      table->sumOver(itr()->table->vars - qvar);  // sum over remaining variables
      table->normalize(qvar);
      if (TRACE) {
	cout << "Result of JT query for " << var << " = " << value << endl;
	cout << *table << endl;
      }
      // the table is guaranteed to have only 2 elements, so we
      // directly index it. 
      float result = table->values[value];
      delete table;
      return result;
    }
  }
  cout << "No node found that contains the query variable" << endl;
}

void junctionTree::save()
{
  listIterator<jtnode *> itr(nodes);
  for (itr.init(); !itr; ++itr) {
    itr()->save();
  }
}

void junctionTree::reset()
{
  listIterator<jtnode *> itr(nodes);
  for (itr.init(); !itr; ++itr) {
    itr()->reset();
  }
}

void junctionTree::fill(float val)
{
  listIterator<jtnode *> itr(nodes);
  for (itr.init(); !itr; ++itr) {
    itr()->fill(val);
  }
}

void junctionTree::print(ostream & str)
{
  // print out each separator node with the two cluster nodes that it
  // joins. 
  listIterator<jtnode *> itr(nodes);
  for (itr.init(); !itr; ++itr) {
    if (!itr()->isClusterNode) {
      str << *(itr()->neighbors[0]) << " --- "
	  << *itr() << " --- "
	  << *(itr()->neighbors[1]) << endl;
    }
  }
}

ostream & operator << (ostream & str, junctionTree & jt)
{
  jt.print(str);
  return str;
}

void junctionTree::printTables(ostream & str)
{
  listIterator<jtnode *> itr(nodes);
  for (itr.init(); !itr; ++itr) {
    str << *itr() << endl;
    str << (*(itr()->table)) << endl;
  }
}

void junctionTree::deleteUselessNodes()
{
  // after a junction tree is built, it often contains useless nodes
  // where a cluster node and one of its separators have identical
  // sets of variables.  In such cases, the separator can be deleted,
  // and the two nodes can be merged.
  listIterator<jtnode *> itr(nodes);
  for (itr.init(); !itr; ++itr) {
    jtnode * mergefrom = 0;
    jtnode * mergeto = 0;
    if (itr()->isClusterNode == 0) {
      jtnode * left = itr()->neighbors[0];
      jtnode * right = itr()->neighbors[1];
      if (left->table->vars == itr()->table->vars) {
	if (TRACE) {
	  cout << "deleting separator: merging left into right "
	       << *left << " --- " << *itr() << " --- " << *right
	       << endl;
	}
	mergefrom = left;
	mergeto = right;
      }
      else if (right->table->vars == itr()->table->vars) {
	if (TRACE) {
	  cout << "deleting separator: merging right into left "
	       << *left << " --- " << *itr() << " --- " << *right
	       << endl;
	}
	mergefrom = right;
	mergeto = left;
      }

      if (mergefrom) {
	// conformal product
	(*mergeto->table) *= (*mergefrom->table);
	// update references
	mergeto->references = mergeto->references | mergefrom->references;
	
	// iterate through the separators incident on mergefrom, and
	// point them to mergeto instead.
	listIterator<jtnode *> sitr(mergefrom->neighbors);
	for (sitr.init(); !sitr; ++sitr) {
	  sitr()->neighbors.remove(mergefrom);
	  sitr()->neighbors.add(mergeto);
	}
	
	// make all of the separators pointed at by mergefrom by pointed
	// at by mergeto instead.  Note that this list will now contain
	// two pointers to us.
	mergeto->neighbors.prepend(mergefrom->neighbors);
	
	// delete both occurrences of us on this list.
	mergeto->neighbors.remove(itr());
	
	// mark mergefrom and us for deletion
	mergefrom->isClusterNode = -1;
	itr()->isClusterNode = -1;
      }
    }
  }
  // now garbage collect the deleted nodes
  for (itr.init(); !itr; ++itr) {
    if (itr()->isClusterNode == -1) {
      delete itr();
      itr.removeCurrent();
    }
  }
}


// ======================================================================
// 
// EM learning algorithm
//
// ======================================================================
void junctionTree::EM(vector< vector<int> > & trainingData,
		      vector< vector<int> > & testData)
{
  // data is a vector of training examples.
  // The value -1 indicates an unobserved value
  double trainingLogLikelihood = - DBL_MAX;
  double testLogLikelihood = - DBL_MAX;
    
  cout << "Iteration    Train Log Likelihood  Test Log Likelihood" << endl;
  for (int iter = 0; ; iter++) {
    EMIteration(trainingData);
    double trainLL = logLikelihood(trainingData);
    double testLL = logLikelihood(testData);
    cout << "   " << iter << "         "
	 << trainLL << "               "
	 << testLL << endl;
    if (fabs(trainingLogLikelihood - trainLL) < 0.001) break;
    trainingLogLikelihood = trainLL;
    testLogLikelihood = testLL;
  }
}

void junctionTree::applyExampleAsEvidence(vector<int> & example)
{
  int N = example.length();
  for (int j = 0; j < N; j++) {
    int val = example[j];
    if (val != -1) {
      // not missing
      observe(j, example[j]);
    }
  }
}

void junctionTree::EMIteration(vector< vector<int> > & data)

  // To implement the EM algorithm, we use the nodes of the belief
  // network (*network) as our "counters".  We place the evidence, e,
  // onto the junction tree and propagate it.  Then, for each table
  // P(x|parents(x)) in the belief network, we find the corresponding
  // table of the junction tree and marginalize it to eliminate any
  // other variables and normalize it.  This then gives
  // P(x|parents(x),e).  If x was an observed variable, this will be 0
  // or 1.  But it if was an unobserved variable, then this will give
  // tilde-P(x), as required by EM.  So the code does some extra
  // work, but it is completely general.
  //
  // To implement the M step, we accumulate these tilde-P values into
  // the belief net tables.  After processing every training example,
  // we normalize each table in the belief network.  Then we
  // re-compute the junction tree from the belief network (using
  // bayesNetwork::updateJT).

{
  int M = data.length();	// number of training examples
  if (M == 0) return;

  network->fill(0.0);		// zero all cells in the belief network
				// note: to get a Dirichlet prior,
				// just change this 0.0 to the desired
				// value of alpha

  for (int i = 0; i < M; i++) {
    // reset the JT tables
    reset();			

    // place the evidence on the junction tree
    applyExampleAsEvidence(data[i]);

    // propagate the evidence
    propagate();

    // read out each of the required PDFs
    network->accumulateEStep();

  }

  // normalize each of the tables in the bayes network
  network->normalizeAllNodes();
  // and save this learned network
  network->save();

  // update the junction tree
  fill(1.0);			// fill all of our tables with 1's
  network->updateJT();
  
  // save the resulting tables
  propagate();			// make it consistent first
  save();
}

double junctionTree::logLikelihood(vector< vector<int> > & data)
{
  // compute the log likelihood of the data
  
  // Because the data set is assumed to be iid, the log likelihood of
  // the entire data set is the sum of the log likelihood of each
  // example.

  int M = data.length();

  double ll = 0.0;
  for (int i = 0; i < M; i++) {
    double thisl = likelihood(data[i]);
    if (thisl == 0.0) thisl = 1e-10; // protect log
    ll = ll + log(thisl);
  }
  return ll;
}

double junctionTree::likelihood(vector<int> & example)
{
  // to compute the likelihood of an example, we apply the example as
  // evidence, perform HuginPropagation, and then choose any node
  // (e.g., the root), and marginalize away ALL of the variables,
  // which gives us P(e), the probability of the evidence.

  reset();

  applyExampleAsEvidence(example);

  jtnode * root = propagate();

  double answer = root->table->sum();
  if (TRACE) cout << "Example likelihood = " << answer << endl;
  return answer;
}

// ======================================================================
//
// Hard EM
//
// Same as EM, except that we compute MPE of the data and count it
// instead.  In a more streamlined implementation, this would be more
// efficient. 
//

void junctionTree::HardEM(vector< vector<int> > & trainingData,
		      vector< vector<int> > & testData)
{
  // CS539 your code here
}

void junctionTree::HardEMIteration(vector< vector<int> > & data)
{
  // CS539 your code here
}
