"""Reader for artifacts/mlx-nvfp4/ produced by scripts/convert.py.

Tensor names are the artifact's own.  Every layer's 128 experts arrive stacked
and gate/up fused, exactly as the model consumes them, so loading copies nothing:

    model.decoder.layers.<i>.experts.gate_up_proj.{weight,scales,global_scale}
    model.decoder.layers.<i>.experts.down_proj.{weight,scales,global_scale}

`weight` is uint32 (8 NVFP4 values per word), `scales` is one E4M3 byte per 16
input values, and `global_scale` is the checkpoint's per-expert float32 scale
(its amax, 448*6*weight_scale_2) which `mx.gather_qmm` takes as an argument.
Everything else is BF16 under its checkpoint name.
"""

from __future__ import annotations

import json
from collections.abc import Iterator, Mapping
from pathlib import Path

import mlx.core as mx

LAYOUT = "stacked-global-scale"
DEFAULT_DIR = Path(__file__).resolve().parent.parent / "artifacts" / "mlx-nvfp4"
if not DEFAULT_DIR.is_dir():
    # the published repo keeps the shards next to the package, not under artifacts/
    DEFAULT_DIR = Path(__file__).resolve().parent.parent


def load_manifest(directory: str | Path = DEFAULT_DIR) -> dict:
    """Quantisation mode, group size, shard list and tensor -> shard map."""
    manifest = json.loads((Path(directory) / "manifest.json").read_text())
    layout = manifest["quantization"].get("layout")
    if layout != LAYOUT:
        raise ValueError(
            f"{directory}: expected a '{LAYOUT}' artifact from scripts/convert.py, got {layout!r}; "
            "the old absolute-scale artifact is ~5% off and must be regenerated"
        )
    return manifest


def load_weights(directory: str | Path = DEFAULT_DIR) -> dict[str, mx.array]:
    """Every tensor in memory (~17.5 GiB) as name -> mx.array, stacked experts included."""
    directory = Path(directory)
    tensors: dict[str, mx.array] = {}
    for shard in load_manifest(directory)["shards"]:
        tensors.update(mx.load(str(directory / shard["file"])))
    return tensors


class LazyWeights(Mapping[str, mx.array]):
    """Name -> tensor mapping that loads a 2 GiB shard on first access and keeps it.

    One shard holds about three layers of stacked experts, so a layer needs one
    shard, not the whole artifact.
    """

    def __init__(self, directory: str | Path = DEFAULT_DIR):
        self.directory = Path(directory)
        self.manifest = load_manifest(self.directory)
        self._weight_map: dict[str, str] = self.manifest["weight_map"]
        self._shards: dict[str, dict[str, mx.array]] = {}

    def __getitem__(self, key: str) -> mx.array:
        shard = self._weight_map[key]
        if shard not in self._shards:
            self._shards[shard] = mx.load(str(self.directory / shard))
        return self._shards[shard][key]

    def __iter__(self) -> Iterator[str]:
        return iter(self._weight_map)

    def __len__(self) -> int:
        return len(self._weight_map)