import numpy as np
import h5py


def load_root_hdf5_file(classlabel, ratio):
    #print("INFO: file name:", filename)

    h5_data= h5py.File('/stage0/gd2392/research/projects/dune/generators/from_root_hdf5_to_dataset/inputs/v2/SN/SN_0.h5', 'r')#1000 images per file.

    extents =  h5_data['/Data/image2d_tpc_group/extents']
    image_extents =  h5_data['/Data/image2d_tpc_group/image_extents']
    image_meta =  h5_data['/Data/image2d_tpc_group/image_meta']
    images = h5_data['/Data/image2d_tpc_group/images']

    image_count = len(extents)
    image_shape = image_meta[0][3] # Get the shape of the first image    

    for i in range(image_count):
        image_shape_i = image_meta[i][3]
        if (image_shape_i[0] != image_shape[0] or image_shape_i[1] != image_shape[1]):
            raise SystemExit("ERROR: the images are not equally shaped in the dataset")

    images = np.asarray(images)
    images = images.reshape(image_count, image_shape[0], image_shape[1])

    image_subset_count = int(image_count * ratio)
    print("INFO: use", image_subset_count, "out of", image_count, "images")
    images = images[0:image_subset_count]

    labels = np.full((image_subset_count,), classlabel)


    return (images, labels)

def dec_to_binary(my_int):
    """   
    Format a number as binary with leading zeros
    """
    if my_int < 4096:
        return "{0:012b}".format(my_int)
    else:
        return "111111111111"



def main():
    images, labels=load_root_hdf5_file(1,0.01)#this willl load 10 images with ratio=0.01
    print(len(images))
    images = np.reshape(images, (1,21542400))
    print(np.shape(images))

    images = images.astype('uint16')
    f=open("binfile.bin","wb")#10 images
    arr=bytearray(images)
    f.write(arr)
    f.close()
    #images_12bit=map(dec_to_binary, images)

    np.savez('images.npz',images)

    np.savetxt("5evts.csv", images_12bit.reshape((5,-1)))

if __name__== '__main__':
    main()