%load_ext cythonmagic
from matplotlib import pyplot as plt

%%cython -lgsl -lgslcblas
'''
Gibbs sampler for function:

f(x,y) = x x^2 \exp(-xy^2 - y^2 + 2y - 4x)

using conditional distributions:

x|y \sim Gamma(3, y^2 +4)
y|x \sim Normal(\frac{1}{1+x}, \frac{1}{2(1+x)})

Original version written by Flavio Coelho.
Tweaked by Chris Fonnesbeck.
Ported to CythonGSL Thomas V. Wiecki.
'''
cimport cython
from cython_gsl cimport *

import numpy as np
cimport numpy as np

from libc.math cimport sqrt

cdef gsl_rng *r = gsl_rng_alloc(gsl_rng_mt19937)

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
def gibbs(int N=20000, int thin=500):
    cdef: 
        double x = 0
        double y = 0
        Py_ssize_t i, j
        np.ndarray[np.float64_t, ndim=2] samples = np.empty((N, 2), dtype=np.float64)

    for i in range(N):
        for j in range(thin):
            x = gsl_ran_gamma(r, 3, 1.0 / (y * y + 4))
            y = gsl_ran_gaussian(r, 1.0 / sqrt(x + 1))
        samples[i, 0] = x
        samples[i, 1] = y
    return samples

posterior = gibbs()

plt.hexbin(posterior[:, 0], posterior[:, 1])

%%cython -l gsl -l gslcblas

cimport cython
from cython_gsl cimport *

import numpy as np
from numpy cimport *

cdef gsl_rng *r = gsl_rng_alloc(gsl_rng_mt19937)

def multinomial(ndarray[double, ndim=1] p, unsigned int N):
    cdef:
       size_t K = p.shape[0]
       ndarray[uint32_t, ndim=1] n = np.empty_like(p, dtype='uint32')
    
    # void gsl_ran_multinomial (const gsl_rng * r, size_t K, unsigned int N, const double p[], unsigned int n[])
    gsl_ran_multinomial(r, K, N, <double*> p.data, <unsigned int *> n.data)
    
    return n

sample = multinomial(np.array([.3, .2, .1, .1, .3], dtype='double'), 500)

plt.bar(range(5), sample, align='center')

import random, math, numpy

def gibbs_python(N=20000, thin=500):
    x = 0
    y = 0
    samples = np.empty((N, 2))
    
    for i in range(N):
        for j in range(thin):
            x = random.gammavariate(3, 1.0 / (y * y + 4))
            y = random.gauss(1.0/(x+1),1.0 / math.sqrt(x + 1))
        samples[i, 0] = x
        samples[j, 1] = y
    return samples

%timeit gibbs_python()

%timeit gibbs()