import numpy as np
import matplotlib.pyplot as plt

def find_phi0(phi, y, n):
    """
    Divide data into phase bins and find phase bin with maximum average y-value.
    
    phi: phase list/array
    y: brightness values list/array
    n: number of bins
    
    returns: value of phi0
    """
    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)]

asu = np.loadtxt('asu.txt')
t = asu[:, 1]   #HJD
y = asu[:, 2]   #ETV
w = np.power(asu[:, 3], -2)   #vahy
W = np.diag(w)
plt.plot(t, y, '.')
M0 = 55950   #hodnota nuloveho HJD; pre pouzivanu fitovaciu funkciu je to hodnota vrcholu (idealne v strednej oblasti HJD)
plt.xlabel('HJD [d]')
plt.ylabel('intensity [mag]')
plt.show()
Sigma = []
Period = []
n = len(t)
for P in np.arange(50, 100, 0.1):   #hodnotu priblizne odhadujeme z predchadzajuceho grafu, mozme menit vzorkovaciu frekvenciu
    phi = abs(np.modf((t - M0)/P)[0])   #fazova funkcia
    #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;
        V = np.dot(X.T,np.multiply(np.transpose(np.vstack((w, w, w))), X))
        #pocetne zjednoduseny vyraz v porovnani s X'*W*X
        U = np.dot(X.T, np.multiply(w, dy))
        db = np.dot(np.linalg.inv(V), U)
        b = db + b
    sigma = np.sqrt(np.dot(np.dot(dy.T, W), dy)/(n - 3)/np.mean(w))   #standardna odchylka
    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('intensity [mag]')
plt.title(P_real)
plt.show()

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