import arviz as az
import matplotlib.pyplot as plt
import numpy as np

from pymc_marketing.mmm import MonthlyFourier
from pymc_marketing.prior import Prior

az.style.use("arviz-white")

seed = sum(map(ord, "Monthly"))
rng = np.random.default_rng(seed)

mu = np.array([0, 0, 0.5, 0])
b = 0.075
dist = Prior("Laplace", mu=mu, b=b, dims="fourier")
monthly = MonthlyFourier(n_order=2, prior=dist)
prior = monthly.sample_prior(samples=100)
curve = monthly.sample_curve(prior)

_, axes = monthly.plot_curve(curve)
axes[0].set(title="Monthly Fourier Seasonality")
plt.show()