# AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/40_Interpolation-gymnastics.ipynb.

# %% auto 0
__all__ = ['MAXRF_SPECTRA', 'MAXRF_CUBE', 'MAXRF_MAX_SPECTRUM', 'MAXRF_SUM_SPECTRUM', 'MAXRF_ENERGIES', 'convert_crono', 'tree',
           'get_array']

# %% ../notebooks/40_Interpolation-gymnastics.ipynb 28
# 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 

# %% ../notebooks/40_Interpolation-gymnastics.ipynb 29
# 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 
