#include <config.h>
#include <sampler/ConjugateGamma.h>
#include <sampler/ConjugateSampler.h>
#include <graph/LogicalNode.h>
#include <graph/StochasticNode.h>
#include <graph/MixtureNode.h>
#include <graph/NodeError.h>

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

#include <Rmath.h>

using std::vector;
using std::set;
using std::sqrt;
using std::invalid_argument;
using std::logic_error;
using std::max;
using std::sort;

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

ConjugateGamma::~ConjugateGamma()
{
}

bool ConjugateGamma::canSample(StochasticNode *snode, Graph const &graph)
{
  switch (getDist(snode)) {
  case GAMMA: case EXP: case CHISQ:
    /* 
       The exponential and chisquare distributions are both special
       cases of the gamma distribution and are handled by the conjugate
       gamma sampler.
    */
    break;
  default:
    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) {
    vector<SArray const *> const &param = stoch_nodes[i]->parameters();
    if (stoch_nodes[i]->isBounded()) {
      return false; //Bounded
    }
    switch(getDist(stoch_nodes[i])) {
    case EXP: case POIS:
      break;
    case GAMMA: case NORM: case DEXP: case WEIB:
      if (paramset.count(param[0])) {
	return false; //non-scale parameter depends on snode
      }
      break;
    default:
      return false;
    }
  }
  
  // Check deterministic descendants
  for (unsigned int j = 0; j < dtrm_nodes.size(); ++j) {
    LogicalNode const *lnode = asLogical(dtrm_nodes[j]);
    if (lnode) {
      unsigned int nfactor = 0;
      vector<SArray const *> const &param = lnode->parameters();
      switch(getOp(lnode)) {
      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!
}

static double getScale(StochasticNode const *snode, ConjugateDist d)
{
  //Get scale parameter of snode
  switch(d) {
  case GAMMA: case NORM: case DEXP: case WEIB:
    return *snode->parameters()[1]->value();
    break;
  case EXP: case POIS:
    return *snode->parameters()[0]->value();
    break;
  default:
    throw NodeError(snode, "Can't get scale parameter: invalid distribution");
  } 
}

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

  double r; // shape
  double mu; // 1/scale
  double xold = *node()->data.value();

  //Prior
  vector<SArray const *> const &param = node()->parameters();
  switch(_target_dist) {
  case GAMMA:
    r = *param[0]->value();
    mu = *param[1]->value();
    break;
  case EXP:
    r = 1;
    mu = *param[0]->value();
    break;
  case CHISQ:
    r = *param[0]->value()/2;
    mu = 1/2;
    break;
  default:
    throw logic_error("invalid distribution in ConjugateGamma sampler");
  }

  /* 
     We know that the scale parameter for each stochastic child is
     either a multiple of x, the value of the sampled node, or a
     constant independent of x (in the case of a mixture model).
     Before calculating the likelihood we need to work out the
     coefficients.
  */
  double *C = new double[nchildren];
  if (deterministicChildren().empty()) {
    // All coefficients are 1
    for (unsigned int i = 0; i < nchildren; ++i) {
      C[i] = 1;
    }
  }
  else {
    for (unsigned int i = 0; i < nchildren; ++i) {
      C[i] = -getScale(stoch_children[i], _child_dist[i]);
    }
    double val = xold + 1;
    setValue(&val, 1);
    for (unsigned int i = 0; i < nchildren; ++i) {
      C[i] += getScale(stoch_children[i], _child_dist[i]);
    }
  }

  // likelihood 
  for (unsigned int i = 0; i < nchildren; ++i) {
    /* 
       C[i] == 0 can only occur when there is a mixture node blocking
       the path from the sampled node to stochastic child i. Such
       children should be ignored.
    */
    if (C[i] > 0) {

      StochasticNode const *schild = stoch_children[i];
      vector<SArray const *> const &cparam = schild->parameters();
      double Y = *schild->data.value();
      double ymean; //normal mean
      switch(_child_dist[i]) {
      case GAMMA:
	r += *cparam[0]->value();
	mu += C[i] * Y;
	break;
      case EXP:
	r += 1;
	mu += C[i] * Y;
	break;
      case NORM:
	r += 0.5;
	ymean = *cparam[0]->value();
	mu += C[i] * (Y - ymean) * (Y - ymean) / 2;
	break;
      case POIS:
	r += Y;
	mu += C[i];
	break;
      case DEXP:
	r += 1;
	ymean = *cparam[0]->value();
	mu += C[i] * fabs(Y - ymean);
	break;
      case WEIB:
	r += 1; 
	mu += C[i] * pow(Y, *cparam[0]->value());
	break;
      default:
	throw logic_error("Invalid distribution in Conjugate Gamma sampler");
      }
    }
  }
  delete [] C;

  // Sample from the posterior
  double xnew;
  if (node()->isBounded()) {
    // Use inversion to get random sample
    double lower = 0;
    SArray const *lb = node()->lowerBound();
    if (lb) {
      lower = max(lower, *lb->value());
    }
    SArray const *ub = node()->upperBound();
    double plower = lb ? pgamma(lower,        r, 1/mu, 1, 0) : 0;
    double pupper = ub ? pgamma(*ub->value(), r, 1/mu, 1, 0) : 1;
    double p = runif(plower, pupper);
    xnew = qgamma(p, r, 1/mu, 1, 0);    
  }
  else {
    xnew = rgamma(r, 1/mu);
  }
  setValue(&xnew, 1);  
}