#include <u.h>
#include <libn.h>

// ----------------------------------------------------------------------------
// Internal structure

uint64 
rol64(uint64 x, int k)
{
	return (x << k) | (x >> (64 - k));
}

typedef struct Rng {
	uint64 s[4];
} Rng;

uint64
xoshiro256ss(Rng *state)
{
	uint64 *s     = state->s;
	uint64 result = rol64(s[1] * 5, 7) * 9;
	uint64 t      = s[1] << 17;

	s[2] ^= s[0];
	s[3] ^= s[1];
	s[1] ^= s[2];
	s[0] ^= s[3];

	s[2] ^= t;
	s[3] = rol64(s[3], 45);

	return result;
}

typedef struct Mix
{
	uint64 s;
} Mix;

uint64 
splitmix64(struct Mix *state) {
	uint64 result = state->s;

	state->s = result + 0x9E3779B97f4A7C15;
	result = (result ^ (result >> 30)) * 0xBF58476D1CE4E5B9;
	result = (result ^ (result >> 27)) * 0x94D049BB133111EB;
	return result ^ (result >> 31);
}

static Rng RNG;

// ----------------------------------------------------------------------------
// Exported functions

/* Initializes the global RNG */

error
rng·init(uint64 seed)
{
	Mix smstate = {seed};

    for (int i=0; i < 4; i++) 
        RNG.s[i] = splitmix64(&smstate);

    return 0;
}

/* Returns a random float64 between 0 and 1 */
double
rng·random(void)
{
    uint64 r = xoshiro256ss(&RNG); 
    return (double)r / (double)UINT64_MAX;
}

double
rng·exponential(double lambda)
{
    double f;

    f = rng·random();
    return -log(1 - f)/lambda;
}

double
rng·normal(void)
{
    double f;
    f = rng·random();

}

/* Returns true or false on success of trial */
bool
rng·bernoulli(double f)
{
    return rng·random() < f;
}

/* Returns a random integer between 0 and max
 * TODO: Modulo throws off uniformity 
 */
uint64
rng·randi(int max)
{
    uint64 r = xoshiro256ss(&RNG); 
    return r % max;
}

/*
 * Ahrens, J. H., & Dieter, U. (1982).
 * Computer Generation of Poisson Deviates from Modified Normal Distributions.
 */
static double factorial[10] = {1., 1., 2., 6., 24., 120., 720., 5040., 40320., 362880.};
static double coeffs[9] = {
    -.500000000, +.333333333, -.249999856,
    +.200011780, -.166684875, +.142187833,
    -.124196313, +.125005956, -.114265030,
};

static inline
double
log1pmx(double x, double off)
{
    int i;
    double r, t;

    if(-0.25 < x && x < 0.25) {
        r = 0;
        t = 1;
        for(i=0;i<arrlen(coeffs);i++) {
            r += coeffs[i]*t;
            t *= x;
        }

        return x*x*r;
    }
    return log(1+x) - off;
}

static inline
double
procf(double mu, double s, int64 K, double *px, double *py, double *fx, double *fy)
{
    double d, V, X;
    double w, b1, b2, c1, c2, c3, c0, c;

    w  = 0.3989422804014327/s;
    b1 = 0.041666666666666664/mu;
    b2 = 0.3*b1*b1;
    c3 = 0.14285714285714285*b1*b2;
    c2 = b2 - 15.*c3;
    c1 = b1 - 6.*b2 + 45.*c3;
    c0 = 1 - b1 + 3.*b2 - 15.*c3;
    c  = .1069/mu;

    if(K < 10) {
        *px = -mu;
        *py = pow(mu,K) / factorial[K];
    }else{
        d = 0.08333333333333333/K;
        d = d - 4.8*d*d*d;
        V = (mu - K) / K;

        *px = K*log1pmx(V,mu-K) - d;
        *py = 0.3989422804014327/sqrt(K);
    }

    X = (K - mu + 0.5)/s;
    *fx = -0.5*X*X;
    *fy = w*(((c3*X*X + c2)*X*X + c1)*X*X + c0);

    return c;
}

static inline
uint64
bigpoisson(double mu)
{
    int64  L,K;
    double G,s,d,U,E,T;
    double px,py,fx,fy,c;

    s = sqrt(mu);
    d = 6*mu*mu;
    L = floor(mu - 1.1484);

stepN:
    G = mu + s*rng·normal();
    K = floor(G);
    if(K<0)
        goto stepP;
stepI:
    if(K>=L)
        return K;
stepS:
    U = rng·random();
    if(d*U >= (mu-K)*(mu-K)*(mu-K))
        return K;
stepP:
    if(G < 0)
        goto stepE;
stepQ:
    c = procf(mu, s, K, &px, &py, &fx, &fy);
stepE:
    E = rng·exponential(1.0);
    U = rng·random();
    U = U + U - 1;
    T = 1.8 + copysign(E,U);
    if(T < 0.6744)
        goto stepE;
    K = floor(mu + s*T);
    c = procf(mu, s, K, &px, &py, &fx, &fy);
stepH:
    if(c*fabs(U) > (py*exp(px + E) - fy*exp(fx + E)))
        goto stepE;
    return K;
}

uint64
rng·poisson(double mean)
{
    int64  n;
    double z;

    if(mean<10.0) {
        for(n=0, z=rng·exponential(1.0); z<mean; ++n, z+=rng·exponential(1.0))
            ;
        return n;
    }

    return bigpoisson(mean);
}