# script nlLCoin.py
''' example to produce random numbers (0/1) 
    as obtained by throwing a coin
.. author:: Guenter Quast <g.quast@kit.edu>

'''

import numpy as np, matplotlib.pyplot as plt, scipy.special as sp

# ------------------------------------------------------

#Binomial distribution
def fBinomial(x, p, N):
#Binomial distribution
  k=np.round(x)
  return sp.binom(N,k) * p**k * (1.-p)**(N-k)

# one possibility to define concrete negLogL function for Binomial distribution
def nlLBinomial(p, k, N):
#neg logl L of Binomial distribution
  return -np.log(fBinomial(k, p, N))
  

# ---------------------------------------------------------------
# general negLogL Function 
def negLogL(p, f, x, *args, **kwargs):
# calulate neg logl L for lists of 
#   - parameters p_i and 
#   - observations x_i
#  for a pdf f(x, p, <other parameters>)
#  
  nlL=np.zeros(len(p))
  for i, p_i in enumerate(p):
    nlL[i]=np.sum(-np.log(f(x, p_i, *args, **kwargs) ) )
  return nlL

# ---------------------------------------------------------------

nthrow=500   # how often to throw

# create figure and plot
fig = plt.figure('throwCoin', figsize=(7.5, 5.) )
axd = fig.add_subplot(1,1,1)

r=[]
Nh=0
for N in range(nthrow):
  if np.random.rand()>0.5:
    Nh+=1
  r.append(float(Nh)/(N+1.))


# plot result ...
x=np.linspace(0., nthrow, nthrow)
axd.plot(x,r,'g.')
# ... and expectation ...
axd.plot(x,0.5*np.ones(nthrow),'r--',lw=2)
# ... and error band
axd.plot(x[1:], 0.5+0.5/np.sqrt(x[1:]),'b--',lw=2)
axd.plot(x[1:], 0.5-0.5/np.sqrt(x[1:]),'b--',lw=2)
axd.set_xlabel('N (number of trials)', size='x-large')
axd.set_ylabel('p ($N_{head} / N$)', size='x-large')

#
# calculate negative log Likelihood L(p|k,N) and plot
npl=4
name = "negLogL Binomial"
fig, axarr = plt.subplots(npl, npl, figsize=(3. * npl, 3.1 * npl), num=name)

pmin, pmax = 0.2, 0.8
nppoints=100
p=np.linspace(pmin, pmax, nppoints, endpoint=True)

Ns = [1, 2, 3, 4, 5, 8, 16, 32, 50, 150, 200, 250, 300, 350, 400, 500]
i=0
j=0
for n in range(0, npl*npl): 
  axl=axarr[i,j]
  N=Ns[n]  
  k=r[N-1]*N

  nL = negLogL(p, fBinomial, k, N) # use general version 
####  nL=nlLBinomial(p, k, N)

  axl.plot(p, nL-min(nL), '-', color='steelblue', linewidth=3)
  axl.text(0.25, 0.9, 'N= %i, k= %i'%(N, k), color='darkblue', 
          transform=axl.transAxes, size='medium')
  axl.set_ylim(0.,5.)
  if(i==(npl-1) and j==0):
    axl.set_xlabel(r'$p$', size='x-large')
    axl.set_ylabel(r'$\Delta\, -{\rm ln}{\cal{L}}(p\,|\, N, k)$', size='large')
  j+=1
  if j==npl: 
    i+=1
    j=0

plt.show()