#include <config.h>
#include <compiler/ParseTree.h>
#include <model/SymTab.h>
#include <compiler/Compiler.h>
#include <graph/LogicalNode.h>
#include <graph/ConstantNode.h>
#include <graph/StochasticNode.h>
#include <graph/MixtureNode.h>
#include <graph/DevianceNode.h>
#include <graph/Graph.h>
#include <distributions/DistTab.h>
#include <functions/FuncTab.h>
#include <sarray/RangeIterator.h>

#include <utility>
#include <vector>
#include <stdexcept>
#include <cfloat>
#include <cmath>
#include <list>
#include <algorithm>
#include <string>
#include <set>

using std::string;
using std::vector;
using std::list;
using std::map;
using std::pair;
using std::invalid_argument;
using std::runtime_error;
using std::logic_error;
using std::ostringstream;
using std::min;
using std::max;
using std::set;
using std::fabs;

//Structure to hold subset indices
struct SSI {
  Node *node;
  long lower;
  long upper;
};

//Structure to hold two Index objects
/*
class Index2 {
public:
  Index *lower;
  Index *upper;
  Index2 (Index const &l, Index const &u) {
    lower = new Index(l);
    upper = new Index(u);
  };
  Index2 (Index2 const &rhs) {
    lower = new Index(*rhs.lower);
    upper = new Index(*rhs.upper);
  }
  ~Index2 () {
    delete lower;
    delete upper;
  }
}
*/

#include <sstream>
template<class T>
std::string ToString(const T& val)
{
    ostringstream strm;
    strm << val;
    return strm.str();
}


/* FIXME: Use this everywhere */
static long asInteger(double fval)
{
  if (fval > LONG_MAX || fval < LONG_MIN) {
    throw runtime_error("double value out of range for conversion to long");
  }
  long ival;
  if (fval > 0) {
    ival = (long) (fval + DBL_EPSILON);
  }
  else {
    ival = (long) (fval - DBL_EPSILON);
  }
  if (fabs(fval - ival) > DBL_EPSILON) {
    throw runtime_error("Invalid integer conversion");
  }
  return ival;
}

double Compiler::constFromTable(ParseTree const *p)
{
  // Try evaluating constant expression from data table
  map<string,SArray>::const_iterator i = _data_table.find(p->name());
  if (i == _data_table.end()) {
    return JAGS_NA;
  }
  Range range = getRange(p->parameters(), i->second.range());
  if (isNULL(range)) {
    return JAGS_NA;
  }
  else {
    // Range expression successfully evaluated
    if (range.length() > 1) {
      throw runtime_error("Vector value in constant expression");
    }
    long offset = i->second.range().leftOffset(range.lower());
    return i->second.value()[offset];
  }
}

double Compiler::constFromNode(ParseTree const*p)
{
  //Evaluate constant expression if it corresponds to a node with
  //a fixed value

  NodeArray *array = _symtab.getVariable(p->name());
  if (array) {
    // We can't call getConstantRange() because we aren't sure that
    // the range expression can be evaluated.
    Range subset_range = getRange(p->parameters(), array->range());
    if (isNULL(subset_range)) {
      return JAGS_NA;
    }
    else if (subset_range.length() > 1) {
      throw runtime_error(string("Vector value ") + p->name() +
			  print(subset_range) + " in constant expression");
    }
    else {
      Node *node = array->getSubset(subset_range);
      if (node && isObserved(node)) {
	return *node->data.value();
      }
      else {
	return JAGS_NA;
      }
    }
  }
  else {
    return JAGS_NA;
  }
}

bool Compiler::constantExpression(ParseTree const *p, double &value)
{
  /* 
     Try to evaluate a constant expression.  An expression can
     only be evaluated if it is a function of
     1) Counters, or
     2) a value given in the data table, or
     3) a fixed Node.
     Functions are not allowed, but inline operators ("*", "/", "+", "-") 
     are.

     If the expression can be evaluated, then we return true and the
     expression value is written to the value argument. If it cannot
     be evaluated (because it depends on missing data) then we return
     false.
  */
  Counter *counter;
  double arg1, arg2;
  std::vector<ParseTree*> const &parameters = p->parameters();

  switch (p->treeClass()) {
  case P_VAR: 
    // Is it a counter?
    counter = _countertab.getCounter(p->name());
    if (counter) {
      value = (*counter)[0];
      return true;
    }
    else {
      // Is it a variable? Try reading value from the data table
      value = constFromTable(p);
      if (value != JAGS_NA) {
	return true;
      }
      else {
	// Failing that, see if it is a node with a fixed value
	value = constFromNode(p);
	return (value != JAGS_NA);
      }
    }
    break;
  case P_VALUE:
    value =  p->value();
    return true;
    break;
  case P_OPERATOR:
    switch(p->getOperator()) {
    case OP_ADD:
      if (parameters.size() < 2) {
	throw logic_error("ADD must have at least two arguments in constant expression");
      }
      arg1 = 0;
      for (unsigned int i = 0; i < parameters.size(); ++i) {
	if (!constantExpression(parameters[i], arg2)) {
	  return false;
	}
	arg1 += arg2;
      }
      value = arg1;
      return true;
      break;
    case OP_SUBTRACT:
      if (parameters.size() != 2) {
	throw logic_error("SUBTRACT must have two arguments in constant expression");
      }
      if (!constantExpression(parameters[0], arg1)) {
	return false;
      }
      if (!constantExpression(parameters[1], arg2)) {
	return false;
      }
      value = arg1 - arg2;
      return true;
      break;
    case OP_MULTIPLY:
      if (parameters.size() < 2) {
	throw logic_error("MULTIPLY must have at least two arguments in constant expression");
      }
      arg1 = 1;
      for (unsigned int i = 0; i < parameters.size(); ++i) {
	if (!constantExpression(parameters[i], arg2)) {
	  return false;
	}
	arg1 *= arg2;
      }
      value = arg1;
      return true;
      break;
    case OP_DIVIDE:
      if (parameters.size() != 2) {
	throw logic_error("DIVIDE must have two arguments in constant expression");
      }
      if (!constantExpression(parameters[0], arg1)) {
	return false;
      }
      if (!constantExpression(parameters[1], arg2)) {
	return false;
      }
      value = arg1/arg2;
      return true;
      break;
    case OP_NEG:
      if (parameters.size() != 1) {
	throw logic_error("NEG must have 1 argument in constant expression");
      }
      if (!constantExpression(parameters[0], arg1)) {
	return false;
      }
      value = -arg1;
      return true;
      break;
    case OP_NONE:
      throw logic_error("Bad constant expression");
      break;
    }
    return false;
    break;
  default:
    throw logic_error("Expected variable, value or expression");
  }
}

bool Compiler::indexExpression(ParseTree const *p, long &value)
{
  /* 
     Evaluate an index expression.
     
     Index expressions occur in three contexts:
     1) In the limits of a "for" loop
     2) On the left hand side of a relation
     3) On the right hand side of a relation
     
     In cases 1) and 2) we expect to be able to evaluate the index
     expression, and an exception will be thrown if it cannot be
     evaluated (strict == true). In case 3) it is not an error if an
     index expression cannot be evaluated (strict == false), but we try
     anyway for efficiency reasons.
  */

  double fvalue;
  if (!constantExpression(p, fvalue)) {
    return false;
  }
  else {
    value = asInteger(fvalue);
    return true;
  }
}

Range Compiler::getRange(vector<ParseTree*> const &range_list, 
			 Range const &default_range)
{
  /* 
     Evaluate a range expression. If successful, it returns the range
     corresponding to the expression.  If unsuccessful (due to missing
     values) returns a null range.
     
     The default_range argument provides default values if the range
     expression is blank: e.g. foo[] or bar[,1].  The default range 
     may be a null range, in which case, missing indices will result in
     failure.
  */
  
  if (range_list.empty()) {
    /* An empty range expression implies the default range, if it exists,
       or a scalar value, if it does not */
    if (isNULL(default_range))
      return Index(1);
    else 
      return default_range;
  }

  // Check size and integrity of range expression
  unsigned int size = range_list.size();
  if (!isNULL(default_range) && size != default_range.ndim(false)) {
    throw logic_error("Default range does not match dimension of range expression");
  }
  for (unsigned int i = 0; i < size; ++i) {
    if (range_list[i]->treeClass() != P_RANGE) {
      throw logic_error("Malformed parse tree. Expected range expression");
    }
  }
  
  // Now step through and evaluate lower and upper index expressions
  Index lower(size), upper(size);
  for (unsigned int i = 0; i < size; i++) {
    switch (range_list[i]->parameters().size()) {
    case 0:
      // Empty index implies default range
      if (isNULL(default_range)) {
	return default_range;
      }
      lower[i] = default_range.lower()[i];
      upper[i] = default_range.upper()[i];
      break;
    case 1:
      // Single index implies lower == upper
      if (!indexExpression(range_list[i]->parameters()[0], lower[i])) {
	return Range();
      }
      else {
	upper[i] = lower[i];
      }
      break;
    case 2:
      if (!indexExpression(range_list[i]->parameters()[0], lower[i]) ||
	  !indexExpression(range_list[i]->parameters()[1], upper[i])) {
	return Range();
      }
      break;
    default:
      throw logic_error("Malformed parse tree in index expression");
    }
  }
  
  //FIXME: Give informative error message if we request invalid range
  return Range(lower, upper);
}

Range Compiler::VariableSubsetRange(ParseTree const *var)
{
  /*
    Get the range of a subset expression of a variable on the LHS of a
    relation.  This means that the subset expression must be constant.
  */
  if (var->treeClass() != P_VAR) {
    throw logic_error("Expecting variable expression");
  }  

  string const &name = var->name();
  if (_countertab.getCounter(name)) {
    throw runtime_error(string("Counter cannot appear on LHS of relation: ")
			+ name);
  }
  NodeArray *array = _symtab.getVariable(name);
  if (array) {
    // It's a declared node
    vector<ParseTree*> const &range_list = var->parameters();
    
    bool ok = true;
    if (range_list.empty()) {
      ok = (array->range().ndim(false) == 1);
    }
    else {
      ok = (range_list.size() == array->range().ndim(false));
    }
    if (!ok) {
      throw runtime_error(string("Dimension mismatch in subset expression ")
			  + "of variable " + name);
    }
    Range range = getRange(var->parameters(), array->range());
    if (isNULL(range)) {
      throw runtime_error(string("Missing values in subset expression ") 
			  + "of variable " + name);
    }
    return range;
  }
  else {
    // Undeclared node
    Range range = getRange(var->parameters(), Range());
    if (isNULL(range)) {
      throw runtime_error(string("Cannot evaluate subset expression for ")
			  + "undeclared variable " + name);
    }
    return range;
  }
}

Range Compiler::CounterRange(ParseTree const *var)
{
  /* The range expression for a counter differs from that of
     a variable in that it is
     1) one-dimensional
     2) may not be empty
     Further, no variables are created for counters in the
     Symbol Table
  */
  if (var->treeClass() != P_COUNTER) {
    throw logic_error("Expecting counter expression");
  }
  if (var->parameters().size() != 1) {
    throw logic_error("Invalid counter expression");
  }
  Range range();
  
  ParseTree const *prange = var->parameters()[0];
  if (prange->treeClass() != P_RANGE) {
    throw logic_error("Expecting range expression");
  }

  unsigned int size = prange->parameters().size();
  if (size < 1 || size > 2) {
    throw logic_error(string("Invalid range expression for counter")
		      + var->name());
  }
  long lower;
  if(!indexExpression(prange->parameters()[0], lower)) {
    throw runtime_error(string("Unable to evaluate lower index of counter ")
			+ var->name());
  }
  long upper;
  if (prange->parameters().size() == 2) {
    if (!indexExpression(prange->parameters()[1], upper)) {
      throw runtime_error(string("Unable to evaluate upper index of counter ")
			  + var->name());
    }
  }
  else {
    upper = lower;
  }
  
  if (lower > upper) {
    return Range();
  }
  else {
    Index ind_lower(1), ind_upper(1);
    ind_lower[0] = lower;
    ind_upper[0] = upper;
    
    return Range(ind_lower, ind_upper);
  }
}

Node* Compiler::VarGetNode(ParseTree const *var)
{
  if (var->treeClass() != P_VAR) {
    throw logic_error("Expecting variable expression");
  }
  NodeArray *array = _symtab.getVariable(var->name());
  Range range = VariableSubsetRange(var);
  Node *node = array->find(range);
  
  return node;
}

Node * Compiler::getSubSetNode(ParseTree const *var)
{
  if (var->treeClass() != P_VAR) {
    throw logic_error("Expecting variable expression");
  }

  NodeArray *array = _symtab.getVariable(var->name());
  if (array == 0) {
    throw runtime_error(string("Unknown variable ") + var->name());
  }
  Range subset_range = getRange(var->parameters(), array->range());
  if (isNULL(subset_range)) {
    return 0;
  }
  if (!array->range().contains(subset_range)) {
    throw runtime_error(string("Subset ") + var->name() + print(subset_range)
			+ " out of range");
  }
  Node *node = array->getSubset(subset_range);
  if (node == 0 && _strict_resolution) {
    throw runtime_error(string("Unable to resolve parameter ") + 
			array->name() + print(subset_range) +
                        " (one of its ancestors may be undefined)");
  }
  return node;
}

static void getMixtureSubsets(vector<pair<Index, Range> > &subsets,  
			      vector<SSI> const &limits,
			      Range const &default_range)
{
  unsigned int ndim = limits.size();
  //Count number of variable indices (nvi)
  int nvi = 0;
  for (unsigned int j = 0; j < ndim; ++j) {
    if (limits[j].node != 0)
      ++nvi;
  }

  // Create upper and lower bounds
  Index variable_offset(nvi), variable_lower(nvi), variable_upper(nvi);
  Index lower_index(ndim), upper_index(ndim);
  int k = 0;
  for (unsigned int j = 0; j < ndim; ++j) {
    if (limits[j].node != 0) {
      variable_offset[k] = j;
      variable_lower[k] = default_range.lower()[j];
      variable_upper[k] = default_range.upper()[j];
      ++k;
    }
    else {
      lower_index[j] = limits[j].lower;
      upper_index[j] = limits[j].upper;
    }
  }
  
  for (RangeIterator i(Range(variable_lower, variable_upper)); !i.atEnd(); 
       i.nextLeft()) {
    for (int k = 0; k < nvi; ++k) {
      lower_index[variable_offset[k]] = i[k];
      upper_index[variable_offset[k]] = i[k];
    }
    subsets.push_back(pair<Index, Range>(i, Range(lower_index, upper_index)));
  }
}

static void getStochasticParents(Node *node, 
				 set<StochasticNode*> &stoch_parents,
				 Graph &visited_nodes)
{
  StochasticNode *snode = dynamic_cast<StochasticNode*>(node);
  if (snode) {
    stoch_parents.insert(snode);
  }
  else {
    set<Node*> const &parents = node->parents();
    for (set<Node*>::iterator p = parents.begin(); p != parents.end(); ++p) {
      getStochasticParents(*p, stoch_parents, visited_nodes);
    }
    visited_nodes.add(node);
  }
}

Node * Compiler::getMixtureNode(ParseTree const *var)
{
  if (var->treeClass() != P_VAR) {
    throw logic_error("Expecting variable expression");
  }

  NodeArray *array = _symtab.getVariable(var->name());
  if (array == 0) {
    throw runtime_error(string("Unknown parameter: ") + var->name());
  }

  vector<ParseTree*> const &range_list = var->parameters();
  vector<SSI> limits;
  unsigned int ndim = array->range().ndim(false);
  if (range_list.size() != ndim) {
    throw runtime_error("Dimension mismatch taking variable subset of " + 
			var->name());
  }
  unsigned int nvi = 0; //Count number of variable indices
  for (unsigned int i = 0; i < ndim; ++i) {
    ParseTree const *range_element = range_list[i];
    if (range_element->treeClass() != P_RANGE) {
      throw runtime_error("Malformed range expression");
    }
    SSI ssi;
    ssi.node = 0;
    ParseTree const *p0, *p1;
    switch(range_element->parameters().size()) {
    case 0:
      // Index is empty, implying the whole range 
      ssi.lower = array->range().lower()[i];
      ssi.upper = array->range().upper()[i];
      break;
    case 1:
      // Single index: upper = lower
      p0 = range_element->parameters()[0];
      if(indexExpression(p0, ssi.lower)) {
	ssi.upper = ssi.lower;
      }
      else {
	ssi.node = getParameter(p0);
	if (ssi.node == 0)
	  return 0;
	else
	  ++nvi;
      }
      break;
    case 2:
      // Upper and lower indices
      p0 = range_element->parameters()[0];
      p1 = range_element->parameters()[1];
      if(indexExpression(p0, ssi.lower)) {
	if (!indexExpression(p1, ssi.upper)) {
	  return 0;
	}
      }
      else {
	ssi.node = getParameter(p0);
	if (getParameter(p1) != ssi.node)
	  return 0;
	else
	  ++nvi;
      }
      break;
    default:
      throw logic_error("Invalid range expression");
    }
    //Check validity of limits
    if (ssi.node == 0) {
      if (ssi.lower < array->range().lower()[i] ||
	  ssi.upper > array->range().upper()[i] ||
	  ssi.upper < ssi.lower)
	{
	  throw runtime_error("Requested invalid variable subset of " + 
			      var->name());
	}
    }
    limits.push_back(ssi);
  }
  
  //Check number of variable indices (nvi)
  if (nvi == 0) {
    throw logic_error("Trivial mixture node");
  }
  
  vector<pair<Index, Range> > ranges;  
  getMixtureSubsets(ranges, limits, array->range());
  vector<pair<Index, Node*> > subsets;  
  for (unsigned int i = 0; i < ranges.size(); ++i) {
    Node *subset_node = array->getSubset(ranges[i].second);
    if (!subset_node)
      return 0;
    subsets.push_back(pair<Index, Node*>(ranges[i].first, subset_node));
  }

  vector<Node*> indices;
  for (unsigned int j = 0; j < ndim; ++j) {
    if(limits[j].node) 
      indices.push_back(limits[j].node);
  }

  return _mixfactory.getMixtureNode(indices, subsets);
}

Node *Compiler::getArraySubset(ParseTree const *p)
{
  Node *node = 0;

  switch(p->treeClass()) {
  case P_VALUE:
    node = _constantfactory.getConstantNode(p->value());
    break;
  case P_VAR:
    {
      Counter *counter = _countertab.getCounter(p->name()); //A counter
      if (counter) {
	node = _constantfactory.getConstantNode((*counter)[0]);
      }
      else {
	NodeArray *array = _symtab.getVariable(p->name());
	if (array == 0) {
	  throw runtime_error(string("Unknown parameter ") + p->name());
	}
	Range subset_range = getRange(p->parameters(), array->range());
	if (isNULL(subset_range)) {
	  node = getMixtureNode(p); //A stochastic subset
	} 
	else {
	  node = getSubSetNode(p); //A fixed subset
	} 
      }
    }
    break;
  default:
    throw logic_error("Expecting value or variable expression");
  }
  return node;
}

static Function const *getLink(ParseTree const *t, FuncTab const &functab)
{
  if (t->treeClass() != P_LINK) {
    throw logic_error("Malformed parse tree: Expected link function");
  }

  Function const *func = functab.findInverse(t->name());
  if (func == 0) {
    string msg("Unable to find inverse of link function ");
    msg.append(t->name());
    throw runtime_error(msg);
  }
  else {
    return func;
  }
}

static Function const *getFunction(ParseTree const *t, FuncTab const &functab)
{
  Function const *func = 0;
  
  switch (t->treeClass()) {
  case P_FUNCTION:
    func = functab.find(t->name());
    if (func == 0) {
      string msg("Unable to find function ");
      msg.append(t->name());
      throw runtime_error(msg);
    }    
    break;
  case P_OPERATOR:
    switch(t->getOperator()) {
    case OP_ADD:
      func = functab.find("+");
      break;
    case OP_SUBTRACT:
      func = functab.find("-");
      break;
    case OP_MULTIPLY:
      func = functab.find("*");
      break;
    case OP_DIVIDE:
      func = functab.find("/");
      break;
    case OP_NEG:
      func = functab.find("NEG");
      break;
    case OP_NONE:
      throw logic_error("Bad operator expression");
      break;
    }
    if (func == 0) {
      throw logic_error("Unable to find operator");
    }
    break;
  default:
    throw logic_error("Malformed parse tree: Expected expression");
  }
  return func;
}

static Distribution const *getDistribution(ParseTree const *pstoch_rel, 
				           DistTab const &table)
{
  //Get the distribution from a stochastic relation
  
  if (pstoch_rel->treeClass() != P_STOCHREL) {
    throw logic_error("Malformed parse tree. Expecting stochastic relation");
  }
  
  // Get the distribution
  ParseTree const *pdist = pstoch_rel->parameters()[1];
  if (pdist->treeClass() != P_DENSITY) {
    throw logic_error("Malformed parse tree. Expected density expression");
  }
  Distribution const *dist = table.find(pdist->name());
  if (dist == 0) {
    throw runtime_error(string("Unknown distribution: ") + pdist->name());
  }
  return dist;
}

Node* Compiler::getParameter(ParseTree const *t)
{
  vector<Node*> parents;
  Node *node = 0;

  switch (t->treeClass()) {
  case P_VALUE:
    node =  _constantfactory.getConstantNode(t->value());
    break;
  case P_VAR:
    node = getArraySubset(t);
    break;
  case P_FUNCTION: case P_OPERATOR:
    if (getLogicalParameterVector(t, parents)) {
      node = _logicalfactory.getLogicalNode(getFunction(t, _functab), parents);
    }
    break;
  default:
    throw  logic_error("Malformed parse tree. Expected value, variable or expression");
    break;
  }

  /* Initialize deterministic nodes now, if they are functions of 
     data, to aid in node recycling */
  if (node && !node->isStochastic() && node->canInitialize() && 
      !node->isInitialized())
    {
      node->initialize();
    }
  return node;
}

bool Compiler::getLogicalParameterVector(ParseTree const *t,
					 vector<Node*> &parents)
{
  if (!parents.empty()) {
    throw logic_error("parent vector must be empty in getLogicalParameterVector");
  }

  switch (t->treeClass()) {
  case P_FUNCTION: case P_LINK: case P_OPERATOR:
    for (unsigned int i = 0; i < t->parameters().size(); ++i) {
      Node *node = getParameter(t->parameters()[i]);
      if (node) {
	parents.push_back(node);
      }
      else {
	parents.clear();
	return false;
      }
			  
    }
    break;
  default:
    throw logic_error("Invalid Parse Tree.  Expected function or operator.");
  }
  return true;
}

void Compiler::setStochasticParameters(ParseTree const *stoch_relation)
{
  if (stoch_relation->treeClass() != P_STOCHREL) {
    return;
  }

  // Get stochastic node
  ParseTree *var = stoch_relation->parameters()[0];
  Node *node = VarGetNode(var);
  StochasticNode *snode = dynamic_cast<StochasticNode*>(node);
  if (!snode) {
    throw logic_error(string("Stochastic node not defined:") + var->name());
  }

  // Set the parameters of the stochastic node
  ParseTree const *distribution = stoch_relation->parameters()[1];
  vector<ParseTree*> const &param_list = distribution->parameters();
  vector<Node*> parameters;

  for (unsigned int i = 0; i < param_list.size(); ++i) {
    Node *param = getParameter(param_list[i]);
    if (param) {
      parameters.push_back(param);
    }
    else {
      string msg = string("Parameter ") + ToString(i+1) + " of node " + 
        // var->name() + print(VariableSubsetRange(var)) + " undefined";
	_symtab.getName(node) + " undefined";
      throw runtime_error(msg);
    }
  }
  snode->setParameters(parameters);

  // Give upper and lower bounds, if truncated
  if (stoch_relation->parameters().size() == 3) {
    ParseTree const *truncated = stoch_relation->parameters()[2];
    Node *lBound = 0, *uBound = 0;
    ParseTree const *ll = truncated->parameters()[0];
    ParseTree const *ul = truncated->parameters()[1];
    if (ll) {
      lBound = getArraySubset(ll);
      if (!lBound) {
	string msg("Unable to resolve lower bound for node ");
	msg.append(_symtab.getName(node));
	throw logic_error(msg);
      }
    }
    if (ul) {
      uBound = getArraySubset(ul);
      if (!uBound) {
	throw logic_error("Unable to resolve upper bound for node " +
			  _symtab.getName(node));
      }
    }
    snode->setBounds(lBound, uBound);
  }
}

void Compiler::allocateStochastic(ParseTree const *stoch_rel)
{
  if (stoch_rel->treeClass() == P_STOCHREL) {
    ParseTree const *var = stoch_rel->parameters()[0];
    NodeArray *array = _symtab.getVariable(var->name());
    Range range = VariableSubsetRange(var);
    Node *node = new StochasticNode(getDistribution(stoch_rel, _disttab),
				    range.dim(false));
    array->insert(node, range);
  }
}

void Compiler::allocateLogical(ParseTree const *dtrm_rel)
{
  if (dtrm_rel->treeClass() != P_DETRMREL || _is_resolved[_nlogical]) {
    return;
  }
  
  ParseTree *expression = dtrm_rel->parameters()[1];
  Node *node = 0;
  vector <Node*> parents;

  switch (expression->treeClass()) {
  case P_VALUE: 
    node = new ConstantNode(expression->value());
    _graph.add(node);
    break;
  case P_VAR: case P_FUNCTION: case P_OPERATOR:
    node = getParameter(expression);
    break;
  case P_LINK:
    if (getLogicalParameterVector(expression, parents)) {
      node = _logicalfactory.getLogicalNode(getLink(expression, _functab), 
					    parents);
    }
    break;
  default:
    throw logic_error("Malformed parse tree");
  }
  if (node) {
    /* Check if a node is already inserted into this range */
    ParseTree *var = dtrm_rel->parameters()[0];
    NodeArray *array = _symtab.getVariable(var->name());
    Range range = VariableSubsetRange(var);
    if (Node const *node2 = array->find(range)) {
      throw runtime_error(string("Attempt to redefine node ") +
			  _symtab.getName(node2));
    }
    array->insert(node, range);
    _nresolved++;
    _is_resolved[_nlogical] = true;
  }
}

void Compiler::getArrayDim(ParseTree const *p)
{
  ParseTree const *var = p->parameters()[0];
  string const &name = var->name();

  if (_symtab.getVariable(name)) {
    return; //Node already declared
  }
  Range new_range = VariableSubsetRange(var);

  map<string, vector<Index> >::iterator i = _ranges.find(name);
  if (i == _ranges.end()) {
    //Create a new entry
    vector<Index> ivec;
    ivec.push_back(new_range.lower());
    ivec.push_back(new_range.upper());
    _ranges.insert(pair<const string, vector<Index> >(name,ivec));
  }
  else {
    //Check against the existing entry, and modify if necessary
    unsigned int ndim = i->second[0].size();
    if (new_range.ndim(false) != ndim) {
      throw runtime_error(string("Inconsistent dimensions for array ") + name);
    }
    else {
	for (unsigned int j = 0; j < ndim; ++j) {
	    i->second[0][j] = min(i->second[0][j], new_range.lower()[j]);
	    i->second[1][j] = max(i->second[1][j], new_range.upper()[j]);
      }
    }
  }
}

void Compiler::writeRelations(ParseTree const *relations)
{
  traverseTree(relations, &Compiler::allocateStochastic);
  _symtab.writeData(_data_table, true);

  // Set up boolean vector for logical nodes to indicate whether they are
  // resolved or not.
  _is_resolved = new bool[_nlogical];
  for (unsigned int i = 0; i < _nlogical; ++i) {
    _is_resolved[i] = false;
  }
  
  for (unsigned long N = _nlogical; N > 0; N -= _nresolved) {
    _nresolved = 0;
    traverseTree(relations, &Compiler::allocateLogical);
    if (_nresolved == 0) {
      // Try again, but this time throw an exception from getSubsetNode
      _strict_resolution = true;
      traverseTree(relations, &Compiler::allocateLogical);
      // If that didn't work (but it should!) just throw a generic message
      throw runtime_error("Unable to resolve logical relations");
    }
  }
  delete _is_resolved; _is_resolved = 0;

  traverseTree(relations, &Compiler::setStochasticParameters);
  collectNodes();
}

void Compiler::traverseTree(ParseTree const *relations, CompilerMemFn fun,
			    bool resetcounter)
{
  /* 
     Traverse parse tree, expanding for loops, applying function fun
     to relations.  
  */

  Counter *counter;
  ParseTree *var;

  if (resetcounter)
    _nlogical = 0;

  vector<ParseTree*> const &relation_list = relations->parameters();
  for (vector<ParseTree*>::const_iterator p = relation_list.begin(); 
       p != relation_list.end(); ++p) 
    {  
      switch ((*p)->treeClass()) {
      case P_FOR:
	var = (*p)->parameters()[0];
	relations = (*p)->parameters()[1];
	if (!isNULL(CounterRange(var))) {
	  Range cr = CounterRange(var);
	  //counter = _countertab.pushCounter(var->name(), CounterRange(var));
	  counter = _countertab.pushCounter(var->name(), cr);
	  // for (counter->reset(); !counter->atEnd(); counter->increment()) {
	  for (; !counter->atEnd(); counter->next()) {
	    traverseTree(relations, fun, false);
	  }
	  _countertab.popCounter();
	}
	break;
      case P_STOCHREL: 
	(this->*fun)(*p);
	break;
      case P_DETRMREL:
	(this->*fun)(*p);
	_nlogical++;
	break;
      default:
	throw logic_error("Malformed parse tree. Expected a relation or for loop");
	break;
      }
    }
}

Compiler::Compiler(Graph &graph, SymTab &symtab,
		   FuncTab const &functab, DistTab const &disttab,
		   map<string, SArray> const &data_table)
    : _graph(graph), _symtab(symtab), _countertab(), _functab(functab), 
      _disttab(disttab), _data_table(data_table), _nresolved(0), 
      _nlogical(0), _is_resolved(0), _strict_resolution(false)
{
  if (graph.size() != 0)
    throw invalid_argument("Non empty graph in Compiler constructor");
  if (symtab.size() != 0)
    throw invalid_argument("Non empty symtab in Compiler constructor");
}

void Compiler::declareVariables(vector<ParseTree*> const &dec_list)
{
  vector<ParseTree*>::const_iterator p;
  for (p = dec_list.begin() ; p != dec_list.end(); ++p) {
    if ((*p)->treeClass() != P_VAR) {
      throw invalid_argument("Expected variable expression");
    }
  }
  
  for (p = dec_list.begin() ; p != dec_list.end(); ++p) {
    ParseTree const *node_dec = *p;
    string const &name = node_dec->name();
    unsigned int ndim = node_dec->parameters().size();
    if (ndim == 0) {
      // Variable is scalar
      _symtab.addVariable(name);
    }
    else {
      // Variable is an array
      Index dim(ndim);
      for (unsigned int i = 0; i < ndim; ++i) {
	if (!indexExpression(node_dec->parameters()[i], dim[i])) {
	  throw runtime_error(string("Unable to calculate dimensions of node ")
			      + name);
	}
      }
      _symtab.addVariable(name, dim);
    }
  }
}

void Compiler::collectNodes()
{
  vector<Node*> nodes;
  _symtab.getNodes(nodes);
  _logicalfactory.graph().getNodes(nodes);
  _constantfactory.graph().getNodes(nodes);
  _mixfactory.graph().getNodes(nodes);

  for (vector<Node*>::iterator p = nodes.begin(); p != nodes.end(); ++p) {
    _graph.add(*p);
  }
}

void Compiler::undeclaredVariables(ParseTree const *prelations)
{
  // Get undeclared variables from data table
  map<string, SArray>::const_iterator p = _data_table.begin();
  for (; p != _data_table.end(); ++p) {
    string const &name = p->first;
    if (!_symtab.getVariable(name)) {
      _symtab.addVariable(name, p->second.dim(false));
    }
  }

  // Infer the dimension of remaining nodes from the relations
  traverseTree(prelations, &Compiler::getArrayDim);
  map<string, vector<Index> >::const_iterator i = _ranges.begin(); 
  for (; i != _ranges.end(); ++i) {
    _symtab.addVariable(i->first, i->second[1]);
  }
  _ranges.clear();
}