#!/usr/bin/env python
'''Writes a NeXus HDF5 file using h5py and numpy'''

import h5py    # HDF5 support
import numpy
import sys

print("Write a NeXus HDF5 file")

if sys.version_info < (3,):
    fileName = "nxdata2d.py2.h5"
else:
    fileName = "nxdata2d.py3.h5"
timestamp = "2010-10-18T17:17:04-0500"

# load data from two column format
data = numpy.arange(100000.)
data.shape = 100, 1000

# create the HDF5 NeXus file
f = h5py.File(fileName, "w")
# point to the default data to be plotted
f.attrs['default']          = u'entry'
# give the HDF5 root some more attributes
f.attrs['file_name']        = fileName
f.attrs['file_time']        = timestamp
f.attrs['creator']          = u'NXdataImage.py'
f.attrs['HDF5_Version']     = h5py.version.hdf5_version
f.attrs['h5py_version']     = h5py.version.version

# create the NXentry group
nxentry = f.create_group('entry')
nxentry.attrs['NX_class'] = 'NXentry'
nxentry.attrs['default'] = u'image_plot'
nxentry.create_dataset('title', data=u'2D Image')

# create the NXdata group
nxdata = nxentry.create_group('image_plot')
nxdata.attrs['NX_class'] = u'NXdata'
nxdata.attrs['signal'] = u'image_data'              # Y axis of default plot
if sys.version_info < (3,):
    string_dtype = h5py.special_dtype(vlen=unicode)
else:
    string_dtype = h5py.special_dtype(vlen=str)
nxdata.attrs['axes'] = numpy.array(['row_name', 'col_name'], dtype=string_dtype) # X axis of default plot

# signal data
ds = nxdata.create_dataset('image_data', data=data)
ds.attrs['interpretation'] = u'image'

data[0:50]= 0
ds = nxdata.create_dataset('image_data2', data=data)
ds.attrs['interpretation'] = u'image'
nxdata.attrs["auxiliary_signals"] = numpy.array(["image_data2"], dtype=string_dtype)


# X axis data
ds = nxdata.create_dataset(u'col_name', data=numpy.arange(data.shape[1]))
ds.attrs['units'] = u'microns'
ds.attrs['long_name'] = u'Pixel Size X (microns)'    # suggested X axis plot label

# Y axis data
ds = nxdata.create_dataset('row_name', data=numpy.arange(data.shape[0]))
ds.attrs['units'] = u'microns'
ds.attrs['long_name'] = u'Pixel Size Y (microns)'    # suggested Y axis plot label

f.close()   # be CERTAIN to close the file

print("wrote file:", fileName)
