# Script to write .dat files for testing
# Gives the opportunity to test the dat files
# with small files

import numpy as np
import tifffile

img = tifffile.imread("AldoJF500k.tiff")

# Filename is name_000000.dat
# Header is frame number and bunch id 8+8 bytes
# Data is 2 bytes per pixel

header = np.zeros(2, dtype=np.uint64)
file_index = 0
frame_number = 1
frames = []
headers = []
for file_index in range(4):
    with open(f"AldoJF500k_{file_index:06d}.dat", "wb") as f:
        for i in range(1,8,1):
            frame = np.random.normal(img, np.sqrt(img))
            frame = np.clip(frame, 0, 65535) #avoid negative values
            frame = frame.astype(np.uint16)
            header[0] = frame_number
            header[1] = frame_number**2
            frame_number += 1
            header.tofile(f)
            frame.tofile(f)
            headers.append(header.copy())
            frames.append(frame.copy( ))

            if file_index == 3 and i == 3:
                break

headers = np.array(headers)
frames = np.array(frames)
print(header.shape, frames.shape)

np.savez('AldoJF500k.npz', headers=headers, frames=frames)