!pip install "welt-client[parquet]"

from welt import Client

with Client() as client:
    client.connect()

import numpy as np
import pandas as pd
from welt import Client, Forecast

# 1. Simulate daily sales for five stores. Replace this with your own history.
HOLIDAYS = {"2024-12-24": "holiday_eve", "2024-12-25": "holiday", "2024-12-26": "holiday",
            "2024-12-31": "holiday_eve", "2025-01-01": "holiday", "2025-04-18": "holiday",
            "2025-04-21": "holiday", "2025-05-01": "holiday", "2025-08-14": "holiday_eve",
            "2025-08-15": "holiday"}
rng = np.random.default_rng(7)
days = pd.date_range("2024-07-01", "2025-08-24", freq="D")
season = np.sin(2 * np.pi * (days.dayofyear - 80) / 365.25)  # peaks in late June
temperature = 13 + 9 * season + rng.normal(0, 2.5, len(days))
holiday = pd.Series(days.strftime("%Y-%m-%d")).map(HOLIDAYS).fillna("regular").to_numpy()
holiday_effect = pd.Series(holiday).map({"regular": 1.0, "holiday_eve": 1.35, "holiday": 0.45}).to_numpy()
stores = {"berlin": (420, 1.25), "hamburg": (310, 1.15), "munich": (380, 1.30),
          "cologne": (240, 1.10), "leipzig": (150, 1.20)}  # base sales, Saturday uplift
frames = []
for store, (base, saturday) in stores.items():
    weekly = np.array([0.85, 0.90, 0.95, 1.00, 1.10, saturday, 0.60])[days.dayofweek]
    promo = (rng.random(len(days)) < 0.1).astype(int)  # planned weeks ahead
    trend = 1 + 0.0004 * np.arange(len(days))
    mean = (base * trend * weekly * (1 + 0.08 * season) * (1 + 0.4 * promo) * holiday_effect
            * (1 + 0.01 * (temperature - 13)))
    frames.append(pd.DataFrame({"date": days, "store": store, "sales": rng.poisson(mean).astype(float),
                                "promo": promo, "holiday": holiday, "temperature": temperature.round(1)}))
data = pd.concat(frames, ignore_index=True)

# 2. Look at the rows before modelling: one row per store and day.
print(data.head())
print("Rows and columns:", data.shape)
print(data.groupby("store")["sales"].describe().round(1))
print(data["holiday"].value_counts())

# 3. Hold out the last four weeks. Split by time, never at random:
#    every held-out day comes after every history day.
HORIZON = 28
cutoff = data["date"].max() - pd.Timedelta(days=HORIZON)
history = data[data["date"] <= cutoff]
test = data[data["date"] > cutoff]
# The future table holds only what you know in advance: no sales, no temperature.
future = test[["date", "store", "promo", "holiday"]]
print("History:", history["date"].min().date(), "to", history["date"].max().date())
print("Held out:", test["date"].min().date(), "to", test["date"].max().date())

# 4. Check the published limits and estimate the cost before sending anything.
with Client() as client:
    chronos = client.model("chronos-2")
variates = 1 + 1 + 2  # target + past covariate + future covariates
longest = history.groupby("store").size().max()
n_series = history["store"].nunique()
fits = any(longest <= c["max_context"] and variates <= c["max_variates"]
           and n_series <= c["max_series"] and HORIZON <= c["max_horizon"]
           for c in chronos["forecast_envelopes"])
print("Published corners:", chronos["forecast_envelopes"])
print("This request fits one:", fits)
tokens = sum((len(rows) + HORIZON) * variates for _, rows in history.groupby("store"))
import math
plain_tokens = sum((len(rows) + HORIZON) * 1 for _, rows in history.groupby("store"))  # step 6 ablation
credits = [max(5, math.ceil(n * 50 / 1_000_000)) for n in (tokens, plain_tokens)]  # 5-credit minimum each
print("Tokens:", tokens, "+", plain_tokens, "· credits:", credits, "=", sum(credits))

# 5. A baseline to beat: repeat each store's last observed week.
parts = []
for store, past in history.groupby("store", sort=False):
    last_week = past.sort_values("date")["sales"].to_numpy()[-7:]
    dates = test.loc[test["store"] == store, "date"].sort_values().to_numpy()
    parts.append(pd.DataFrame({"store": store, "date": dates, "baseline": np.resize(last_week, len(dates))}))
baseline = pd.concat(parts, ignore_index=True)
print(baseline.head(8))

# 6. Forecast every store in one request, with the known future.
model = Forecast(model="chronos-2").fit(
    history, time="date", id="store", target=["sales"],
    past_covariates=["temperature"],
    future_covariates=["promo", "holiday"],
    categorical=["holiday"])
result = model.forecast(HORIZON, future=future, quantile_levels=[0.1, 0.5, 0.9])
print("Job:", result.job_id, "· sorted quantile steps:", result.rearranged_steps)
frame = result.to_frame().rename(columns={"id": "store"})
frame["date"] = frame["timestamp"].dt.tz_localize(None)
print(frame.head())

# The same forecast without covariates: the difference is what the drivers add.
plain = Forecast(model="chronos-2").fit(history, time="date", id="store", target=["sales"])
plain_frame = plain.forecast(HORIZON).to_frame().rename(columns={"id": "store", "median": "plain"})
plain_frame["date"] = plain_frame["timestamp"].dt.tz_localize(None)

# 7. Score the three forecasts on the held-out weeks.
#    MASE divides the mean absolute error by the in-sample error of a one-week-back
#    forecast, so 1.0 means "as good as last week's value on the history".
def mase(actual, forecast, past, m=7):
    return np.mean(np.abs(actual - forecast)) / np.mean(np.abs(past[m:] - past[:-m]))

scored = (test.merge(baseline, on=["store", "date"])
          .merge(frame[["store", "date", "median", "q_0.1", "q_0.9"]], on=["store", "date"])
          .merge(plain_frame[["store", "date", "plain"]], on=["store", "date"]))
rows = []
for store, part in scored.groupby("store", sort=False):
    past = history.loc[history["store"] == store, "sales"].to_numpy()
    inside = (part["sales"] >= part["q_0.1"]) & (part["sales"] <= part["q_0.9"])
    rows.append({"store": store, "seasonal naive": mase(part["sales"], part["baseline"], past),
                 "Chronos-2, no covariates": mase(part["sales"], part["plain"], past),
                 "Chronos-2": mase(part["sales"], part["median"], past), "80% coverage": inside.mean()})
per_store = pd.DataFrame(rows).set_index("store")
scores = per_store.mean()
print(per_store.round(3))
print(scores.round(3))

# 8. Plot the last eight weeks of history, the held-out weeks and both forecasts.
import matplotlib.pyplot as plt

fig, axes = plt.subplots(len(stores), 1, figsize=(10, 2.4 * len(stores)), sharex=True)
for ax, store in zip(axes, stores):
    recent = history[history["store"] == store].tail(56)
    part = scored[scored["store"] == store]
    ax.plot(recent["date"], recent["sales"], color="0.6", label="history")
    ax.plot(part["date"], part["sales"], color="black", label="actual")
    ax.plot(part["date"], part["baseline"], linestyle="--", label="seasonal naive")
    ax.plot(part["date"], part["median"], label="Chronos-2 median")
    ax.fill_between(part["date"], part["q_0.1"], part["q_0.9"], alpha=0.25, label="Chronos-2 10–90%")
    ax.set_title(store, loc="left")
axes[0].legend(ncol=5, fontsize="small", loc="upper left")
fig.tight_layout()
plt.show()

print("Pinned model:", model.model_version_, "· configuration:", model.configuration_version_)
print("Series:", model.series_ids_, "· frequency:", model.frequency_)