#!/usr/bin/python3

import h5py

from os import scandir
import re

import tifffile

def pilatus_get_tifflist(basepath, set_index=None):
    '''
    Returns a sorted list of Pilatus TIFF files (full path!)
    from the base pilatus dir `basepath` the data index `index`.
    '''

    if set_index:
        tiffdir = "%s/pilatus/S%.5d" % (basepath, set_index)
    else:
        tiffdir = basepath

    try:

        # The TIFF files generated by Pilatus usually have some format of the
        # form <mumble>_xxx.<extension>, with 'xxx' being an index with a variable
        # number of digits. <mumble> is usually the same everywhere, and <ext>
        # is usually "tif". This algorithm doesn't exacly rely on that being true,
        # but it's a good idea to check, just in case.
        #
        # In any case, we need to reorder the files by the 'xxx' sequence, which
        # means that we have to parse the filename first, sort by 'xxx', then
        # rebuild the path again.
        
        regex = re.compile(r"(.*)([^\d])(\d+).([a-zA-Z]*)")
        tuples = []
        
        for f in scandir(tiffdir):
            m = regex.match(f.path)
            tuples.append((m.group(1), m.group(2), m.group(3), m.group(4)))

        files_sorted = [ f[0]+f[1]+f[2]+"."+f[3]
            for f in sorted(tuples, key=lambda t: int(t[2]))]

        return files_sorted
        
    except FileNotFoundError as e:
        #print("pilatus_get_tifflist: path error:", e)
        raise e

    except IndexError as e:
        #print("pilatus_get_rifflist: file name parse error:", e)
        raise e    
    

if __name__ == "__main__":

    # Trying to open a Nexus HDF5 file and read some UDKM data

    hdfile  =  "/var/home/florin/Documents/sync/Uni/Analysis/221-11141-vonreppert/tmp.hdf5"
    database = "/var/home/florin/Documents/sync/Uni/DataCache/xpp-data/221-11141-vonreppert"

    # name of the TIFF data within the /Sxxx/data/ section of the Nexus file
    tiff_group_name = u'maps'
    
    f = h5py.File(hdfile, "a")
    
    for item in f.attrs.keys():
        print (item + ":", f.attrs[item])

    print ("%d items in file: %r" % (len(f.keys()), f.keys()))

    
    for key in sorted(f.keys(), key=lambda k: int(k[1:])):

        dataset = int(key[1:])

        #print ("Dataset %d" % dataset)
        
        datafile = "%s/pilatus/S%.5d" % (database, dataset)
        setpath  = "/S%d" % (dataset)
        datapath = "/S%d/data" % (dataset)

        #print ("%d / HDF5 set path: %s" % (dataset, datapath))
        #print ("%d / TIFF file path: %s" % (dataset, datafile))
        
        setobj = f[setpath]
        

        try:
            tiffs = pilatus_get_tifflist(database, dataset)

            #mapsobj = setobj.require_group(tiff_group_name)
            #mapsobj.attrs[u'NX_class'] = u'NXentry'
            #mapsobj.attrs[u'set_index'] = dataset

            setobj.attrs[u'default'] = tiff_group_name

            maps_parent = setobj

            dataobj = maps_parent.require_group(tiff_group_name)
            dataobj.attrs['NX_class'] = 'NXdata'
            dataobj.attrs['axes'] = ['x', 'y']            
            maps_parent.attrs[u'default'] = tiff_group_name

            for i in range(len(tiffs)):
                
                cnts = tifffile.imread(tiffs[i])
               
                data_name = "map-%d" % i

                dataobj.attrs['signal'] = data_name
                dataobj.attrs['title'] = tiffs[i]

                # Just a convenience -- put the last succesfully loaded data
                # map as the default map in this map
                
                dcnt = dataobj.require_dataset(name=data_name,
                                               data=cnts,
                                               shape=cnts.shape,
                                               dtype=cnts.dtype,
                                               exact=True)

                dx = dataobj.require_dataset(name='x',
                                             data=[i for i in range(cnts.shape[0])],
                                             shape=(cnts.shape[0]),
                                             dtype=int)
            
                dy = dataobj.require_dataset(name='y',
                                             data=[i for i in range(cnts.shape[1])],
                                             shape=(cnts.shape[1]),
                                             dtype=int)
                
                
            print ("%d / %d TIFF files" % (dataset, len(tiffs)))
            
        except FileNotFoundError:
            print ("%d / No TIFF data" % dataset)
            
        
        #for f in tiffs:
        #    print (f)
        #
        #    print (tiffs[5])
        #    data = tifffile.imread(tiffs[10])
        #    print (data, data.max())
