#!/usr/bin/env python3
"""TinyLunarPolicy — a 32-parameter single-layer neural network that lands LunarLander-v3.

    logits = Linear(8 -> 4, bias=False)(obs)     # 32 parameters
    action = argmax(logits)

This is the *minimal* policy that solves Gymnasium `LunarLander-v3`:
mean return ~231.2, 99.7% landings (0.25% crashes), 60/60 evaluation seeds solved (see README).
It is exactly a single linear layer + argmax, and is mathematically identical to a
single-centroid nearest-neighbour policy in the C2 space (s (x) onehot(a)).

Dependencies: torch, safetensors (numpy optional, only for `.act`).
"""
from __future__ import annotations

import numpy as np
import torch
import torch.nn as nn

OBS_DIM = 8
N_ACTIONS = 4

# Gymnasium LunarLander-v3 observation layout (for interpretability only).
OBS_NAMES = [
    "x", "y", "vx", "vy", "angle", "angular_velocity",
    "leg_left_contact", "leg_right_contact",
]


class TinyLunarPolicy(nn.Module):
    """Linear 8->4 policy head; the action is the argmax of the four logits."""

    def __init__(self, obs_dim: int = OBS_DIM, n_actions: int = N_ACTIONS,
                 bias: bool = False) -> None:
        super().__init__()
        self.obs_dim = int(obs_dim)
        self.n_actions = int(n_actions)
        self.linear = nn.Linear(self.obs_dim, self.n_actions, bias=bias)

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

    def weights(self) -> np.ndarray:
        """Return W of shape (obs_dim, n_actions) such that logits = obs @ W."""
        return self.linear.weight.detach().cpu().numpy().T  # (8, 4)

    # ── policy ───────────────────────────────────────────────────────────────
    def forward(self, obs: torch.Tensor) -> torch.Tensor:
        """obs: (..., 8) -> logits (..., 4)."""
        return self.linear(obs)

    @torch.no_grad()
    def logits(self, obs) -> torch.Tensor:
        o = torch.as_tensor(np.asarray(obs, dtype=np.float32))
        return self.forward(o)

    @torch.no_grad()
    def act(self, obs) -> int:
        """Greedy action for a single observation (array-like of length 8)."""
        return int(torch.argmax(self.logits(obs)).item())

    # ── loading ──────────────────────────────────────────────────────────────
    @classmethod
    def from_safetensors(cls, path: str, **kwargs) -> "TinyLunarPolicy":
        """Load from a local `.safetensors` file."""
        from safetensors.torch import load_file
        model = cls(**kwargs)
        model.load_state_dict(load_file(path))
        model.eval()
        return model

    @classmethod
    def from_pretrained(cls, repo_id_or_path: str = "Dimitrius174/lunarlander-lin32",
                        filename: str = "model.safetensors", **kwargs) -> "TinyLunarPolicy":
        """Load from the Hugging Face Hub (or a local directory)."""
        import os
        if os.path.isdir(repo_id_or_path):
            return cls.from_safetensors(os.path.join(repo_id_or_path, filename), **kwargs)
        from huggingface_hub import hf_hub_download
        return cls.from_safetensors(hf_hub_download(repo_id_or_path, filename), **kwargs)