import matplotlib as matplotlib
import matplotlib.pyplot as plt
import numpy as np
import torch
from chronos import BaseChronosPipeline

pipeline = BaseChronosPipeline.from_pretrained(
    # "amazon/chronos-bolt-base",
    "amazon/chronos-t5-base",
    device_map="cpu",
    cache_dir="./"
)

# generate gaussian noise
data = np.random.standard_normal(256+64)
train_data, test_data = data[:256], data[256:]

# predict
pred, _ = pipeline.predict_quantiles(
    context=torch.tensor(train_data), 
    quantile_levels=[0.1, 0.5, 0.9],
    prediction_length=len(test_data),
)
low, med, high = pred[0, :, 0], pred[0, :, 1], pred[0, :, 2]

# draw graph
plt.figure(figsize=(10, 5))
plt.plot(data, label="ground truth")
plt.plot(list(range(256, 256+64)), med, linestyle="--", label="median forecast")
plt.fill_between(list(range(256, 256+64)), low, high, color="tomato", alpha=0.2, label="80% prediction interval")

plt.grid()
plt.legend()
plt.savefig(f"gaussian_noise_test.png")