import numpy	# modul pro matematické funkce a práci s poli
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit

lines = numpy.loadtxt("cv12/upf_cv12_covid.txt")	# nacti soubor do pole N x 3

x = lines[:,0]
y = lines[:,1]
sigma = lines[:,2]
N = numpy.size(x)

fig, ax = plt.subplots(1,3,figsize=(16,4))

ax[0].errorbar(x,y,sigma,marker='o',ls='None', ms=3)
ax[0].set_xlabel("$x$")
ax[0].set_ylabel("$y$")

# v logaritmicke skale:
ax[1].semilogy()
ax[1].errorbar(x,y,sigma,marker='o',ls='None', ms=3)
ax[1].set_xlabel("$x$")
ax[1].set_ylabel("$y$")

############ fitovani obecne funkce numericky: ##############################

def fit_func(x,a,b):
	return a * numpy.exp(b*x)

M = 2	# pocet parametru
p = curve_fit(fit_func, x, y, sigma=sigma)

yModel = numpy.empty(N)
for i in range(0,N):
	yModel[i] = fit_func(x[i], p[0][0], p[0][1])
ax[0].plot(x,yModel)
ax[1].plot(x,yModel)

a, b = p[0]		# p[0] je pole nafitovanych parametru
sigma_a = numpy.sqrt(p[1][0][0])	# p[1] je matice kovariance
sigma_b = numpy.sqrt(p[1][1][1])

print("\nNelineární fit:")
print("a = {:.0f} +/- {:.0f}".format(a,sigma_a))
print("b = {:.3f} +/- {:.3f}".format(b,sigma_b))


############ linearizace a fitovani polynomu: ###############################
yLin = numpy.log(y)				# linearizovane pocty
sigmaLin = 1.0 / y * sigma		# odpovodajici chyby (prenosem chyb)
ax[2].set_xlabel("$x$")
ax[2].set_ylabel("ln($y$)")

ax[2].errorbar(x,yLin,sigmaLin,marker='o',ls='None', ms=3)

invSigmaLin = 1.0 / sigmaLin
pLin = numpy.polyfit(x, yLin, deg=(M-1), w=invSigmaLin, cov=True)

bLin, lnaLin = pLin[0]		# p[0] je pole nafitovanych parametru
sigma_bLin = numpy.sqrt(pLin[1][1][1])	# p[1] je matice kovariance
sigma_lnaLin = numpy.sqrt(pLin[1][0][0])

aLin = numpy.exp(lnaLin)
sigma_aLin = aLin * sigma_lnaLin

print("\nLinearizovaný fit:")
print("ln(a) = {:.3f} +/- {:.3f}".format(lnaLin,sigma_lnaLin))
print("a = {:.0f} +/- {:.0f}".format(aLin,sigma_aLin))
print("b = {:.3f} +/- {:.3f}".format(bLin,sigma_bLin))

yModel = numpy.polyval(pLin[0], x)
ax[2].plot(x,yModel)

plt.show()