# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/load_tiff_images.ipynb (unless otherwise specified).

__all__ = ['load_image', 'load_slice']

# Cell
# export
import pandas
import dotenv
import os
import matplotlib.pyplot as plt
dotenv.load_dotenv()
from PIL import Image
import numpy

# Cell

def load_image(path):
    """
    path - Path to the multipage-tiff file
    Returns all slices
    """
    img = Image.open(path)
    images = []
    for i in range(img.n_frames):
        img.seek(i)
        images.append(numpy.array(img))
    return numpy.array(images)

# Cell

def load_slice(c=None, z=0, t=0, series=None, path=None):
    """
    Return one slice from a multi-page tiff
    """

    n_slices = 3
    n_channels = 7

    if type(z) is not str:
        assert z in range(n_slices), f"There are only {n_slices} Z slices."
    assert c in range(n_channels), f"There are only {n_channels} channels."

    if series is None:
        img_path = path
    else:
        assert os.path.isdir(path), "Path is not found."
        img_path = os.path.join(path, f"series_{series}.ome.tiff")

    assert os.path.isfile(img_path), "Path is not found."

    with Image.open(img_path) as img:
        if (type(z) is str) and (z == "all"):
            images = []
            for z in range(n_slices):
                s = c+z*n_channels
                img.seek(s)
                images.append(numpy.array(img))
            return numpy.array(images, dtype=numpy.uint16)
        else:
            s = c + z*n_channels
            img.seek(s)
            return numpy.array(img, dtype=numpy.uint16)