import numpy as np
from timesfm3 import TimesFM3Evaluator, ModelConfig


# ============================================================
# 1. 加载 TimesFM 3.0
# ============================================================

config = ModelConfig(
    checkpoint_path="google/timesfm-3.0-pytorch",
    device="cpu",
)

forecaster = TimesFM3Evaluator(config)


# ============================================================
# 2. 单变量预测
# ============================================================

print("=" * 60)
print("单变量预测")
print("=" * 60)

context = np.sin(
    np.linspace(0, 20, 256)
).astype(np.float32)

horizon = 24

out = forecaster.predict(
    context,
    horizon=horizon,
    return_quantiles=True,
)

print("forecast shape :", out.forecast.shape)
print("quantiles shape:", out.quantiles.shape)

print("\n未来 24 步预测值：")
print(out.forecast)


# ============================================================
# 3. 批量预测
# ============================================================

print("\n" + "=" * 60)
print("批量预测")
print("=" * 60)


# 这里就是原来缺失的 ts1 / ts2 / ts3
ts1 = np.sin(np.linspace(0, 20, 256)).astype(np.float32)
ts2 = np.sin(np.linspace(2, 22, 256)).astype(np.float32)
ts3 = np.sin(np.linspace(4, 24, 256)).astype(np.float32)

contexts = [ts1, ts2, ts3]

outs = list(
    forecaster.predict_batch(
        contexts,
        horizon=24,
        return_quantiles=True,
    )
)


for i, result in enumerate(outs, start=1):
    print(f"\n序列 {i}:")
    print("forecast shape :", result.forecast.shape)
    print("quantiles shape:", result.quantiles.shape)
    print("forecast:")
    print(result.forecast)


# ============================================================
# 4. 多变量 + Covariates
# ============================================================

print("\n" + "=" * 60)
print("多变量 + Covariates")
print("=" * 60)


context_len = 128
horizon = 24


# ------------------------------------------------------------
# 多变量目标
# shape = (num_variates, context_len)
# ------------------------------------------------------------

target = np.stack([
    np.sin(np.linspace(0, 24, context_len)),
    np.sin(np.linspace(1, 26, context_len)),
]).astype(np.float32)


# ------------------------------------------------------------
# Past-only covariate
#
# 只有历史数据，没有未来数据
#
# shape = (num_covariates, context_len)
# ------------------------------------------------------------

past_only_cov = np.random.randn(
    1,
    context_len
).astype(np.float32)


# ------------------------------------------------------------
# Past-future covariate
#
# 历史 + 未来都已知
#
# shape = (num_covariates, context_len + horizon)
# ------------------------------------------------------------

past_future_cov = np.sin(
    np.linspace(
        0,
        30,
        context_len + horizon
    )
)[None, :].astype(np.float32)


# ------------------------------------------------------------
# 预测
# ------------------------------------------------------------

out = forecaster.predict(
    target,
    horizon=horizon,
    past_only_covariates=past_only_cov,
    past_future_covariates=past_future_cov,
    return_quantiles=True,
)


print("forecast shape :", out.forecast.shape)
print("quantiles shape:", out.quantiles.shape)


# ============================================================
# 5. 打印多变量预测结果
# ============================================================

print("\n未来 24 步预测：")

for variable_index in range(out.forecast.shape[0]):

    print(
        f"\n变量 {variable_index + 1}:"
    )

    print(
        out.forecast[variable_index]
    )


print("\n测试完成。")