from Datasource import Datasource
import numpy as np
import h5py
import json
import os

class HDF5(Datasource):
    """ Loads images from hdf5 files

    Attributes
    ------------
    _read: list
        All load functions for :data:`_meta_files`
    """

    # All readers for _meta_files
    _read = [json.load]
    #: All extensions of files pointing to h5
    _meta_files = ['.json']

    @staticmethod
    def dtype(n):
        only_pos = dict(zip(
            map(np.dtype, ('int64', 'int32', 'int16', 'int8')),
            map(np.dtype, ('uint64', 'uint32', 'uint16', 'uint8'))
        ))
        d = n.dtype
        return only_pos.get(d,d)

    @staticmethod
    def load_tile(t_query):
        """load a single tile (image)

        Gets the image path from the \
:data:`TileQuery.RUNTIME`. ``IMAGE`` attribute.

        Gets the position of the image with the whole \
volume from :meth:`TileQuery.all_scales`, \
:meth:`TileQuery.tile_origin`, and \
:meth:`TileQuery.blocksize`.

        Arguments
        -----------
        t_query: :class:`TileQuery`
            With file path and image position

        Returns
        -----------
        np.ndarray
            An image array that may be as large \
as an entire full resolution slice of \
the whole hdf5 volume. Based on the value \
of :meth:`TileQuery.all_scales`, this array \
will likely be downsampled by to a small fraction \
of the full tile resolution.
        """
        # call superclass
        Datasource.load_tile(t_query)
        # Load data for all the h5 files
        h5_files = t_query.RUNTIME.IMAGE.SOURCE.HDF5.VALUE
        # Get all the z indices and coordinates
        z_stops = list(enumerate(zip(*h5_files)[-1]))
        z_starts = z_stops[::-1]

        # Find the region to crop
        sk,sj,si = t_query.all_scales
        [z0,y0,x0],[z1,y1,x1] = t_query.source_tile_bounds
        # Get the scaled blocksize for the output array
        zb,yb,xb = t_query.blocksize

        # get the right h5 files for the current z index
        start_z = next((i for i, z in z_starts if z <= z0), 0)
        stop_z = next((i for i, z in z_stops if z >= z1), len(z_stops))
        needed_files = [h5_files[zi] for zi in range(start_z, stop_z)]

        ####
        # Load from all needed files
        ####
        dtype = getattr(np, t_query.OUTPUT.INFO.TYPE.VALUE)
        # Make the full volume for all needed file volumes
        full_vol = np.zeros([zb, yb, xb], dtype = dtype)

        # Get the first offset
        offset_0 = needed_files[0][-1]

        # Loop through all needed h5 files
        for h5_file in needed_files:
            # Offset for this file
            z_offset = h5_file[-1]
            # Get input and output start
            iz0 = max(z0 - z_offset, 0)
            # Scale output bounds by z-scale
            oz0 = (z_offset - offset_0) // sk

            # Load the image region from the h5 file
            with h5py.File(h5_file[0]) as fd:
                # read from one file
                vol = fd[h5_file[1]]
                # Get the input and output end-bounds
                iz1 = min(z1 - z_offset, vol.shape[0])
                # Scale the output bounds by the z-scale
                dz = iz1 - iz0
                oz1 = oz0 + dz // sk
                # Get the volume from one file
                file_vol = vol[iz0:iz1:sk, y0:y1:sj, x0:x1:si]
                yf, xf = file_vol.shape[1:]
                # Add the volume to the full volume
                full_vol[oz0:oz1,:yf,:xf] = file_vol

        # Combined from all files
        return full_vol

    @staticmethod
    def load_file(h5_file):
        """ Load the needed volume from a single h5 File

        Arguments
        -----------
        t_query: :class:`TileQuery`
            With file path and image position

        """
    @staticmethod
    def preload_source(t_query):
        """load info from example tile (image)

        Calls :meth:`valid_path` to get filename and \
inner dataset path for the full h5 image volume.

        Then gets three needed values from the given \
path from the :class:`TileQuery` t_query

        Arguments
        -----------
        t_query: :class:`TileQuery`
            Only the file path is needed

        Returns
        --------
        dict
            Will be empty if :meth:`valid_path` finds\
this filname to not give a valid h5 volume.

            * :class:`RUNTIME` ``.IMAGE.BLOCK.NAME``
                (numpy.ndarray) -- 3x1 for any give tile shape
            * :class:`OUTPUT` ``.INFO.TYPE.NAME``
                (str) -- numpy dtype of any given tile
            * :class:`OUTPUT` ``.INFO.SIZE.NAME``
                (numpy.ndarray) -- 3x1 for full volume shape
        """
        # Keyword names
        output = t_query.OUTPUT.INFO
        runtime = t_query.RUNTIME.IMAGE
        k_h5 = runtime.SOURCE.HDF5.NAME
        # Get the max block size in bytes for a single tile
        max_bytes = t_query.RUNTIME.CACHE.MAX_BLOCK.VALUE