"""
    Utility functions for data extraction and processing

    Authors: Martin Horvat, January 2026
"""

import numpy as np
import pandas as pd

"""
  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)



def  prepare_data(x0, y0, drop_mask):
    """
        Prepare data for plotting and fitting by creating a DataFrame with different subsets and scales.

        Parameters:
        - x0: Original feature vector (e.g., SUV percentiles).
        - y0: Original target vector (e.g., patient states).
        - drop_mask: Boolean array indicating which samples to drop for the "TRIM" dataset.

        Returns:
        dataframe, scales, datasets:
        - df_data: A pandas DataFrame containing the prepared data with columns ['X', 'Y', 'scale', 'dataset'].
        - scales: List of scale names.
        - datasets: List of dataset names.
    """
    
    datasets = ["FULL", "TRIM"]
    scales = ["plain", "log"]

    lst = []
    for dataset in datasets:
        for scale in scales:

            if dataset == "TRIM":
                y_tmp = y0[~drop_mask]
                x_tmp = x0[~drop_mask]
            else:
                y_tmp = y0[:]
                x_tmp = x0[:]
            
            if scale == "log":
                # remove non-positive values for log scale
                valid_mask = (x_tmp > 0)
                y_tmp = y_tmp[valid_mask]
                x_tmp = x_tmp[valid_mask]
                x_tmp = np.log(x_tmp)

            df_tmp = pd.DataFrame({
                "X" : x_tmp,
                "Y" : y_tmp ,
                "scale" : scale,
                "dataset": dataset})
            lst.append(df_tmp)

    return pd.concat(lst, ignore_index=True)