import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import chi2

theta_true=([0.3,-5,1.5]) #skutečné hodnoty parametrů

Nsim=1000 #počet simulací
chi2_exp=np.empty(Nsim)

Ndata=1000
x=np.linspace(0,20,Ndata)
y_true=np.polyval(theta_true,x)
ey=np.random.uniform(0,5,Ndata)

fig,ax=plt.subplots(2,1,height_ratios=(0.3,1))
ax[0].set_ylim(-5,5)
ax[0].set_ylabel("rezidua")
ax[1].set_xlabel("x")
ax[1].set_ylabel("y")
ax[0].plot([0,20],[0,0],c='black')
for i in range(Nsim):
    y=y_true+np.random.normal(0,ey,Ndata)
    pars=np.polyfit(x,y,2,w=1/ey)
    residuals=(y-np.polyval(pars,x))/ey
    #ax[1].errorbar(x,y,ey,lw=0,elinewidth=1,capsize=5,markersize=1,marker='o')
    ax[1].scatter(x,y,s=1)
    ax[1].plot(x,np.polyval(pars,x),c='black')
    ax[0].scatter(x,residuals,s=1)
    chi2_exp[i]=np.sum(((y-np.polyval(pars,x))/ey)**2)
plt.show()

fig,ax=plt.subplots()
ax.hist(chi2_exp,bins=20,density=True,label='histogram chi2 hodnot')
x_plot=np.linspace(Ndata-5*np.sqrt(Ndata),Ndata+5*np.sqrt(Ndata),1000)
ax.plot(x_plot,chi2.pdf(x_plot,Ndata),c='red',lw=2,label='chi2 rozdělení')
ax.set_xlabel("chi2")
ax.set_ylabel("pdf")
plt.legend()
plt.show()