import numpy as np
import matplotlib.pyplot as plt
import astropy.units as u

from beast.physicsmodel.priormodel import PriorDustModel

fig, ax = plt.subplots()

# distance grid with linear spacing
d1, d2 = (50.e3, 70.e3)
dists = np.linspace(d1, d2, num=100)
rv1, rv2 = (2.0, 6.0)
rvs = np.arange(rv1, rv2, 0.05)
distim, rvim = np.meshgrid(dists, rvs)

dustmod = {
    "name": "step",
    "dist0": 60 * u.kpc,
    "amp1": 3.1,
    "damp2": 1.4,
    "lgsigma1": 0.01,
    "lgsigma2": 0.01}

dustprior = PriorDustModel(dustmod)
probim = dustprior(rvim, y=distim)

ax.imshow(
    probim, origin="lower", aspect="auto", extent=[d1, d2, rv1, rv2], norm="log"
)

ax.set_ylabel("R(V)")
ax.set_xlabel("distance [kpc]")
ax.set_title("step")
plt.tight_layout()
plt.show()