from __future__ import division

import sympy as S
x=S.symbols('x')
F = sum([S.Heaviside(x-i+1) for i in range(1,7)])/6 # offset to satisfy sympy definition
S.plot(F,(x,0,6),ylabel='CDF(x)',xlabel='x');

invF=S.Piecewise((1,0<x<=1/6),    # if uniform sample between 0 and 1/6, choose die-side labeled 1
                 (2,1/6<x<=2/6),
                 (3,2/6<x<=3/6),
                 (4,3/6<x<=4/6),
                 (5,4/6<x<=5/6),
                 (6,5/6<x<=1))

samples=array([invF.subs(x,i) for i in  rand(500)])
hist(samples,bins=[1,2,3,4,5,6,7],normed=1,align='left');
axis(ymax=2/6.)
title('Estimated PMF of Fair Six-Sided Die');

bar(-.5+arange(1,7),cumsum([mean(samples==i) for i in range(1,7)]),alpha=.3,label='ECDF')   # estimated CDF
bar(-.25+arange(1,7),[F.subs(x,i).evalf() for i in range(1,7)],fc='g',alpha=.3,width=.5,label='CDF') # target CDF
legend(loc=0)
title('Estimated CDF for fair six-sided die',fontsize=18)
ylabel('Frequency',fontsize=18)
xlabel('Die face index',fontsize=18)
axis(ymax=1.1)

p = array([1,2,9,2,1,1]) # the faces 2,4 and more probable than 1,5,6 and 3 is the most probable
p = p/sum(p) # normalize to 1
Fu = sum([S.Heaviside(x-i+1)*p[i-1] for i in range(1,7)])
S.plot(Fu,(x,0,6),ylabel='CDF(x)',xlabel='x',title='Unfair 6-sided Die');

cp=cumsum([0]+p.tolist()) # need to find edges on vertical axis, add [0] to get left edge
invFu=S.Piecewise(*[(i,j<x<=k) for i,j,k in zip(range(1,7),cp[:-1],cp[1:])])
hist([invFu.subs(x,i) for i in  rand(1000)],bins=[1,2,3,4,5,6,7],normed=1,align='left');
axis(ymax=1.)
title('Estimated PMF of Unfair Six-Sided Die');

import scipy.stats

alpha = 1.
nsamp = 1000

u=scipy.stats.uniform(0,1)
x=scipy.stats.expon(alpha)

Finv=lambda u: 1/alpha*log(1/(1-u))
#Finv=lambda u: 1+1/alpha*log(1/(1-u)) # shift over to correct
    
fig,ax = subplots()

xhat = array(map(Finv,u.rvs(nsamp)))
# exponential distrib samples by inverse method
ax.hist(xhat,normed=1,bins=30,alpha=0.3,label='inverse method') 
# exponential distrib samples by scipy.stats
xrvs = x.rvs(nsamp)
xe = linspace(0,xrvs.max(),100)
ax.hist(xrvs,bins=30,normed=1,alpha=0.3,label='scipy.stats')
# exponential theoretical density function
ax.plot(xe,x.pdf(xe),'r-',label='theoretical')
ax.set_xticks(range(int(xrvs.max())))
ax.axis(ymax=1.05)
ax.set_title('Inverse CDF for Exponential Distribution')
ax.set_ylabel('pdf(x)',fontsize=18)
ax.set_xlabel('x',fontsize=18)
ax.legend();

x = linspace(0.001,15,100)
fig,ax=subplots()
f= lambda x: exp(-(x-1)**2/2./x)*(x+1)/12.
fx = f(x)
ax.plot(x,fx,label='$f(x)$')
ax.plot(x,cumsum(fx)*diff(x)[0],'g',label='$F(x)$')
ax.legend(loc=0,fontsize=16);

M=.3 # scale factor
u1 = rand(10000)*15 # uniform random samples scaled out
u2 = rand(10000)    # uniform random samples
idx=where(u2<=f(u1)/M)[0] # rejection criterion
v = u1[idx]

fig,ax=subplots()
ax.hist(v,normed=1,bins=40,alpha=.3)
ax.plot(x,fx,'r',lw=3.,label='$f(x)$')
ax.set_title('Estimated Efficency=%3.1f%%'%(100*len(v)/len(u1)))
ax.legend(fontsize=18)

fig,ax=subplots()
ax.plot(u1,u2,'.',label='rejected',alpha=.3)
ax.plot(u1[idx],u2[idx],'g.',label='accepted',alpha=.3)
ax.legend(fontsize=16)

ch=scipy.stats.chi2(4) # chi-squared
h = lambda x: f(x)/ch.pdf(x) # h-function

fig,axs=subplots(1,2)
fig.set_size_inches(12,4)
axs[0].plot(x,fx,label='$f(x)$')
axs[0].plot(x,ch.pdf(x),label='$g(x)$')
axs[0].legend(loc=0,fontsize=18)
axs[1].plot(x,h(x))
axs[1].set_title('$h(x)=f(x)/g(x)$',fontsize=18)

hmax=h(x).max()
u1 = ch.rvs(5000) # samples from chi-square distribution
u2 = rand(5000)   # uniform random samples
idx = (u2 <= h(u1)/hmax)  # Rejection criterion

v = u1[idx]  # keep these only

fig,ax=subplots()
ax.hist(v,normed=1,bins=40,alpha=.3)
ax.plot(x,fx,'r',lw=3.,label='$f(x)$')
ax.set_title('Estimated Efficency=%3.1f%%'%(100*len(v)/len(u1)))
ax.legend(fontsize=18)

fig,ax=subplots()
ax.plot(u1,u2,'.',label='rejected',alpha=.3)
ax.plot(u1[idx],u2[idx],'g.',label='accepted',alpha=.3)
ax.plot(x,h(x)/hmax,'r',lw=3.,label='$h(x)$')
ax.legend(fontsize=16)