import matplotlib.pyplot as plt

from pymc_marketing.mmm import YearlyFourier
from pymc_extras.prior import Prior

# "fourier" is the default prefix!
prior = Prior(
    "Laplace",
    mu=Prior("Normal", dims="fourier"),
    b=Prior("HalfNormal", sigma=0.1, dims="fourier"),
    dims=("fourier", "hierarchy"),
)
yearly = YearlyFourier(n_order=3, prior=prior)

coords = {"hierarchy": ["A", "B", "C"]}
prior = yearly.sample_prior(coords=coords)
curve = yearly.sample_curve(prior)
yearly.plot_curve(curve)
plt.show()