//   -*-  C++  -*-
//
//======================================================================
//  File:          bayes.c
//  Author:        Thomas G. Dietterich
//  Description:   Implementation of bayesNet class and related
//                 classes.
//  Copyright (c) 1995 by Thomas G. Dietterich.  All Rights Reserved.
//       may be reproduced by any non-commercial purpose.
//======================================================================
//
//  The most complex part of this code is the proper maintenance of
//  the references and vars fields of nodes and factors.  These permit quick
//  checks, during the greedy algorithm, of which sums can be
//  performed. 
//
//  Whenever a network is created, its reset() function should be
//  called.  One of the things this does is to compute the references
//  field of each node.  This field tells each node (i.e., each
//  variable), which probability tables (including its own) refer to
//  it.  A conformal product (i.e., a factor) must contain all such
//  tables before this variable can be summed away.
//
//  Each factor also contains a references field that has one bit set
//  for each probability table contained in that factor.  These
//  references fields are initialized by the factor constructors. 
//
//  The function factor::computeSums takes the information from these
//  two fields are computes the sums that are permitted to be
//  performed on the factor.
//
//  However, we don't want to perform all possible sums, because (a)
//  we don't want to sum over output variables or observed variables
//  and (b) we don't want to sum over variables we've already summed
//  over.  Here is how this is handled.
//
//  First, at the start of query(), we compute the value of
//  desiredSums to be all variables except those that have been
//  observed or the query variable.  
//
//  Second, each factor also keeps track of the variables that would
//  result from taking the conformal product and doing any sums.
//  These are stored on the vars member.
//
//  So, to compute the set of permitted sums, we take
//
//        factor.sums & desiredSums & factor.vars
//
//  These are the variables we sum over.  We then compute the new vars
//  by removing these summed-over variables.
//

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

extern int TRACE;

//----------------------------------------------------------------------
//
// class node implementation
//

node::node(const node & old):
  parents(old.parents), children(old.children), references(old.references),
  name(old.name), valueNames(old.valueNames), flag(old.flag),
  index(old.index), jtcontainer(old.jtcontainer)
  
// do a deep copy of a node.
{
  // duplicate the ptables
  if(old.initialTable) {
    initialTable = old.initialTable->duplicate();
  }
  else initialTable = 0;

  if(old.currentTable) {
    currentTable = old.currentTable->duplicate();
  }
  else currentTable = 0;
}

node::~node()
{
  if(initialTable) delete initialTable;
  if(currentTable) delete currentTable;
}

node & node::operator = (const node & other)
{
  // assignment does a deep copy too
  initialTable = other.initialTable->duplicate();
  currentTable = other.currentTable->duplicate();
  parents = other.parents;
  children = other.children;
  references = other.references;
  name = other.name;
  valueNames = other.valueNames;
  flag = other.flag;
  index = other.index;
  return *this;
}

void node::save ()
{
  if (initialTable) delete initialTable;
  initialTable = currentTable->duplicate();
}

void node::reset ()
{
  if (currentTable) delete currentTable;
  currentTable = initialTable->duplicate();
  references = (1 << index);
}

void node::randomize(clrandom & rng)
{
  assert(currentTable);
  currentTable->randomize(rng, index);
}

void node::observe (int var, int value)
// remove var from the ptable by selecting only rows having its value
{
  currentTable->project(var, value);
}

istream & operator >> (istream & str, node & thisNode)
{
  // format
  // nodenumber (parent indexes) (child indexes) ptable
  thisNode.flag = 0;
  str >> thisNode.index;
  str >> thisNode.name >> thisNode.valueNames;
  str >> thisNode.parents >> thisNode.children;
  // compute the set of all variables in this ptable
  // and allocate it
  list<int> temp(thisNode.parents);
  temp.add(thisNode.index);

  probabilityTable * ptable = new probabilityTable(temp);
  // read the table
  str >> *ptable;
  thisNode.initialTable = ptable;
  thisNode.currentTable = 0;
  thisNode.reset();

  return str;
}

ostream & operator << (ostream & str, node & thisNode)
{
  if (TRACE) {
    if (thisNode.jtcontainer) {
      str << "jtcontainer = " << *(thisNode.jtcontainer) << endl;
    }
  }
  str << thisNode.index << " " << thisNode.name << " (" <<
    thisNode.valueNames[0] << " " << thisNode.valueNames[1] << " ) "
      << thisNode.parents << " " <<
	thisNode.children << endl << *(thisNode.currentTable);
  return str;
}  



//----------------------------------------------------------------------
// implementation of factor class
//

factor::factor(node * nodeptr)
// construct a factor from a node
{
  leaf = nodeptr;
  left = 0;
  right = 0;
  vars = nodeptr->currentTable->vars;
  sharedVariables = 0;
  references = (1 << nodeptr->variable());
  jtReferences = (1 << nodeptr->variable());
  sums = 0;
  ptable = nodeptr->currentTable->duplicate();
}

factor::factor(factor * leftFactor, factor * rightFactor)
{
  leaf = 0;
  left = leftFactor;
  right = rightFactor;
  vars = left->vars | right->vars;
  sharedVariables = left->vars & right->vars;
  references = leftFactor->references | rightFactor->references;
  sums = 0;
  jtReferences = leftFactor->jtReferences | rightFactor->jtReferences;

  // take the union of the separators from the left and right
  // factors. 
  list<jtnode *> * leftSeparators = leftFactor->separators.duplicate();
  list<jtnode *> * rightSeparators = rightFactor->separators.duplicate();
  separators.prepend(*leftSeparators);
  separators.prepend(*rightSeparators);
  delete leftSeparators;
  delete rightSeparators;
}

factor::~factor()
// reclaim the memory for the ptable
// each ptable is allocated fresh, including primitive ones.
// We don't delete separators because it is not a pointer.  The ~list
// routine will clean up the buddlinks.
{
  if (ptable) delete ptable;
}

void factor::computeSums(bayesNet & net) 
// determine which sums could be performed on this factor
{
  // this->references is a varset indicating all nodes whose tables
  // have been multiplied together to create this factor.  For
  // example, if we multiplied together P(A|B,C) and P(C|D),
  // this->references would be {A,C}.
  //
  // net.nodes[i]->references tells all of the variables whose tables
  // mention this variable.  These are all of node i's children plus
  // itself.  We can sum over variable i if this->references contains
  // all of the tables that involved node i = all of the bits
  // mentioned in net.nodes[i]->references

  for (int i = 0; i < net.nvars; i++) {
    if( (net.nodes[i]->references - references).bits == 0) {
      sums = sums | variableSet(1 << i);
    }
  }
}

ostream & operator << (ostream & str, factor & f)
{
  str << "Factor";
  if (f.leaf) {
    str << "[" << f.leaf->variable()
	<< "]";
  }
  else {
    str << "[" << f.left->references << f.right->references
	<< "]";
  }
  return str;
}

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

bayesNet::bayesNet (int n):
  nvars(n), nodes(n,(node *) 0), isObserved(n,0),
  observedValue(n,-1), observedVars(0)
{
  // compute a mask that has the low-order nvars bits set
  Mask m = 0;
  for (int i = 0; i < nvars; i++) {
    m = (m << 1) | 1;
  }
  vars = variableSet(m);
}

bayesNet::bayesNet (char filename[]):
  nvars(0), nodes(0,(node *) 0), isObserved(0,0),
  observedValue(0,-1), observedVars(0)
{
  ifstream str(filename);
  str >> *this;
}

// copy constructor.  We make a deep copy.
bayesNet::bayesNet (const bayesNet &  other): 
  nodes(other.nvars),
  nvars(other.nvars),
  vars(other.vars),
  isObserved(other.isObserved),
  observedVars(other.observedVars),
  observedValue(other.observedValue)
{
  // we have initialized the nodes to be of the right size.  Now
  // duplicate each node
  for(int i = 0; i < nvars; i++) {
    nodes[i] = new node(*(other.nodes[i]));
  }
}

// assignment operator.  We make a deep copy here too.
bayesNet & bayesNet::operator = (const bayesNet & other)
{
  nvars = other.nvars;
  nodes.setSize(nvars);
  for (int i = 0; i < nvars; i++) {
    nodes[i] = new node(*(other.nodes[i]));
  }
  vars = other.vars;
  isObserved = other.isObserved;
  observedVars = other.observedVars;
  observedValue = other.observedValue;
  return *this;
}

// destructor.  Release those pointers to nodes.
bayesNet::~bayesNet()
{
  for(int i = 0; i < nvars; i++) {
    delete nodes[i];
  }
}

void bayesNet::observe(int var, int value)
{
  // loop through all of the nodes and observe var.
  vectorIterator<node *> itr(nodes);
  for (itr.init(); !itr; ++itr) {
    itr()->observe(var, value);
  }
  isObserved[var] = 1;
  observedValue[var] = value;
  observedVars =  observedVars | variableSet(1 << var);
}

void bayesNet::addFactor(factor * newFactor)
  // newFactor is a new factor that has been defined.  (a) Update the
  // candidates list to remove old factors that are invalidated by the
  // new one. (b) Create new candidates by considering all pairs of
  // existing factors with the new factor.  (c) Add the new factor to
  // the list of factors.
{
  listIterator<factor *> fitr(factors);

  if (newFactor->left && newFactor->right) {
    // delete the old factors matching the left and right children of
    // newFactor.
    for(fitr.init(); !fitr; ++fitr) {
      if ((fitr() == newFactor->left) ||
	  (fitr() == newFactor->right)) {
	fitr.removeCurrent();
      }
    }
    // delete all candidates that involve any of the tables included in
    // this factor.
    listIterator<factor *> citr(candidates);
    for(citr.init(); !citr; ++citr) {
      if (citr()->references.bits & newFactor->references.bits) {
	citr.removeCurrent();
      }
    }
  }
  // compute new candidates
  for(fitr.init(); !fitr; ++fitr) {
    factor * newCandidate = new factor(newFactor, fitr());
    newCandidate->computeSums(*this);
    candidates.add(newCandidate);
  }
  // add this factor to the factors.
  factors.add(newFactor);
}

float bayesNet::query(int var, int value)
// compute the marginal probability that var = value given all
// observations so far.  Uses the greedy set factoring algorithm of Li
// and D'Ambrosio.
{
  // first check if we've already observed this variable
  if (isObserved[var]) {
    return ((observedValue[var] == value) ? 1.0 : 0.0);
  }

  return queryInternal(variableSet(1 << var), variableSet(value << var));
}

float bayesNet::query(list<int> & queryVars, list<int> & queryValues)
{
  // check each of the variables to see if we have already observed
  // its other value.
  variableSet vars(0);
  variableSet values(0);

  listIterator<int> varitr(queryVars);
  listIterator<int> valitr(queryValues);

  for(varitr.init(), valitr.init();
      !varitr && !valitr;
      ++varitr, ++valitr) {
    if (isObserved[varitr()]) {
      if (observedValue[varitr()] != valitr()) return 0.0;
    }
    else {
      vars = vars | (variableSet(1 << varitr()));
      values = values | (variableSet(valitr() << varitr()));
    }
  }
  return queryInternal(vars, values);
}


//
// This function implements the SPI algorithm of Li and D'Ambrosio
//
float bayesNet::queryInternal(variableSet queryVars, variableSet queryValues)
// queryVars is a bit map of the variables of interest
// queryValues is the corresponding set of values for those variables. 
{
  // variables we desire to sum over.  These include all variables
  // except the final answer and any variables already observed.
  desiredSums = (vars - queryVars) - observedVars;
  factors.deleteAllValues();
  candidates.deleteAllValues();

  // first, add each node to factors.
  vectorIterator<node *> itr(nodes);
  for (itr.init(); !itr; ++itr) {
    factor * primitive = new factor(itr());
    if (TRACE) {
      cout << "Primitive ptable is " << endl;
      primitive->ptable->print(cout);
    }
    // can we do any sums immediately?
    primitive->computeSums(*this);
    variableSet sums = primitive->sums & desiredSums & primitive->vars;
    if(sums.bits) {
      if (TRACE) cout << "summing over " << sums << endl;
      primitive->ptable->sumOver(sums);
      primitive->vars = primitive->vars - sums;
    }
    addFactor(primitive);
  }

  // now enter main loop
  while (factors.length() > 1) {
    if (TRACE) {
      cout << "Factors: ";
      printFactors(factors);
      cout << endl;
    }
    factor * bestCandidate = 0;
    int bestScore = nvars + 10;  // 10 is just for insurance. number
				 // of variables in the resulting ptable
    listIterator<factor *> canditr(candidates);

    for(canditr.init(); !canditr; ++canditr) {
      // compute number of unobserved variables that will be output from this product
      // form a variableSet starting with all variables mentioned in
      // this subtree, remove any variables that are observed or that
      // can be summed over.
      variableSet outputs = (canditr()->vars -
			     (canditr()->sums & desiredSums))
	                    - observedVars;
      int score = outputs.size();
      if (TRACE) cout << "Candidate: " << *(canditr()) << " outputs = "
	<< outputs << " score = " << score << endl;
      if (score < bestScore) {
	bestCandidate = canditr();
	bestScore = score;
      }
      // I am skipping the tie-breaking heuristic of Li and Dambrosio.
      // They skip ties by preferring the factor that will involve the
      // fewest multiplications.
    }
    // perform the conformal product
    assert(bestCandidate);
    bestCandidate->ptable = *(bestCandidate->left->ptable) *
                            *(bestCandidate->right->ptable);
    if (TRACE) {
      cout << "Multiplying bestCandidate: " << *bestCandidate << endl;
      cout << "Resulting ptable is " << endl;
      bestCandidate->ptable->print(cout);
    }
    // now sum over all summable variables
    variableSet sums = bestCandidate->sums & desiredSums & bestCandidate->vars;
    if(sums.bits) {
      if (TRACE) cout << "Summing over " << sums << endl;
      bestCandidate->ptable->sumOver(sums);
      bestCandidate->vars = bestCandidate->vars - sums;
    }
    addFactor(bestCandidate);
  }
  // the single element remaining in factors is the solution ptable.
  // It will involve only our queryVariables
  probabilityTable & solution = *(factors.firstElement()->ptable);
  solution.normalize(queryVars);
  return solution.values[compressUnderMask(queryValues.bits, queryVars.bits)];
}

float bayesNet::fullQuery(int var, int value)
// compute the marginal probability that var = value
// for debugging purposes, we implement this "guaranteed perfect"
// routine to compute the full joint distribution and then marginalize
// out all of the other variables
{
  variableSet  thisVar (1 << var);

  // an initial table with no variables and probability 1.
  probabilityTable * joint = new probabilityTable(0);
  joint->values[0] = 1;
  cout << *joint;

  vectorIterator<node *> itr(nodes);
  for(itr.init(); !itr; ++itr) {
    joint = ((*joint) * (*itr()->currentTable));
    cout << *joint;
  }

  variableSet  sum = joint->vars - thisVar;
  joint->sumOver(sum);
  cout << "Marginalized table: " << *joint;
  // now normalize
  joint->normalize(thisVar);
  return joint->values[value];
}

void bayesNet::reset()
{
  // reset each node
  int i;
  for (i = 0; i < nvars; i++) {
    nodes[i]->reset();
  }
  // now update the references fields of each node
  for (i = 0; i < nvars; i++) {
    node & thisNode = *(nodes[i]);
    for (int j = 0; j < nvars; j++) {
      if (thisNode.currentTable->vars.bits & (1 << j)) {
	// node i mentions variable j.  Add node i to j's references
	nodes[j]->references = nodes[j]->references | variableSet(1 << i);
      }
    }
  }
  isObserved.fill(0);
  observedValue.fill(-1);
  observedVars = 0;
}


void bayesNet::save()
{
  // save each node
  int i;
  for (i = 0; i < nvars; i++) {
    nodes[i]->save();
  }
}

void bayesNet::randomize (clrandom & rng)
{
  for (int i = 0; i < nodes.length(); i++) nodes[i]->randomize(rng);
}

void bayesNet::printFactors(list <factor *> f)
{
  listIterator<factor *> itr(f);
  cout << "(";
  for(itr.init(); !itr; ++itr) {
    cout << *(itr()) << " ";
  }
  cout << ")";
}

istream & operator >> (istream & str, bayesNet & net)
{
  int nvars;
  str >> nvars;
  bayesNet  temp(nvars);
  for (int i = 0; i < nvars; i++) {
    node * thisNode = new node(i);
    str >> *thisNode;
    temp.nodes[i] = thisNode;
  }
  temp.reset();  // compute references fields, set up currentTable's
  net = temp;
  return str;
}

ostream & operator << (ostream & str, bayesNet & net)
{
  str << net.nvars << endl;
  for (int i = 0; i < net.nvars; i++) {
    str << *(net.nodes[i]) << endl;
  }
  return str;
}

void bayesNet::createClusterNode(junctionTree & jt,
				 factor * fact,
				 variableSet sums)
  // jt = junction tree we are building
  // fact = factor we are marginalizing to create cluster node
  // sums = variables we are marginalizing away
{
  if (TRACE) {
    cout << "Creating a cluster node over variables "
	 << fact->vars
	 << " and marginalizing over "
	 << sums
	 << endl;
  }

  // allocate cluster node and install its ptable
  jtnode * cluster = new jtnode(fact->ptable, 1, fact->jtReferences);
  jt.nodes.add(cluster);
  fact->ptable = 0;


  // now link all of fact's separators to this new node.
  listIterator<jtnode *> itr(fact->separators);
  for (itr.init(); !itr; ++itr) {
    if (TRACE) {
      cout << "linking separator " << itr() << endl;
    }
    cluster->neighbors.add(itr());  // cluster --> separator
    itr()->neighbors.add(cluster);  // separator --> cluster
    itr.removeCurrent(); 
  }

  // now allocate a new separator with a ptable of all 1's. 

  // allocate a separator table of 1's 
  variableSet separatorVars = fact->vars - sums;
  variableSet emptySet(0); 
  if (separatorVars.bits) {
    jtnode * separator = new jtnode(new probabilityTable(separatorVars),
				    0,
				    emptySet);
    jt.nodes.add(separator);
    
    if (TRACE) {
      cout << "creating new separator " << *separator << separator << endl;
    }

    // link new separator 
    cluster->neighbors.add(separator);
    separator->neighbors.add(cluster);

    // update the factor fact to include this new separator and new status
    fact->separators.add(separator);
    fact->ptable = new probabilityTable(separatorVars);
  }
  else {
    // no remaining variables.  Fix the factor to be empty.
    fact->vars = separatorVars;
  }
}

void CandidateConformalProduct(factor * bestCandidate)
  // bestCandidate has been chosen for a conformal product and
  // (potential) marginalization step. 
{
  bestCandidate->ptable = (*(bestCandidate->left->ptable) *
			   *(bestCandidate->right->ptable));
}

// ======================================================================
// The following algorithm builds a junction tree by performing the
// SPI algorithm and extracting the cluster nodes as we go.

void bayesNet::buildJT(junctionTree & jt)
{
  // variables we desire to sum over.  To build the junction tree, we
  // request to sum over all variables.
  desiredSums = vars;
  factors.deleteAllValues();
  candidates.deleteAllValues();

  // first, add each node to factors and see if we can sum over any of
  // them.
  vectorIterator<node *> itr(nodes);
  for (itr.init(); !itr; ++itr) {
    factor * primitive = new factor(itr());
    if (TRACE) {
      cout << "Primitive ptable is " << endl;
      primitive->ptable->print(cout);
    }
    // can we do any sums immediately?
    primitive->computeSums(*this);
    variableSet sums = primitive->sums & desiredSums & primitive->vars;
    if (sums.bits) {
      // yes, we can sum over it.  Create a cluster node and a
      // separator node.  The factor modified to be the ptable from
      // the separator node.
      createClusterNode(jt, primitive, sums);
      primitive->vars = primitive->vars - sums;
      primitive->jtReferences = variableSet(0);
    }
    addFactor(primitive);
  }

  // now enter main loop
  // This is the same basic loop as in Li and D'Ambrosio, except that
  // whenever we marginalize away a set of variables, this defines a
  // cluster node.  That cluster node is then linked to all of the
  // separators that were multiplied together to create the factor
  // that is being marginalized.

  while (factors.length() > 1) {
    if (TRACE) {
      cout << "Factors: ";
      printFactors(factors);
      cout << endl;
    }
    factor * bestCandidate = 0;
    int bestScore = nvars + 10;  // 10 is just for insurance. number
				 // of variables in the resulting ptable
    int bestShared = 0;          // candidate with most shared variables

    listIterator<factor *> canditr(candidates);

    for (canditr.init(); !canditr; ++canditr) {
      // compute number of unobserved variables that will be output
      // from this product form a variableSet starting with all
      // variables mentioned in this subtree, remove any variables
      // that are observed or that can be summed over.
      variableSet outputs = (canditr()->vars -
			     (canditr()->sums & desiredSums));

      int score = outputs.size();
      if (TRACE) cout << "Candidate: " << *(canditr()) << " outputs = "
	<< outputs << " score = " << score << endl;
      if (score < bestScore) {
	bestCandidate = canditr();
	bestScore = score;
	bestShared = canditr()->sharedVariables.size();
      }
      else if (score == bestScore) {
	// break ties in favor of shared variables
	int nshared = canditr()->sharedVariables.size();
	if (nshared > bestShared) {
	  bestShared = nshared;
	  bestCandidate = canditr();
	}
      }
    }
    
    // perform the conformal product

    // We only take the product of the primitive factors, since these
    // will become the ptables of the cluster nodes in the JT.
    assert(bestCandidate);
    // if this is an empty separator factor, then ignore it
    if (bestCandidate->vars.isEmpty()) continue;

    CandidateConformalProduct(bestCandidate);

    if (TRACE) {
      cout << "Multiplying bestCandidate: " << *bestCandidate << endl;
      cout << "Resulting ptable is " << endl;
      bestCandidate->ptable->print(cout);
    }
    // now sum over all summable variables
    variableSet sums = bestCandidate->sums & desiredSums & bestCandidate->vars;
    /*
    if (TRACE) {
      cout << "bestCandidate->sums = " << bestCandidate->sums << endl;
      cout << "desiredSums = " << desiredSums << endl;
      cout << "bestCandidate->vars = " << bestCandidate->vars << endl;
      cout << "sums = " << sums << endl;
    }
    */
    if (sums.bits) {
      if (TRACE) cout << "Summing over " << sums << endl;
      createClusterNode(jt, bestCandidate, sums);
      bestCandidate->vars = bestCandidate->vars - sums;
      bestCandidate->jtReferences = variableSet(0);
    }
    addFactor(bestCandidate);
  }
  jt.deleteUselessNodes();

  // now establish pointers from each node of the original belief
  // network to the corresponding node in the junction tree (jtcontainer)
  listIterator<jtnode *> jtitr(jt.nodes);
  for (jtitr.init(); !jtitr; ++jtitr) {
    list<int> * varlist = jtitr()->references.variableList();
    listIterator<int> nitr(*varlist);
    for (nitr.init(); !nitr; ++nitr) {
      nodes[nitr()]->jtcontainer = jtitr();
    }
  }
}

void bayesNet::normalizeAllNodes()
{
  // loop through all of the nodes of this belief network and
  // normalize them appropriately.
  int len = nodes.length();
  for (int i = 0; i < len; i++) {
    node & thisnode = *(nodes[i]);
    thisnode.currentTable->normalize(variableSet(1 << i));
  }
}

void bayesNet::updateJT()
{
  // for each node
  //   find corresponding cluster node in the JT and multiply
  //   node.table into it.

  int n = nodes.length();
  for (int i = 0; i < n; i++) {
    node & thisnode = *(nodes[i]);
    *(thisnode.jtcontainer->table) *= *(thisnode.currentTable);
  }
}

void bayesNet::accumulateEStep()
{
  int n = nodes.length();
  for (int i = 0; i < n; i++) {
    node & thisnode = *(nodes[i]);
    //   cout << "AccumulateEStep: node: " << i << *(thisnode.currentTable);
    probabilityTable * jtTable = thisnode.jtcontainer->table->duplicate();
    variableSet marginalizers = jtTable->vars - thisnode.currentTable->vars;
    //    cout << "AccumulateEStep: jtnode: " << *thisnode.jtcontainer
    //	 << *jtTable << endl;
    //    cout << "summing over " << marginalizers << endl;
    jtTable->sumOver(marginalizers);
    jtTable->normalize(jtTable->vars);
    //    cout << "result is " << *jtTable << endl;
    *(thisnode.currentTable) += *jtTable;
    //    cout << "resulting node: " << *(thisnode.currentTable) << endl;
    delete jtTable;
  }
}


void bayesNet::accumulateHardEStep()
{
  // CS539 your code here
}
