import numpy as np
import matplotlib.pyplot as plt

def get_X(phi, N):
    """
    Vypocita maticu X pre harmonicky polynom N-teho radu.
    
    Args:
        phi: Faza.
        N: Stupen harmonickeho polynomu.
    
    Returns: Matica parcialnych derivacii.
    """
    X = np.ones((len(phi), 2*N + 1))
    for i in range(1, N + 1):
        X[:, 2*i - 1] = np.cos(2*i*np.pi*phi)
        X[:, 2*i] = np.sin(2*i*np.pi*phi)
    return X

def fit(Y, E, N):
    """
    Vypocita parametre fitu a data prelozi krivkou.
    
    Args:
        Y: Hodnoty merani.
        E: Chyby merani.
        N: Stupen harmonickeho polynomu.
    """
    X = get_X(phase, N)
    W = np.diag(np.power(E, -2))
    V = np.dot(np.dot(X.T, W), X)   # .T transponuje pole
    U = np.dot(np.dot(X.T, W), Y)
    b = np.dot(np.linalg.inv(V), U)
    
    t = np.arange(0, 1, 0.01)
    y_p = np.dot(get_X(t, N), b)
    plt.plot(t, y_p)
    return max(abs(y_p))


binary = np.loadtxt('binary.txt')
phase = binary[:, 7]
RV1 = binary[:, 1]
RV2 = binary[:, 4]
err1 = binary[:, 2]
err2 = binary[:, 5]
plt.errorbar(phase, RV1, yerr=err1, ls='', marker='.', ms=4)
plt.errorbar(phase, RV2, yerr=err2, ls='', marker='.', ms=4)
plt.xlabel('faza')
plt.ylabel('radialna rychlost [km/s]')
maxA = fit(RV1, err1, 4)
maxB = fit(RV2, err2, 4)

# a)
# Ide o stvorhviezdny system, ktory sa sklada z dvoch spektroskopickych
# dvojhviezd: Aa+Ab, Ba+Bb. PhasA a PhasB oznacuju fazy jednotlivych dvojic.

# b)
# Z ReadMe pripojeneho k datam zistime periodu P = 57.32 dna. Z pomeru
# maximalnych radialnych rychlosti urcime pomer hmotnosti M1/M2 = v1/v2.
print('Hmotnost mensej zlozky je {} hmotnosti Slnka.'
      .format(np.round(maxA/maxB, 2)))
# Z tretieho Keplerovho zakona urcime velku polos: a**3 = (M1 + M2)*P**2.
P = 57.32/365.25
a = ((1 + maxA/maxB)*P**2)**(1/3)
print('Velka polos sustavy je {} AU.'.format(np.round(a, 2)))
