# pyhton2 - python3 compatibility
from __future__ import division
from __future__ import absolute_import
from __future__ import print_function
from __future__ import unicode_literals

# script nlLExp.py
''' example of negative Log Likelihood
         for exponential distribution
.. author:: Guenter Quast <g.quast@kit.edu>

'''

import sys, numpy as np, matplotlib.pyplot as plt

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

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

def fPoisson(x, lam):
# Poisson distribution
  k=np.round(x)
  return (lam**k)/np.exp(lam)/sp.gamma(k+1.)

def fGauss(x, mu=0., sigma=1.):
# Gauss distribution
    return (np.exp(-(x-mu)**2/2./sigma**2)/np.sqrt(2.*np.pi)/sigma)

def fExp(x, tau=1):
# exponential distribution
   return np.exp(-x/tau)/tau

# ---- end distributions ------

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

# 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

# ---------------------------------------------------------------
##### ---- main Program starts here -----
if __name__ == "__main__":

  print('*==* script ' + sys.argv[0]+ ' executing \n')

# get a random sample
  Ns=9      # sample size               
  tau = 1
  # draw exponentially distributed numbers  
  m = -tau * np.log(np.random.rand(Ns))  

# set parameter range to be considered:
  pmin, pmax = 0.1, 3. 
  parname=r'$\tau$'
  name = "negLogL Exponential"

# construct a subplot array, max. 5*5
  npl=np.int32(np.sqrt(Ns))
  npl=min(npl,25)
  fig, axarr = plt.subplots(npl, npl, figsize=(3. * npl, 3.1 * npl), num=name)

#
# plot likelihood function for 1, ..., Ns measurements  
  nppoints=250  # number of points for plot of parameter dependence
  p=np.linspace(pmin, pmax, nppoints, endpoint=True)
  print('*==*: sample of Gaussian random numbers \n', m) 
  i=0
  j=0
  for n in range(0, npl*npl): 
    axl=axarr[i,j]
    nL=negLogL(p, fExp, m[:n+1]) # first n measurements

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

  # analyse nlL for all measurements
  idx_min = np.argmin(nL)
  tau_hat = p[idx_min]
  print('Mimimum of Likelihood: %.4g'%tau_hat)
  print('Mean of Mesurements: %.4g'%m.mean())

  plt.show()