import numpy as np
import matplotlib.pyplot as plt

from beast.physicsmodel.prior_weights_dust import PriorWeightsDust

fig, ax = plt.subplots()

# rv grid with linear spacing
rvs = np.linspace(2.0, 6.0, num=200)

dust_prior_models = [
    {"name": "flat"},
    {"name": "lognormal", "max_pos": 3.1, "sigma": 0.25},
    {
        "name": "two_lognormal",
        "max_pos1": 3.1,
        "max_pos2": 4.5,
        "sigma1": 0.1,
        "sigma2": 0.2,
        "N1_to_N2": 2.0 / 5.0
    }
]

for dmod in dust_prior_models:
    dmodel = PriorWeightsDust(
        [1.0], {"name": "flat"}, rvs, dmod, [1.0], {"name": "flat"}
    )

    ax.plot(rvs, dmodel.rv_priors, label=dmod["name"])

ax.set_ylabel("probability")
ax.set_xlabel("R(V)")
ax.legend(loc="best")
plt.tight_layout()
plt.show()