"""Dump every hand-written number in handweights.py to weights.safetensors.

The safetensors format is an 8-byte little-endian header length, a JSON header,
then raw tensor bytes, so it is written here with numpy alone.
"""
import json
import struct
from pathlib import Path

import numpy as np

from handweights import (CLR, COL, DEF_COLOR, DEF_SIZE, DEF_Y, DIM, EXP, EXPRS, FACES, GND, GROUNDS, HAND, OBJ, POS,
                         PARTS, PRED, PROPS, REF, SHAPES, SIZ, SKIES, SKY, VOCAB, count_params)

LAYOUT = {"shape_onehot": [0, OBJ], "is_object": OBJ, "pronoun(1=singular,2=plural)": REF,
          "rgb_has_color": [COL, COL + 4], "dx_dy_has_x_has_y": [POS, POS + 4], "size_has_size": [SIZ, SIZ + 2],
          f"expression_onehot({','.join(EXPRS)})": [EXP, PRED], "copula": PRED, "conjunction": CLR,
          f"sky_onehot({','.join(SKIES)})": [SKY, GND], f"ground_onehot({','.join(GROUNDS)})": [GND, DIM]}


def leaves(g, name, out):
    """Flatten an AND/OR/MINUS tree into named plane tensors plus a structure string."""
    op, arg = g
    if op == "and":
        out[name] = arg
        return name
    return f"{op}(" + ", ".join(leaves(c, f"{name}.{i}", out) for i, c in enumerate(arg)) + ")"


def collect():
    words = list(VOCAB)
    tensors = {
        "encoder.embedding": np.stack([VOCAB[w] for w in words]),
        "decoder.default_color": DEF_COLOR,
        "decoder.default_size": DEF_SIZE,
        "decoder.default_y": np.array([y if isinstance(y, float) else np.nan for y in DEF_Y], float),
        "layout.prop_grip_scale": np.array([[*g, sc] for _, g, sc in PROPS.values()], float),
        "layout.hand": np.array(list(HAND.values()), float),
    }
    structure = {}
    groups = list(PARTS.items()) + [(f"face.{k}.{e}", parts) for (k, e), parts in sorted(FACES.items())]
    for name, parts in groups:
        for i, (geom, col, alpha, shaded, outline) in enumerate(parts):
            key = f"renderer.{name}.part{i}"
            structure[key] = dict(geometry=leaves(geom, key, tensors), color=col if isinstance(col, str) else list(col),
                                  alpha=alpha, shaded=shaded, outline=outline)
    meta = {"vocab": json.dumps(words), "shapes": json.dumps(SHAPES), "embedding_layout": json.dumps(LAYOUT),
            "props": json.dumps(list(PROPS)), "default_y_special": "NaN = stands on the ground line",
            "renderer_structure": json.dumps(structure),
            "plane_rows": "each row (nx, ny, d) is a half-plane neuron firing where nx*px + ny*py <= d; a leaf = min-pool over its rows",
            "parameter_count": str(count_params()), "trained": "false"}
    return tensors, meta


def save_safetensors(tensors, meta, path):
    header, blobs, offset = {"__metadata__": meta}, [], 0
    for name, t in tensors.items():
        b = np.ascontiguousarray(t, dtype="<f4").tobytes()
        header[name] = {"dtype": "F32", "shape": list(t.shape), "data_offsets": [offset, offset + len(b)]}
        blobs.append(b)
        offset += len(b)
    h = json.dumps(header, separators=(",", ":")).encode()
    h += b" " * (-len(h) % 8)
    Path(path).write_bytes(struct.pack("<Q", len(h)) + h + b"".join(blobs))


def load_safetensors(path):
    raw = Path(path).read_bytes()
    n = struct.unpack("<Q", raw[:8])[0]
    header = json.loads(raw[8:8 + n])
    meta = header.pop("__metadata__")
    data = raw[8 + n:]
    return {k: np.frombuffer(data[v["data_offsets"][0]:v["data_offsets"][1]], "<f4").reshape(v["shape"])
            for k, v in header.items()}, meta


if __name__ == "__main__":
    out = Path(__file__).parent / "weights.safetensors"
    tensors, meta = collect()
    save_safetensors(tensors, meta, out)
    back, _ = load_safetensors(out)
    assert all(np.allclose(back[k], np.asarray(v, "<f4"), equal_nan=True) for k, v in tensors.items())
    print(f"{out.name}: {len(tensors)} tensors, {sum(t.size for t in tensors.values()):,} numbers, round-trip ok")