reconutils.py

import ctypes
import logging
import os
import pathlib
import queue
import threading

import ismrmrd
import numpy as np
import numpy.typing as npt
import rtoml
import sigpy as sp
from scipy.io import loadmat

import coils
import connection
import GIRF
import mrdhelper

# Waveform IDs for different physiological signals
ECG_WAVEFORM_ID = 0
PULSEOX_WAVEFORM_ID = 1
RESP_WAVEFORM_ID = 2
EXT1_WAVEFORM_ID = 3
EXT2_WAVEFORM_ID = 4
RESPPT_WAVEFORM_ID = 12

def data_acquisition_worker(conn: connection.Connection, data_queue: queue.Queue, stop_event: threading.Event):
    """Worker function for data acquisition using concurrent.futures."""
    try:
        for arm in conn:
            if arm is None or stop_event.is_set():
                data_queue.put(None)
                logging.info("Acquisition worker has stopped.")
                break
            
            data_queue.put(arm)
            
            # Yield control to avoid hogging CPU
            for _ in range(500):
                continue
            
    except Exception as e:
        logging.error(f"Error in data acquisition worker: {e}")
        data_queue.put(None)  # Ensure main thread doesn't hang
    finally:
        logging.info("Data acquisition worker finished.")

def data_acquisition_with_save_worker(conn: connection.Connection, data_queue: queue.Queue, stop_event: threading.Event, output_file_path: str, metadata: ismrmrd.xsd.ismrmrdHeader):
    """Worker function for data acquisition that also saves the raw data using concurrent.futures."""
    
    with ismrmrd.Dataset(output_file_path, create_if_needed=True) as dset:
        dset.write_xml_header(ismrmrd.xsd.ToXML(metadata))

        try:
            for arm in conn:
                if arm is None or stop_event.is_set():
                    data_queue.put(None)
                    logging.info("Acquisition worker has stopped.")
                    break
                
                data_queue.put(arm)
                if type(arm) is ismrmrd.Acquisition:
                    dset.append_acquisition(arm)
                elif type(arm) is ismrmrd.Waveform:
                    dset.append_waveform(arm)

                # Yield control to avoid hogging CPU
                # time.sleep(0.001)
                
        except Exception as e:
            logging.error(f"Error in data acquisition worker: {e}")
            data_queue.put(None)  # Ensure main thread doesn't hang
        finally:
            logging.info("Data acquisition worker finished.")

def girf_calibration(g_nom: np.ndarray, patient_position: str, arm: ismrmrd.Acquisition, dt: float, msize: int, girf_file: str) -> np.ndarray:

    r_PCS2DCS = GIRF.calculate_matrix_pcs_to_dcs(patient_position)
    r_GCS2RCS = np.array(  [[0,    1,   0],  # [PE]   [0 1 0] * [r]
                        [1,    0,   0],  # [RO] = [1 0 0] * [c]
                        [0,    0,   1]]) # [SL]   [0 0 1] * [s]
    r_GCS2PCS = np.array([arm.phase_dir, -np.array(arm.read_dir), arm.slice_dir])
    r_GCS2DCS = r_PCS2DCS.dot(r_GCS2PCS)
    sR = r_GCS2DCS.dot(r_GCS2RCS)
    tRR = 3*1e-6/dt

    k_pred, _ = GIRF.apply_GIRF(g_nom, dt, sR, girf_file=girf_file, tRR=tRR)
    # k_pred = np.flip(k_pred[:,:,0:2], axis=2) # Drop the z
    k_pred = k_pred[:,:,0:2] # Drop the z

    kmax = np.max(np.abs(k_pred[:,:,0] + 1j * k_pred[:,:,1]))
    k_pred = np.swapaxes(k_pred, 0, 1)
    k_pred = 0.5 * (k_pred / kmax) * msize

    return k_pred

def process_csm(frames):
    data = np.zeros(frames[0].shape, dtype=np.complex128)
    for g in frames:
        data += sp.to_device(g)
    (csm_est, rho) = coils.calculate_csm_inati_iter(data, smoothing=32)

    return csm_est


def process_group(group, frames: list, sens: npt.NDArray | None, device, rep, image_idx, config, metadata):
    xp = device.xp
    with device:
        data = xp.zeros(frames[0].shape, dtype=np.complex128)
        for g in frames:
            data += g

        if sens is None:
            # Sum of squares coil combination
            data = np.abs(np.flip(data, axis=(1,)))
            data = np.square(data)
            data = np.sum(data, axis=0)
            data = np.sqrt(data)
        else:
            # Coil combine
            data = np.flip(np.abs(np.sum(np.conj(sens) * data, axis=0)), axis=(0,))
            
        # Determine max value (12 or 16 bit)
        BitsStored = 12
        if (mrdhelper.get_userParameterLong_value(metadata, "BitsStored") is not None):
            BitsStored = mrdhelper.get_userParameterLong_value(metadata, "BitsStored")
        maxVal = 2**BitsStored - 1

        # Normalize and convert to int16
        dscale = maxVal/data.max()
        data *= dscale
        data = np.around(data)
        data = data.astype(np.int16)

    data = sp.to_device(data)
    data_prct = np.percentile(data, [1,99])
    wdw_w = data_prct[1] - data_prct[0]
    wdw_c = (data_prct[1] + data_prct[0])//2

    # Format as ISMRMRD image data
    # data has shape [RO PE], i.e. [x y].
    # from_array() should be called with 'transpose=False' to avoid warnings, and when called
    # with this option, can take input as: [cha z y x], [z y x], or [y x]
    image = ismrmrd.Image.from_array(data.transpose(), acquisition=group, transpose=False)

    image.image_index = image_idx
    image.repetition = rep

    # Set field of view
    image.field_of_view = (ctypes.c_float(metadata.encoding[0].reconSpace.fieldOfView_mm.x), 
                            ctypes.c_float(metadata.encoding[0].reconSpace.fieldOfView_mm.y), 
                            ctypes.c_float(metadata.encoding[0].reconSpace.fieldOfView_mm.z))

    # Set ISMRMRD Meta Attributes
    meta = ismrmrd.Meta({'DataRole':               'Image',
                         'ImageProcessingHistory': ['FIRE', 'PYTHON', 'simplenufft1arm'],
                         'WindowCenter':           str(wdw_c),
                         'WindowWidth':            str(wdw_w),
                         'NumArmsPerFrame':        str(config['reconstruction']['arms_per_frame']),
                         'GriddingWindowShift':    str(config['reconstruction']['window_shift']), 
                         'ImageScaleFactor':       str(dscale)
                         })

    
    # Set slice position
    # Center of the excited volume, in (left, posterior, superior) (LPS) coordinates relative to isocenter in millimeters. 
    # NB this is different than DICOM’s ImageOrientationPatient, which defines the center of the first (typically top-left) voxel.
    # TODO: This assumes slice direction is along "AP" axis. Can be generalized.
    Nslc = metadata.encoding[0].encodingLimits.slice.maximum+1
    slice_shift = [0, 0, metadata.encoding[0].reconSpace.fieldOfView_mm.z * (image.slice-Nslc//2)]
    #del_img_pos = r_GCS2PCS @ slice_shift

    #
    r_GCS2RCS = np.array(  [[0,    1,   0],  # [PE]   [0 1 0] * [r]
                                        [1,    0,   0],  # [RO] = [1 0 0] * [c]
                                        [0,    0,   1]]) # [SL]   [0 0 1] * [s]
    r_GCS2PCS = np.array([group.phase_dir, -np.array(group.read_dir), group.slice_dir])
    r_PCS2DCS = GIRF.calculate_matrix_pcs_to_dcs(metadata.measurementInformation.patientPosition.value)
    r_GCS2DCS = r_PCS2DCS.dot(r_GCS2PCS)
    sR = r_GCS2DCS.dot(r_GCS2RCS)
    del_img_pos = r_GCS2PCS.T @ slice_shift

    # for ii in range(3):
        # image.position[ii] = ctypes.c_float(image.position[ii] + del_img_pos[ii])
    if rep == 0:
        print(f"Slice position: {image.position[:]}, delta: {del_img_pos}")

    # Add image orientation directions to MetaAttributes if not already present
    if meta.get('ImageRowDir') is None:
        meta['ImageRowDir'] = ["{:.18f}".format(image.getHead().read_dir[0]), "{:.18f}".format(image.getHead().read_dir[1]), "{:.18f}".format(image.getHead().read_dir[2])]

    if meta.get('ImageColumnDir') is None:
        meta['ImageColumnDir'] = ["{:.18f}".format(image.getHead().phase_dir[0]), "{:.18f}".format(image.getHead().phase_dir[1]), "{:.18f}".format(image.getHead().phase_dir[2])]

    xml = meta.serialize()
    logging.debug("Image MetaAttributes: %s", xml)
    logging.debug("Image data has %d elements", image.data.size)

    image.attribute_string = xml
    return image


def process_frame_complex(group, frames: list, sens: npt.NDArray | None, device, rep, image_idx, config, metadata):
    xp = device.xp
    with device:
        data = xp.zeros(frames[0].shape, dtype=np.complex128)
        for g in frames:
            data += g

        if sens is None:
            logging.error("No coil sensitivity maps found. Cannot perform coil combination.")
            # Sum of squares coil combination
            data = np.abs(np.flip(data, axis=(1,)))
            data = np.square(data)
            data = np.sum(data, axis=0)
            data = np.sqrt(data)
        else:
            # Coil combine
            data = np.flip(np.sum(np.conj(sens) * data, axis=0), axis=(0,))

    data = sp.to_device(data)

    # Format as ISMRMRD image data
    # data has shape [RO PE], i.e. [x y].
    # from_array() should be called with 'transpose=False' to avoid warnings, and when called
    # with this option, can take input as: [cha z y x], [z y x], or [y x]
    image = ismrmrd.Image.from_array(data.transpose(), acquisition=group, transpose=False)

    image.image_index = image_idx
    image.repetition = rep

    # Set field of view
    image.field_of_view = (ctypes.c_float(metadata.encoding[0].reconSpace.fieldOfView_mm.x), 
                            ctypes.c_float(metadata.encoding[0].reconSpace.fieldOfView_mm.y), 
                            ctypes.c_float(metadata.encoding[0].reconSpace.fieldOfView_mm.z))

    # Set ISMRMRD Meta Attributes
    meta = ismrmrd.Meta({'DataRole':               'Image',
                         'ImageProcessingHistory': ['FIRE', 'PYTHON', 'simplenufft1arm'],
                         'WindowCenter':           str((data.max()+1)/2),
                         'WindowWidth':            str((data.max()+1)),
                         'NumArmsPerFrame':        str(config['reconstruction']['arms_per_frame']),
                         'GriddingWindowShift':    str(config['reconstruction']['window_shift'])})

    # Add image orientation directions to MetaAttributes if not already present
    if meta.get('ImageRowDir') is None:
        meta['ImageRowDir'] = ["{:.18f}".format(image.getHead().read_dir[0]), "{:.18f}".format(image.getHead().read_dir[1]), "{:.18f}".format(image.getHead().read_dir[2])]

    if meta.get('ImageColumnDir') is None:
        meta['ImageColumnDir'] = ["{:.18f}".format(image.getHead().phase_dir[0]), "{:.18f}".format(image.getHead().phase_dir[1]), "{:.18f}".format(image.getHead().phase_dir[2])]

    xml = meta.serialize()
    logging.debug("Image MetaAttributes: %s", xml)
    logging.debug("Image data has %d elements", image.data.size)

    image.attribute_string = xml
    return image

def process_waveforms(wf_list: list[ismrmrd.Waveform], acq_init_timestamp: float) -> tuple[npt.NDArray, npt.NDArray,  npt.NDArray,  npt.NDArray]:
    ecg = []
    resp_pt = []
    ext1 = []
    pulseox = []
    t_init_pox = 0
    pulseox_sample_time = 0
    t_init_ecg = 0
    ecg_sample_time = 0
    t_init_resp = 0
    resp_sample_time = 0
    t_init_ext1 = 0
    ext1_sample_time = 0

    for arm in wf_list:
        if arm.waveform_id == ECG_WAVEFORM_ID:
            ecg.append(arm.data)
            if t_init_ecg == 0:
                t_init_ecg = arm.time_stamp*2.5e-3
                ecg_sample_time = arm.sample_time_us*1e-6
        elif arm.waveform_id == PULSEOX_WAVEFORM_ID:
            pulseox.append(arm.data)
            if t_init_pox == 0:
                t_init_pox = arm.time_stamp*2.5e-3
                pulseox_sample_time = arm.sample_time_us*1e-6
        elif arm.waveform_id == RESPPT_WAVEFORM_ID:
            resp_pt.append(arm.data)
            if t_init_resp == 0:
                t_init_resp = arm.time_stamp*2.5e-3
                resp_sample_time = arm.sample_time_us*1e-6
        elif arm.waveform_id == EXT1_WAVEFORM_ID:
            ext1.append(arm.data)
            if t_init_ext1 == 0:
                t_init_ext1 = arm.time_stamp*2.5e-3
                ext1_sample_time = arm.sample_time_us*1e-6

    ecg = np.concatenate(ecg, axis=1).T if len(ecg) > 0 else np.array([])
    if len(ecg) > 0:
        if np.isnan(ecg).all():
            t_ecg = np.array([])
        else:
            ecg = ecg[:, 0].astype(np.float32)
            ecg -= np.percentile(ecg, 5, axis=0)
            ecg /= np.max(np.abs(ecg), axis=0, keepdims=True)
            t_ecg = np.arange(ecg.shape[0])*ecg_sample_time + t_init_ecg - acq_init_timestamp
    else:
        t_ecg = np.array([])

    pulseox = np.concatenate(pulseox, axis=1).T if len(pulseox) > 0 else np.array([])
    if len(pulseox) > 0:
        if not np.isnan(pulseox).all():
            pulseox = pulseox[:, 0].astype(np.float32)
            pulseox -= np.percentile(pulseox, 5, axis=0)
            pulseox /= np.max(np.abs(pulseox), axis=0, keepdims=True)
            t_pox = np.arange(pulseox.shape[0])*pulseox_sample_time + t_init_pox - acq_init_timestamp
        else:
            t_pox = np.array([])
    else:
        t_pox = np.array([])

    resp_pt = np.concatenate(resp_pt, axis=1).T if len(resp_pt) > 0 else np.array([])
    if len(resp_pt) > 0:
        resp_pt = resp_pt[:, 0].astype(np.float32)
        resp_pt -= np.mean(resp_pt, axis=0, keepdims=True)
        resp_pt /= np.max(np.abs(resp_pt), axis=0, keepdims=True)
        t_resp = np.arange(resp_pt.shape[0])*resp_sample_time + t_init_resp - acq_init_timestamp
    else:
        t_resp = np.array([])

    ext1 = np.concatenate(ext1, axis=1).T if len(ext1) > 0 else np.array([])
    if len(ext1) > 0:
        ext1 = ext1[:, 1].astype(np.float32)
        ext1[ext1 > 0] = 1
        t_ext1 = np.arange(ext1.shape[0])*ext1_sample_time + t_init_ext1 - acq_init_timestamp
    else:
        t_ext1 = np.array([])

    resp = resp_pt
    card = ecg if (len(ecg) > 0 and not np.isnan(ecg).all()) else (pulseox if len(pulseox) > 0 else ext1)
    t_card = t_ecg if (len(ecg) > 0 and not np.isnan(ecg).all()) else (t_pox if len(pulseox) > 0 else t_ext1)
    return t_resp, resp, t_card, card

def waveforms_asarray2(wf_list: list[ismrmrd.Waveform]) -> dict:
    '''An alternative function to sort and convert a list of waveforms to numpy arrays.
        Parameters
        ----------
        wf_list : list[ismrmrd.Waveform]
            List of waveforms.
        
        Returns
        -------
        waveform_dict : dict
            Dictionary of waveforms. Possible keys are 'ecg', 'pulseox', 'resp', 'ext1'. 
            Keys will only be present if the corresponding waveform is found.
    '''
    ecg = []
    resp_pt = []
    ext1 = []
    t_ext1 = []
    pulseox = []
    t_init_pox = 0
    pulseox_sample_time = 0
    t_init_ecg = 0
    ecg_sample_time = 0
    t_init_resp = 0
    resp_sample_time = 0
    t_init_ext1 = 0
    ext1_sample_time = 0

    for wf in wf_list:
        if wf.waveform_id == ECG_WAVEFORM_ID:
            ecg.append(wf.data)
            if t_init_ecg == 0:
                t_init_ecg = wf.time_stamp*2.5e-3
                ecg_sample_time = wf.sample_time_us*1e-6
        elif wf.waveform_id == PULSEOX_WAVEFORM_ID:
            pulseox.append(wf.data)
            if t_init_pox == 0:
                t_init_pox = wf.time_stamp*2.5e-3
                pulseox_sample_time = wf.sample_time_us*1e-6
        elif wf.waveform_id == RESPPT_WAVEFORM_ID:
            resp_pt.append(wf.data)
            if t_init_resp == 0:
                t_init_resp = wf.time_stamp*2.5e-3
                resp_sample_time = wf.sample_time_us*1e-6
        elif wf.waveform_id == EXT1_WAVEFORM_ID:
            ext1.append(wf.data)
            if t_init_ext1 == 0:
                t_init_ext1 = wf.time_stamp*2.5e-3
                ext1_sample_time = wf.sample_time_us*1e-6
            t_ext1.extend(wf.time_stamp*2.5e-3 + np.arange(wf.data.shape[1])*ext1_sample_time)


    waveform_dict = {}
    ecg = np.concatenate(ecg, axis=1).T if len(ecg) > 0 else np.array([])
    if len(ecg) > 0:
        is_flat = np.all(ecg[:, 0] == ecg[0, 0], axis=0) or \
            np.all(ecg[:, 1] == ecg[0, 1], axis=0) or \
            np.all(ecg[:, 2] == ecg[0, 2], axis=0) or \
            np.all(ecg[:, 3] == ecg[0, 3], axis=0)                  # Check if the waveform is flat

        if not np.isnan(ecg).all() and not is_flat:
            ecg = ecg[:, :].astype(np.float32)
            ecg -= np.percentile(ecg, 5, axis=0)
            ecg /= np.max(np.abs(ecg), axis=0, keepdims=True)
            t_ecg = np.arange(ecg.shape[0])*ecg_sample_time + t_init_ecg
            waveform_dict['ecg'] = (t_ecg, ecg)

    pulseox = np.concatenate(pulseox, axis=1).T if len(pulseox) > 0 else np.array([])
    if len(pulseox) > 0:
        is_flat = np.all(pulseox[:, 0] == pulseox[0, 0], axis=0) # Check if the waveform is flat
        if not np.isnan(pulseox).all() and not is_flat and not np.all(pulseox[:,1]):
            pulseox_trigs = pulseox[:, 1].astype(np.int32)
            pulseox_trigs[pulseox_trigs > 0] = 1
            pulseox = pulseox[:, 0].astype(np.float32)
            pulseox -= np.percentile(pulseox, 5, axis=0)
            pulseox /= np.max(np.abs(pulseox), axis=0, keepdims=True)
            t_pox = np.arange(pulseox.shape[0])*pulseox_sample_time + t_init_pox
            waveform_dict['pulseox'] = (t_pox, pulseox, pulseox_trigs)

    resp_pt = np.concatenate(resp_pt, axis=1).T if len(resp_pt) > 0 else np.array([])
    if len(resp_pt) > 0:
        resp_pt = resp_pt[:, 0].astype(np.float32)
        resp_pt -= np.mean(resp_pt, axis=0, keepdims=True)
        resp_pt /= np.max(np.abs(resp_pt), axis=0, keepdims=True)
        t_resp = np.arange(resp_pt.shape[0])*resp_sample_time + t_init_resp
        waveform_dict['resp'] = (t_resp, resp_pt)

    ext1 = np.concatenate(ext1, axis=1).T if len(ext1) > 0 else np.array([])
    if len(ext1) > 0:
        is_flat = np.all(ext1[:, 0] == ext1[0, 0], axis=0) or np.all(ext1[:, 1] == ext1[0, 1], axis=0)
        if not np.isnan(ext1).all() and not is_flat:
            ext1 = ext1[:, 1].astype(np.float32)
            ext1[ext1 > 0] = 1
            # t_ext1 = np.arange(ext1.shape[0])*ext1_sample_time + t_init_ext1
            t_ext1 = np.array(t_ext1)
            waveform_dict['ext1'] = (t_ext1, ext1)

    return waveform_dict

def load_trajectory(metadata, metafile_paths: list[str]) -> dict | None:
    # get the k-space trajectory based on the metadata hash.
    for str_param in metadata.userParameters.userParameterString:
        if str_param.name == "tSequenceVariant":
            traj_name = str_param.value[:32] # Get first 32 chars, because a bug sometimes causes this field to have /OSP added to the end.
            break
    else:
        logging.error("Sequence hash is not found in metadata user parameters.")
        return None

    # load the .mat file containing the trajectory
    # Search for the file in the metafile_paths
    for path in metafile_paths:
        metafile_fullpath = os.path.join(path, traj_name + ".mat")
        if os.path.isfile(metafile_fullpath):
            logging.info(f"Loading metafile {traj_name} from {path}...")
            traj = loadmat(metafile_fullpath)
            return traj
    else:
        logging.error(f"Trajectory file {traj_name}.mat not found in specified paths.")
        return None
    
def load_config(config_file_name: str) -> dict | None:
    path_list = [
        pathlib.Path('/tmp/share/configs'),
        pathlib.Path(__file__).parent / 'configs', 
        ]
    for path in path_list:
        cfg_fullpath = path / config_file_name
        logging.debug(f"Checking for config file at: {cfg_fullpath}")
        if os.path.isfile(cfg_fullpath):
            # We now read these parameters from toml file, so that we won't have to keep restarting the server when we change them.
            logging.info(f"Using the configuration at {cfg_fullpath}")
            try:
                with open(cfg_fullpath) as jf:
                    cfg = rtoml.load(jf)
                return cfg
            except Exception as e:
                logging.error(f"Error loading configuration file {config_file_name}: {e}")
                return None
    logging.error(f"Configuration file {config_file_name} not found in specified paths: {path_list}.")
    return None