# AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/30_parsing-bytes.ipynb.

# %% auto 0
__all__ = ['PDZ_25_STRUCTURE_DICT', 'extract_spectra', 'multiparse', 'prefix', 'file_to_bytes', 'parse', 'read_strings',
           'skip_bytes', 'read_table', 'read_counts', 'get_block_at', 'get_blocks', 'get_blocktypes']

# %% ../notebooks/30_parsing-bytes.ipynb 47
import struct 
import numpy as np 
import os 
import re
import matplotlib.pyplot as plt
import pandas as pd 
from IPython.display import display

# %% ../notebooks/30_parsing-bytes.ipynb 48
PDZ_25_STRUCTURE_DICT = {
    25:  {'xformat': 'hi-10X-i', 
          'param_keys': ['pdz_type', 'block_size', 'FileFormatString?', '??']}, 
    1:   {'xformat': 'hi-2S-6s-2S-h-S-T', 
          'param_keys': ['block_type', 'block_size', '??', 'SerialString', '??', '??', '??', '??', '??', '??', 
                         '??', '??', '??', '??', '??', '??', '??']}, 
    2:   {'xformat': 'hi3i8f-*X', 
          'param_keys': ['block_type', 'block_size', '??', 'RawCounts', 'ValidCounts', '??', '??', 
                         '??', 'ActiveTimeInSeconds', 'DeadTimeInSeconds', 'ResetTimeInSeconds', 
                         'LiveTimeInSeconds', 'TotalElapsedTimeInSeconds', '??']}, 
    3:   {'xformat': 'hi-3i9f7hfhfhfhf8hfhi-S-h-Z', 
          'param_keys': ['block_type', 'block_size', '??', 'RawCounts', 'ValidCounts', 
                         '??', '??', '??', 'ActiveTimeInSeconds', 'DeadTimeInSeconds', 
                         'ResetTimeInSeconds', 'LiveTimeInSeconds', 'XrayVoltageInkV', 'XrayFilamentCurrentInMicroAmps', 
                         'Filter1ElementAtomicNum', 'Filter1Thickness', 'Filter2ElementAtomicNum', 'Filter2Thickness', 
                         'Filter3ElementAtomicNum', 'Filter3Thickness', '??', 'DetectorTempInC', '??', 
                         '??', '??', 'eVPerChannel', '??', 'eVStart', 
                         'Year', 'Month', 'AM/PM code?', 'Day', 'Hour', 'Minutes', 'Seconds', 
                         '??', 'NosePressureInMilliBars', 'NumberOfChannels', 'NoseTemperatureInC', 
                         'TubeSpec?', '??', '2048 counts']}}


def extract_spectra(pdz_file, to_csv=True, verbose=True): 
    '''Directly extract spectral data from `pdz_file`. 
    
    Robust extraction of spectrum counts and energy calibration. Other meta data is ignored. '''

    pdz_bytes = file_to_bytes(pdz_file)
    block_list = get_blocks(pdz_bytes, verbose=False) 
    
    # select type 3 blocks
    b3_list = [b for b in block_list if b['block_type'] == 3] 
    n_spectra = len(b3_list) 

    # parsing spectrum parameters 
    #from first block to compute channel energies (keV) 
    # (assuming that these are similar for all spectra in the pdz file) 
    
    arr = b3_list[0]['bytes'] # only using first spectrum!  
    spectrum_params = multiparse('hi-3i9f7hfhfhfhf8hfhi-S-h', arr, verbose=False)[0]  
    tube_keV = spectrum_params[12] # FYI 
    delta_keV = spectrum_params[25] / 1000
    start_keV = spectrum_params[27] / 1000 
    n_channels = spectrum_params[37] 
    if n_channels != 2048: 
        print(f'Found unexpected number of channels in pdz metadata: {n_channels}')
        
    stop_keV = start_keV + delta_keV * (n_channels -1)
    x_keV = np.linspace(start_keV, stop_keV, num=n_channels) 

    # initialize array 
    
    spectra_df = pd.DataFrame(index=x_keV)

    # TODO: Test if computed channel energies `x_keV` are reasonable, 
    # Parsing of spectrum parameters might fail 
    # if Bruker messes with file format.   
    # otherwise fall back on 0-40 keV range  

    spectrum_list = [] 

    for i, b3 in enumerate(b3_list):  
        arr = b3['bytes'] 
        counts = np.array(parse(f'{n_channels}i', arr[-4*n_channels:], verbose=False)[0])  

        spectra_df[f'spectrum #{i+1}'] = counts

    if to_csv == True: 
        csv_file = f'{pdz_file}.csv' 
        print(f'Saving spectral data to: {csv_file}')
        spectra_df.to_csv(csv_file, float_format='%10.5f')

    return spectra_df
           

def multiparse(xformat, arr, param_keys=None, verbose=True): 
    '''Parse segments in extendend format string `xformat` e.g. '<i5f-2S-T-3S-S-f' '''

    
    parts = re.split('-', xformat) 

    result_list = []
    for p in parts: 
        if 'S' in p:
            result, arr = read_strings(p, arr, verbose=False) 
        elif p == 'T': 
            result, arr = read_table(p, arr, verbose=False) 
        elif 'X' in p: 
            result, arr = skip_bytes(p, arr, verbose=False) 
        elif 'Z' in p: 
            result, arr = read_counts(p, arr, verbose=False)
            result = [result]
        elif 'z' in p: 
            result, arr = read_counts(p, arr, n_channels=1028, verbose=False)
            result = [result]    
        
        else: 
            result, arr = parse(p, arr, verbose=False) 
            
        result_list.extend(result)   

    if verbose: 
        if param_keys == None: 
            result_df = pd.DataFrame({'values': result_list})
        else: 
            result_df = pd.DataFrame({'values': result_list, 'param_keys': param_keys})      
        display(result_df)

    return result_list, arr   

def prefix(format): 
    '''Prefix little endian byte order (<) to struct type format string if missing. '''

    if not format.startswith('<'): 
        format = f'<{format}'
    
    return format    

def file_to_bytes(pdz_file): 
    '''Read all bytes from filepath `pdz_file` into a byte array. 
    
    Returns: `pdz_arr` (numpy array of bytes)
    '''

    with open(pdz_file, 'rb') as fh: 
        blob = fh.read() 
        
    pdz_arr = np.array([v[0] for v in struct.iter_unpack('c', blob)])
    #pdz_arr = bytearray(blob)

    return pdz_arr 
    

def parse(format, arr, verbose=True): 
    '''Parse first bytes from bytes array `arr` into human readable text according to `format` string. 
    
    See struct library for format string specification. For example, '<ff' would result 
    the first 8 bytes to be converted into two Little-Endian floats. 
    
    Returns: `parsed` list and remaining bytes array of `tail_arr` unprocessed `values`.'''

    format = prefix(format) 
    size = struct.calcsize(format)
    buffer = arr[0:size]
    tail_arr = arr[size:]

    parsed = list(struct.unpack(format, buffer)) 

    if verbose: 
        print(parsed)
    
    return parsed, tail_arr

def read_strings(xformat, arr, verbose=True): 
    '''Parse `n` variable length character strings preceded by a length integer. 
     
    '''
    # get multiplier  
    if xformat == 'S': 
        n = 1
    else: 
        n = int(re.sub('(^\d+)S', '\g<1>', xformat))

    # parse strings 
    string_list = [] 
    while n > 0: 
        [length], arr = parse('<i', arr, verbose=False) # read length 
        n_bytes = 2 * length 

        # do some testing 
        assert (n_bytes > 1) and (type(n_bytes) is int), f'{n_bytes} is invalid string length' 
        
        char_list, arr = parse(f'<{n_bytes}c', arr, verbose=False) 
        string = b''.join(char_list).decode(encoding='utf-16') 
        string_list.append(string) 
        n = n -1 
        
    if verbose: 
        print(string_list) 

    return string_list, arr


def skip_bytes(xformat, arr, verbose=True): 
    '''Skip a number of bytes as specified in `xformat` string. 
    
    If multiplier is `*` then skip all. 
    '''

    # get multiplier  
    if xformat == 'X': 
        n_bytes = 1
    elif xformat == '*X': 
        n_bytes = len(arr)    
    else: 
        n_bytes = int(re.sub('(^\d+)X', '\g<1>', xformat))

    skipped = [b''.join(arr[0:n_bytes])] 

    arr = arr[n_bytes:] 

    if verbose: 
        print(skipped)

    return skipped, arr 
    

def read_table(xformat, arr, verbose=True): 
    '''Extract numbered table'''

    assert xformat == 'T', 'Incorrect format string'

    [table_length], arr = parse('<i', arr, verbose=False) 

    table = []
    for i in range(table_length): 
        [num], arr = parse('<h', arr, verbose=False)  
        [string], arr = read_strings('S', arr, verbose=False)
        table.append([f'#{num}', string]) 
        
    if verbose: 
        print(table)

    return table, arr


def read_counts(xformat, arr, n_channels=2048, verbose=True): 
    '''Extract counts. '''

    assert xformat == 'Z'  or xformat == 'z', 'Incorrect format string'

    format = f'<{n_channels}i'

    counts, arr = parse(format, arr, verbose=False) 
    counts = np.array(counts)
        
    if verbose: 
        print(counts)

    return counts, arr
    

def get_block_at(pdz_arr, start): 
    '''Read first data block from bytes array `pdz_arr` from index position `start`. 

    Assumes that first 4 bytes are (block type and size)
    
    Returns: `block_dict`, `block` 
    '''

    file_size = len(pdz_arr)
    
    [block_type, block_size], arr = parse('hi', pdz_arr[start:], verbose=False)
    
    stop = start + block_size + 6 # four bytes extra due to `dtype` and `size` shorts plus two empty pad bytes? 

    # read block bytes
    arr = pdz_arr[start:stop] 
    
    block_dict = {'block_type': block_type, 'block_size': block_size, 'start': start, 'stop': stop, 
                  'file_size': file_size, 'bytes': arr}

    return block_dict


def get_blocks(pdz_bytes, verbose=True): 
    '''Parse `pdz_byte_array` into consequtive blocks. ''' 

    if verbose: 
        print('Detecting block sequence...')
    
    start = 0 
    total_size = len(pdz_bytes) 

    block_list = [] 
    while start < total_size: 
        block_dict = get_block_at(pdz_bytes, start)
        start = block_dict['stop']
        block_list.append(block_dict)

    else:
        if start == total_size: 
            if verbose: 
                print('Ok!')
        else: 
            print('Error while reading last block! ')
            print(f'Stop index: {start} does not match total file size: {total_size}')
        
    return block_list


def get_blocktypes(block_list): 
    '''Extract `block_types` list from `block_list`. '''
    
    block_types = []

    for block_dict in block_list: 
        t = block_dict['block_type']
        block_types.append(t)

    return block_types

