from astropy.coordinates import EarthLocation, GeocentricMeanEcliptic, SkyCoord
from astropy.time import Time
from astropy import units as u
from matplotlib import pyplot as plt
from m4opt.dynamics import nominal_roll
import numpy as np

# Place the observer at the center of the Earth.
# Note that for observers in Earth orbit, the impact of the orbital phase
# is negligible because the orbit is much smaller than the Earth-Sun distance.
observer_location = EarthLocation.from_geocentric(0, 0, 0, unit=u.m)
obstime = Time('2024-01-1') + np.linspace(0, 1, 1000) * u.year

fig, ax = plt.subplots()
for lat in np.arange(0, 100, 10) * u.deg:
    target_coord = SkyCoord(0 * u.deg, lat, frame=GeocentricMeanEcliptic)
    roll = nominal_roll(observer_location, target_coord, obstime)
    ax.plot(
        obstime.to_datetime(),
        np.unwrap(roll, period=360 * u.deg),
        label=lat.to_string(format='latex'))
ax.set_xlabel('Date')
ax.set_ylabel('Nominal roll angle (°)')
ax.legend(title='Ecliptic latitude')