from __future__ import annotations

try:
    import torch
    from torch import nn
except ImportError:  # pragma: no cover
    torch = None
    nn = None


class ChartEncoder(nn.Module if nn is not None else object):
    """Encode a chart summary into a compact CTT chart token.

    The first implementation uses exported numeric chart summaries such as the
    base action chunk. When visual-language embeddings are exported, the same
    module can consume them without changing the CTT operator interface.
    """

    def __init__(self, input_dim: int, hidden_dim: int = 128, output_dim: int = 64) -> None:
        if nn is None:  # pragma: no cover
            raise ImportError("ChartEncoder requires torch")
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim),
            nn.LayerNorm(output_dim),
        )

    def forward(self, chart_features: "torch.Tensor") -> "torch.Tensor":
        return self.net(chart_features)