import numpy as np
import matplotlib.pyplot as plt

P = 1.087
matrix = np.loadtxt('V2509_Sgr.txt', usecols=range(11))
#loadtxt nevie citat zmiesane datove typy, preto specifikujeme vyber prvych 11 stlpcov
HJD = matrix[:, 0]
M2 = matrix[:, 4]   #vyberame stlpec MAG_2 s najpresnejsimi datami
E2 = matrix[:, 10]

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.plot(HJD, M2, '*')
plt.xlabel('HJD [d]')
plt.ylabel('MAG_2 [mag]')
plt.gca().invert_yaxis()
plt.show()

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))

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

for i in range(5):
    plt.plot(fi, M2, '*')
    plt.xlim([-0.2, 1.2])
    plt.gca().invert_yaxis()
    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]')
    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)
    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.show()
    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)
    print(b, db)
    fi = np.delete(fi, errors)
    M2 = np.delete(M2, errors)
    E2 = np.delete(E2, errors)