"""Numerical primitive parameters stored in ``primitives.npz``."""

from __future__ import annotations

from collections.abc import Sequence
from pathlib import Path

import numpy as np

from .configuration import ConfigError

PRIMITIVES_NAME = "primitives.npz"


def _key(dimension: int) -> str:
    return f"dimension_{dimension:04d}"


class PrimitiveStore:
    """Immutable numerical primitive axes in global action-dimension order."""

    def __init__(self, values: Sequence[np.ndarray | Sequence[float]]) -> None:
        if not values:
            raise ConfigError("primitives must contain at least one dimension")
        arrays = []
        for dimension, raw in enumerate(values):
            array = np.asarray(raw)
            if array.dtype.hasobject:
                raise ConfigError(f"primitive dimension {dimension} must not use object dtype")
            if array.ndim != 1 or not len(array):
                raise ConfigError(f"primitive dimension {dimension} must be a non-empty vector")
            real_numeric = np.issubdtype(array.dtype, np.integer) or np.issubdtype(
                array.dtype, np.floating
            )
            if not real_numeric or not np.all(np.isfinite(array)):
                raise ConfigError(f"primitive dimension {dimension} must contain finite numbers")
            arrays.append(np.array(array, copy=True))
        self.values = tuple(arrays)
        self.bins = tuple(len(array) for array in arrays)

    @property
    def dimension(self) -> int:
        return len(self.values)

    @classmethod
    def from_pretrained(cls, directory: str | Path, *, input_dim: int) -> PrimitiveStore:
        """Load exactly one safe numeric vector per expected global dimension."""
        path = Path(directory) / PRIMITIVES_NAME
        with np.load(path, allow_pickle=False) as payload:
            expected = {_key(index) for index in range(input_dim)}
            actual = set(payload.files)
            missing = sorted(expected - actual)
            unknown = sorted(actual - expected)
            if missing:
                raise ConfigError(f"{PRIMITIVES_NAME} is missing array(s): {', '.join(missing)}")
            if unknown:
                raise ConfigError(
                    f"{PRIMITIVES_NAME} contains unknown array(s): {', '.join(unknown)}"
                )
            values = [payload[_key(index)] for index in range(input_dim)]
        return cls(values)

    def save_pretrained(self, directory: str | Path) -> None:
        """Write ``primitives.npz`` without pickled objects or path metadata."""
        path = Path(directory) / PRIMITIVES_NAME
        path.parent.mkdir(parents=True, exist_ok=True)
        np.savez(path, **{_key(index): value for index, value in enumerate(self.values)})