import argparse
import os
from concurrent import futures
from glob import glob

import h5py
import vigra
import z5py
from tqdm import tqdm

#
# A conda environment with all necessary dependencies can be installed with the following
# command:
# conda create -c conda-forge -n env-name h5py z5py vigra tqdm
#


N_SLICES = 11416
H5_EXTS = ('.h5', '.hdf5', '.hdf')
N5_EXTS = ('.n5',)


def open_file(path, mode='r'):
    ext = os.path.splitext(path)[1]
    if ext in H5_EXTS:
        return h5py.File(path, mode)
    elif ext in N5_EXTS:
        return z5py.File(path, mode)
    else:
        raise ValueError("Invalid file extension %s" % ext)


def to_key(path):
    ext = os.path.splitext(path)[1]
    key = 'setup0/timepoint0/s0' if ext in N5_EXTS else 't00000/s00/0/cells'
    return key


def to_tiff(in_path, out_folder, in_key=None, n_threads=16):
    """ Convert input hdf5 or n5 file to tiff.
    """

    os.makedirs(out_folder, exist_ok=True)

    with open_file(in_path, 'r') as f:

        # if the in_key is None, decuce it from the extension
        if in_key is None:
            in_key = to_key(in_path)

        ds = f[in_key]
        n_slices = ds.shape[0]

        def save_slice(z):
            im = ds[z]
            out_path = os.path.join(out_folder, 'z%05i.tif' % z)
            vigra.impex.writeImage(im, out_path, compression="DEFLATE")

        with futures.ThreadPoolExecutor(n_threads) as tp:
            list(tqdm(tp.map(save_slice, range(n_slices))))


def to_container(in_folder, out_path, out_key=None, n_threads=16, chunk_size=(32, 128, 128)):
    """ Convert input folder with tifs to hdf5 or n5 container.
    """
    if not os.path.isdir(in_folder):
        raise ValueError("Input needs to be a folder with tiff files")
    n_files = glob(os.path.join(in_folder, 'z*.tif'))
    if len(n_files) != N_SLICES:
        raise ValueError("Only found %i slices, expect %i" % (n_files, N_SLICES))

    def load_im(z):
        pass

    im0 = load_im(0)
    shape = (N_SLICES,) + im0.shape

    if out_key is None:
        out_key = to_key(out_path)

    with open_file(out_path, 'a') as f:
        ds = f.create_dataset(out_key, shape=shape, compression='gzip', dtype='uint8',
                              chunks=chunk_size)

        def write_slice(im, z):
            if (im != 0).sum() == 0:
                return
            ds[z] = im

        write_slice(im0, 0)

        def save_slice(z):
            im = load_im(z)
            write_slice(im, z)

        with futures.ThreadPoolExecutor(n_threads) as tp:
            list(tqdm(tp.map(save_slice, range(N_SLICES))))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('input', type=str)
    parser.add_argument('output', type=str)
    args = parser.parse_args()

    in_ext = os.path.split(args.input)
    if in_ext in N5_EXTS or in_ext in H5_EXTS:
        to_tiff(args.input, args.output)
    else:
        to_container(args.input, args.output)