#!/usr/bin/env python3
"""TinyLunarPolicy - a 14-parameter neural network that lands Gymnasium LunarLander-v3.

    logits = obs @ W          # obs: (8,), W: (8, 4)
    action = argmax(logits)   # action in {0,1,2,3}

W is not free: it is built from **14 trainable scalars** by hard-coding the exact
mirror symmetry of the environment (x -> -x, vx -> -vx, angle -> -angle, w -> -w,
legs swap, action left <-> right; actions 0 "nothing" and 2 "main engine" are
self-paired). So the model has 14 parameters.

Dependencies: torch, safetensors, numpy.
"""
from __future__ import annotations

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

OBS_DIM = 8
N_ACTIONS = 4
N_PARAMS = 14

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


class TinyLunarPolicy(nn.Module):
    """Linear policy 8 -> 4 built from 14 parameters via the mirror symmetry."""

    def __init__(self) -> None:
        super().__init__()
        self.sym_theta = nn.Parameter(torch.zeros(N_PARAMS))

    @property
    def n_params(self) -> int:
        return int(self.sym_theta.numel())

    def weights(self) -> torch.Tensor:
        """Materialise W (8, 4) from the 14 parameters; logits = obs @ W."""
        f = self.sym_theta
        W = torch.zeros(OBS_DIM, N_ACTIONS, dtype=f.dtype)
        # actions 0 and 2 are self-paired under the mirror: only y, vy and leg contact survive
        W[1, 0], W[3, 0], W[6, 0], W[7, 0] = f[0], f[1], f[2], f[2]
        W[1, 2], W[3, 2], W[6, 2], W[7, 2] = f[3], f[4], f[5], f[5]
        # actions 1 and 3 form the chiral pair: W[:, 3] = M @ W[:, 1]
        W[:, 1] = f[6:14]
        W[0, 3], W[1, 3], W[2, 3], W[3, 3] = -W[0, 1], W[1, 1], -W[2, 1], W[3, 1]
        W[4, 3], W[5, 3], W[6, 3], W[7, 3] = -W[4, 1], -W[5, 1], W[7, 1], W[6, 1]
        return W

    def numpy_weights(self) -> np.ndarray:
        return self.weights().detach().cpu().numpy()

    def forward(self, obs: torch.Tensor) -> torch.Tensor:
        """obs: (..., 8) -> logits (..., 4)."""
        return obs @ self.weights()

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

    # -- loading -------------------------------------------------------------
    @classmethod
    def from_safetensors(cls, path: str) -> "TinyLunarPolicy":
        from safetensors.torch import load_file
        model = cls()
        model.load_state_dict(load_file(path))
        model.eval()
        return model

    @classmethod
    def from_pretrained(cls, repo_id_or_path: str = "Dimitrius174/lunarlander-14params",
                        filename: str = "model.safetensors") -> "TinyLunarPolicy":
        import os
        if os.path.isdir(repo_id_or_path):
            return cls.from_safetensors(os.path.join(repo_id_or_path, filename))
        from huggingface_hub import hf_hub_download
        return cls.from_safetensors(hf_hub_download(repo_id_or_path, filename))