# A simple top-view of the inner Solar system
# (c) SsODNet, B.Carry & J.Berthier, Lagrange/OCA & LTE/Obs. Paris

import rocks
import matplotlib.pyplot as plt

# Read SsODNet BFT: only dynamical class and cartesian positions
bft = rocks.load_bft(columns=["class", "position"])

# User choices (color...) for selected populations
populations = {
    "NEA": {"class": {"mode": "contains", "key": "NEA"}, "color": "red", "zorder": 100},
    "Main Belt": {
        "class": {"mode": "contains", "key": "MB"},
        "color": "grey",
        "zorder": -100,
    },
    "Hilda": {
        "class": {"mode": "contains", "key": "Hilda"},
        "color": "black",
        "zorder": 100,
    },
    "L4 Trojans": {
        "class": {"mode": "is", "key": "Trojan>L4"},
        "color": "green",
        "zorder": 50,
    },
    "L5 Trojans": {
        "class": {"mode": "is", "key": "Trojan>L5"},
        "color": "cornflowerblue",
        "zorder": 50,
    },
}

fig, ax = plt.subplots(figsize=(6, 6))
for c in populations:
    # Fine-selection of dynamical class
    if populations[c]["class"]["mode"] == "contains":
        cond = bft["class"].str.contains(populations[c]["class"]["key"])
    else:
        cond = bft["class"] == populations[c]["class"]["key"]

    # Simple scatter plot
    ax.scatter(
        bft.loc[cond, "position"].str[0],
        bft.loc[cond, "position"].str[1],
        s=1,
        alpha=0.4,
        color=populations[c]["color"],
        zorder=populations[c]["zorder"],
        label=c,
    )

# Tweaking the axes
ax.set_aspect("equal")
ax.set_xlabel("x / au")
ax.set_ylabel("y / au")
ax.set_xlim(-6, 6)
ax.set_ylim(-6, 6)
leg = ax.legend(loc="lower center", ncol=3)
for lh in leg.legend_handles:
    lh.set_alpha(1)

# Export to disk
fig.tight_layout()
fig.savefig("ISS-DynClasses.png", transparent=True, dpi=180)