# Code adapted from example on Starplot documentation website
from skyfield.api import load
from skyfield.data import mpc
from skyfield.constants import GM_SUN_Pitjeva_2005_km3_s2 as GM_SUN

from starplot import MapPlot, Projection, Star, _
from starplot.styles import PlotStyle, extensions

# First, we use Skyfield to get comet data
# Code adapted from: https://rhodesmill.org/skyfield/kepler-orbits.html#comets
with load.open(mpc.COMET_URL) as f:
    comets = mpc.load_comets_dataframe(f)

# Keep only the most recent orbit for each comet,
# and index by designation for fast lookup.
comets = (
    comets.sort_values("reference")
    .groupby("designation", as_index=False)
    .last()
    .set_index("designation", drop=False)
)

# Find comet
row = comets.loc["3I/ATLAS"]
print(row)

ts = load.timescale()
eph = load("de440.bsp")
sun, earth = eph["sun"], eph["earth"]
comet = sun + mpc.comet_orbit(row, ts, GM_SUN)

# Find the RA/DEC of comet for every 8 days starting on October 2, 2024
radecs = []
for day in range(-1000, 2000, 10):
    t = ts.utc(2024, 9 , 28 + day)
    ra, dec, distance = earth.at(t).observe(comet).radec()
    radecs.append((t, ra.hours * 15, dec.degrees))

# Now let's plot the data on a map!
style = PlotStyle().extend(
    extensions.GRAYSCALE_DARK,
    {"figure_background_color": "hsl(136, 0%, 20%)"},
    extensions.MAP,
    {
        "star": {
            "label": {
                "font_weight": "normal",
            }
        },
        "legend": {
            "___location": "lower center",
        },
        "ecliptic": {
            "label": {
                "font_size": 25}},
        "constellation_labels": {
            "font_size": 20,
            "font_weight": "heavy",}
    },
)
style.legend.___location = "lower center"

p = MapPlot(
    projection=Projection.MERCATOR,
    ra_min=6 * 15,
    ra_max=20 * 15,
    dec_min=-32,
    dec_max=32,
    style=style,
    resolution=2500,
    scale=0.4,
)

# Plot the comet markers first, to ensure their labels are plotted
for t, ra, dec in radecs:
    if (t.tt>=2460793) and (t.tt<=2461120):
        if (t.utc.month==1 and t.utc.day<=10) or t.tt <=2460793+10:
            label=f"{t.utc_strftime('%Y %b %-d')}"
        else:
            label = f"{t.utc_strftime('%b %-d')}"
        p.marker(
        ra=ra,
        dec=dec,
        style={
            "marker": {
                "size": 10,
                "symbol": "circle",
                "fill": "none",
                "color": "hsl(0,100%,50%)",
                "edge_color": "hsl(0,100%,50%)",
                "alpha": 1,
                "zorder": 4096,
            },
            "label": {
                "font_size": 21,
                "font_weight": "bold",
                "font_color": "hsl(60, 70%, 72%)",
                "zorder": 4096,
                "offset_x": "auto",
                "offset_y": "auto",
            },
        },
        label=label,
        legend_label="3I/ATLAS",
    )
    else:
        label = None
        p.marker(
        ra=ra,
        dec=dec,
        style={
            "marker": {
                "size": 2,
                "symbol": "circle",
                "fill": "none",
                "color": "hsl(0,100%,50%)",
                "edge_color": "hsl(0,100%,50%)",
                "alpha": 1,
                "zorder": 2048,
            },
        },
        legend_label="3I/ATLAS",
    )

p.gridlines(labels=False)
p.constellations()
p.constellation_borders()
p.stars(where=[_.magnitude < 6], where_labels=[_.magnitude < 2])
p.constellation_labels()
p.milky_way()
p.ecliptic()
p.title('Path of 3I/ATLAS', style={
    "font_size":58,
    "font_weight":"bold",
    "font_color":"hsl(136, 0%, 90%)",
    "anchor_point":'center',
})

p.export("map_comet_3i.svg",format='svg', padding=0)