import zmat, time

data = bytes(range(256)) * 4000   # ~1 MB, moderately compressible

methods = ['zlib', 'gzip', 'lzma', 'lzip', 'lz4', 'lz4hc', 'zstd']

print(f"{'Method':<10} {'Size':>8}  {'Ratio':>6}  {'Comp':>8}  {'Decomp':>8}")
print("-" * 50)

for m in methods:
    t0 = time.perf_counter()
    c  = zmat.compress(data, method=m)
    tc = time.perf_counter() - t0

    t0 = time.perf_counter()
    d  = zmat.decompress(c, method=m)
    td = time.perf_counter() - t0

    assert d == data, ff"{m} round-trip failed!"
    print(ff"{m:<10} {len(c):>8}  {len(c)/len(data):>6.3f}  {tc:>7.4f}s  {td:>7.4f}s")

# ── Method names are case-insensitive ───────────────────────────
c1 = zmat.compress(data, method='ZLIB')
c2 = zmat.compress(data, method='zlib')
assert c1 == c2

import zmat

# ── Blosc2 meta-compressor (ideal for numeric arrays) ───────────
import struct
# Simulate a 3-D uint16 volume (128×128×64)
vol = bytes(range(256)) * (128*128*64 * 2 // 256)

# All five blosc2 variants
for m in ['blosc2blosclz','blosc2lz4','blosc2lz4hc','blosc2zlib','blosc2zstd']:
    c = zmat.compress(vol, method=m)
    d = zmat.decompress(c, method=m)
    assert d == vol
    print(ff"{m:<18}  {len(c):>8} B  ratio={len(c)/len(vol):.3f}")

# ── Advanced: multi-threading + byte-shuffle ────────────────────
# typesize=2 → uint16 elements; shuffle=1 → byte-shuffle filter
c_adv = zmat.zmat(vol, iscompress=-6, method='blosc2zstd',
                  nthread=4, shuffle=1, typesize=2)
d_adv = zmat.zmat(c_adv, iscompress=0, method='blosc2zstd')
assert d_adv == vol
print(ff"Advanced blosc2zstd: {len(c_adv)} B")

import numpy as np
import zmat

# ── Compress / restore a NumPy array ────────────────────────────
arr = np.random.rand(1000, 1000)           # 8 MB float64
raw = arr.tobytes()                          # zmat works on bytes

c = zmat.compress(raw, method='lz4')
print(f"lz4: {len(raw):,} → {len(c):,} B")

restored = np.frombuffer(
    zmat.decompress(c, method='lz4'),
    dtype=arr.dtype
).reshape(arr.shape)
assert np.array_equal(arr, restored)

# ── Blosc2 + shuffle for integer arrays ─────────────────────────
vol   = np.arange(128*128*64, dtype=np.uint16).reshape(128,128,64)
cbest = zmat.zmat(vol.tobytes(), iscompress=1, method='blosc2zstd',
                  typesize=2, shuffle=1)    # typesize = itemsize
dbest = zmat.zmat(cbest, iscompress=0, method='blosc2zstd')
vrest = np.frombuffer(dbest, dtype=np.uint16).reshape(vol.shape)
assert np.array_equal(vol, vrest)

# ── Structured dtype round-trip ─────────────────────────────────
dt  = np.dtype([('x', np.float32), ('y', np.float32), ('z', np.float32)])
pts = np.zeros(100000, dtype=dt)
c2  = zmat.compress(pts.tobytes(), method='zstd')
print(f"Struct array: {pts.nbytes:,} → {len(c2):,} B")

import zmat, zlib, gzip, base64, io

# ── Base64 encode / decode ──────────────────────────────────────
data    = b"ZMat Python base64 test"
encoded = zmat.encode(data, method='base64')
decoded = zmat.decode(encoded, method='base64')
assert decoded == data

# zmat base64 matches Python stdlib (strip newlines for comparison)
assert encoded.replace(b'\n', b'') == base64.b64encode(data)

# Chain: compress → base64  (useful for JSON / XML embedding)
payload = b"scientific data" * 5000
c   = zmat.compress(payload, method='zlib')
b64 = zmat.encode(c, method='base64')
print(f"base64 string: {len(b64)} chars")

# Reverse
c2  = zmat.decode(b64, method='base64')
out = zmat.decompress(c2, method='zlib')
assert out == payload

# ── stdlib interoperability ─────────────────────────────────────
text = b"interoperability test" * 500

# zmat zlib ↔ Python zlib
assert zlib.decompress(zmat.compress(text, method='zlib')) == text
assert zmat.decompress(zlib.compress(text), method='zlib')  == text

# zmat gzip ↔ Python gzip
cg = zmat.compress(text, method='gzip')
with gzip.open(io.BytesIO(cg), 'rb') as f:
    assert f.read() == text

! test_zmat.f90 – Fortran90 binding via zmat_module
program zmat_demo
    use iso_c_binding
    use zmat_module
    implicit none

    character(len=40)         :: input_str = "ZMat Fortran demo"
    integer(c_size_t)          :: input_len, output_len, decoded_len
    integer(c_int)             :: zipid, status, clevel
    type(c_ptr)                :: out_ptr, dec_ptr
    character(kind=c_char), pointer :: out_arr(:), dec_arr(:)
    integer :: i

    input_len = len_trim(input_str)
    zipid  = 0    ! 0=zlib, 5=lz4, 7=zstd …
    clevel = -6   ! compression level

    ! Compress
    call zmat_run(input_len, input_str, output_len, out_ptr, zipid, status, clevel)
    call c_f_pointer(out_ptr, out_arr, [output_len])
    print *, "Compressed:", input_len, "→", output_len, "bytes"

    ! Decompress (clevel=0)
    call zmat_run(output_len, out_arr, decoded_len, dec_ptr, zipid, status, 0)
    call c_f_pointer(dec_ptr, dec_arr, [decoded_len])
    print *, "Decompressed:", (dec_arr(i), i=1,decoded_len)
end program zmat_demo