//   -*- C++  -*-
//
//======================================================================
//  File:          ptable.c
//  Author:        Thomas G. Dietterich
//  Description:   Implements classes for probability table operations
//                 including conformal product, marginalizing,
//                 normalizing, and observing as well as reading and
//                 writing. 
//                 
//  Copyright (c) 1995 by Thomas G. Dietterich.  All Rights Reserved.
//       may be reproduced by any non-commercial purpose.
//======================================================================

#include "ptable.h"

variableSet::variableSet(list<int> & varList)
// construct a variable set from a list of variable indexes
{
  bits = 0;
  listIterator<int> itr(varList);
  for(itr.init(); !itr; ++itr) {
    bits |= (1 << itr());
  }
}

int variableSet::size()
// how many bits are set?
{
  Mask set = bits;
  int answer = 0;
  while (set) {
    answer += set & 1;
    set = (set >> 1);
  }
  return answer;
}

list<int> * variableSet::variableList()
{
  list<int> * answer = new list<int>;
  Mask temp = bits;
  int count = 0;
  while(temp) {
    if (temp & 1) {
      answer->add(count);
    }
    count++;
    temp = (temp >> 1);
  }
  return answer;
}


istream & operator >> (istream & str, variableSet & vs) 
{
  // format is a list of non-negative integers
  list<int> members;
  str >> members;
  variableSet newMembers(members);
  vs = newMembers;
  return str;
}

ostream & operator << (ostream & str, variableSet & vs)
{
  return str << *(vs.variableList());
}

Mask compressUnderMask(Mask item, Mask mask)
// Create a new mask whose bits are the bits from item selected by mask
// All bits that are zero in mask are deleted
// Example:  item = 11001100  
//           mask = 10101010
// Result:  00001010
{
//  cout << "compress(" << hex << item << "," << mask << ") = ";
  Mask stack = 0;
  Mask answer = 0;
  int count = 0;
  while (mask) {
    if (mask & 1) {
      // if this bit is set, push corresponding bit from item onto stack
      stack = (stack << 1) | (item & 1);
      count++;
    }
    item = (item >> 1);
    mask = (mask >> 1);
  }
  // now reverse the stack by popping it onto answer
  for (int i = 0; i < count; i++) {
    answer = (answer << 1) | (stack & 1);
    stack = (stack >> 1);
  }
//  cout << answer << endl << dec;
  return answer;
}

Mask expandUnderMask(Mask item, Mask mask)
// create a new mask whose bits are the bits from item interpolated
// with the zero bits from the mask.
// Example:  item = 00001010
//           mask = 10101010
// result:          10001000
//                  ^ ^ ^ ^  these 4 bits came from the item.
//
{
//  cout << "expand(" << hex << item << "," << mask << ") = ";
  Mask stack = 0;
  Mask answer = 0;
  int count = 0;
  while (mask) {
    count++;
    if (mask & 1) {
      // copy next bit from the item
      stack = (stack << 1) | (item & 1);
      item = (item >> 1);
    }
    else {
      stack = (stack << 1);  // shift the zero from the mask
    }
    mask = (mask >> 1);
  }
  // now reverse the stack, popping into answer
  for (int i = 0; i < count; i++) {
    answer = (answer << 1) | (stack & 1);
    stack = (stack >> 1);
  }
//  cout << answer << endl << dec;
  return answer;
}  

probabilityTable * probabilityTable::operator * (probabilityTable & other)
// compute the conformal product of two probability tables
{
  variableSet newvars = (vars | other.vars);
  probabilityTable * answer = new probabilityTable(newvars);
  int n = answer->values.length();
  Mask left = compressUnderMask(vars.bits, newvars.bits);
  Mask right = compressUnderMask(other.vars.bits, newvars.bits);
  for(int i = 0; i < n; i++) {
    answer->values[i] = values[compressUnderMask(i,left)] *
                        other.values[compressUnderMask(i,right)];
  }
  return answer;
}

probabilityTable * probabilityTable::operator *= (probabilityTable & other)
// update our table to be the conformal product of itself with other
{
  variableSet newvars = (vars | other.vars);
  probabilityTable * answer = new probabilityTable(newvars);
  int n = answer->values.length();
  Mask left = compressUnderMask(vars.bits, newvars.bits);
  Mask right = compressUnderMask(other.vars.bits, newvars.bits);
  for(int i = 0; i < n; i++) {
    answer->values[i] = values[compressUnderMask(i,left)] *
                        other.values[compressUnderMask(i,right)];
  }
  vars = newvars;
  values = answer->values;
  return answer;
}

probabilityTable * probabilityTable::operator / (probabilityTable & other)
// compute the conformal quotient of two probability tables
{
  variableSet newvars = vars;   // result table variables
  probabilityTable * answer = new probabilityTable(newvars);
  int n = answer->values.length();
  Mask left = compressUnderMask(vars.bits, newvars.bits);
  Mask right = compressUnderMask(other.vars.bits, newvars.bits);
  for(int i = 0; i < n; i++) {
    double numerator = values[compressUnderMask(i,left)];
    double denominator = other.values[compressUnderMask(i,right)];
    if (denominator == 0.0) {
      if (numerator != 0.0) {
	cerr << "Attempt to divide by zero in conformal quotient.  i =  " 
	     << i << endl;
	abort();
      }
      else {
	answer->values[i] = 0.0;
      }
    }
    else {
	answer->values[i] = numerator / denominator;
    }
  }
  return answer;
}

void probabilityTable::operator += (probabilityTable & other)
// update our table to be the element-wise sum of itself with other
{
  assert(vars == other.vars);
  values += other.values;
}


void probabilityTable::sumOver (variableSet vs)
// sum over all of the variables mentioned in vs
{
  variableSet  keep = vars - vs;
  Mask keepMask = compressUnderMask(keep.bits, vars.bits);
  probabilityTable  answer(keep);
  int nkeep = keep.size();
  answer.values.setSize((1 << nkeep));   // 2**nkeep
  answer.values.fill(0);
  int n = values.length();

  for (int i = 0; i < n; i++) {
    answer.values[compressUnderMask(i, keepMask)] += values[i];
  }
  *this = answer;  // destructively overwrite current
}

double probabilityTable::sum()
{
  // sum all elements in this table
  int n = values.length();
  double result = 0.0;
  for (int i = 0; i < n; i++) {
    result += values[i];
  }
  return result;
}

void probabilityTable::maxOver (variableSet vs)
// max over all of the variables mentioned in vs
{
  variableSet  keep = vars - vs;
  Mask keepMask = compressUnderMask(keep.bits, vars.bits);
  probabilityTable  answer(keep);
  int nkeep = keep.size();
  answer.values.setSize((1 << nkeep));   // 2**nkeep
  answer.values.fill(0);
  int n = values.length();

  for (int i = 0; i < n; i++) {
    int index = compressUnderMask(i, keepMask);
    if (values[i] > answer.values[index]) {
      answer.values[index] = values[i];
    }
  }
  *this = answer;  // destructively overwrite current
}


probabilityTable * probabilityTable::duplicate()
// duplicate this probability table
{
  probabilityTable * answer = new probabilityTable(vars);
  answer->values = values;
  return answer;
}

void probabilityTable::project(int var, int value)
// modify the current ptable to remove var by selecting all rows in
// which var has the indicated value.  Does nothing if the table does
// not contain var
{
  Mask  varBit = (1 << var);
  if (varBit & vars.bits) {
    // this table contains var.
    Mask newmask = compressUnderMask(vars.bits, vars.bits);
    Mask ourbit = compressUnderMask(varBit, vars.bits);
    // If the vars.bits for this table were 10110 and varBit is 00100,
    // then newmask = 111 and ourbit = 010
    newmask = newmask & (~ourbit);
    // now newmask = 101.
    probabilityTable  result(vars - variableSet(varBit));
    result.values.fill(0.0);
    int n = result.values.length();
    for(int i = 0; i < n; i++) {
      int oldindex = expandUnderMask(i, newmask);
      // if the index is 2 = 10, then oldindex = 100
      if (value) oldindex = oldindex | ourbit;
      // if the value was 1, then oldindex becomes 110.  Note this is
      // where we assume that the tables are organized with 0 values
      // before 1 values.
      result.values[i] = values[oldindex];
    }
    // now change ourselves.  
    (*this) = result;
  }
}

void probabilityTable::print(ostream & str)
// print as a joint probability table
{
  int n = vars.size();
  str << vars << " Probability" << endl;
  for (int i = 0; i < values.length(); i++) {
    for (int iv = n - 1; iv >= 0; iv--) {
      str << " " << ((i & (1 << iv)) != 0);
    }
    str << "    " << values[i] << endl;
  }
}
  
void probabilityTable::normalize(variableSet vs)
// normalize so that the sum of the values in the table specified by
// the value set vs is = 1.  This normalization is done WITHIN each
// setting of any OTHER variables in the table
{
  variableSet others = vars - vs;
  // compute masks over the index bits
  Mask otherMask = compressUnderMask(others.bits, vars.bits);
  Mask normMask = compressUnderMask(vs.bits, vars.bits);

  int n = (1 << vs.size());
  int nothers = (1 << others.size());
  for (int iother = 0; iother < nothers; iother++) {
    Mask thisOtherMask = expandUnderMask(iother, otherMask);

    float sum = 0.0;  // sum of values
    for (int i = 0; i < n; i++) {
      // construct index into table
      int index = expandUnderMask(i, normMask) | thisOtherMask;
      sum += values[index];
    }
    if (sum > 0) {
      for (int i = 0; i < n; i++) {
	int index = expandUnderMask(i, normMask) | thisOtherMask;
	values[index] /= sum;
      }
    }
  }
}

void probabilityTable::randomize(clrandom & rng, int var)
{
  // for each setting of the variables besides var
  // generate a random value between 0 and 1, and its complement, and
  // put them into the appropriate cells.
  //
  // Example:  vars = 1011100
  // suppose var = 3
  // 0001000
  Mask mybit = (1 << var);
  // 1011100 => 1111
  Mask allBits = compressUnderMask(vars.bits, vars.bits);
  // 0010
  Mask myCBit = compressUnderMask(mybit, vars.bits);

  int n = (1 << vars.size());
  for (int i = 0; i < n; i++) {
    // In this ordering we will always visit a value for i where
    // i & myCBit == 0 first.  When this happens, we will generate a
    // random value.
    // When we encounter i & myCBit != 1, then we compute the
    // corresponding index and compute 1 - the random value
    if (i & myCBit) {
      // second entry for this variable
      int zeroindex = i & (~myCBit);
      values[i] = 1.0 - values[zeroindex];
    }
    else {
      values[i] = rng.between(0.0, 1.0);
    }
  }
}

void probabilityTable::maximize(variableSet vs)
  // modify the table so that WITHIN each configuration of the
  // variables not mentioned in vs, the configuration of maximum
  // probability is set to 1.0 and the rest to 0.0.
  // Specifically, if you are interested in a CPT P(a|b,c,d), you
  // should invoke maximize({a}), and for each configuration of
  // {b,c,d}, the value maximizing {a} will be set to 1.0.
{
  variableSet others = vars - vs;
  // compute masks over the index bits
  Mask otherMask = compressUnderMask(others.bits, vars.bits);
  Mask normMask = compressUnderMask(vs.bits, vars.bits);

  int n = (1 << vs.size());
  int nothers = (1 << others.size());
  for (int iother = 0; iother < nothers; iother++) {
    Mask thisOtherMask = expandUnderMask(iother, otherMask);

    float maxval = 0.0;  // maximum of values
    int maxindex = -1;
    for (int i = 0; i < n; i++) {
      // construct index into table
      int index = expandUnderMask(i, normMask) | thisOtherMask;
      if (values[index] > maxval) {
	maxval = values[index];
	maxindex = index;
      }
      values[index] = 0.0;  // set all entries to zero
    }
    values[maxindex] = 1.0;  // set the highest entry to 1.0
  }
}


probabilityTable & probabilityTable::operator = (probabilityTable & other)
{
  vars = other.vars;
  values = other.values;
  return *this;
}


istream & operator >> (istream & str, probabilityTable & table)
{
  // ascii format of table:
  // (var1 var2 var3 ... vark)
  // val1 ... val2^k
  str >> table.vars;
  int n = (1 << table.vars.size());
  table.values.setSize(n);
  for (int i = 0; i < n; i++) {
    str >> table.values[i];
  }
  return str;
}


ostream & operator << (ostream & str, probabilityTable & table)
{
  str << table.vars << endl;
  int n = (1 << table.vars.size());
  for (int i = 0; i < n; i++) {
    str << table.values[i] << " ";
  }
  str << endl;
  return str;
}
  
