# script animated_nlLCoin.py
''' example to produce random numbers (0/1) 
    as obtained by throwing a coin
    animated fraction of heads vs. number of throws
    shown is the likelihood function after each trial
.. author:: Guenter Quast <g.quast@kit.edu>

'''
import numpy as np, matplotlib.pyplot as plt, scipy.special as sp
import matplotlib.animation as anim
import sys

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

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
# ---------------------------------------------------------------

# generate the results of the coin flipping experiment
nthrow=250    # how often to throw
r=[]
Nh=0
for N in range(nthrow):
  if np.random.rand()>0.5:
    Nh+=1
  r.append(float(Nh)/(N+1.))

#create a figure for the experiment
fig=plt.figure('throwCoin', figsize=(7.5, 10.))
ax=fig.add_subplot(2,1,1)
ax.grid(True)
ax.set_xlabel('$N$ (number of trials)', size='large')
ax.set_ylabel('p ($N_{head} / N$)', size='large')
# plot expectation ...
x=np.linspace(0.,nthrow,nthrow)
ax.plot(x,0.5*np.ones(nthrow),'r--',lw=2)
# ... and error band
ax.plot(x[1:],0.5+0.5/np.sqrt(x[1:]),'b--',lw=2)
ax.plot(x[1:],0.5-0.5/np.sqrt(x[1:]),'b--',lw=2)
txt1 = ax.text(0.81, 0.93, ' ', color='darkblue',
            transform=ax.transAxes, size='medium')

# figure for negative log L
axl=fig.add_subplot(2,1,2)
axl.grid(True)
axl.plot([], [] , '-', color='steelblue', linewidth=3)
axl.set_ylim(0.,5.)
axl.set_xlabel(r'$p$', size='x-large')
axl.set_ylabel(r'$\Delta\, -{\rm ln}{\cal{L}}(p\,|\, N, k)$', size='large')
txtl=axl.text(0.75, 0.95, '', color='darkblue', 
           transform=axl.transAxes, size='medium')
pmin, pmax = 0.2, 0.8
nppoints=100
p=np.linspace(pmin, pmax, nppoints, endpoint=True)

def animate(i):
  # plot experiment and negLogL
  graphc, = ax.plot(range(i), r[:i], 'g.')
  txt1.set_text('N = %i' %(i+1))
  N=i+1
  k=r[i]*(i+1.)

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

  graphl,= axl.plot(p, nL-min(nL), '-', color='steelblue', linewidth=3)
  txtl.set_text('N= %i, k= %i'%(N, k))
  return graphc, graphl, txt1, txtl

# show as animated (=updating) graph
print('\n*==* script ' + sys.argv[0]+' executing')

ani=anim.FuncAnimation(fig, animate, nthrow, interval=75, blit=True,
                         init_func=None, fargs=None, repeat=False)
plt.show()
# save as movie
#ani.save('nlLCoin.mp4')