import numpy as np
import matplotlib.pyplot as plt

cepheids = np.loadtxt('cepheids.txt')
P = cepheids[:, 7]   # perioda
m_V = cepheids[:, 8]   # zdanliva hviezdna velkost vo filtri V
E = cepheids[:, 6]   # chyba

# Odstranime data, ktore maju spatne hodnoty periody.
errors = np.where(P <= 0)
P = np.delete(P, errors)
m_V = np.delete(m_V, errors)
E = np.delete(E, errors)

M_V = m_V + 5 - 5*np.log10(50000)
X = np.ones((len(P), 2))
X[:, 1] = np.log10(P)
W = np.diag(np.power(E, -2))
V = np.dot(np.dot(X.T, W), X)
U = np.dot(np.dot(X.T, W), M_V)
b = np.dot(np.linalg.inv(V), U)

t = np.arange(0.1, 140, 0.1)
y = b[0] + b[1]*np.log10(t)

plt.plot(P, M_V, '.', t, y)
plt.gca().invert_yaxis()
plt.ylabel('absolutna hviezdna velkost [mag]')
plt.xlabel('P [d]')
print('Pre cefeidy plati: M_V = {} + {}*log(P).'.format(np.round(b[0], 2),
                                                        np.round(b[1], 2)))
