# Histograms of albedo with Gaussian fit for C, S, and M taxonomies
# (c) SsODNet, B.Carry & J.Berthier, Lagrange/OCA & LTE/Obs. Paris

import rocks
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit


# Define a Gaussian function
def gauss(x, mu, sigma, peak):
    return peak * np.exp(-0.5 * ((x - mu) / sigma) ** 2)


# Load BFT with necessary columns only for speedup
bft = rocks.load_bft(columns=["albedo.value", "taxonomy.complex", "taxonomy.technique"])

# Choice of albedo range
range = [np.log10(0.02), np.log10(0.5)]
xx = np.linspace(range[0], range[1], num=100)

fig, ax = plt.subplots()
for i, cl in enumerate(["C", "S", "M"]):
    # Spectrally classified objects
    cond = (bft["taxonomy.complex"] == cl) & (bft["taxonomy.technique"] == "Spec")
    n, h, _ = ax.hist(
        np.log10(bft.loc[cond, "albedo.value"]),
        range=range,
        bins=50,
        density=True,
        alpha=0.6,
        color=f"C{i}",
        label=f"{cl} type",
    )

    # Gaussian fit
    param = curve_fit(gauss, h[:-1], n)
    ax.plot(xx, gauss(xx, *param[0]), color=f"C{i}")
    ax.axvline(param[0][0], linestyle="--", color=f"C{i}")

# Axe
ax.legend(loc="upper right")
ax.set_xlabel("Albedo")
ax.set_ylabel("Density count")
xticks = [0.02, 0.05, 0.07, 0.1, 0.2, 0.3, 0.4, 0.5]
ax.set_xticks([np.log10(x) for x in xticks], labels=xticks)
ax.set_xlim(range)

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