"""Hugging Face `PreTrainedModel` wrapper for TinyLunarPolicy (auto_map / trust_remote_code).

Usage:
    from transformers import AutoModel
    model = AutoModel.from_pretrained("<repo>", trust_remote_code=True)
    logits = model(torch.tensor(obs, dtype=torch.float32))   # (4,)
    action = int(logits.argmax())
"""
import torch
import torch.nn as nn
from transformers import PreTrainedModel

from configuration_tiny import TinyLunarConfig


class TinyLunarPolicy(PreTrainedModel):
    config_class = TinyLunarConfig
    # no tied weights; these two forms keep compatibility across transformers v4 and v5
    _tied_weights_keys = []
    all_tied_weights_keys = {}

    def __init__(self, config: TinyLunarConfig):
        super().__init__(config)
        self.linear = nn.Linear(config.obs_dim, config.n_actions, bias=config.bias)

    @property
    def n_params(self) -> int:
        return sum(p.numel() for p in self.parameters())

    def forward(self, obs: torch.Tensor) -> torch.Tensor:
        """obs: (..., obs_dim) -> logits (..., n_actions)."""
        return self.linear(obs)

    @torch.no_grad()
    def act(self, obs) -> int:
        o = torch.as_tensor(obs, dtype=torch.float32)
        return int(self.forward(o).argmax().item())