#!/usr/bin/env python
# ---- script bootstrapErrors.py
''' resample measurements many times (bootstrap method) to determine 
    - uncertainties
    - correlations
    of fit parameters
   
.. author:: Guenter Quast <g.quast@kit.edu>
'''
# --------------------------------------------------------------------------

# pyhton2 - python3 compatibility
from __future__ import division
from __future__ import absolute_import
from __future__ import print_function
from __future__ import unicode_literals
#
import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import curve_fit

import Histogram as h

# set data
xm = np.array([.05,0.36,0.68,0.80,1.09,1.46,1.71,1.83,2.44,2.09,3.72,4.36,4.60])
ym = np.array([0.35,0.26,0.52,0.44,0.48,0.55,0.66,0.48,0.75,0.70,0.75,0.80,0.90])
ye = np.array([0.06,0.07,0.05,0.05,0.07,0.07,0.09,0.1,0.11,0.1,0.11,0.12,0.1])

# define fit function
def poly2(x, a=1.0, b=0.0, c=0.0):
    return a * x**2 + b * x + c  

if __name__ == '__main__':  # --------------  execution starts here  ----------
# fit model parameters to data
  par, cov= curve_fit( poly2, xm, ym, sigma=ye, absolute_sigma=True )
  print("Fit parameters:\n", par)
  print("Covariance matrix:\n", cov)
  npar=len(par)

# show fit results
#  xp = np.linspace( 0., 5., 100 )
#  ffit = poly2( xp, par[0], par[1], par[2] )
#  plt.errorbar( xm, ym, yerr=ye, fmt='o' )
#  plt.plot( xp, ffit, '-' )
#  plt.xlim( 0, 5 )
#  plt.ylim( 0, 1 )
#  plt.show()

# "bootstrap method": repeat fit many times
#                     with random samplings of input data
  nsamples=10000
  pvec=[]
  for i in range(nsamples):
    ytmp = ym + np.random.randn(len(ye))*ye
    p, c = curve_fit( poly2, xm, ytmp, sigma=ye, absolute_sigma=True )
    pvec.append(p)
# turn into numpy ndarray
  pvec=np.array(pvec)

# plot results as array of sub-figures
  fig, axarr = plt.subplots(npar, npar, figsize=(3. * npar, 3.1 * npar))
  fig.tight_layout()
  fig.subplots_adjust(top=0.92, bottom=0.1, left=0.1, right=0.95,
                        wspace=0.2, hspace=0.3)  

  nb1=50 
  nb2=50
  ip = -1
  for i in range(0, npar):
    ip += 1
    jp = -1
    for j in range(0, npar):
      jp += 1
      if ip > jp:
        axarr[jp, ip].axis('off')      # set empty space
      elif ip == jp:
        ax=axarr[jp, ip]
        bc, be, _ = ax.hist(pvec[:, ip], nb1) # plot 1d-histogram
        ax.set_xlabel('par%i' %ip)
        ax.locator_params(nbins=5) # number of axis labels
        m, s = h.histstat(bc, be, False)  # calculate statistics
        ax.text(0.66, 0.85, '$\\mu=$%f \n $\\sigma=$%f'%(m, s),
                    transform=ax.transAxes,
                    backgroundcolor='white')
      else:
         # 2d-plto to visualize correlations
        ax=axarr[jp, ip]
        H, xe, ye, _ = ax.hist2d(pvec[:, jp], pvec[:, ip], nb2, cmap='Blues')
        ax.set_xlabel('par%i' %jp)
        ax.set_ylabel('par%i' %ip)
        ax.locator_params(nbins=5) # number of axis labels
        mx, my, vx, vy, cov, cor =h.hist2dstat(H,xe,ye,False)
        ax.text(0.33, 0.85, '$\\rho$=%.2f' %cor,
                    transform=ax.transAxes,
                    backgroundcolor='white')

  plt.show()