from vidata.io import (
    load_image, save_image,
    load_tif, save_tif,
    load_sitk, save_sitk,
    load_nib, save_nib,
    load_blosc2, save_blosc2,
    load_blosc2pkl, save_blosc2pkl,
    load_npy, save_npy,
    load_npz, save_npz,
)

# Standard image formats (PNG, JPG, BMP, …)
img, meta = load_image("example.png")
save_image(img, "out.png", meta)

# TIFF (supports multipage / high bit depth)
img, meta = load_tif("example.tif")
save_tif(img, "out.tif", meta)

# Medical imaging (NIfTI, MHA, NRRD) with SimpleITK
vol, meta = load_sitk("example.nii.gz")
save_sitk(vol, "out_sitk.nii.gz", meta)

# Medical imaging with Nibabel
vol, meta = load_nib("example.nii.gz")
save_nib(vol, "out_nib.nii.gz", meta)

# Blosc2 compressed array
arr, meta = load_blosc2("example.b2nd")
save_blosc2(arr, "out.b2nd", meta)

# Blosc2 compressed array but with metadata in a separate pickle file (with the same name)
arr, meta = load_blosc2pkl("example.b2nd")
save_blosc2pkl(arr, "out.b2nd", meta)

# NumPy arrays
arr, _ = load_npy("example.npy")
save_npy(arr, "out.npy")

arrs, _ = load_npz("example.npz")
save_npz(arrs, "out.npz")

from vidata.io import (
    load_json, save_json,
    load_yaml, save_yaml,
    load_pickle, save_pickle,
    load_txt, save_txt,
)

# JSON
obj = load_json("config.json")
save_json(obj, "out.json")

# YAML
obj = load_yaml("config.yaml")
save_yaml(obj, "out.yaml")

# Pickle (Python objects)
obj = load_pickle("config.pkl")
save_pickle(obj, "out.pkl")

# Plain text
obj = load_txt("config.txt")
save_txt(obj, "out.txt")

# custom_io_template.py  — fill in the TODOs and import this module somewhere at startup.
import numpy as np
from typing import Tuple, Dict, List

# TODO: import your backend library (e.g., imageio, tifffile, nibabel, SimpleITK, ...)
# import imageio.v3 as iio

from vidata.registry import register_loader, register_writer

# --------------------------- READER ------------------------------------------
# Replace file extension and backend name to your custom function
@register_loader("image", ".png", ".jpg", ".jpeg", ".bmp", backend="imageio")  # To Register Image Loading
@register_loader("mask", ".png", ".bmp", backend="imageio") # To Register Label Loading
def load_custom(file: str) -> tuple[np.ndarray, dict]:
    """
    Load a file and return (data, metadata).
    metadata can be empty or include keys like: spacing, origin, direction, shear, dtype, etc.
    """
    # data = iio.imread(file)  # example for imageio
    data = ...  # TODO: replace
    meta = {}   # TODO: replace
    return data, meta
    # return meta # for metadata files like json or yaml

# --------------------------- WRITER ------------------------------------------
# Replace file extension and backend name to your custom function
@register_writer("image", ".png", ".jpg", ".jpeg", ".bmp", backend="imageio") # To Register Image Writer
@register_writer("mask", ".png", ".bmp", backend="imageio") # To Register Label Writer
def save_custom(data: np.ndarray, file: str) -> list[str]:
    """
    Save array to `file`. Return all created paths (include sidecars if any).
    """
    # TODO: write using your backend
    # iio.imwrite(file, data)
    return [file]

from vidata.loaders import ImageLoader, SemSegLoader, MultilabelLoader
from vidata.writers import ImageWriter, SemSegWriter, MultilabelWriter

# Image Loader/Writer
loader = ImageLoader(ftype=".png")
writer = ImageWriter(ftype=".png")
image, meta = loader.load("path/to/file.png")
writer.save(image,"path/to/output.png",meta)

# Semantic Segmentation Loader/Writer
loader = SemSegLoader(ftype=".png")
writer = SemSegWriter(ftype=".png")
mask, meta = loader.load("path/to/file.png")
writer.save(mask,"path/to/output.png",meta)

# Multilabel Segmentation Loader/Writer
loader = MultilabelLoader(ftype=".png")
writer = MultilabelWriter(ftype=".png")
mask, meta = loader.load("path/to/file.png")
writer.save(mask,"path/to/output.png",meta)

from vidata.loaders import ImageStackLoader, MultilabelStackedLoader
from vidata.writers import ImageStackWriter, MultilabelStackedWriter

# Image Loader
#       expects multiple files with numeric suffixes: file_0000.png, file_0001.png, file_0002.png
loader = ImageStackLoader(ftype=".png", channels=3)           # channels=3 --> 3 files are expected
writer = ImageStackWriter(ftype=".png", channels=3)           # channels=3 --> 3 files are expected
image, meta = loader.load("path/to/file")                     # Pass only the base path without suffix and file_type
writer.save(image,"path/to/output.png",meta)                  # each channel will be saved in a separate file

# Multilabel Segmentation Loader
#       expects file names: file_0000.png,... , file_0006.png
loader = MultilabelStackedLoader(ftype=".png", classes=7)     # classes=7 --> 7 files are expected
writer = MultilabelStackedWriter(ftype=".png", classes=7)     # classes=7 --> 7 files are expected
mask, meta = loader.load("path/to/file")                      # Pass only the base path without suffix and file_type
writer.save(mask,"path/to/output.png",meta)                   # each class will be saved in a separate file

from vidata.loaders import ImageLoader

# Explicit backend (SimpleITK)
loader=ImageLoader(ftype=".nii.gz",channels=1, backend="sitk")
image, meta = loader.load("path/to/file.nii.gz")

# Explicit backend (Nibabel)
loader=ImageLoader(ftype=".nii.gz",channels=1, backend="nibabel")
image, meta = loader.load("path/to/file.nii.gz")

# If not specified the first one from the table below is chosen (SimpleITK in this case)
loader=ImageLoader(ftype=".nii.gz",channels=1)
image, meta = loader.load("path/to/file.nii.gz")

from vidata.task_manager import (
    SemanticSegmentationManager,
    MultiLabelSegmentationManager,
)

# Semantic segmentation (2D example)
ssm = SemanticSegmentationManager()

mask = ssm.random(size, num_classes)       # shape: (H, W) or (D, H, W), dtype=int, range: [0:num_classes]
empty = ssm.empty(size, num_classes)       # zeros by default
ids = ssm.class_ids(mask)                  # sorted unique class IDs (e.g., [0,1,2,3,4])
count_bg = ssm.class_count(mask, 1)        # number of pixels/voxels with class 1
coords = ssm.class_location(mask, 3)       # indices (tuple of arrays) where class==3
spatial = ssm.spatial_dims(mask.shape)     # (H, W) or (D, H, W)
has_bg = ssm.has_background()              # True (class 0 is background)

# Multilabel segmentation (3D example)
mlm = MultiLabelSegmentationManager()
ml_mask = mlm.random(size, num_classes)    # shape: (N, H, W) or (N, D, H, W), channel-first one-hot
ml_empty = mlm.empty(size, num_classes)    # all zeros
ml_ids = mlm.class_ids(ml_mask)            # sorted unique class IDs (e.g., [0,1,2,3,4])
ml_count = mlm.class_count(ml_mask, 2)     # number of pixels/voxels where channel 2 is active
ml_coords = mlm.class_location(ml_mask, 6) # indices (tuple of arrays) where channel 6 is active
ml_spatial = mlm.spatial_dims(ml_mask.shape) # (H, W) or (D, H, W)
mlm.has_background()                       # False (class 0 is not background)

from vidata.file_manager import FileManager

# Single-file data (e.g., images or masks stored one-per-file)
fm = FileManager(
    path="data/images",
    file_type=".png",
    pattern="*_image",            # optional; glob-like,
    include_names=["case_","sample_"], # optional; keep only if any token is in file name
    exclude_names=["corrupt"],   # optional; drop if any token is in file name
)
print(len(fm), "files")           # > 30 files | how many files are found
print(fm[0])                      # > data/images/case_1_image.png

from vidata.file_manager import FileManagerStacked

# Single-file data (e.g., images or masks stored one-per-file)
fm = FileManagerStacked(
    path="data/stacked_images",
    file_type=".png",
    pattern="*_image",            # optional; glob-like,
    include_names=["case_","sample_"], # optional; keep only if any token is in file name
    exclude_names=["corrupt"],   # optional; drop if any token is in file name
)
print(len(fm), "files")           # > 30 files | how many files are found (number of base path)
print(fm[0])                      # > data/stacked_images/case_1_image | gives the base path without suffix and dtype

from vidata.config_manager import ConfigManager
from vidata.io import load_yaml

# Load a dataset config
cfg = load_yaml("dataset.yaml")
cm = ConfigManager(cfg)

# Access layers
print(cm.layer_names())                 # ['SomeImages', 'SomeLabels']
image_layer = cm.layer("SomeImages")

# Get a file manager for train split
fm = image_layer.file_manager(split="train", fold=0)
print(len(fm), "files")

# Build a loader
loader = image_layer.data_loader()
arr, meta = loader.load(str(fm[0]))

# Get the task manager for labels
label_layer = cm.layer("SomeLabels")
tm = label_layer.task_manager()
print("Classes in labels:", tm.class_ids(arr))