"""Check sign/quantile symmetry with random weights, not forecast accuracy.
Install pinned TimesFM source, torch==2.8.0 (CPU), numpy==2.2.6.
"""
import json
import numpy as np
import torch
from timesfm import configs
from timesfm.timesfm_2p5.timesfm_2p5_torch import TimesFM_2p5_200M_torch
torch.manual_seed(0)
torch.set_num_threads(4)
tfm = TimesFM_2p5_200M_torch(torch_compile=False)
tfm.model.eval()
tfm.compile(configs.ForecastConfig(max_context=256, max_horizon=256,
 force_flip_invariance=True, use_continuous_quantile_head=False,
 infer_is_positive=False, normalize_inputs=False,
 fix_quantile_crossing=False, return_backcast=False))
x = np.random.RandomState(0).standard_normal((1, 256)).astype(np.float32)
masks = np.zeros((1, 256), dtype=bool)
with torch.no_grad():
    _, pos = tfm.compiled_decode(256, x, masks)
    _, neg = tfm.compiled_decode(256, -x, masks)
pos, neg = np.asarray(pos), np.asarray(neg)
expected = -np.concatenate([pos[..., :1], pos[..., 1:][..., ::-1]], axis=-1)
error = np.abs(neg - expected)
np.testing.assert_allclose(neg, expected, atol=1e-4, rtol=1e-4)
print(json.dumps({'torch': torch.__version__, 'shape': list(pos.shape),
 'random_weights': True, 'max_absolute_error_first_128': float(error[:, :128].max()),
 'max_absolute_error_next_128': float(error[:, 128:].max())}, indent=2))
