#!/usr/bin/env python3
"""Reproducibly build `model.safetensors` from the source centroid vector.

The source policy is a single 32-dimensional centroid `c` (C2 formulation). It is
identical to the linear layer:   a = argmax(obs @ W),   W[:, a] = c[a*8:(a+1)*8].

  python export_weights.py /path/to/lunar_centroid_12.npy
"""
import argparse
import numpy as np
import torch
from safetensors.torch import save_file

OBS_DIM, N_ACTIONS = 8, 4


def centroid_to_W(c: np.ndarray, obs_dim: int = OBS_DIM, n_actions: int = N_ACTIONS) -> np.ndarray:
    c = np.asarray(c, dtype=np.float32).ravel()
    assert c.size == obs_dim * n_actions
    return np.column_stack([c[a * obs_dim:(a + 1) * obs_dim] for a in range(n_actions)])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("centroid", help="path to the 32-d centroid .npy")
    ap.add_argument("--out", default="model.safetensors")
    a = ap.parse_args()

    c = np.load(a.centroid)
    W = centroid_to_W(c)                          # (8, 4)
    weight = torch.tensor(W.T.copy())             # nn.Linear.weight is (out, in) = (4, 8)
    assert weight.shape == (N_ACTIONS, OBS_DIM)
    save_file({"linear.weight": weight.contiguous()}, a.out)
    print(f"wrote {a.out}  |  {weight.numel()} params  |  W=\n{np.round(W.T, 4)}")


if __name__ == "__main__":
    main()