import numpy as np
import matplotlib.pyplot as plt

from beast.physicsmodel.prior_weights_stars import compute_mass_prior_weights

fig, ax = plt.subplots()

# mass grid from 0.01 to 100 solar masses (log spacing)
masses = np.logspace(-2.0, 2.0)

mass_prior_models = [
    {"name": "kroupa"},
    {"name": "salpeter"},
    {"name": "flat"}
]

for mp_mod in mass_prior_models:
    ax.plot(masses, compute_mass_prior_weights(masses, mp_mod), label=mp_mod["name"])

ax.set_ylabel("probability")
ax.set_xlabel("mass")
ax.set_yscale("log")
ax.set_xscale("log")
ax.legend(loc="best")
plt.tight_layout()
plt.show()