import matplotlib.pyplot as plt
from matplotlib.dates import HourLocator, DateFormatter

from sunpy import timeseries as ts
from sunpy.net import Fido
from sunpy.net import attrs as a
from sunpy.time import parse_time

import numpy as np

tr = a.Time('2023-12-14 00:00', '2023-12-15 00:00')
results = Fido.search(tr, a.Instrument.xrs & a.goes.SatelliteNumber(16) & a.Resolution("avg1m") | a.hek.FL & (a.hek.FRM.Name == 'SWPC'))

file = Fido.fetch(results)[0]
# Note: first file's data ends at 2023-12-14 23:59
goes = ts.TimeSeries(file).to_dataframe()
hek_results = results['hek']


fig, ax = plt.subplots()
fig.suptitle('GOES-16 XRS X-Ray Flux (1-minute data)')

ax.plot(goes.index, goes['xrsa'], '-', label='0.05–0.4 nm', color='#005AB5', lw=2)
ax.plot(goes.index, goes['xrsb'], '-', label='0.1–0.8 nm', color='#DC3220', lw=2)

# SXR classification on right side
labels = ['A', 'B', 'C', 'M', 'X']
centers = np.logspace(-7.6, -3.6, len(labels))
for value, label in zip(centers, labels):
    ax.text(1.025, value, label,
            transform=ax.get_yaxis_transform(), horizontalalignment='center', size='large', color='#DC3220')

# Label major HEK events
for flare in hek_results:
    goes_cls = flare['fl_goescls']
    if 'M' in goes_cls or 'X' in goes_cls:
        peak_time = parse_time(flare['event_peaktime']).datetime
        peak_flux = goes['xrsb'][peak_time]
        # ax.axvline(peak_time, ls='-', lw=1, alpha=0.8, color='gray')
        ax.annotate(goes_cls, (peak_time, peak_flux), (-3, 6), textcoords='offset points')

ax.legend(loc='upper center', bbox_to_anchor=(0.5, 1.05), ncol=2, fancybox=False, shadow=True)

ax.set_yscale("log")
ax.set_ylim(1e-9, 1e-2)
ax.set_ylabel("Watts m$^{-2}$")

ax.set_xlim(tr.start.to_datetime(), tr.end.to_datetime())
ax.grid(True, which='major', axis='y', alpha=1)
ax.grid(True, which='minor', axis='y', alpha=0.4)

ax.xaxis.set_major_locator(HourLocator(byhour=range(0, 24, 4)))
ax.xaxis.set_major_formatter(DateFormatter('%H:%M'))
ax.set_xlabel('2023 Dec 14 (UTC)')

"""
# Alternative date formatting
import matplotlib.dates as mdates
locator = mdates.AutoDateLocator(minticks=3, maxticks=11)
formatter = mdates.ConciseDateFormatter(locator)
ax.xaxis.set_major_locator(locator)
ax.xaxis.set_major_formatter(formatter)
"""

plt.savefig("GOES-16_X-Ray_Flux_2023-12-14.svg")