from bioio import BioImage
import bioio_ome_zarr

img = BioImage("my_file.zarr", reader=bioio_ome_zarr.Reader)
img.data

from bioio import BioImage
path = "https://allencell.s3.amazonaws.com/aics/nuc-morph-dataset/hipsc_fov_nuclei_timelapse_dataset/hipsc_fov_nuclei_timelapse_data_used_for_analysis/baseline_colonies_fov_timelapse_dataset/20200323_09_small/raw.ome.zarr"
image = BioImage(path)
print(image.get_image_dask_data())

from bioio_ome_zarr.writers import OMEZarrWriter
import numpy as np

# Minimal 2D example (Y, X)
data = np.random.randint(0, 255, size=(64, 64), dtype=np.uint8)

writer = OMEZarrWriter(
    store="basic.zarr",
    shape=data.shape,   # (Y, X)
    dtype=data.dtype,
)

# Write the data to the store
writer.write_full_volume(data)

from bioio_ome_zarr.writers import OMEZarrWriter, Channel
import numpy as np

shape = (2, 3, 4, 8, 8)  # (T, C, Z, Y, X)
data = np.random.randint(0, 255, size=shape, dtype=np.uint8)

channels = [Channel(label=f"c{i}", color="FF0000") for i in range(shape[1])]

writer = OMEZarrWriter(
    store="output.zarr",
    shape=shape,
    dtype=data.dtype,
    zarr_format=3,  # 2 for Zarr v2
    scale=((1, 1, 0.5, 0.5, 0.5),),  # add level at half Z/Y/X
    channels=channels,
    axes_names=["t", "c", "z", "y", "x"],
    axes_types=["time", "channel", "space", "space", "space"],
    axes_units=[None, None, "micrometer", "micrometer", "micrometer"],
)

# Write the entire volume
writer.write_full_volume(data)

# NumPy array is accepted and will be wrapped into Dask automatically
writer.write_full_volume(data)

# Or with a Dask array if you already have one
import dask.array as da
writer.write_full_volume(da.from_array(data, chunks=(1, 1, 1, 8, 8)))

# Suppose your writer axes include "T"; write timepoints in flexible batches
from bioio import BioImage
import dask.array as da

bioimg = BioImage("/path/to/any/bioimage")
data = bioimg.get_image_dask_data()

# Write the entire timeseries at once
writer.write_timepoints(data)

# Write in 5-timepoint batches
for t in range(0, data.shape[0], 5):
    writer.write_timepoints(
    data,
    start_T_src=t,
    start_T_dest=t,
    total_T=5,
    )

# Write source timepoints [10:20] into destination positions [50:60]
writer.write_timepoints(
    data,
    start_T_src=10,
    start_T_dest=50,
    total_T=10,
)

# Provide one chunk shape per level; must match ndim
chunk_shape = (
    (1, 1, 1, 64, 64),  # level 0
    (1, 1, 1, 32, 32),  # level 1
)
writer = OMEZarrWriter(
    store="custom_chunks.zarr",
    shape=(1, 1, 2, 256, 256),
    dtype="uint16",
    zarr_format=3,
    scale=((1, 1, 0.5, 0.5, 0.5),),
    chunk_shape=chunk_shape,
)

# Example data matching the declared shape
arr = np.random.randint(0, 65535, size=(1, 1, 2, 256, 256), dtype=np.uint16)
writer.write_full_volume(arr)

from zarr.codecs import BloscCodec, BloscShuffle

writer = OMEZarrWriter(
    store="sharded_v3.zarr",
    shape=(1, 1, 16, 1024, 1024),
    dtype="uint8",
    zarr_format=3,
    scale=((1, 1, 0.5, 0.5, 0.5),),
    shard_factor=(1, 1, 2, 2, 2),  # multiply chunks to form shard sizes
    compressor=BloscCodec(cname="zstd", clevel=5, shuffle=BloscShuffle.bitshuffle),
)

writer.write_full_volume(
    np.random.randint(0, 255, size=(1, 1, 16, 1024, 1024), dtype=np.uint8)
)

import numcodecs

writer = OMEZarrWriter(
    store="target_v2.zarr",
    shape=(2, 1, 4, 256, 256),
    dtype="uint8",
    zarr_format=2,  # write NGFF 0.4
    scale=((1, 1, 0.5, 0.5, 0.5), (1, 1, 0.25, 0.25, 0.25)),
    compressor=numcodecs.Blosc(cname="zstd", clevel=3, shuffle=numcodecs.Blosc.BITSHUFFLE),
)

writer.write_full_volume(
    np.random.randint(0, 255, size=(2, 1, 4, 256, 256), dtype=np.uint8)
)

# Requires creds for private buckets; public can be anonymous
writer = OMEZarrWriter(
    store="s3://my-bucket/path/to/out.zarr",
    shape=(1, 2, 8, 2048, 2048),
    dtype="uint16",
    zarr_format=3,
)

writer.write_full_volume(
    np.random.randint(0, 65535, size=(1, 2, 8, 2048, 2048), dtype=np.uint16)
)

writer = OMEZarrWriter(
    store="with_transform.zarr",
    shape=(1, 1, 1, 128, 128),
    dtype="uint8",
    zarr_format=3,
    root_transform={"type": "scale", "scale": [1.0, 1.0, 1.0, 0.1, 0.1]},
)
# preview metadata (no disk write)
md = writer.preview_metadata()
print(md)

# Write file 
writer.write_full_volume(
    np.random.randint(0, 255, size=(1, 1, 1, 128, 128), dtype=np.uint8)
)

from bioio_ome_zarr.writers import chunk_size_from_memory_target
chunk_size_from_memory_target((1, 1, 16, 1024, 1024), "uint16", 16 * 1024 * 1024) # 16 mbs
# Returns: (1, 1, 4, 512, 512)

from bioio_ome_zarr.writers import add_zarr_level
add_zarr_level(
    "my_existing.zarr",
    scale_factors=(1, 1, 0.5, 0.5, 0.5),
    compressor=numcodecs.Blosc(cname="zstd", clevel=3, shuffle=numcodecs.Blosc.BITSHUFFLE)
)