#include "BinaryLogit.h"
#include "Classify.h"

#include "KS.h"

#include <graph/StochasticNode.h>
#include <rng/TruncatedNormal.h>
#include <rng/RNG.h>

#include <cmath>

using std::exp;
using std::log;
using std::sqrt;

#define REG_PENALTY 0.001

//Sample from the left-truncated logistic distribution
static double llogit(double left, jags::RNG *rng, double mu)
{
    double u = rng->uniform();
    if (left - mu <= 3) {
	/* Sampling from the body. For left -> -Inf this tends to the
	   unconstrained sampling algorithm used by the R math
	   library */
	return mu + log((u + exp(left - mu))/(1.0 - u));
    }
    else {
	/* Sampling from the right tail. This algorithm is more stable
	   for left - mu >> 1. It's not obvious what the optimal
	   cutoff is but we use this branch for exp(left - mu) > 20 */
	return left + log1p(u * exp(mu - left)) - log(1 - u);
    }
}

//Right truncated logistic distribution
static double rlogit(double right, jags::RNG *rng, double mu)
{
    return - llogit(-right, rng, -mu);
}

namespace jags {
namespace glm {

    BinaryLogit::BinaryLogit(StochasticNode const *snode, unsigned int chain)
	: Outcome(snode, chain), _y(snode->value(chain)[0]), _z(0), _tau(1), _sigma2(1)
    {
	//fixme: sanity checks on snode
    }

    double BinaryLogit::value() const 
    {
	return _z;
    }

    double BinaryLogit::precision() const 
    {
	return _tau;
    }
    
    void BinaryLogit::update(RNG *rng)
    {
	/* Albert Chib update */
	if (_y) {
	    _z = llogit(0, rng, _lp);
	}
	else {
	    _z = rlogit(0, rng, _lp);
	}
	_sigma2 = sample_lambda(_z - _lp, rng);
	_tau = REG_PENALTY + 1/_sigma2;

    }

    void BinaryLogit::update(double mean, double var, RNG *rng)
    {
	/* Holmes-Held update */
	if (_y) {
	    _z = lnormal(0, rng, mean, sqrt(var + _sigma2));
	}
	else {
	    _z = rnormal(0, rng, mean, sqrt(var + _sigma2));
	}
    }

    bool BinaryLogit::canRepresent(StochasticNode const *snode)
    {
	//FIXME: Big overlap with BinaryProbit::canRepresent

	Node const *N = nullptr;

	switch(getFamily(snode)) {
	case GLM_BERNOULLI:
            break;
	case GLM_BINOMIAL:
            // We can also model binomial distribution if the denominator is fixed to be 1.                                                                                    
            N = snode->parents()[1];
            if (N->length() != 1)
		return false;
            if (!N->isFixed())
		return false;
            if (N->value(0)[0] != 1)
		return false;
            break;
	default:
            return false;
	}

	return getLink(snode) == LNK_LOGIT;
    }

}}
