# A quick view of NEA source region probabilities
# (c) SsODNet, B.Carry & J.Berthier, Lagrange/OCA & LTE/Obs. Paris

import rocks
import matplotlib.pyplot as plt

# Read the full SsODNet BFT
bft = rocks.load_bft(full=True)
cond = bft['class'].str.contains('NEA')

fig, ax = plt.subplots(2, 1, sharex=True, gridspec_kw={"hspace": 0.05})
vmin, vmax = 0.0, 0.8

# First source region: Hungaria
im1 = ax[0].scatter(
    bft.loc[cond,"orbital_elements.semi_major_axis.value"],
    bft.loc[cond,"orbital_elements.eccentricity.value"],
    s=2,
    alpha=0.1,
    c=bft.loc[cond,"source_regions.hun.value"],
    vmin=vmin,
    vmax=vmax,
)

# Second source region: 3:1 mean-motion resonance
im2 = ax[1].scatter(
    bft.loc[cond,"orbital_elements.semi_major_axis.value"],
    bft.loc[cond,"orbital_elements.eccentricity.value"],
    s=2,
    alpha=0.1,
    c=bft.loc[cond,"source_regions.mm31.value"],
    vmin=vmin,
    vmax=vmax,
)

# Add colorbars
for axis, image, label in zip(ax, (im1, im2), ("Hungaria", "3:1")):
    cax = axis.inset_axes([0.55, 0.25, 0.4, 0.08])
    cbar = fig.colorbar(
        image, cax=cax, orientation="horizontal", label=f"Probability from {label}"
    )
    cbar.ax.tick_params(labelsize=8)
    cbar.solids.set_alpha(1)

# Axes
ax[0].set_xlim(0.5, 3.5)
ax[1].set_xlabel("Semi-major axis (au)")
ax[0].set_ylabel("Eccentricity")
ax[1].set_ylabel("Eccentricity")

# To disk
fig.tight_layout()
fig.savefig("NEA-Probability.png", transparent=True, dpi=180)