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=2020,12,11 
ey,em,ed=2020,12,13 
sdate,edate=dt.datetime(sy,sm,sd), dt.datetime(ey,em,ed)
unit="he1"
pha_buffer=9    # see bottom of document

# load data
data = EPDData(sdate,edate)
pha = data.het.science.pha(unit, calibrated=True)
events = pha[pha['buffer'].isin([pha_buffer,pha_buffer])]

# calculate total energy in A1, A2, B1, B2 (inner and outer segment) and C
a1 = np.maximum(np.array(events['A_A11']), np.array(events['A_A12']))/1e3
a2 = np.maximum(np.array(events['A_A21']), np.array(events['A_A22']))/1e3
b1 = np.maximum(np.array(events['A_B11']), np.array(events['A_B12']))/1e3
b2 = np.maximum(np.array(events['A_B21']), np.array(events['A_B22']))/1e3
c = np.array(events['A_C'])/1e3
mask=events["trig_4"]   # trigger 'trig_4' is for sunward stopping
a,b,c=a1[mask],b1[mask],c[mask] # now we have energy losses in a,b,c in MeV :)

# calculate PHA weights: number of recorded events divided by number of events sent down
weights = events['nevents'][mask] / events['nwords'][mask]

# make hist
nbins = 200  # number of bins
xbins = np.logspace(-0.3,2.3, nbins)
ybins = np.logspace(-0.3,1, nbins)
bins=[xbins,ybins]
x,y=c,a+b
h,xedges,yedges=np.histogram2d(x,y,bins=bins,weights=weights)

# plot hist
plt.figure()
p=plt.pcolormesh(xedges,yedges,h.T,norm=LogNorm(),rasterized=True) 
cbar=plt.colorbar()
plt.xscale("log")
plt.yscale("log")
plt.show()



""" pha classes:
######################### PHA classes #############################
ept_pha_invalid      := 0   # ept trigger invalid event
ept_pha_pen          := 1   # penetrating event
ept_pha_stop         := 2   # stopping event

het_pha_pen_invalid  := 3   # penetrating trigger invalid event
het_pha_pen_gcr      := het_pha_pen_invalid # gcr trigger event
het_pha_pen          := 4   # penetrating event
het_pha_pen_highZ    := 5   # penetrating event with C>600MeV

het_pha_highZ_ABC    := 6   # heavy ion (S to Fe) stopping in C
het_pha_ion_ABC      := 7   # light ion (Li to Si) stopping in C
het_pha_a_ABC        := 8   # helium stopping in C
het_pha_p_ABC        := 9   # proton stopping in C
het_pha_e_ABC        := 10  # electron stopping in C

het_pha_highZ_AB     := 11  # heavy ion (S to Fe) stopping in B
het_pha_ion_AB       := 12  # light ion (Li to Si) stopping in B
het_pha_a_AB         := 13  # helium stopping in B
het_pha_p_AB         := 14  # proton stopping in B
het_pha_e_AB         := 15  # electron stopping in B
het_pha_stop_invalid := het_pha_pen_invalid  # stopping trigger invalid event
"""