import numpy as np
import matplotlib.pyplot as plt

def find_phi0(phi, y, n):
    """
    Rozdeli data do fazovych binov a najde bin s najvyssou priemernou y-ovou hodnotou.
    
    Args:
        phi: Faza.
        y: Zdanliva hviezdna velkost.
        n: Pocet binov.
    
    returns: Hodnota nulovej fazy.
    """
    averages = []
    bin_centers = np.linspace(0.5/n, 1, n)
    for bin_center in bin_centers:
        averages.append(np.mean(y[np.where(abs(phi - bin_center) < 0.5/n)]))
    return bin_centers[np.where(max(averages) == averages)]

# a)
asu = np.loadtxt('asu.txt')

# b)
t = asu[:, 1]   # HJD
y = asu[:, 2]   # ETV
w = np.power(asu[:, 3], -2)   # vahy
W = np.diag(w)
plt.plot(t, y, '.')
# M0 je hodnota referencneho HJD. Pre pouzivanu fitovaciu funkciu je to hodnota
# vrcholu (idealne v strednej oblasti HJD).
M0 = 55950
plt.xlabel('HJD [d]')
plt.ylabel('ETV [s]')
plt.show()
Sigma = []
Period = []
n = len(t)
# Hodnotu periody priblizne odhadujeme z predchadzajuceho grafu, mozme menit
# vzorkovaciu frekvenciu.
for P in np.arange(50, 100, 0.1):
    phi = abs(np.modf((t - M0)/P)[0])   # fazova funkcia
    
    # c)
    # Funkciu find_phi0 pouzivame, aby sme co najpresnejsie stanovili
    # pociatocnu polohu parametra phi0. V pripade, ze ju nahradime konstantnou
    # hodnotou (povedzme 0), fit zdiverguje a v grafe Sigma/Period sa objavia
    # oscilacie.
    b = [0, 50, find_phi0(phi, y, 20)]   # pociatocne hodnoty parametrov
    for i in range(5):   # pocet cyklov nelinearnej regresie
        f0 = b[0]
        A = b[1]
        phi0 = b[2]
        X = np.ones((n, 3))
        X[:, 1] = np.cos(2*np.pi*(phi - phi0))
        X[:, 2] = A*2*np.pi*np.sin(2*np.pi*(phi - phi0))
        f = f0 + A*np.cos(2*np.pi*(phi - phi0))   # funkcne hodnoty fitu
        dy = y - f
        # Pocetne zjednoduseny vyraz v porovnani s X'*W*X.
        V = np.dot(X.T, np.multiply(np.transpose(np.vstack((w, w, w))), X))
        U = np.dot(X.T, np.multiply(w, dy))
        db = np.dot(np.linalg.inv(V), U)
        b = db + b
    # Standardna odchylka.
    sigma = np.sqrt(np.dot(np.dot(dy.T, W), dy)/(n - 3)/np.mean(w))
    Sigma.append(sigma)
    Period.append(P)
    if sigma <= min(Sigma):   # hladame minimalnu periodu
        P_real = P
        phi_real = phi
        f_real = f      

plt.figure()
plt.plot(phi_real, y, '.', phi_real, f_real, 'r.')
plt.xlabel('phase')
plt.ylabel('ETV [s]')
plt.title('P = {} d.'.format(np.round(P_real, 2)))
plt.show()

plt.figure()
plt.plot(Period, Sigma)
plt.xlabel('perioda [d]')
plt.ylabel('sigma')
plt.title('periodogram')
plt.show()

# V sustave je ocividne dalsie teleso, ktore obieha s periodou 71.1 dna.