import numpy as np
import matplotlib.pyplot as plt

# a)
# Zo stranky http://www.astrouw.edu.pl/asas/?page=acvs sme ziskali data
# pre vybranu hviezdu a v datovom subore sme odstranili vsetky komentare.
# Funkcia loadtxt nevie citat zmiesane datove typy, preto specifikujeme vyber
# prvych 11 stlpcov.
matrix = np.loadtxt('V2509_Sgr.txt', usecols=range(11))

P = 1.087
HJD = matrix[:, 0]
M2 = matrix[:, 4]   # vyberame stlpec MAG_2 s najpresnejsimi datami
E2 = matrix[:, 10]

plt.figure()
plt.plot(HJD, M2, '*')
plt.xlabel('HJD [d]')
plt.ylabel('MAG_2 [mag]')
plt.title('surove data')
# Obratime y-ovu os, aby zodpovedala vzrastajucej hviezdnej velkosti.
plt.gca().invert_yaxis()
plt.show()
errors = np.where(M2 > 20)[0]   #zistujeme pozicie vsetkych odlahlych bodov
HJD = np.delete(HJD, errors)
M2 = np.delete(M2, errors)
E2 = np.delete(E2, errors)

plt.figure()
plt.plot(HJD, M2, '*')
plt.xlabel('HJD [d]')
plt.ylabel('MAG_2 [mag]')
plt.title('data po odstraneni chyb')
plt.gca().invert_yaxis()

fi = np.modf(HJD/P)[0]   # modf()[0] vrati desatinnu cast
M2 = np.concatenate((M2, M2, M2))   # nasklada vektory vedla seba
E2 = np.concatenate((E2, E2, E2))
fi = np.concatenate((fi - 1, fi, fi + 1))

# Funkcia argsort() vrati indexy prvkov zoradeneho pola fi podla poradia,
# v akom sa vyskytuju v povodnom poli fi.
sort_seq = fi.argsort()
fi = fi[sort_seq]   # presklada prvky vzostupne
M2 = M2[sort_seq]
E2 = E2[sort_seq]

plt.figure()
for i in range(5):
    plt.clf()
    plt.gca().invert_yaxis()
    plt.plot(fi, M2, '.')
    plt.xlim([-0.2, 1.2])
    
    # b)
    n = len(fi)
    X = np.ones((n, 7))
    X[:, 1] = np.cos(2*np.pi*fi)
    X[:, 2] = np.sin(2*np.pi*fi)
    X[:, 3] = np.cos(4*np.pi*fi)
    X[:, 4] = np.sin(4*np.pi*fi)
    X[:, 5] = np.cos(6*np.pi*fi)
    X[:, 6] = np.sin(6*np.pi*fi)
    W = np.diag(np.power(E2, -2))
    V = np.dot(np.dot(X.T, W), X)   # .T transponuje pole
    U = np.dot(np.dot(X.T, W), M2)
    b = np.dot(np.linalg.inv(V), U)
    y_p = np.dot(X, b)
    plt.plot(fi, y_p, 'r')
    plt.xlabel('phase')
    plt.ylabel('V [mag]')
    plt.title('fazova krivka po odstraneni odlahlych bodov')
    
    # c)
    # Chi kvadrat ma teoreticky zodpovedat hodnote n - g, co je v nasom pripade
    # pribliznee 4000. Modifikovany chi kvadrat ma byt priblizne rovny jednej.
    chi2 = np.dot(np.dot(M2.T, W), M2) - np.dot(b.T, U)
    g = len(X[0])   # pocet stupnov volnosti
    chi2_mu = chi2/(n - g)
    
    # d)
    sigma = np.sqrt(chi2_mu/np.average(np.diag(W)))
    db = np.sqrt(chi2_mu**2*np.diag(np.linalg.inv(V)))
    dy = np.sqrt(chi2_mu**2*np.diag(np.dot(np.dot(X, np.linalg.inv(V)), X.T)))
    plt.plot(fi, y_p - 10*dy, 'g')    
    plt.plot(fi, y_p + 10*dy,'g')  
    plt.pause(1)
    plt.show()
    
    # e)
    errors = np.where(np.abs(M2 - np.dot(X, b)) > 3.5*sigma)[0]
    print('pocet merani:', n)
    print('chi2: ', chi2)
    print('chi2_mu: ', chi2_mu)
    print('sigma: ', sigma)
    fi = np.delete(fi, errors)
    M2 = np.delete(M2, errors)
    E2 = np.delete(E2, errors)
