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

from pymc_marketing.mmm import WeeklyFourier
from pymc_marketing.prior import Prior

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

seed = sum(map(ord, "Weekly"))
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")
weekly = WeeklyFourier(n_order=2, prior=dist)
prior = weekly.sample_prior(samples=100)
curve = weekly.sample_curve(prior)

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