# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/40_Interpolation-gymnastics.ipynb (unless otherwise specified).

__all__ = ['convert_crono', 'tree', 'get_array', 'MAXRF_SPECTRA', 'MAXRF_CUBE', 'MAXRF_MAX_SPECTRUM',
           'MAXRF_SUM_SPECTRUM', 'MAXRF_ENERGIES']

# Cell

# python package for reading hdf5 files
import h5py

# python package for processing too-big-for-memory data
import dask
import dask.array as da
import dask_ndfilters
from dask.diagnostics import ProgressBar

# python package for a new data container file format called z-arrays
import zarr

# array computions
import numpy as np

# regular expressions
import re

# plotting and printing
import matplotlib.pyplot as plt
from IPython.display import HTML

# standard datapath locations for zarr zipstore
MAXRF_SPECTRA = 'maxrf_all_spectra'
MAXRF_CUBE = 'maxrf_cube'
MAXRF_MAX_SPECTRUM = 'maxrf_max_spectrum'
MAXRF_SUM_SPECTRUM = 'maxrf_sum_spectrum'
MAXRF_ENERGIES = 'maxrf_energies'


# FUNCTION TO READ AND CONVERT CRONO HDF5 FILE

def convert_crono(crono_filename):
    '''Read spectral data from Crono maxrf hdf5 file and convert to zarr zipstore format.

    Compose (smoothed) spectral image data cube from selected spectra. Furthermore,
    calculate smoothed sum and max spectra. '''


    with h5py.File(crono_filename, mode='r') as fh:

        # STEP 1: SCHEDULE COMPUTIONS WITH (LAZY) DASK ARRAYS

        # read spectral data cube from hdf5 dataset into dask array
        dataset = fh['/XRF/Spectra']
        arr = da.from_array(dataset)
        arr = arr.astype(np.float32)

        # schedule spectral dimension gaussian smoothing computation
        spectra_arr = dask_ndfilters.gaussian_filter(arr, (0, 7))

        n_spectra, n_channels = spectra_arr.shape

        # READ SCAN POSITIONS AND SHUFFLE SPECTRA INTO CUBE

        # read and convert position indices to numpy array and remove inner bracket
        spectra_indices = fh['/XRF/SpectraSelectedIndex'][:,:,0]

        # spectral image height and width
        h, w = spectra_indices.shape

        # flatten indices table into list
        spectra_indices = spectra_indices.flatten()

        # shuffle spectrum positions according to indices
        cube_arr = spectra_arr[spectra_indices].reshape([h, w, n_channels])

        # read energy (keV) calibration values of channels
        energy_vector = fh['/XRF/EnergyVector'][:]
        # find channel index of our favorite element iron
        FeKa_i = np.argmin((energy_vector - 6.402)**2)

        # STEP 2: EXECUTE DASK ARRAY COMPUTIONS AND WRITE TO ZIPSTORE

        # create/overwrite and open an empty zarr zipstore file for writing
        zs_filename = re.sub('\.[^.]*$', '.zipstore', crono_filename)
        zs = zarr.ZipStore(zs_filename, mode='w')

        # compute and write maxrf data to datapath in zipstore

        print('Computing (smoothed) spectra...')
        with ProgressBar():
            spectra_arr.to_zarr(zs, component=MAXRF_SPECTRA)

        print('Computing (smoothed) spectral data cube from selection...')
        with ProgressBar():
            cube_arr.to_zarr(zs, component=MAXRF_CUBE)

        print('Computing (smoothed) max spectrum...')
        with ProgressBar():
            (cube_arr.reshape([h*w, n_channels])).max(axis=0).to_zarr(zs, component=MAXRF_MAX_SPECTRUM)

        print('Computing (smoothed) sum spectrum...')
        with ProgressBar():
            ((cube_arr.reshape([h*w, n_channels])).sum(axis=0) / (h*w) ).to_zarr(zs, component=MAXRF_SUM_SPECTRUM)

        print('Writing channel energies...')
        with ProgressBar():
            da.from_array(energy_vector).to_zarr(zs, component=MAXRF_ENERGIES)

        zs.close()

        return zs_filename


# FUNCTIONS FOR USING ZARR ZIPSTORE FILES

def tree(zs_filename, show_arrays=False):
    '''Prints content tree of *zipstore_file*'''

    with zarr.ZipStore(zs_filename, mode='r') as zs:
        root = zarr.group(store=zs)
        tree = root.tree(expand=True).__repr__()
        print(f'\n{zs_filename}:\n\n{tree}')

        if show_arrays:
            datasets = sorted(root)
            arrays_html = ''

            for ds in datasets:
                arr = da.from_array(root[ds])
                html = arr._repr_html_()
                arrays_html = f'{arrays_html}- Dataset: <h style="color:brown">{ds}</h>{html}'

            return HTML(arrays_html)


def get_array(zipstore_filename, datapath, compute=True):
    '''Open zipstore *zipstore_filename* and return dataset from *datapath*.

    For large files that do not fit into memory it is advised to set option compute=False. '''

    # open existing zipstore filehandle
    zs = zarr.ZipStore(zipstore_filename, mode='r')
    root = zarr.group(store=zs)

    # initialize dask array
    arr = da.from_array(root[datapath])

    # return numpy array and close,
    # otherwise return dask array and do not close
    if compute:
        arr = arr.compute()
        zs.close()

    return arr