"""Load the exported weights (models/maya2, written by export/export_weights.py)."""
import json
import os

import mlx.core as mx

MODELS = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "models")
# the dynamic 8-bit weights (the published model); the 16-bit export (models/maya2, made by export_weights.py) when
# only that one is present
DEFAULT_DIR = next((os.path.join(MODELS, d) for d in ("maya2-8bit", "maya2") if os.path.isdir(os.path.join(MODELS, d))),
                   os.path.join(MODELS, "maya2-8bit"))
RELEASE_16BIT = os.path.join(MODELS, "maya2")


def load_meta(path=DEFAULT_DIR):
    return json.load(open(os.path.join(path, "meta.json")))


def load(names, path=DEFAULT_DIR):
    """Merge the given safetensors files into one dict of mx arrays (dtypes as stored)."""
    out = {}
    for n in names:
        part = mx.load(os.path.join(path, n))
        dup = set(part) & set(out)
        if dup:
            raise RuntimeError(f"{n}: duplicate tensors {sorted(dup)[:4]}")
        out.update(part)
    return out


def backbone_files(meta):
    return [f for f in meta["files"] if f.startswith("backbone-")]