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


"""
    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.
"""
def prepare_datasets(x0, y0, drop_mask):

    x0 = np.asarray(x0)
    y0 = np.asarray(y0)
    drop_mask = np.asarray(drop_mask, dtype=bool)

    if not (len(x0) == len(y0) == len(drop_mask)):
        raise ValueError(
            "x0, y0, and drop_mask must have the same length. "
            f"Got len(x0)={len(x0)}, len(y0)={len(y0)}, "
            f"len(drop_mask)={len(drop_mask)}."
        )

    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.copy()
                x_tmp = x0.copy()

            if scale == "log":
                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)

# ==============================================================================
# Using dataframe to store data
# ==============================================================================

from pathlib import Path

import numpy as np
import pandas as pd
import scipy.io


FLAG_INDEX = {
    "bowel": 1,
    "lung": 3,
    "thyroid": 5,
}


def load_mat_data_to_dataframe(
    suv_filename: str | Path,
    flags_filename: str | Path,
    nr_patients: int = 58,
) -> pd.DataFrame:
    """
    Load SUV percentile and AE flag data from MATLAB files and convert them
    to a long-format pandas DataFrame.

    One row corresponds to one combination of:
        patient, organ, measurement/visit, percentile

    Parameters
    ----------
    suv_filename
        Path to the MATLAB file containing arrays such as
        'lung_SUVperc_COMBINED'.

    flags_filename
        Path to the MATLAB file containing the 'flags' array.

    nr_patients
        Maximum number of patients to include.

    Returns
    -------
    pd.DataFrame
        Columns:
        patient_id, organ, measurement, percentile, suv, y
    """
    suv_filename = Path(suv_filename)
    flags_filename = Path(flags_filename)

    if not suv_filename.is_file():
        raise FileNotFoundError(f"SUV file does not exist: {suv_filename}")

    if not flags_filename.is_file():
        raise FileNotFoundError(f"Flags file does not exist: {flags_filename}")

    suv_dict = scipy.io.loadmat(suv_filename)
    flags_dict = scipy.io.loadmat(flags_filename)

    if "flags" not in flags_dict:
        raise KeyError(
            f"'flags' was not found in {flags_filename}. "
            f"Available keys: {list(flags_dict)}"
        )

    flags_array = np.asarray(flags_dict["flags"])

    frames = []

    for organ, flag_idx in FLAG_INDEX.items():
        suv_key = f"{organ}_SUVperc_COMBINED"

        if suv_key not in suv_dict:
            print(f"Skipping {organ!r}: key {suv_key!r} was not found.")
            continue

        suv = np.asarray(suv_dict[suv_key])

        if suv.ndim != 3:
            raise ValueError(
                f"{suv_key!r} must be a 3D array, but has shape {suv.shape}."
            )

        n_available = min(
            nr_patients,
            suv.shape[0],
            flags_array.shape[0],
        )

        suv = suv[:n_available]
        flags = flags_array[:n_available, flag_idx].squeeze()

        patient_idx, measurement_idx, percentile_idx = np.indices(suv.shape)

        frame = pd.DataFrame(
            {
                "patient_id": patient_idx.ravel() + 1,
                "organ": organ,
                "measurement": measurement_idx.ravel(),
                "percentile": percentile_idx.ravel() + 1,
                "suv": suv.ravel(),
            }
        )

        frame["y"] = flags[patient_idx.ravel()]

        frames.append(frame)

    if not frames:
        raise ValueError(
            "No supported SUV arrays were found in the SUV MATLAB file."
        )

    return pd.concat(frames, ignore_index=True)

def get_data_from_dataframe(
    df: pd.DataFrame,
    organ: str,
    perc: int,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Select one organ and percentile and compute the maximum SUV value
    across measurements for each patient.
    """
    selected = (
        df.loc[
            (df["organ"] == organ)
            & (df["percentile"] == perc)
        ]
        .groupby("patient_id", as_index=False)
        .agg(
            x=("suv", "max"),
            y=("y", "first"),
        )
        .sort_values("patient_id")
    )

    return (
        selected["x"].to_numpy(),
        selected["y"].to_numpy(),
    )
