from astropy.coordinates import EarthLocation, get_body, ICRS
from astropy.time import Time
from astropy import units as u
from astropy_healpix import HEALPix
from matplotlib import pyplot as plt
from matplotlib.colors import LogNorm
import numpy as np
import ligo.skymap.plot

from m4opt.synphot.background import ZodiacalBackground
from m4opt.synphot import observing

wave = 10000 * u.angstrom
hpx = HEALPix(nside=512, frame=ICRS())
loc = EarthLocation.from_geocentric(0 * u.m, 0 * u.m, 0 * u.m)
coord = hpx.healpix_to_skycoord(np.arange(hpx.npix))
obstime = Time('2023-06-30T00:00:00')
zodi = ZodiacalBackground()
with observing(observer_location=loc, target_coord=coord, obstime=obstime):
    surf = zodi(wave, flux_unit=u.ABmag)

fig = plt.figure()
ax = fig.add_subplot(projection='astro hours mollweide')
im = ax.imshow_hpx(surf.value, cmap='viridis')
fig.colorbar(im, extend='both', orientation='horizontal').set_label(surf.unit)

sun = get_body('sun', obstime)
transform = ax.get_transform('world')
ax.plot(
    sun.ra, sun.dec, 'or', ms=1,
    transform=transform)
ax.plot(
    sun.ra, sun.dec, 'or', mfc='none',
    transform=transform)
ax.text(sun.ra, sun.dec, '  Sun', color='red', transform=transform)

ax.grid()