import numpy as np
from matplotlib import pyplot as plt


def read_data(file):
    with open(file, 'r') as f:
        data = f.readlines()
    k = 0
    while data[k][0] == '#':
        k += 1
    data = [list(float(s.replace('\n', '').strip()) for s in row.split(',')) for row in data[k:]]
    return np.array(data)

E_proton = [0.1, 0.2, 0.4, 0.6, 0.8, 1, 2, 4, 6, 8, 10, 15, 20, 30, 50, 60, 80, 100, 150, 200, 300, 400, 700, 1200, 2000]
E_electron = [0.04, 0.07, 0.1, 0.25, 0.5, 0.75, 1, 1.5, 2, 2.5, 3, 3.5, 4, 4.5, 5, 5.5, 6, 6.5, 7, 8.5, 10]

data = read_data('assets/AP9_mean_flux.txt')
proton = [np.mean(data[:, k]) for k in range(4, data.shape[1])]

data = read_data('assets/AE9_mean_flux.txt')
electron = [np.mean(data[:, k]) for k in range(4, data.shape[1])]

_, ax = plt.subplots(1, 2, dpi=150, figsize=(16, 6))

ax1, ax2 = ax[0], ax[1]

ax1.plot(E_proton, proton)
ax1.set_ylim([1e-3, 5 * 1e2])
ax1.set_yscale('log')
ax1.set_ylabel(r'$Differential\ Flux\ (\mathrm{cm^{-2}s^{-1}})$')
ax1.set_title('Average spectra - Proton')
ax1.set_xlabel(r'$Energy\ (\mathrm{MeV})$')
ax1.set_xlim(1e-1, 2000)
ax1.set_xscale('log')

ax2.plot(E_electron, electron)
ax2.set_ylim([1e-3, 1e6])
ax2.set_yscale('log')
ax2.set_ylabel(r'$Differential\ Flux\ (\mathrm{cm^{-2}s^{-1}MeV^{-1}})$')
ax2.set_title('Average spectra - Electron')
ax2.set_xlabel(r'$Energy\ (\mathrm{MeV})$')
ax2.set_xlim(0.04, 10)
ax2.set_xscale('log')

plt.tight_layout()
plt.savefig('docs/spectra-ep-9')