import numpy as np
import matplotlib.pyplot as plt

matrix = np.loadtxt('cont0109a.txt')
x = matrix[:, 0]   # vlnova dlzka
y = matrix[:, 1]   # intenzita
n = len(x)

b0 = [1.0, -0.6, 4340.0, 5.0]   # pociatocny odhad hodnot
for i in range(5):
    G0 = b0[0]   # intenzita kontinua
    A = b0[1]   # amplituda
    x0 = b0[2]   # stred ciary
    sigma = b0[3]   # polosirka ciary

    exp = np.exp(-np.power(x - x0, 2)/2/sigma**2)
    # Prekladame Gaussovou funkciou.
    G = G0 + A*exp
    X = np.ones((n, 4))
    X[:, 1] = exp   # parcialna derivacia podla A
    X[:, 2] = A*exp*(x - x0)/sigma**2   # parcialna derivacia podla x0
    # Parcialna derivacia podla sigma.
    X[:, 3] = A*exp*np.power(x - x0, 2)/sigma**3
    dY = y - G
    V = np.dot(X.T, X)
    U = np.dot(X.T, dY)
    db = np.dot(np.linalg.inv(V), U)
    b0 = db + b0
    plt.clf()
    plt.plot(x, y, 'b')
    plt.plot(x, G, 'r')
    plt.show()
    plt.pause(1)
    chi2_G = np.dot(dY.T, dY)
    b_G = b0

b0 = [1.0, -3.0, 4340.0, 5.0]   # pociatocny odhad hodnot
for i in range(5):
    L0 = b0[0]   # intenzita kontinua
    A = b0[1]   # amplituda
    x0 = b0[2]   # stred ciary
    gamma = b0[3]   # polosirka ciary
    Q = np.power(x - x0, 2) + gamma**2
    L = L0 + A*gamma/Q   # prekladana Lorentzova funkcia
    X = np.ones((n, 4))
    X[:, 1] = gamma/Q   # derivacia podla A
    X[:, 2] = 2*A*gamma*(x - x0)/np.power(Q, 2)   # derivacia podla x0
    # Derivacia podla gamma.
    X[:, 3] = A*(np.power(x - x0, 2) - gamma**2)/np.power(Q, 2)
    dY = y - L
    V = np.dot(X.T, X)
    U = np.dot(X.T, dY)
    db = np.dot(np.linalg.inv(V), U)
    b0 = db + b0
    plt.clf()
    plt.plot(x, y, 'b')
    plt.plot(x, L, 'g')
    plt.show()
    plt.pause(1)
    chi2_L = np.dot(dY.T, dY)
    b_L = b0

plt.clf()
plt.plot(x, y, 'b')
plt.plot(x, L, 'g')
plt.plot(x, G, 'r')
plt.show()

print('Centrum čiary použitím Gaussovej funkcie má vlnovú dĺžku {} A.'
      .format(np.round(b_G[2], 2)))
print('Centrum čiary použitím Lorentzovej funkcie má vlnovú dĺžku {} A.'
      .format(np.round(b_L[2],2)))
print('Chi2 pre Gaussovu funkciu je {}.'.format(np.round(chi2_G, 2)))
print('Chi2 pre Lorentzovu funkciu je {}.'.format(np.round(chi2_L, 2)))
