from matplotlib.colors import LogNorm
from solo_loader.epd import EPDData
import datetime as dt
import matplotlib.pyplot as plt
import numpy as np
import os

# set startdate, enddate, unit and pha_buffer
sy,sm,sd=2021,10,28
ey,em,ed=2021,10,31
sdate,edate=dt.datetime(sy,sm,sd), dt.datetime(ey,em,ed)
unit,direc="he1","sun"
energy_scaler=1/1e3
flux_scaler=1e5

# load data
data = EPDData(sdate,edate)
species="p"
ranges="B" # can be B, C or P for stopping in B, C and penetrating respectively
hetdata = data.het.science.nominal(species, ranges)[direc]
flux=hetdata.to_flux() * flux_scaler # now in units of (cm$^2$ sr s MeV/nuc)$^{-1}$
eff_energies = hetdata.attrs['eff_energies'] * energy_scaler # now in units of MeV

# other counters have to be loaded slightly differently, i.e.
gcr_counter=data.het.science.gcr()
single_counter=data.het.housekeeping.l1_counters(unit)



# consider resampling
flux=flux.resample("10Min").mean()


# plot timeseries
plt.figure()
plt.plot(flux)
plt.yscale("log")
plt.ylabel("Flux / (cm$^2$ sr s MeV/nuc)$^{-1}$")
plt.show()



# plot spectrum
plt.figure()
plt.plot(eff_energies,flux.mean(),"o")
plt.xscale("log")
plt.xlabel("Energy / MeV")
plt.yscale("log")
plt.ylabel("Flux / (cm$^2$ sr s MeV/nuc)$^{-1}$")
plt.show()