# A simple (a,e) (a,i) figure of the albedo distribution in the asteroid bel
# (c) SsODNet, B.Carry & J.Berthier, Lagrange/OCA & LTE/Obs. Paris

import numpy as np
import rocks
import matplotlib.pyplot as plt


# Read SsODNet BFT: only dynamical class, proper a/e/i, and albedo
bft = rocks.load_bft(
    columns=[
        "class",
        "proper_elements.proper_semi_major_axis.value",
        "proper_elements.proper_eccentricity.value",
        "proper_elements.proper_inclination.value",
        "albedo.value"
    ]
)
# Select Main-belt, Hungaria and Phocaea populations
cond = (
    (bft["class"].str.contains("MB"))
    | bft["class"].str.contains("Hungaria")
    | bft["class"].str.contains("Phocaea")
) & (bft["albedo.value"].notna())
df = bft[cond].sort_values("albedo.value")


fig, ax = plt.subplots(
    2,
    1,
    figsize=(6, 6),
    sharex=True,
    gridspec_kw={"hspace": 0.02, "left": 0.12, "right": 0.98, "bottom": 0.08, "top": 0.98},
)

# Define the color scale
vmin, vmax, c = np.log10(0.02), np.log10(0.4), np.log10(df["albedo.value"])

# a/e scatter plot
ax[0].scatter(
    df["proper_elements.proper_semi_major_axis.value"],
    df["proper_elements.proper_eccentricity.value"],
    s=1,
    alpha=0.1,
    c=c,
    vmin=vmin,
    vmax=vmax,
)

# a/i scatter plot
im = ax[1].scatter(
    df["proper_elements.proper_semi_major_axis.value"],
    df["proper_elements.proper_inclination.value"],
    s=1,
    alpha=0.1,
    c=c,
    vmin=vmin,
    vmax=vmax,
)

# Color bar
cbar = fig.colorbar(im, ax=ax.ravel().tolist(), shrink=0.6, label="Albedo")
cbar.set_ticks([np.log10(x) for x in [0.02, 0.05, 0.1, 0.2, 0.4]])
cbar.set_ticklabels([0.02, 0.05, 0.1, 0.2, 0.4])
cbar.solids.set_alpha(1)


# Axes
ax[1].set_xlabel("Proper semi-major axis / au")
ax[0].set_ylabel("Proper eccentricity")
ax[1].set_ylabel("Proper inclination / deg.")
ax[0].set_xlim(1.78, 3.75)
ax[0].set_ylim(0, 0.4)
ax[1].set_ylim(0, 35)

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