"""Non-IEEE floats (bfloat16, FP8 E4M3/E5M2, FP6 E2M3/E3M2, FP4 E2M1) read
as h5py 3.16 reads them: as the narrowest IEEE float that holds every value
(bfloat16 as float32, the 1-byte formats as float16), in the file's byte
order, with the values libhdf5 converts them to.

The fixture was written by libhdf5 2.2.0; the JSON next to it holds what
libhdf5 2.2.0 itself returns for every element (see gen_mx_floats.py)."""

import json
import os
import shutil

import numpy as np
import pytest

import clawhdf5

FIXTURES = os.path.join(os.path.dirname(__file__), "..", "..", "clawhdf5", "tests", "fixtures")
FILE = os.path.join(FIXTURES, "mx_floats_hdf5_2_2.h5")
REFERENCE = json.load(open(os.path.join(FIXTURES, "mx_floats_hdf5_2_2.json")))["objects"]

# What h5py 3.16 reports for each (checked 2026-09-28).
DTYPES = {
    "bf16le": "<f4",
    "bf16be": ">f4",
    "f8e4m3": "<f2",
    "f8e5m2": "<f2",
    "f6e2m3": "<f2",
    "f6e2m3_pad": "<f2",
    "f6e3m2": "<f2",
    "f6e3m2_pad": "<f2",
    "f4e2m1": "<f2",
    "f4e2m1_pad": "<f2",
}


def _expected(name, dtype):
    """The libhdf5 2.2.0 values as `dtype`, NaN bits included (sign kept,
    every mantissa bit set)."""
    out = []
    for text in REFERENCE[name]["f64"]:
        if text.startswith("nan:"):
            negative = int(text[4:], 16) >> 63
            bits = dtype.itemsize * 8
            word = (negative << (bits - 1)) | ((1 << (bits - 1)) - 1)
            out.append(np.frombuffer(word.to_bytes(dtype.itemsize, "little"), dtype.newbyteorder("<"))[0])
        else:
            out.append(float(text))
    return np.array(out, dtype=dtype.newbyteorder("<")).astype(dtype)


def test_values_match_libhdf5_2_2():
    assert set(DTYPES) == set(REFERENCE)
    with clawhdf5.File(FILE, "r") as f:
        for name, dtype in DTYPES.items():
            ds = f[name]
            assert ds.dtype == np.dtype(dtype), name
            assert ds.dtype.str == dtype, name
            want = _expected(name, np.dtype(dtype))
            got = ds[()]
            assert got.dtype.str == dtype
            assert got.tobytes() == want.tobytes(), name
            # Selections convert the same way.
            assert ds[3:9].tobytes() == want[3:9].tobytes(), name
            assert ds[[0, 5, 7]].tobytes() == want[[0, 5, 7]].tobytes(), name
            assert ds[5].tobytes() == want[5].tobytes(), name
            if REFERENCE[name]["attribute"]:
                attr = f.attrs[name]
                assert attr.dtype.str == dtype
                assert attr.tobytes() == want.tobytes(), name


def test_matches_h5py(h5py):
    with clawhdf5.File(FILE, "r") as ours, h5py.File(FILE, "r") as theirs:
        for name in DTYPES:
            assert ours[name].dtype == theirs[name].dtype, name
            assert ours[name][()].tobytes() == theirs[name][()].tobytes(), name
            if name in theirs.attrs:
                assert ours.attrs[name].dtype == theirs.attrs[name].dtype
                assert ours.attrs[name].tobytes() == theirs.attrs[name].tobytes()


def test_writing_is_refused(tmp_path):
    path = tmp_path / "copy.h5"
    shutil.copyfile(FILE, path)
    with clawhdf5.File(str(path), "r+") as f:
        with pytest.raises(NotImplementedError, match="non-IEEE"):
            f["bf16le"][0] = 1.0
        with pytest.raises(NotImplementedError, match="non-IEEE"):
            f["f4e2m1"][:] = np.zeros(16)
    with open(path, "rb") as a, open(FILE, "rb") as b:
        assert a.read() == b.read()