#include <config.h>
#include <model/Model.h>
#include <model/TraceMonitor.h>
#include <sampler/Sampler.h>
#include <sampler/SamplerFactory.h>
#include <graph/GraphMarks.h>
#include <graph/StochasticNode.h>
#include <graph/NodeError.h>
#include <graph/Node.h>

#include <fstream>
#include <sstream>
#include <set>
#include <stdexcept>
#include <string>
#include <algorithm>
#include <functional>
#include <map>

#include <Rmath.h>

using std::map;
using std::pair;
using std::binary_function;
using std::sort;
using std::vector;
using std::list;
using std::set;
using std::ofstream;
using std::logic_error;
using std::runtime_error;
using std::string;
using std::ostringstream;

Model::Model()
{
  _iteration = 0;
  _initialized = false;
  _cansample = false;
}

Model::~Model()
{
  while(!_samplers.empty()) {
    Sampler *sampler = _samplers.back();
    delete sampler;
    _samplers.pop_back();
  }

  while(!_monitors.empty()) {
    TraceMonitor *monitor = _monitors.back();
    delete monitor;
    _monitors.pop_back();
  }
}

Graph &Model::graph() 
{
  return _graph;
}

bool Model::isInitialized()
{
  return _initialized;
}

void Model::initialize()
{
  if (_initialized)
    throw logic_error("Model already initialized");
  if (!_cansample)
    throw logic_error("Graph not checked yet");

  /* Count the uninitialized nodes */
  unsigned int n = _nodes.size();
  for (vector<Node*>::iterator i = _nodes.begin(); i != _nodes.end(); i++)
    {
      if ((*i)->isInitialized())
	n--;
    }

  for (vector<Node*>::iterator i = _nodes.begin(); i != _nodes.end(); i++)
    {
      Node *node = *i;
      if (!node->isInitialized()) {
	node->initialize();
	if (!node->isInitialized()) {
	  throw NodeError(node, "Initialization failure");
	}
	n--;
      }
    }

  if (n != 0) {
    throw logic_error("Model::initialize failed");
  }
  _initialized = true;
}



void Model::checkGraph()
{
  if (_cansample) {
    throw logic_error("Already checked ability to sample");
  }
  if (!_graph.isClosed())
    throw runtime_error("Graph not closed");
  if (_graph.hasCycle()) 
    throw runtime_error("Directed cycle in graph");
  if (_nodes.empty()) {
    _graph.getSortedNodes(_nodes);
  }

  _cansample = true;
}

struct less_sampler {  
  /* 
     Comparison operator for Samplers which sorts them according to
     the partial ordering defined by the DAG, using the first node
     in the vector of nodes sampled by the sampler
  */
  map<Node const*, int> const &_node_map;

  less_sampler(map<Node const*, int> const &node_map) : _node_map(node_map) {};
  
  bool operator()(Sampler *x, Sampler *y) const {
    int indx = _node_map.find(x->nodes()[0])->second;
    int indy = _node_map.find(y->nodes()[0])->second;
    //return indx < indy; //Forward sampling order
    return indx > indy; //Backward sampling order
  };
};

void Model::chooseSamplers(vector<SamplerFactory const *> const &samplers)
{
  if (!_samplers.empty())
    throw logic_error("Samplers already chosen");
  if (!_cansample) 
    throw logic_error("Graph not checked");

  // Mark observed nodes
  GraphMarks marks(_graph);
  vector<Node*>::reverse_iterator i = _nodes.rbegin();
  for (; i != _nodes.rend(); ++i) {
    if (isObserved(*i)) {
      marks.mark(*i,2);
    }
  }
    
  // Now mark ancestors of observed nodes
  for (i = _nodes.rbegin(); i != _nodes.rend(); ++i) {
    if (marks.mark(*i) != 2) {
      for (set<Node*>::const_iterator ch = (*i)->children().begin(); 
	   ch != (*i)->children().end(); ++ch) 
	{
	  if (marks.mark(*ch) != 0) {
	    marks.mark(*i,1);
	    break;
	  }
	}
    }
  }

  // Create set of unobserved stochastic nodes, for which we need
  // to find a sampler,  a graph within which sampling will take
  // place (excluding uninformative nodes), and a set of "extra"
  // uninformative nodes that will be updated by the model at the
  // end of every iteration.
  set<StochasticNode*> stochastic_nodes;
  Graph graph;
  for (vector<Node*>::iterator i = _nodes.begin(); i != _nodes.end(); i++) {
    switch(marks.mark(*i)) {
    case 0:
      _extra_nodes.insert(*i);
      break;
    case 1:
      graph.add(*i);
      if ((*i)->isStochastic()) {
	stochastic_nodes.insert(dynamic_cast<StochasticNode*>(*i));
      }
      break;
    case 2:
      graph.add(*i);
      break;
    }
  }
      
  // Traverse the list of samplers, selecting nodes that can be sampled
  for (vector<SamplerFactory const *>::const_iterator p = samplers.begin();
       p != samplers.end(); ++p)
    {
      (*p)->makeSampler(stochastic_nodes, graph, _samplers);
    }
    
  // Make sure we found a sampler for all the nodes
  if (!stochastic_nodes.empty()) {
    throw NodeError(*stochastic_nodes.begin(),
		    "Unable to find appropriate sampler");
  }

  /**
   * Now sort the samplers in order
   *
   * The map node_map associates each node in the graph with its index
   * in the vector of sorted nodes. This is used by the comparison
   * operator less_sampler.
   */
  static map <Node const *, int> node_map;
  int index = 0;
  for (vector<Node*>::iterator i = _nodes.begin(); i != _nodes.end(); i++)
    {
      node_map.insert(pair<Node const*, int>(*i, index));
      index++;
    }
  sort(_samplers.begin(), _samplers.end(), less_sampler(node_map));
}

void Model::update(long niter)
{
  if (!_initialized) {
    throw logic_error("Attempt to update uninitialized model");
  }

  for (int iter = 0; iter < niter; ++iter) {
    for (vector<Sampler*>::iterator i = _samplers.begin(); 
         i != _samplers.end(); ++i) 
      {
	(*i)->update();
      }
    for (vector<Node*>::const_iterator k = _sampled_extra.begin();
	 k != _sampled_extra.end(); ++k)
      {
	(*k)->forwardSample();
      }
    _iteration++;
    for (list<TraceMonitor*>::iterator k = _monitors.begin(); 
	 k != _monitors.end(); k++) 
      {
	(*k)->update(_iteration);
      }
  }
}

long Model::iteration() const
{
  return _iteration;
}

static void addAncestors(Node *node, Graph &to, set<Node*> const &from)
{
  // Take ancestors of "node" belonging to set "from" and add them to
  // graph "to", along with "node" itself.
  if (from.count(node) == 0 || to.contains(node)) {
    return;
  }
  to.add(node);
  for (set<Node*>::iterator p = node->parents().begin(); 
       p != node->parents().end(); ++p) 
    {
      addAncestors(*p, to, from);
    }
}

TraceMonitor const * Model::setMonitor(Node *node, int thin)
{
  if (_monitored_nodes.count(node))
    throw NodeError(node, "Node already being monitored");

  if (_monitors.empty()) {
    // The first monitor: turn off burnin mode.
    for (vector<Sampler*>::iterator p = _samplers.begin();
	 p != _samplers.end(); ++p)
      {
	(*p)->burninOff();
      }
  }

  TraceMonitor *monitor = new TraceMonitor(node, _iteration + 1, thin);
  _monitors.push_back(monitor);
  _monitored_nodes.insert(node);

  // Recalculate the vector of uninformative nodes that need sampling
  Graph egraph;
  for (set<Node*>::const_iterator j = _monitored_nodes.begin();
       j != _monitored_nodes.end(); ++j)
    {
      addAncestors(*j, egraph, _extra_nodes);
    }
  _sampled_extra.clear();
  egraph.getSortedNodes(_sampled_extra);

  return monitor;
}

void Model::clearMonitor(Node const *node)
{
  for (list<TraceMonitor*>::iterator j = _monitors.begin();
       j != _monitors.end(); j++) 
    {
      if ((*j)->node() == node) {
	_monitors.erase(j);
	return;
      }
    }
}

static string printIndex (string const &name, Index const &index)
{
  ostringstream ostr;
  ostr << name << "[";
  for (unsigned int i = 0; i < index.size(); ++i) {
    if (i > 0) 
      ostr << ",";
    ostr << index[i];
  }
  ostr << "]";
  return ostr.str();
}



list<TraceMonitor*> const &Model::monitors() const
{
  return _monitors;
}

void Model::addExtraNode(Node *node)
{
  if (!_initialized) 
    throw logic_error("Attempt to add extra node to uninitialized model");
  if (isObserved(node)) {
    throw logic_error("Cannot add observed node to initialized model");
  }
  if (!node->children().empty()) {
    throw logic_error("Cannot add extra node with children");
  }
  if (_graph.contains(node)) {
    throw logic_error("Extra node already in model");
  }
  for (set<Node*>::const_iterator p = node->parents().begin(); 
       p != node->parents().end(); ++p)
    {
      if (!_graph.contains(*p)) {
	throw logic_error("Extra node has parents not in model");
      }
    }

  _extra_nodes.insert(node);
  _graph.add(node);
}