import numpy as np
import pandas as pd

import matplotlib.pyplot as plt

from pymc_marketing.mmm import HSGP

seed = sum(map(ord, "Change of covariance function"))
rng = np.random.default_rng(seed)

n = 52
X = np.arange(n)

hsgp = HSGP.parameterize_from_data(
    X=X,
    cov_func="matern32",
    dims="time",
)

dates = pd.date_range("2022-01-01", periods=n, freq="W-MON")
coords = {
    "time": dates,
}
prior = hsgp.sample_prior(coords=coords, random_seed=rng)
curve = prior["f"]
hsgp.plot_curve(curve, random_seed=rng)
plt.show()