# data.py
import numpy as np

"""
LUNG data extraction (with NaN filtering)

This file does one job:
1) Read lung SUV percentile data from `suv_dict`
2) Build one feature per patient:
      x[i] = max value (ignoring NaNs) of the chosen percentile for patient i
3) Read lung labels from `flags_dict`:
      y[i] = 0 (NC) or 1 (AE)
4) Remove patients where x is NaN (and keep y aligned)

Important:
- Logistic regression (sklearn) cannot use NaNs in X.
- So we filter NaNs here to always return clean data.

Inputs:
- perc: percentile number (1..100), e.g. 95
- suv_dict: dict loaded from the SUV .mat file
- flags_dict: dict loaded from the flags .mat file
- nr_patient: number of patients to use (default 58)

Outputs:
- x: 1D numpy array (NaNs removed)
- y: 1D numpy array (aligned with x)
"""


def get_data(perc, suv_dict, flags_dict):
    """
    Extract SUV percentile feature and AE labels.

    perc = 95 -> uses column index 94
    """

    # SUV matrix: (patients × percentiles × organs)
    suv = suv_dict["lung_SUVperc_COMBINED"]

    # Flags: AE indicator
    flags = flags_dict["flags"]

    # We only use first 58 patients and column 3 for AE
    y = flags[:58, 3].astype(int).ravel()

    # Percentile index (95th -> column 94)
    p_idx = perc - 1

    # Extract feature: max over organs at that percentile
    x = np.nanmax(suv[:58, :, p_idx], axis=1).astype(float)

    # Remove NaNs
    mask = np.isfinite(x)
    return x[mask], y[mask]