import numpy as np
from astropy import units as u
from astropy.coordinates import (
    SkyCoord,
    CartesianRepresentation,
    get_body,
)
from astropy.time import Time
from m4opt.missions import uvex as mission, uvex_downlink_orientation
from matplotlib import pyplot as plt

time = Time("2025-01-01") + np.linspace(0, 1, 1000) * u.year
target, roll = uvex_downlink_orientation(time)

observer_location = mission.observer_location(time)
sun = get_body("sun", time, observer_location)
earth = get_body("earth", time, observer_location)
spacecraft_frame = target.skyoffset_frame(roll)
antenna = SkyCoord(
    CartesianRepresentation(-1 / np.sqrt(2), 0, 1 / np.sqrt(2)), frame=spacecraft_frame
)
solar_array = SkyCoord(CartesianRepresentation(0, 1, 0), frame=spacecraft_frame)

fig = plt.figure()
ax = fig.add_subplot()
ax.yaxis.set_major_locator(plt.MultipleLocator(45))
ax.yaxis.set_major_formatter(plt.FormatStrFormatter("%g°"))
dt = time.datetime
ax.plot(dt, target.separation(earth), label=r"Target $\leftrightarrow$ Earth")
ax.plot(dt, target.separation(sun), label=r"Target $\leftrightarrow$ Sun")
ax.plot(dt, solar_array.separation(sun), label=r"Solar axis $\leftrightarrow$ Sun")
ax.plot(dt, antenna.separation(earth), label=r"Target $\leftrightarrow$ antenna")
ax.legend(title="Separation", loc="center right", bbox_to_anchor=(0.975, 0.25))