"""
A small, patch-based transformer for zero-shot time series forecasting —
a toy-scale version of the same general recipe used by real time-series
foundation models (Chronos, TimesFM, t0-alpha, etc.):

  1. Per-series normalization: each input window is normalized by its
     OWN mean and std before entering the model, and de-normalized
     using those same stats afterward. This is the single most
     important design choice for generalization — it means the model
     learns SHAPE (trend, seasonality, noise structure), not absolute
     numbers, so a model trained on synthetic series scaled 0.5-500 can
     meaningfully forecast a real series scaled in the thousands.

  2. Patching: rather than feeding every single timestep as a separate
     token (expensive, and treats each point independently), the input
     window is split into small chunks ("patches") and each patch is
     projected into the model's embedding space as one token. This is
     the same idea PatchTST and most modern time-series transformers
     use — it's both more efficient and tends to generalize better
     than point-wise tokenization.

  3. A standard transformer encoder processes the patch tokens, and a
     linear head maps the final representation directly to the
     forecast horizon.

CONTEXT_LEN and HORIZON are fixed (96 and 42) to match the evaluation
setup already used for SARIMA/Prophet/statsforecast/t0-alpha in the
sibling demand-forecasting project, so results are directly comparable.
"""
import torch
import torch.nn as nn

CONTEXT_LEN = 96
HORIZON = 42
PATCH_SIZE = 8
N_PATCHES = CONTEXT_LEN // PATCH_SIZE  # 12
D_MODEL = 32
N_HEADS = 4
N_LAYERS = 2
DIM_FEEDFORWARD = 64


class TinyForecaster(nn.Module):
    def __init__(self):
        super().__init__()
        self.patch_embed = nn.Linear(PATCH_SIZE, D_MODEL)
        self.pos_embed = nn.Parameter(torch.randn(1, N_PATCHES, D_MODEL) * 0.02)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=D_MODEL, nhead=N_HEADS, dim_feedforward=DIM_FEEDFORWARD,
            batch_first=True, dropout=0.1,
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=N_LAYERS)

        self.head = nn.Sequential(
            nn.Linear(D_MODEL * N_PATCHES, 128),
            nn.ReLU(),
            nn.Linear(128, HORIZON),
        )

    def forward(self, context_raw: torch.Tensor) -> torch.Tensor:
        """
        context_raw: (batch, CONTEXT_LEN) raw (unnormalized) values.
        Returns: (batch, HORIZON) forecast, in the ORIGINAL scale.
        """
        mean = context_raw.mean(dim=1, keepdim=True)
        std = context_raw.std(dim=1, keepdim=True).clamp(min=1e-3)
        normed = (context_raw - mean) / std

        patches = normed.view(-1, N_PATCHES, PATCH_SIZE)
        tokens = self.patch_embed(patches) + self.pos_embed
        encoded = self.encoder(tokens)
        flat = encoded.reshape(encoded.size(0), -1)
        forecast_normed = self.head(flat)

        forecast = forecast_normed * std + mean
        return forecast