"""
    Utility functions for data extraction and processing

    Authors: Martin Horvat, January 2026
"""

import numpy as np

"""
  Extract data fom dictionaries for specific organ

  Input: 
    organ: string in ["lung", "bowel", "thyroid"]
    perc: int, percentiles [1, ... , 100]
    suv_dict: dict containing percentiles of suv
    flags_dict: dict containing states
"""

def get_data(organ, perc, suv_dict, flags_dict, nr_patient = 58):
    fname = "get_data"

    # get concrete data set
    suv = suv_dict[organ + '_SUVperc_COMBINED'][:nr_patient,:,:]  # suv percentiles
    
    # index of percentile
    perc_idx = perc -1

    # computing max SUV percentile per patient, ignoring nans
    x = np.nanmax(suv[:,:,perc_idx], axis = 1)

    # determining index in the flags based on organ
    match organ:
        case "lung":
            flags_idx = 3
        case "bowel":
            flags_idx = 1
        case "thyroid":
            flags_idx = 5
        case _:
            assert False, f"{fname}::this organ {organ = } is not supported"

    # getting state of patients: 0 == NC, 1  == AE
    y = flags_dict['flags'][:nr_patient, flags_idx]  

    return x, y


"""
    Check if a vector lies within the specified bounds for each dimension.

    Parameters:
    - vector (np.ndarray): 1D array representing the point to check. Shape: (n,)
    - bounds (np.ndarray): 2D array of shape (n, 2), where each row is (min, max) for a dimension.

    Returns:
    - bool: True if the vector is within bounds in all dimensions, False otherwise.
"""
def within_bounds(vector: np.ndarray, bounds: np.ndarray) -> bool:

    if vector.shape[-1] != bounds.shape[0]:
        raise ValueError("Dimension mismatch: vector length and bounds rows must be equal.")
    
    return np.apply_along_axis(lambda x: np.all((x >= bounds[:, 0]) & (x <= bounds[:, 1])), -1, vector)