import numpy as np
import matplotlib.pyplot as plt
filename = 'data.txt'

# a)
def fit(x, y, b0, show=False):
    """
    Prelozi data fitovanou funkciou, vrati parametre fitu aj s chybou a
    pripadne data vykresli do grafu.
    
    Args:
        x: x-ove hodnoty vo tvare array.
        y: y-ove hodnoty vo tvare array.
        b0: Pociatocny odhad vektora parametrov fitu.
        show: Urcuje ci sa maju data prelozene fitom vykreslit.
        
    Returns: Vektor parametrov fitu aj s ich chybou.
    """
    for i in range(5):
        (f0, A, x0, sigma, gamma) = tuple(b0)
        g = 1.0 - np.exp(-(x - x0)**2/2/sigma**2) 
        f = f0 + A*g**gamma
        if show:
            plt.clf()
            plt.plot(x, y, '.')
            plt.plot(x, f)
            plt.gca().invert_yaxis()
            plt.xlabel('MJD [d]')
            plt.ylabel('zdanliva hviezdna velkost [mag]')
            plt.show()
            plt.pause(1)
            
        X = np.ones((len(x), 5))
        X[:, 1] = g**gamma    # derivacia podla A
        # Derivacia podla x0.
        X[:, 2] = A*gamma*g**(gamma - 1)*(g - 1)*(x - x0)/sigma**2
        X[:, 3] = X[:, 2]*(x - x0)/sigma   # derivacia podla sigma
        X[:, 4] = A*g**gamma*np.log(g)    # derivacia podla gamma
        dY=y - f
        V = np.dot(np.transpose(X), X)
        U = np.dot(np.transpose(X), dY)
        b0 += np.dot(np.linalg.inv(V), U)
        s =  np.sqrt(np.dot(dY.T, dY)/(len(X) - len(X[0])))
        db = s*np.sqrt(np.diag(np.linalg.inv(V)))
        
    return(b0, db)

data = np.loadtxt('data.txt')
# Data usporiadame podla MJD, aby sa pri fite vykreslila spojita krivka.
data = data[np.argsort(data[:, 0]), :]
x = data[:, 0]
y = data[:, 1]
b0 = [9.38, -0.03, 25348.91, 0.05, 5.0]
(b, db) = fit(x, y, b0, show=True)
print('Chyby parametrov db:', db)

# b)
b_list = []
for j in range(1000):
    ind = np.random.choice(range(len(x)), len(x))
    x = np.loadtxt('data.txt')[ind, 0]
    y = np.loadtxt('data.txt')[ind, 1]
    (b, db) = fit(x, y, b0)
    b_list.append(b)
b_list = np.array(b_list)
print('Chyby parametrov db:', np.std(b_list, axis = 0))

# Porovnanim hodnot chyb parametrov z casti a) a b) vidime, ze bootstrap
# systematicky znizuje chybu urcenia linearnych parametrov, ale zvysuje tuto
# hodnotu pre nelinearne parametre. To poukazuje na zaver, ze hodnoty chyb
# nelinearnych parametrov pouzitim klasickej statistiky su znacne podhodnotene!
