#include <config.h>
#include <sampler/ConjugateNormal.h>
#include <sampler/ConjugateSampler.h>
#include <graph/AggNode.h>
#include <graph/MixtureNode.h>
#include <graph/LogicalNode.h>
#include <graph/StochasticNode.h>

#include <set>
#include <stdexcept>
#include <vector>
#include <cmath>

#include <Rmath.h>

using std::vector;
using std::set;
using std::sqrt;
using std::invalid_argument;

ConjugateNormal::ConjugateNormal(StochasticNode *snode, Graph const &graph)
  : ConjugateSampler(snode, graph)
{
  if (!canSample(snode, graph)) {
    throw invalid_argument("Can't construct ConjugateNormal sampler");
  }
}

ConjugateNormal::~ConjugateNormal()
{
}

bool ConjugateNormal::canSample(StochasticNode *snode, Graph const &graph)
{
  /*
    Normal Conjugate samplers are more complex than other samplers
    because conjugacy is preserved by linear functions, and we want
    to be able to exploit this. 

    1) Stochastic children of snode must be normal, and depend on
    snode only via the mean parameter
    2) The mean parameter must be a linear function of snode. This
    means that any path from snode to the stochastic child must 
    consist of logical nodes representing functions +,-,*,/ and
    the identity function.
  */

  if (getDist(snode) != NORM)
    return false;

  vector<StochasticNode const*> stoch_nodes;
  vector<Node*> dtrm_nodes;
  classifyChildren(vector<StochasticNode*>(1,snode), 
		   graph, stoch_nodes, dtrm_nodes);

  /* 
     Create a set of nodes containing snode and its deterministic
     descendants for the checks below.
  */
  set<SArray const *> paramset;
  paramset.insert(&snode->data);
  for (unsigned int j = 0; j < dtrm_nodes.size(); ++j) {
    paramset.insert(&dtrm_nodes[j]->data);
  }

  // Check stochastic children
  for (unsigned int i = 0; i < stoch_nodes.size(); ++i) {
    if (getDist(stoch_nodes[i]) != NORM) {
      return false; //Not normal
    }
    if (stoch_nodes[i]->isBounded()) {
      return false;
    }
    vector<SArray const *> const &param = stoch_nodes[i]->parameters();
    if (paramset.count(param[1])) {
      return false; //Precision depends on snode
    }
  }
  
  // Check deterministic descendants
  for (unsigned int j = 0; j < dtrm_nodes.size(); ++j) {
    if (isLogical(dtrm_nodes[j])) {
      LogicalNode const *lnode = asLogical(dtrm_nodes[j]);
      vector<SArray const *> const &param = lnode->parameters();
      unsigned int nfactor = 0;
      switch(getOp(lnode)) {
      case ADD: case SUBTRACT: case NEG:
	break;
      case DIVIDE:
	if (paramset.count(param[1]))
	  return false; //reciprocal term
	break;
      case MULTIPLY:
	for (unsigned int k = 0; k < param.size(); ++k) {
	  nfactor += paramset.count(param[k]);
	}
	if (nfactor != 1)
	  return false; //quadratic or higher term
	break;
      default:
	return false;
      }
    }
    else if (isMixture(dtrm_nodes[j])) {
      // Check that indices do not depend on snode
      vector<Node *> const &index = asMixture(dtrm_nodes[j])->index();
      for (unsigned int i = 0; i < index.size(); ++i) {
	if (paramset.count(&index[i]->data))
	  return false;
      }
    }
    else {
      return false;
    }
  }

  return true; //We made it!
}

void ConjugateNormal::update()
{
  vector<StochasticNode const*> const &stoch_children = stochasticChildren();
  unsigned int nchildren = stoch_children.size();

  /* For convenience in the following computations, we shift the 
     origin to xold, the previous value of the node */
  const double xold = *node()->data.value();
  const double priormean = *node()->parameters()[0]->value() - xold; 
  const double priorprec = *node()->parameters()[1]->value(); 

  double A = priormean * priorprec; //Weighted sum of means
  double B = priorprec; //Sum of weights

  if (deterministicChildren().empty()) {
    /* Stochastic children are direct children of snode */
    for (unsigned int i = 0; i < nchildren; ++i) {
      double Y = *stoch_children[i]->data.value();
      double tau = *stoch_children[i]->parameters()[1]->value();
      A += (Y - xold) * tau;
      B += tau;
    }
  }
  else {
    /* The prior mean of any stochastic child depends on the value of
       snode (x say) via a linear function. We can write this as y =
       alpha + beta * (x - x.cur), where x.cur is the currrent value
       of snode.  In order to calculate the value of the coefficients
       alpha and beta, we evaluate the prior mean twice: once at x.cur
       and once at x.cur + 1.
    */
    double *alpha = new double[nchildren];
    for (unsigned int i = 0; i < nchildren; ++i) {
      alpha[i] = *stoch_children[i]->parameters()[0]->value();
    }
    double xnew = xold + 1;
    setValue(&xnew, 1);
    for (unsigned int i = 0; i < nchildren; ++i) {
      double Y = *stoch_children[i]->data.value();
      double tau = *stoch_children[i]->parameters()[1]->value();
      double beta = *(stoch_children[i]->parameters()[0]->value()) - alpha[i];
      A += (Y - alpha[i]) * tau * beta;
      B += tau * beta * beta;
    }

    // Tidy up    
    delete [] alpha;
  }
  
  // Draw the sample
  double postmean = xold + A/B;
  double postsd = sqrt(1/B);
  double xnew;

  if (node()->isBounded()) {
    SArray const *lb = node()->lowerBound();
    SArray const *ub = node()->upperBound();
    double plower = lb ? pnorm(*lb->value(), postmean, postsd, 1, 0) : 0;
    double pupper = ub ? pnorm(*ub->value(), postmean, postsd, 1, 0) : 1;
    double p = runif(plower, pupper);
    xnew = qnorm(p, postmean, postsd, 1, 0);
  }
  else {
    xnew = rnorm(postmean, postsd);  
  }
  setValue(&xnew, 1);
}