import numpy as np
import matplotlib.pyplot as plt


# Ukol 1
with open('eclipse.txt', 'r') as f:
    lines = f.read().split('\n')

dates = []
eclipse_numbers = []   # P = polotienove (1), U = uplne (2), C = ciastocne (3)
for line in lines:
    line_list = line.split()
    year = line_list[0]
    month = line_list[1]
    eclipe_type = line_list[2]
    dates.append(int(year) + (int(month) - 1)/12.)
    if eclipe_type == 'P':
        eclipse_numbers.append(1)
    elif eclipe_type == 'Ú':
        eclipse_numbers.append(2)
    elif eclipe_type == 'Č':
        eclipse_numbers.append(3)

plt.figure(figsize=(9, 3))
plt.plot(dates, eclipse_numbers, 'o')
plt.xlabel('cas [rok]')
plt.ylabel('typ zatmenia')
plt.yticks([1, 2, 3], ['P', 'U', 'C'])
plt.show()

time1 = [dates[i] for i in range(len(dates)) if (dates[i] > 1910 and
         (eclipse_numbers[i] == 3))][0]
time2 = [dates[i] for i in range(len(dates)) if (dates[i] > 1985 and
         (eclipse_numbers[i] == 3))][0]

saros = (time2 - time1)/4
print('Perioda sarosu je {} rokov.'.format(saros))


# Ukol 2
data = np.loadtxt('SN_m_tot_V2.0.txt')
x = data[:,2]
y = data[:,3]
plt.figure(figsize=(15, 3))
plt.xticks(range(1700,2020,10))
plt.plot(x, y, 'o')
plt.xlabel('cas [rok]')
plt.ylabel('Wolfovo cislo')
plt.show()
time2 = 2014.
time1 = 1762
n = 24
print('Hlavny slnecny cyklus ma periodu {} rokov.'.format((time2 - time1)/n))


# Ukol 3
with open('asu-txt.txt', 'r') as f:
    lines = f.read().split('\n')

u_U = []
T_U = []
u_V = []
T_V = []
u_B = []
T_B = []

for line in lines:
    line_list = line.split()
    if line_list[5] is 'U':
        u_U.append(float(line_list[4]))
        T_U.append(float(line_list[1]))
    elif line_list[5] is 'B':
        u_B.append(float(line_list[4]))
        T_B.append(float(line_list[1]))
    elif line_list[5] is 'V':
        u_V.append(float(line_list[4]))
        T_V.append(float(line_list[1]))
 
plt.figure(figsize=(15, 10))
plt.scatter(T_U, u_U, s=2, c='m', label='U')
plt.scatter(T_B, u_B, s=2, c='b', label='B')
plt.scatter(T_V, u_V, s=2, c='g', label='V')
plt.xlabel('T [K]')
plt.ylabel('u')
plt.legend()
plt.show()