import numpy as np
import matplotlib.pyplot as plt

data = np.loadtxt('hd81009.txt')
P = 33.978462
phi = np.modf(data[:, 1]/P)[0]   # modf vrati desatinnu cast cisla
indexes = np.argsort(phi)   # zoradime fazu
phi = phi[indexes]
m = data[indexes, 2]   # zdanliva hviezdna velkost
ID = data[indexes, 5]   # ID pozorovatela

# Pri vykresleni grafu chceme vidiet ako svetelna krivka pokracuje aj mimo
# intervalu [0, 1], preto si data multiplikujeme.
phi = np.concatenate((phi - 1, phi, phi + 1))
m = np.concatenate((m, m, m))
ID = np.concatenate((ID, ID, ID))
# Kvoli lepsiemu farebnemu rozliseniu pozorovatelov upravime ID na color.
color = 1*np.equal(ID, 2) + 2*np.equal(ID, 27) + 3*np.equal(ID, 37) + \
        4*np.equal(ID, 110)
n = len(phi)
b0 = [0.07, 0.07, 0.07, 0.07, 0.06, 0.7, 0.15]
for i in range(5):
    (f0, f1, f2, f3, A, phi0, s0) = tuple(b0)
    z = phi - phi0 - np.round(phi - phi0)
    Q = A*np.exp(1 - np.cosh(z/s0))
    f = f0*np.equal(ID, 2) + f1*np.equal(ID, 27) + f2*np.equal(ID, 37) + \
        f3*np.equal(ID, 110) + Q
    X = np.zeros((len(phi), 7))
    X[:, 0] = np.equal(ID, 2)
    X[:, 1] = np.equal(ID, 27)
    X[:, 2] = np.equal(ID, 37)
    X[:, 3] = np.equal(ID, 110)
    X[:, 4] = Q/A
    X[:, 5] = Q*np.sinh(z/s0)/s0
    X[:, 6] = Q*np.sinh(z/s0)*z/s0**2
    Y = m - f
    V = np.dot(X.T, X)
    U = np.dot(X.T, Y)
    db = np.dot(np.linalg.inv(V), U)
    b0 = db + b0
    plt.clf()
    plt.scatter(phi, m - f + Q, 5, color)
    plt.plot(phi, Q, 'r')
    plt.xlabel('faza')
    plt.ylabel('relativna zmena hviezdnej velkosti [mag]')
    plt.title('HD 81009')
    plt.ylim([-0.02, 0.07])
    plt.xlim([-0.2, 1.2])
    plt.gca().invert_yaxis()
    plt.pause(1)
