from astropy.coordinates import EarthLocation, Galactic, SkyCoord
from astropy.time import Time
from astropy import units as u
from matplotlib import pyplot as plt
from m4opt.synphot.background import GalacticBackground
from m4opt.synphot.background._core import BACKGROUND_SOLID_ANGLE
from m4opt.synphot import observing
import numpy as np

lat = np.linspace(-90, 90, 100)
wavelengths = [1539, 2316] * u.angstrom
spectrum = GalacticBackground()

target_coord = SkyCoord(0 * u.deg, lat * u.deg, frame=Galactic())
# Observer location and obstime are arbitrary
observer_location = EarthLocation(0 * u.m, 0 * u.m, 0 * u.m)
obstime = Time("2024-01-01")

with observing(observer_location, target_coord, obstime):
    flux_density = (spectrum(wavelengths[:, np.newaxis]) / BACKGROUND_SOLID_ANGLE).to(
        u.photon * u.cm**-2 * u.s**-1 * u.sr**-1 * u.angstrom**-1
    )

ax = plt.axes()
ax.set_xlim(-90, 90)
ax.set_ylim(0, 5000)
ax.set_xlabel("Galactic latitude (°)")
ax.set_ylabel(f"UV background ({flux_density.unit:latex_inline}")
for y, wavelength in zip(flux_density, wavelengths):
    ax.plot(lat, y, label=f"{wavelength:latex_inline}")
ax.legend()