//  -*- C++  -*- 
//
//======================================================================
//  File:          jtdriver.c
//  Author:        Thomas G. Dietterich
//  Description:   This program reads in a bayes net, converts it to a
//                 junction tree, and then permits 
//                 the user to pose queries and make observations.  I
//                 used it for debugging the jt code.  It prints
//                 the entire jt after each query, which makes it
//                 hard to use for large nets.
//                 
//  Copyright (c) 1995, 2000 by Thomas G. Dietterich.  All Rights Reserved.
//       may be reproduced by any non-commercial purpose.
//======================================================================

#include "bayes.h"
#include "jt.h"
#include <iostream.h>
#include <fstream.h>
#include "clrandom.h"

int TRACE = 0;

int main(int argc, char * argv[])
{
  extern char getopt(int, char *[], char *);
  extern char * optarg;
  char optc;
  int trainingDataPresent = 0;
  vector< vector<int> > trainingData(0);
  int testDataPresent = 0;
  vector< vector<int> > testData(0);
  int netFilePresent = 0;
  bayesNet network;
  clrandom rng(3748374);

  while (( optc = getopt(argc, argv, "vl:t:n:")) != -1 ) {
    switch (optc) {
    case 'v':
      TRACE = 1;
      break;
    case 'l':
      trainingDataPresent = 1;
      {
	ifstream dstr(optarg);
	dstr >> trainingData;
	cout << "Read " << trainingData.length() << " training examples from "
	     << optarg << endl;
      }
      break;
    case 't':
      testDataPresent = 1;
      {
	ifstream dstr(optarg);
	dstr >> testData;
	cout << "Read " << testData.length() << " test examples from "
	     << optarg << endl;
      }
      break;
    case 'n':
      netFilePresent = 1;
      {
	ifstream str(optarg);
	str >> network;
      }
      break;
    default:
      cout << "Usage: jtdriver [-v] [-l trainingdata] [-t testdata] -n networkfile" 
	   << endl;
      cout << " -v                       turn on trace output" << endl;
      cout << " -l trainingdata          file of training examples" << endl;
      cout << " -t testdata              file of test examples" << endl;
      cout << " -n networkfile              file describing belief network" << endl;
      return 1;
    }
  }

  if (!netFilePresent) {
    cout << "You must specify a belief network file with -n" << endl;
    return 1;
  }

  // create the junction tree
  junctionTree * jt = new junctionTree(&network);
  jt->save();

  int var;
  int value;

  while (1) {
    cout << "Junction tree: " << endl;
    cout << *jt << endl;
    jt->printTables(cout);
    cout << "Network: " << endl << network << endl;
    cout << "Observe (o), query (q), MPE (m), reset (r), \n";
    cout << "learn (EM) (l), learn (hard EM) (h),\n";
    cout << "trace (t), untrace(u), randomize (z)  or exit (x): ";
    char command;
    cin >> command;
    switch (command) {
    case 'm':
      jt->maxPropagate();
      break;
    case 'z':
      network.randomize(rng);
      network.save();
      jt->fill(1.0);
      network.updateJT();  // recompute JT tables
      jt->save();
      break;
    case 'o':
      cout << "Variable to observe: ";
      cin >> var;
      cout << "Value observed: ";
      cin >> value;
      network.observe(var, value);
      jt->observe(var, value);
      break;
    case 'q':
      cout << "Enter variable to query: ";
      cin >> var;
      cout << "Enter value to query: ";
      cin >> value;
      cout << "BN: P(var " << var << " = " << value << ") = " <<
	network.query(var, value) << endl;
      cout << "JT: P(var " << var << " = " << value << ") = " <<
	jt->query(var, value) << endl;
      break;
    case 'r':
      network.reset();
      jt->reset();
      break;
    case 'l':
      jt->EM(trainingData, testData);
      cout << "Training set likelihood = " <<
	jt->logLikelihood(trainingData) << endl;
      cout << "Test set likelihood = " <<
	jt->logLikelihood(testData) << endl;
      break;
    case 'h':
      jt->HardEM(trainingData, testData);
      cout << "Training set likelihood = " <<
	jt->logLikelihood(trainingData) << endl;
      cout << "Test set likelihood = " <<
	jt->logLikelihood(testData) << endl;
      break;
    case 't':
      TRACE = 1;
      break;
    case 'u':
      TRACE = 0;
      break;
    case 'x':
      return 0;
    default:
      cout << "Unknown command.  Try again." << endl;
    }
  }
  return -1;
}


  
