import numpy as np
import matplotlib.pyplot as plt
import mplhep as hep_style

hep_style.style.use("CMS")

label_font  = dict(family='serif', color='darkred', weight='normal', size=20)
annot_font  = dict(family='serif', color='black',   weight='normal', size=13)
annot_font2 = dict(family='serif', color='#555555', weight='normal', size=12)

# ── total fluxes (BS2021 SSM, cm⁻²s⁻¹) ──────────────────────────
PP_TOTAL  = 5.99e10
N13_TOTAL = 3.07e8
O15_TOTAL = 2.33e8
F17_TOTAL = 5.84e6
B8_TOTAL  = 5.69e6
HEP_TOTAL = 7.93e3
BE7_TOTAL = 4.84e9
PEP_TOTAL = 1.40e8

# ── load spectral data ────────────────────────────────────────────
pp  = np.loadtxt('pp.txt')
n13 = np.loadtxt('n13.dat')
o15 = np.loadtxt('o15.dat')
f17 = np.loadtxt('f17.dat')
b8  = np.loadtxt('b8.txt')
hep = np.loadtxt('hep.dat')

# ── normalize to absolute flux ────────────────────────────────────
pp_flux  = pp[:,1]  * PP_TOTAL
n13_flux = n13[:,1] * N13_TOTAL
o15_flux = o15[:,1] * O15_TOTAL
f17_flux = f17[:,1] * F17_TOTAL
b8_flux  = b8[:,1]  * B8_TOTAL
hep_flux = hep[:,1] * HEP_TOTAL

COLOR_PP  = '#F2C4AE'
COLOR_B8  = '#C5DDB5'
COLOR_HEP = '#C0B6D4'

fig, ax = plt.subplots(figsize=(8, 8), facecolor='#fafafa')
ax.set_facecolor('#fafafa')

# ── filled spectra (layered — B8 and hep start from bottom) ──────
ax.fill_between(pp[:,0],  pp_flux,  where=pp_flux  > 0, color=COLOR_PP,  alpha=0.85, zorder=1)
ax.fill_between(b8[:,0],  b8_flux,  where=b8_flux  > 0, color=COLOR_B8,  alpha=0.80, zorder=2)
ax.fill_between(hep[:,0], hep_flux, where=hep_flux > 0, color=COLOR_HEP, alpha=0.80, zorder=3)

# ── contour lines ─────────────────────────────────────────────────
kw_cont = dict(color='black', linewidth=1.2, zorder=5)
kw_dash = dict(color='black', linewidth=0.9, linestyle=':', zorder=5)

ax.plot(pp[:,0],  pp_flux,  **kw_cont)
ax.plot(n13[:,0], n13_flux, **kw_dash)
ax.plot(o15[:,0], o15_flux, **kw_dash)
ax.plot(f17[:,0], f17_flux, **kw_dash)
ax.plot(b8[:,0],  b8_flux,  **kw_cont)
ax.plot(hep[:,0], hep_flux, **kw_cont)

# ── monochromatic lines: ⁷Be and pep ─────────────────────────────
be7_energies  = [0.3843, 0.8613]
be7_fractions = [0.1050, 0.8950]
be7_fluxes    = [BE7_TOTAL * f for f in be7_fractions]

for E, F in zip(be7_energies, be7_fluxes):
    ax.vlines(E, 1e-2, F, color='black', linewidth=2.5, zorder=6)

ax.vlines(1.442, 1e-2, PEP_TOTAL, color='black', linewidth=2.5, zorder=6)

# ── labels ───────────────────────────────────────────────────────
ax.text(0.225, 5e10, 'pp',        **annot_font)
ax.text(0.118, 5.4e7, r'$^{13}$N', **annot_font2)
ax.text(0.118, 3.5e6, r'$^{15}$O', **annot_font2)
ax.text(0.118, 3.5e5, r'$^{17}$F', **annot_font2)
ax.text(7.5,   7.8e4, r'$^8$B',    **annot_font)
ax.text(10.0,  95,    'hep',       **annot_font)

ax.annotate('', xy=(0.862, 2.5e3), xytext=(0.384, 2.5e3),
            arrowprops=dict(arrowstyle='<->', color='black', lw=1.2))
ax.text(0.50, 5e2, r'$^7$Be', **annot_font)

ax.annotate('pep', xy=(1.442, 1.5e7), xytext=(1.442+0.35, 1.5e7),
            fontsize=13, fontfamily='serif',
            arrowprops=dict(arrowstyle='->', color='black', lw=1.2),
            ha='left', va='center')

# ── axes ─────────────────────────────────────────────────────────
ax.set_xscale('log', base=10)
ax.set_yscale('log', base=10)
ax.set_xlim(0.1, 25)
ax.set_ylim(1e-2, 1e15)

x_ticks = [0.1, 0.2, 0.5, 1, 2, 5, 10, 20]
ax.set_xticks(x_ticks)
ax.set_xticklabels([str(v) for v in x_ticks])
ax.set_yticks([10**i for i in range(0, 16, 2)])

ax.grid(True, which='major', linestyle='--', linewidth=0.4, color='gray', alpha=0.4)
ax.grid(True, which='minor', linestyle=':', linewidth=0.2, color='gray', alpha=0.2)

ax.set_xlabel(r'$E_\nu\ \mathrm{[MeV]}$', fontdict=label_font)
ax.set_ylabel(
    r'Flux at 1 AU $(\mathrm{cm}^{-2}\,\mathrm{s}^{-1}\,\mathrm{MeV}^{-1})$'
    r'[for lines, $\mathrm{cm}^{-2}\,\mathrm{s}^{-1}$]',
    fontdict=label_font
)

plt.savefig('solar_neutrino_flux.pdf', dpi=200, bbox_inches='tight')
plt.savefig('solar_neutrino_flux.png', dpi=200, bbox_inches='tight')
print("Saved.")