"""Numeric and log state encoders.

NumericStateEncoder: a fixed-length feature vector (metrics, sensor readings, a
tabular row) -> one token per feature (value * feature embedding + feature id),
so the core can attend to individual features. Missing values are NaN -> masked.

LogStateEncoder: a text state encoder for logs. Logs are text, so it reuses the
text encoder; it only normalises volatile tokens (hex ids, numbers, timestamps)
and keeps the last `max_lines` lines before tokenisation (see minijev.text_task_utils).
"""
from __future__ import annotations

import torch
from torch import nn


class NumericStateEncoder(nn.Module):
    uses_backbone = False

    def __init__(self, d: int, n_features: int, **_):
        super().__init__()
        self.feat = nn.Parameter(torch.zeros(n_features, d))
        self.scale = nn.Parameter(torch.zeros(n_features, d))
        nn.init.normal_(self.feat, std=0.02)
        nn.init.normal_(self.scale, std=0.02)
        self.mlp = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, d), nn.GELU(), nn.Linear(d, d))

    def forward(self, x):  # (B, F) float, NaN = missing
        valid = ~torch.isnan(x)
        v = torch.nan_to_num(x).unsqueeze(-1)
        tok = self.feat + v * self.scale
        tok = tok + self.mlp(tok)
        return tok, valid