import matplotlib.pyplot as plt
import numpy as np
import arviz as az
from pymc_marketing.mmm.transformers import WeibullType, weibull_adstock
plt.style.use('arviz-darkgrid')

spend = np.zeros(50)
spend[0] = 1

shapes = [0.5, 1., 1.5, 5.]
scales = [10, 20, 40]
modes = [WeibullType.PDF, WeibullType.CDF]

fig, axes = plt.subplots(
    len(shapes), len(modes), figsize=(12, 8), sharex=True, sharey=True
)
fig.suptitle("Effect of Changing Weibull Adstock Parameters", fontsize=16)

for m, mode in enumerate(modes):
    axes[0, m].set_title(f"Mode: {mode.value}")

    for i, shape in enumerate(shapes):
        for j, scale in enumerate(scales):
            adstock = weibull_adstock(
                spend, lam=scale, k=shape, type=mode, l_max=len(spend)
            ).eval()

            axes[i, m].plot(
                np.arange(len(spend)),
                adstock,
                label=f"Scale={scale}",
                linestyle="-",
            )

fig.legend(
    *axes[0, 0].get_legend_handles_labels(),
    loc="center right",
    bbox_to_anchor=(1.2, 0.85),
)

plt.tight_layout(rect=[0, 0, 0.9, 1])
plt.show()