# -*- coding: utf-8 -*-
#      diese Zeile legt die Codierung von Umlauten fest
##################################################################

# script PlotAverage-withChi2.py
''' 
   Mittelwert mit Chi2-Methode

.. author:: Günter Quast <g.quast@kit.edu> 
     für den Kurs Computergestützte Datenauswertung (CgDA)
'''
#--------------------------------------------------------------

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

FitFuncName='Const $c$'
Npar=1
def FitFunc(x, av):
    return np.ones(len(x)) * av

def calcS(f, x, y, sigy, par):
  return np.sum( ( (y-f(x, par) ) /sigy )**2)

##### ---- main Program starts here -----

# genearate "data"
N=10
av=10.
Ym=av + np.random.randn(N)
Yerr=1.
Xm=np.array(range(N)) + 1.
## Ein numpy-array, das die Stuetzstellen enthaelt (np.pi ist pi=3.141..)
xmin, xmax=0.75, N+0.25
npoints=100
X=np.linspace( xmin, xmax, npoints, endpoint=True)

fig=plt.figure(figsize=(7.5, 10))
ax=fig.add_subplot(2,1,1)

## Plots der Funktion generieren
for par in (9., 9.5, 10., 10.5, 11.):
  ax.plot( X, FitFunc(X, par), linewidth=2, label=FitFuncName+"="+str(par))

## Plot noch etwas polieren
ax.set_title('Constant Function')
ax.set_xlabel('number of measurement')
ax.set_ylabel('Measurement $m$')
ax.set_xlim(0.5, 10.5)
ax.set_ylim(0., 15.)
ax.legend(loc='lower right')
ax.grid(True)

# generierte Daten einzeichnen
ax.errorbar(Xm, Ym, xerr=0., yerr=Yerr, fmt='ro') 

ax2=fig.add_subplot(2,1,2)
AvVals=np.linspace(9., 11., 50)
S=[]
for pval in AvVals:
  S.append(calcS(FitFunc, Xm, Ym, Yerr, pval) )

ax2.plot(AvVals, S, 'g-', linewidth=2)
ax2.set_title('Summe der Fehlerquadrate $S$' )
ax2.set_xlabel('constant $c$', size='x-large')
ax2.set_ylabel('$S(c)$', size='x-large')
#ax2.text(0.3, 0.9, 
#  r'$S(D)=\sum \frac{  \left(A_i - A(\eta_i,D)\right)^2}{\sigma_i^2}$',
#  transform=ax2.transAxes,size='x-large')

# find and print minimum
idx=S.index(min(S))
chi2prb=1.- stats.chi2.cdf(S[idx], N-Npar)
txt_result=\
'Minimum bei %.3f\n    $\chi^2$= %.3f\n    $\chi^2$-prb.= %.2f'\
% (AvVals[idx], S[idx], chi2prb)


print('\n *==*' + txt_result + '\n')
ax2.text(0.3, 0.4, txt_result, transform=ax2.transAxes,size='large') 



# more space between subplots than default
fig.subplots_adjust(top=0.92, bottom=0.1, left=0.1, right=0.95,
                        wspace=0.2, hspace=0.3)
plt.show()