"""History-only regime conditioning: trailing volatility / relative volume statistics -> embedding added to the
decoder queries (option 'stats'); 'learned' = a small codebook selected by a soft classifier over the same stats;
'none' = disabled (ablation H3)."""
from __future__ import annotations

import torch
import torch.nn as nn


class RegimeEmbedding(nn.Module):
    def __init__(self, f_global: int, d_model: int, mode: str = "stats", n_codes: int = 8):
        super().__init__()
        self.mode = mode
        if mode == "stats":
            self.mlp = nn.Sequential(nn.Linear(f_global, d_model), nn.GELU(), nn.Linear(d_model, d_model))
        elif mode == "learned":
            self.cls = nn.Sequential(nn.Linear(f_global, d_model), nn.GELU(), nn.Linear(d_model, n_codes))
            self.codes = nn.Embedding(n_codes, d_model)
        elif mode != "none":
            raise ValueError(mode)

    def forward(self, glob: torch.Tensor) -> torch.Tensor | None:
        if self.mode == "none":
            return None
        if self.mode == "stats":
            return self.mlp(glob)
        w = torch.softmax(self.cls(glob), dim=-1)
        return w @ self.codes.weight