import numpy as np
import matplotlib.pyplot as plt

# a)
data = np.loadtxt('SQ.dat')
MJD = data[:, 0]   # modifikovane julianske datum
Y = data[:, 1]   # intenzita vo filtri V
plt.plot(MJD, Y, '.')

# b)
t0 = [0]   # pociatocna hodnota odhadu polohy centra zakrytu
for i in range(len(MJD)):
    if Y[i] > 8.4 and abs(MJD[i] - t0[-1]) > 10:
        t0.append(MJD[i])   # odhadovane polohy centier vsetkych zakrytov
t0 = t0[1:]   # umelo dodanu prvu hodnotu musime odstranit
f0 = 8.15
A = 0.3
s = 0.7
b = [f0, A, s]   # pociatocne parametre modelovej funkcie
b.extend(t0)
# Vytvorime spojity casovy vektor pre vykreslenie fitovanej funkcie ft.
t = np.arange(min(MJD), max(MJD), 0.1)
for j in range(5):
    f0 = b[0]
    A = b[1]
    s = b[2]
    t0 = b[3:]
    X1 = np.ones(len(MJD))   # derivacia podla f0
    X2 = 0
    X3 = 0
    X4 = []
    f = f0
    ft = f0
    for i in range(len(t0)):
        exp = A*np.exp(-(MJD - t0[i])**2/2/s**2)
        X2 = X2 + exp/A   # derivacia podla A
        X3 = X3 + exp*(MJD - t0[i])**2/s**3   # derivacia podla s
        X4.append(exp*(MJD - t0[i])/s**2)   # derivacie podla t0i
        f = f + exp   # funkcna hodnota v casoch MJD
        # Funkcna hodnota v casoch t.
        ft = ft + A*np.exp(-(t - t0[i])**2/2/s**2)
    plt.clf()
    plt.plot(MJD, Y, '.', t, ft, 'r')
    plt.gca().invert_yaxis()
    plt.xlabel('MJD [d]')
    plt.ylabel('zdanliva hviezdna velkost [mag]')
    plt.title('SQ Tau lightcurve')
    plt.show()
    plt.pause(1)
    X = np.array([X1, X2, X3])
    X = np.concatenate((X, X4))
    X = X.T
    dY = Y - f
    V = np.dot(X.T, X)
    U = np.dot(X.T, dY)
    db = np.dot(np.linalg.inv(V), U)
    Yp = np.dot(X, b)
    b = b + db
    sigma = np.sqrt(np.dot(dY.T, dY)/(len(X) - len(X[0])))
    # Chyba urcenia parametra b.
    db = np.sqrt(sigma**2*np.diag(np.linalg.inv(V)))

# c)
M0 = 42446.22;
P = 230.01323;
minima = M0 + np.arange(-100*P, 55000, P)
OC = np.array([])
for i in range(len(t0)):
    print(t0[i] - minima[np.where(abs(t0[i] - minima) < 10)])
    OC = np.append(OC, t0[i] - minima[np.where(abs(t0[i] - minima) < 10)])

E = np.round((t0 - M0)/P)
plt.figure()
plt.errorbar(E, OC, db[3:], ls='', marker='.')
tt = np.arange(-90, 60, 1)
XX = np.ones((len(t0), 3))
XX[:, 1] = E
XX[:, 2] = E**2
WW = np.diag(db[3:]**(-2))
VV = np.dot(np.dot(XX.T, WW), XX)
UU = np.dot(np.dot(XX.T, WW), OC)
bb = np.dot(np.linalg.inv(VV), UU)
ff = bb[0] + bb[1]*tt + bb[2]*tt**2
plt.plot(tt, ff, 'r')
plt.xlabel('epocha')
plt.ylabel('O-C [s]')

# Parabolicky tvar O-C diagramu poukazuje na to, ze perioda sa linearne
# zvysuje. To je mozne pozorovat napriklad pri pretoku latky tesnych dvojhviezd.

# Hviezda SQ Tau v skutocnosti neexistuje, taketo oznacenie ani nie je mozne.
# Data boli vygenerovane, nie su to realne merania.