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

import numpy as np
import tifffile

def write_dat_file_module():
    # 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)
    with np.load('AldoJF500k.npz') as f:
        headers = f["headers"]
        frames = f["frames"]

    new_header = np.zeros(headers.shape[0], dtype = [('framenum', np.uint64), ('bunchid', np.uint64)])
    for i in range(headers.shape[0]):
        new_header[i][0] = headers[i][0]
        new_header[i][1] = headers[i][1]

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

def write_dat_file_half_module():
    # img = tifffile.imread("AldoJF500k.tiff")
    # img = img[0:256,:]

    # header = np.zeros(2, dtype=np.uint64)
    # file_index = 0
    # frame_number = 1
    # frames = []
    # headers = []
    # for file_index in range(7):
    #     with open(f"AldoJF250k_{file_index:06d}.dat", "wb") as f:
    #         for i in range(1,9,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 == 6 and i == 5:
    #                 break

    # headers = np.array(headers)
    # frames = np.array(frames)
    # print(header.shape, frames.shape)
    with np.load('AldoJF250k.npz') as f:
        headers = f["headers"]
        frames = f["frames"]

    new_header = np.zeros(headers.shape[0], dtype = [('framenum', np.uint64), ('bunchid', np.uint64)])
    for i in range(headers.shape[0]):
        new_header[i][0] = headers[i][0]
        new_header[i][1] = headers[i][1]
    np.savez('AldoJF250k.npz', headers=new_header, frames=frames)


def write_dat_file_single_chip():
    # img = tifffile.imread("AldoJF500k.tiff")
    # x = 580
    # y = 220
    # img = img[y:y+256, x:x+256]
    # header = np.zeros(2, dtype=np.uint64)
    # file_index = 0
    # frame_number = 1
    # frames = []
    # headers = []
    # for file_index in range(7):
    #     with open(f"AldoJF65k_{file_index:06d}.dat", "wb") as f:
    #         for i in range(1,19,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 == 6 and i == 5:
    #                 break

    # headers = np.array(headers)
    # frames = np.array(frames)
    # print(header.shape, frames.shape)
    with np.load('AldoJF65k.npz') as f:
        headers = f["headers"]
        frames = f["frames"]

    new_header = np.zeros(headers.shape[0], dtype = [('framenum', np.uint64), ('bunchid', np.uint64)])
    for i in range(headers.shape[0]):
        new_header[i][0] = headers[i][0]
        new_header[i][1] = headers[i][1]
    np.savez('AldoJF65k.npz', headers=new_header, frames=frames)


if __name__ == "__main__":
    # write_dat_file_module()
    write_dat_file_half_module()
    write_dat_file_single_chip()