"""internal utilities for working with gamma software"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/API/gamma_.ipynb.

# %% auto 0
__all__ = ['read_gamma_image', 'read_gamma_pdata', 'write_gamma_image', 'read_gamma_plist', 'write_gamma_plist']

# %% ../nbs/API/gamma_.ipynb 4
import numpy as np
import os

# %% ../nbs/API/gamma_.ipynb 5
def read_gamma_image(imag:str, # gamma raster data
                     width:int, # data width
                     dtype:str='float', # data format, only 'float', 'fcomplex', 'double' and 'int' are supported.
                     y0:int=0, # line number to start reading
                     ny:int=None, # number of lines to read, default: to the last line
                    ):
    # zero value in the image is transformed to 
    # gpu:bool=False, # return a cupy array if true 
    # !! nvidia gpu do not support big endian data, so no way to directly read it
    '''read gamma image into numpy array.'''
    if dtype == 'float':
        dt = '>f4'
    elif dtype == 'fcomplex':
        dt = '>c8'
    elif dtype == 'int':
        dt = '>i4'
    elif dtype =='double':
        dt = '>f8'
    else: raise ValueError('Unsupported data type')

    with open(imag,"rb") as datf:
        datf.seek(0,os.SEEK_END)
        size = datf.tell()
    n_byte = int(dt[-1]) # number of byte per data point
    nlines = int(size / n_byte / width)
    assert y0 <= nlines, 'y0 is larger than height.'
    if ny is None: ny = nlines-y0
    assert y0+ny <= nlines, 'y0+ny is larger than height.'
    offset = y0*width*n_byte
    count = ny*width*n_byte
    # if gpu:
    #     import cupy as cp
    #     import kvikio
    #     import kvikio.zarr
    #     with kvikio.CuFile(imag,'rb') as cufile:
    #         gpu_dt = dt.replace('>','=')
    #         data = cp.empty((ny,width),dtype=dt)
    #         cufile.read(data,size=count,file_offset=offset)
    #     return data
    # else:
    with open(imag,'rb') as datf:
        data = np.fromfile(datf,dtype=dt,offset=offset,count=ny*width)
    if dt != '>i4':
        mask = data==0
        if dt == '>c8':
            data[mask] = np.nan+1j*np.nan
        else:
            data[mask] = np.nan
    data = data.astype(data.dtype.newbyteorder('native'))
    return data.reshape(-1,width)

# %% ../nbs/API/gamma_.ipynb 6
def read_gamma_pdata(
    imag:str, # gamma point cloud data
    dtype:str='float', # data format, only 'float', 'fcomplex', 'double' and 'int' are supported.
):
    # zero value in the data is transformed to nan
    '''read gamma image into numpy array.'''
    if dtype == 'float':
        dt = '>f4'
    elif dtype == 'fcomplex':
        dt = '>c8'
    elif dtype == 'int':
        dt = '>i4'
    elif dtype =='double':
        dt = '>f8'
    else: raise ValueError('Unsupported data type')

    data = np.fromfile(imag,dtype=dt)
    if dt != '>i4':
        mask = data==0
        if dt == '>c8':
            data[mask] = np.nan+1j*np.nan
        else:
            data[mask] = np.nan
    data = data.astype(data.dtype.newbyteorder('native'))
    return data

# %% ../nbs/API/gamma_.ipynb 7
def write_gamma_image(imag,path):
    # support both raster and point cloud
    imag = imag.astype(imag.dtype.newbyteorder('big'))
    imag.tofile(path)

# %% ../nbs/API/gamma_.ipynb 8
def read_gamma_plist(plist:str,dtype='int'):
    return read_gamma_image(plist,width=2,dtype=dtype)

# %% ../nbs/API/gamma_.ipynb 9
def write_gamma_plist(imag,path):
    return write_gamma_image(imag,path)
