# AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/40_pdz-archeology.ipynb.

# %% auto 0
__all__ = ['PDZ11_STRUCT_DICT', 'check_pdz_type', 'extract_spectrum']

# %% ../notebooks/40_pdz-archeology.ipynb 20
from . import file_to_bytes, get_blocks, multiparse, parse, PDZ_25_STRUCTURE_DICT
import struct
import numpy as np
import pandas as pd
import os 

# %% ../notebooks/40_pdz-archeology.ipynb 21
PDZ11_STRUCT_DICT = {
    'pdz11_2048_channels' : {'xformat': '2X-i-h-34X-2d-86X-2i-10X-2f-188X-Z-*X', 
                          'param_keys': ['pdz-version', '??', 'NumberOfChannels', '??', '??', 
                                         'eVPerChannel', '??', 'RawCounts', 'ValidCounts', '??',  'XrayVoltageInkV', 
                                         'XrayFilamentCurrentInMicroAmps', '??', 'Intensity (2048 channels)', '??']}, 
    'pdz11_1024_channels' : {'xformat': '2X-i-h-34X-2d-86X-2i-10X-2f-24X-z-*X', 
                          'param_keys': ['pdz-version', '??', 'NumberOfChannels', '??', '??', 
                                         'eVPerChannel', '??', 'RawCounts', 'ValidCounts', '??',  'XrayVoltageInkV', 
                                         'XrayFilamentCurrentInMicroAmps', '??', 'Intensity (1024 channels)', '??']} 
}


def check_pdz_type(pdz_file, verbose=True): 
    '''Read first two bytes and for legacy pdz files number of detector channels to check pdz file type.'''

    pdz_bytes = file_to_bytes(pdz_file) 
    file_size = len(pdz_bytes)
    first_two_bytes = struct.unpack('<h', pdz_bytes[0:2])[0] 

    if first_two_bytes == 25: 
        pdz_type = 'pdz25' 
    elif first_two_bytes == 257: 
        n_channels = struct.unpack('<h', pdz_bytes[6:8])[0] 
        if n_channels == 1024:
            pdz_type = 'pdz11_1024_channels'
        elif n_channels == 2048:
            pdz_type = 'pdz11_2048_channels'  
        else: 
            pdz_type = f'pdz11_with_unexpected_number_of_{n_channels}_channels'
    else:
        pdz_type = f'pdz_type_unknown:{first_two_bytes}'
                
    return pdz_type 


def extract_spectrum(pdz_file, to_csv=True, verbose=True): 
    '''Direct extraction of spectral data from `pdz_file`. 
    
    Should now work for both pdz25 and both pdz11 file  formats. 
    ''' 
    
    pdz_type = check_pdz_type(pdz_file)
    pdz_bytes = file_to_bytes(pdz_file) 

    # get spectral data xformat string for pdz_type 
    if pdz_type == 'pdz25': 
        xformat = PDZ_25_STRUCTURE_DICT[3]['xformat']      
    elif pdz_type == 'pdz11_2048_channels': 
        xformat = PDZ11_STRUCT_DICT['pdz11_2048_channels']['xformat']
    elif pdz_type == 'pdz11_1024_channels': 
        xformat = PDZ11_STRUCT_DICT['pdz11_1024_channels']['xformat'] 
    else: 
        return f'Unknown pdz type: {pdz_type}'
 

    # pdz25 
    if pdz_type == 'pdz25': 
        # extracting spectral data in block 
        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) 
        if verbose and n_spectra > 1: 
            print('Found multiple spectral data blocks. Only parsing first spectrum. ')
        arr = b3_list[0]['bytes'] 
    
        # parsing spectrum parameters and data 
        parsed, tail = multiparse(xformat, arr, verbose=False)
        tube_keV = parsed[12] # FYI 
        delta_keV = parsed[25] / 1000
        start_keV = parsed[27] / 1000 
        n_channels = parsed[37]
        intensity = parsed[-1]
        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) 
        
        spectrum_df = pd.DataFrame(index=x_keV)
        basename = os.path.basename(pdz_file)
        spectrum_df[basename] = intensity 
        
        # 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??  

    # pdz11_2048_channels 
    elif pdz_type == 'pdz11_2048_channels': 
        
        xformat = PDZ11_STRUCT_DICT['pdz11_2048_channels']['xformat']
        parsed, tail = multiparse(xformat, pdz_bytes, verbose=False) 

        # The problem with legacy files is that we do not know if and at which position 
        # the energy offset is stored. So we need to set start_keV=0  
        # parsing spectrum parameters and data 
        parsed, tail = multiparse(xformat, pdz_bytes, verbose=False)
        
        tube_keV = parsed[10] # FYI 
        delta_keV = parsed[5] / 1000
        start_keV = 0 
        n_channels = parsed[2]
        intensity = parsed[13]
        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) 
        
        spectrum_df = pd.DataFrame(index=x_keV)
        basename = os.path.basename(pdz_file)
        spectrum_df[basename] = intensity 
        
    
    # pdz11_1024_channels
    elif pdz_type == 'pdz11_1024_channels': 
        
        xformat = PDZ11_STRUCT_DICT['pdz11_1024_channels']['xformat']
        parsed, tail = multiparse(xformat, pdz_bytes, verbose=False) 

        # The problem with legacy files is that we do not know if and at which position 
        # the energy offset is stored. So we need to set start_keV=0  
        # parsing spectrum parameters and data 
        parsed, tail = multiparse(xformat, pdz_bytes, verbose=False)
        
        tube_keV = parsed[10] # FYI 
        delta_keV = parsed[5] / 1000
        start_keV = 0 
        n_channels = parsed[2]
        intensity = parsed[13]
        if n_channels != 1024: 
            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) 
        
        spectrum_df = pd.DataFrame(index=x_keV)
        basename = os.path.basename(pdz_file)
        spectrum_df[basename] = intensity 
    
    else: 
        print(f'pdz_type: {pdz_type} Sorry, this specific pdz type is not yet implemented...')
        return 

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