"""Simple SUV distribution summaries."""

from __future__ import annotations

import numpy as np
import pandas as pd
import nibabel as nib


def get_suv_percentiles(
    image: nib.Nifti1Image,
    percentiles: tuple[float, ...] = (25, 50, 75),
    min_suv: float = 0.0,
) -> pd.DataFrame:
    """Compute SUV percentile thresholds and voxel counts.

    Parameters
    ----------
    image:
        3D NIfTI image containing SUV values.

    percentiles:
        Percentiles to compute.

    min_suv:
        Only finite voxels with ``SUV > min_suv`` are used. For processed images
        where background is set to zero, the default ``min_suv=0`` uses only the
        non-background region.

    Returns
    -------
    pandas.DataFrame
        Columns: ``percentile``, ``suv_threshold``,
        ``n_voxels_ge_threshold`` and ``n_voxels_lt_threshold``.
    """

    if not isinstance(image, nib.Nifti1Image):
        raise TypeError("image must be a nib.Nifti1Image.")

    data = image.get_fdata(dtype=np.float64)

    if data.ndim != 3:
        raise ValueError("Input image must be 3D.")

    suv_values = data[np.isfinite(data) & (data > min_suv)]

    if suv_values.size == 0:
        return pd.DataFrame(
            {
                "percentile": percentiles,
                "suv_threshold": [np.nan] * len(percentiles),
                "n_voxels_ge_threshold": [0] * len(percentiles),
                "n_voxels_lt_threshold": [0] * len(percentiles),
            }
        )

    percentile_values = np.percentile(suv_values, percentiles)

    return pd.DataFrame(
        {
            "percentile": percentiles,
            "suv_threshold": percentile_values,
            "n_voxels_ge_threshold": [
                int(np.sum(suv_values >= val)) for val in percentile_values
            ],
            "n_voxels_lt_threshold": [
                int(np.sum(suv_values < val)) for val in percentile_values
            ],
        }
    )