from astropy import units as u
from astropy.visualization import quantity_support
from matplotlib import pyplot as plt
import numpy as np

from m4opt.synphot.background import SkyBackground

quantity_support()

wave = np.linspace(3750, 6868) * u.angstrom
ax = plt.axes()
for key in ['veryhigh', 'high', 'medium', 'low']:
    surf = getattr(SkyBackground, key)()(wave, flux_unit=u.ABmag)
    ax.plot(wave, surf, label=f'SkyBackground.{key}')
ax.legend()
ax.invert_yaxis()