import os
import logging
from tqdm import tqdm
import numpy as np
import csv
import nibabel as nib
from PIL import Image
from typing import List, Optional, Tuple

from ._helpers import (
    validate_nifti_path,
    list_nifti_files,
    ensure_dir,
    strip_nifti_ext,
    validate_view,
)

logger = logging.getLogger("nidataset")


# Slices

def extract_slices(nii_path: str,
                   output_path: str,
                   view: str = "axial",
                   target_size: Optional[Tuple[int, int]] = None,
                   pad_value: float = 0.0,
                   debug: bool = False) -> List[str]:
    """
    Extracts slices from a NIfTI file and saves them as images .tif, following the structure
        <NIFTI FILENAME>_<VIEW>_<PROGRESSIVE SLICE NUMBER>.tif

    :param nii_path:
        Path to the input .nii.gz file with shape (X, Y, Z).

    :param output_path:
        Path where the extracted slices will be saved.

    :param view:
        Anatomical view for slice extraction:

        - ``"axial"`` → extracts along the Z-axis.
        - ``"coronal"`` → extracts along the Y-axis.
        - ``"sagittal"`` → extracts along the X-axis.

    :param target_size:
        Optional target dimensions (height, width) for the output slices. If specified,
        slices will be padded symmetrically to reach the target size. Padding is applied
        equally on both sides when possible; if odd padding is needed, the extra pixel
        is added to the right/bottom. If ``None``, slices are saved at their original size.

        Example: ``target_size=(512, 512)``

    :param pad_value:
        Value used for padding when ``target_size`` is specified. Default is ``0.0``.

    :param debug:
        Verbose print about the total number of slices extracted. Default is ``False``.

    :returns:
        List of saved slice file paths.

    :raises FileNotFoundError:
        If the input NIfTI file does not exist.

    :raises ValueError:
        If the NIfTI file is empty, has invalid extension, or invalid dimensions.
        If the view is not 'axial', 'coronal', or 'sagittal'.
        If ``target_size`` is smaller than the original slice dimensions.

    Example
    -------
    >>> from nidataset.slices import extract_slices
    >>>
    >>> # define paths
    >>> nii_path = "path/to/input_image.nii.gz"
    >>> output_path = "path/to/output_directory"
    >>>
    >>> # choose the anatomical view ('axial', 'coronal', or 'sagittal')
    >>> view = "axial"
    >>>
    >>> # extract slices without padding
    >>> extract_slices(nii_path=nii_path,
    ...                output_path=output_path,
    ...                view=view,
    ...                debug=True)
    >>>
    >>> # extract slices with padding to 512x512
    >>> extract_slices(nii_path=nii_path,
    ...                output_path=output_path,
    ...                view=view,
    ...                target_size=(512, 512),
    ...                pad_value=0.0,
    ...                debug=True)
    """

    # check if the input file exists and has a valid NIfTI extension
    validate_nifti_path(nii_path)

    # validate the view parameter
    validate_view(view)

    # create output dir if it does not exist
    ensure_dir(output_path)

    # load the NIfTI file
    nii_img = nib.load(nii_path)
    nii_data = nii_img.get_fdata()

    # validate NIfTI data dimensions
    if nii_data.ndim != 3:
        raise ValueError(f"Error: expected a 3D NIfTI file. Got shape '{nii_data.shape}' instead.")

    # mapping of views to slicing axes
    view_mapping = {
        "axial": (2, lambda data, i: data[:, :, i]),       # Z-axis
        "coronal": (1, lambda data, i: data[:, i, :]),     # Y-axis
        "sagittal": (0, lambda data, i: data[i, :, :])     # X-axis
    }

    # get axis and slicing function
    axis, slice_func = view_mapping[view]

    # get number of slices along the selected axis
    num_slices = nii_data.shape[axis]

    # check if the dimension is not zero
    if num_slices == 0:
        raise ValueError("Error: the NIfTI file contains no slices (empty volume).")

    # define prefix as the nii.gz filename
    prefix = strip_nifti_ext(os.path.basename(nii_path))

    # collect saved file paths
    saved_paths = []

    # iterate over slices and save as images
    for i in tqdm(range(num_slices), desc=f"Processing {prefix} ({view})", unit="slice"):
        # extract the slice using the dynamic function
        slice_data = slice_func(nii_data, i)

        # apply padding if target_size is specified
        if target_size is not None:
            target_h, target_w = target_size
            current_h, current_w = slice_data.shape

            # validate that target size is not smaller than current size
            if target_h < current_h or target_w < current_w:
                raise ValueError(
                    f"Error: target_size {target_size} is smaller than slice dimensions "
                    f"({current_h}, {current_w}). Target size must be >= slice dimensions."
                )

            # calculate padding amounts
            pad_h = target_h - current_h
            pad_w = target_w - current_w

            # distribute padding symmetrically (extra pixel goes to right/bottom if odd)
            pad_top = pad_h // 2
            pad_bottom = pad_h - pad_top
            pad_left = pad_w // 2
            pad_right = pad_w - pad_left

            # apply padding
            slice_data = np.pad(
                slice_data,
                ((pad_top, pad_bottom), (pad_left, pad_right)),
                mode='constant',
                constant_values=pad_value
            )

        # construct filename with zero-padded slice index
        slice_filename = f"{prefix}_{view}_{str(i).zfill(3)}.tif"
        slice_path = os.path.join(output_path, slice_filename)

        # save slice as an image
        slice_to_save = Image.fromarray(slice_data)
        slice_to_save.save(slice_path)
        saved_paths.append(slice_path)

    # debug verbose logging
    if debug:
        padding_info = f"\nPadding applied: target size {target_size}" if target_size else "\nNo padding applied"
        logger.info(f"\nInput file: '{nii_path}'\nOutput path: '{output_path}'"
              f"{padding_info}\nTotal {view} slices extracted: {num_slices}")

    return saved_paths


def extract_slices_dataset(nii_folder: str,
                           output_path: str,
                           view: str = "axial",
                           saving_mode: str = "case",
                           target_size: Optional[Tuple[int, int]] = None,
                           pad_value: float = 0.0,
                           save_stats: bool = False) -> int:
    """
    Extracts slices from all NIfTI files in a dataset folder and saves them as images .tif, following the structure

        <NIFTI FILENAME>_<VIEW>_<PROGRESSIVE SLICE NUMBER>.tif

    :param nii_folder:
        Path to the folder containing all .nii.gz files with shape (X, Y, Z).

    :param output_path:
        Path where the extracted slices will be saved.

    :param view:
        Anatomical view for slice extraction:

        - ``"axial"`` → extracts along the Z-axis.
        - ``"coronal"`` → extracts along the Y-axis.
        - ``"sagittal"`` → extracts along the X-axis.

    :param saving_mode:
        - ``"case"`` → creates a folder for each case.
        - ``"view"`` → saves all slices inside a single view folder.

    :param target_size:
        Optional target dimensions (height, width) for the output slices. If specified,
        slices will be padded symmetrically to reach the target size. If ``None``,
        slices are saved at their original size.

        Example: ``target_size=(512, 512)``

    :param pad_value:
        Value used for padding when ``target_size`` is specified. Default is ``0.0``.

    :param save_stats:
        If ``True``, saves a CSV file with FILENAME and NUM_SLICES information per case
        as ``<VIEW>_slices_stats.csv``.

    :returns:
        Total slice count across all NIfTI files.

    :raises FileNotFoundError:
        If the dataset folder does not exist or contains no .nii.gz files.

    :raises ValueError:
        If an invalid view or saving_mode is provided.

    Example
    -------
    >>> from nidataset.slices import extract_slices_dataset
    >>>
    >>> # define paths
    >>> nii_folder = "path/to/dataset"
    >>> output_path = "path/to/output_directory"
    >>>
    >>> # choose the anatomical view ('axial', 'coronal', or 'sagittal')
    >>> view = "axial"
    >>>
    >>> # extract slices without padding
    >>> extract_slices_dataset(nii_folder=nii_folder,
    ...                        output_path=output_path,
    ...                        view=view,
    ...                        saving_mode="view",
    ...                        save_stats=True)
    >>>
    >>> # extract slices with padding to 512x512
    >>> extract_slices_dataset(nii_folder=nii_folder,
    ...                        output_path=output_path,
    ...                        view=view,
    ...                        saving_mode="case",
    ...                        target_size=(512, 512),
    ...                        pad_value=0.0,
    ...                        save_stats=True)
    """

    # get all NIfTI files in the dataset folder (sorted; raises if missing)
    nii_files = list_nifti_files(nii_folder)

    # validate the view parameter
    validate_view(view)

    # validate input parameters
    if saving_mode not in ["case", "view"]:
        raise ValueError(f"Error: saving_mode must be either 'case' or 'view'. Got '{saving_mode}' instead.")

    # create output dir if it does not exist
    ensure_dir(output_path)

    # create a single folder if using "view" mode
    if saving_mode == "view":
        view_output_dir = os.path.join(output_path, view)
        ensure_dir(view_output_dir)

    # initialize statistics tracking
    stats = []
    total_slices = 0
    stats_file = os.path.join(output_path, f"{view}_slices_stats.csv") if save_stats else None

    # iterate over nii.gz files with tqdm progress bar
    for nii_file in tqdm(nii_files, desc=f"Processing {view} slices", unit="file"):
        # nii.gz file path
        nii_path = os.path.join(nii_folder, nii_file)

        # extract the filename prefix (case ID)
        prefix = strip_nifti_ext(os.path.basename(nii_path))

        # update tqdm description with the current file prefix
        tqdm.write(f"Processing: {prefix}")

        # determine number of slices **before** calling extract_slices
        try:
            nii_img = nib.load(nii_path)
            nii_data = nii_img.get_fdata()
            num_slices = nii_data.shape[{"axial": 2, "coronal": 1, "sagittal": 0}[view]]
        except Exception as e:
            tqdm.write(f"Error processing {nii_file} for statistical analysis: {e}")
            continue  # skip this file if an error occurs

        # keep track of the total number of slices
        total_slices += num_slices
        if save_stats:
            stats.append([nii_file, num_slices])

        # determine the appropriate output folder
        if saving_mode == "case":
            case_output_dir = os.path.join(output_path, prefix, view)
            ensure_dir(case_output_dir)
            extract_slices(nii_path, case_output_dir, view,
                          target_size=target_size, pad_value=pad_value, debug=False)
        else:
            extract_slices(nii_path, view_output_dir, view,
                          target_size=target_size, pad_value=pad_value, debug=False)

    # save statistics if enabled
    if save_stats:
        with open(stats_file, mode="w", newline="") as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(["FILENAME", "NUM_SLICES"])
            writer.writerows(stats)
            writer.writerow(["TOTAL_SLICES", total_slices])

        padding_info = f" (with padding to {target_size})" if target_size else ""
        logger.info(f"\nSlice extraction statistics saved in: '{stats_file}'{padding_info}")

    return total_slices


# Annotations

def extract_annotations(nii_path: str,
                        output_path: str,
                        view: str = "axial",
                        saving_mode: str = "slice",
                        data_mode: str = "center",
                        target_size: Optional[Tuple[int, int]] = None,
                        debug: bool = False) -> List[str]:
    """
    Extracts annotations from a NIfTI annotation file and saves them as CSV, based on the selected view and named with:

        <NIFTI FILENAME>_<VIEW>_<PROGRESSIVE SLICE NUMBER>.csv

    or

        <NIFTI FILENAME>.csv

    :param nii_path:
        Path to the input .nii.gz file with shape (X, Y, Z).

    :param output_path:
        Path where the CSV annotations will be saved.

    :param view:
        Anatomical view for annotation extraction:

        - ``"axial"`` → extracts along the Z-axis (2D: X=Y_3D, Y=X_3D).
        - ``"coronal"`` → extracts along the Y-axis (2D: X=Z_3D, Y=X_3D).
        - ``"sagittal"`` → extracts along the X-axis (2D: X=Z_3D, Y=Y_3D).

    :param saving_mode:
        - ``"slice"`` → generates a CSV per slice.
        - ``"volume"`` → generates a single CSV for the whole volume.

    :param data_mode:
        - ``"center"`` → saves the center of the bounding box.
        - ``"box"`` → saves the bounding box coordinates.
        - ``"radius"`` → saves the center and radius (from center to bounding box border).

    :param target_size:
        Optional target dimensions (height, width) for coordinate adjustment. If specified,
        coordinates will be adjusted to account for padding that would be applied to match
        the target size. This ensures annotations align with padded images from ``extract_slices``.
        If ``None``, coordinates are saved at their original values. Note that this is applied only if
        the saving_mode is ``"slice"``.

        Example: ``target_size=(512, 512)``

    :param debug:
        If ``True``, prints additional information about the extraction.

    :returns:
        List of saved CSV file paths.

    :raises FileNotFoundError:
        If the input NIfTI file does not exist.

    :raises ValueError:
        If the NIfTI file is empty, has invalid extension, or invalid parameters.
        If the view is not 'axial', 'coronal', or 'sagittal'.
        If the saving_mode and data_mode are not correct.

    Example
    -------
    >>> from nidataset.slices import extract_annotations
    >>>
    >>> # define paths
    >>> nii_path = "path/to/input_image.nii.gz"
    >>> output_path = "path/to/output_directory"
    >>>
    >>> # choose the anatomical view ('axial', 'coronal', or 'sagittal')
    >>> view = "axial"
    >>>
    >>> # extract annotations without padding adjustment
    >>> extract_annotations(nii_path=nii_path,
    ...                     output_path=output_path,
    ...                     view=view,
    ...                     saving_mode="slice",
    ...                     data_mode="center",
    ...                     debug=True)
    >>>
    >>> # extract annotations with padding adjustment to match 512x512 images
    >>> extract_annotations(nii_path=nii_path,
    ...                     output_path=output_path,
    ...                     view=view,
    ...                     saving_mode="slice",
    ...                     data_mode="box",
    ...                     target_size=(512, 512),
    ...                     debug=True)
    >>>
    >>> # extract annotations with center and radius
    >>> extract_annotations(nii_path=nii_path,
    ...                     output_path=output_path,
    ...                     view=view,
    ...                     saving_mode="volume",
    ...                     data_mode="radius",
    ...                     debug=True)
    """

    # check if the input file exists and has a valid NIfTI extension
    validate_nifti_path(nii_path)

    # validate the view parameter
    validate_view(view)

    # create output dir if it does not exist
    ensure_dir(output_path)

    # load the NIfTI file
    nii_img = nib.load(nii_path)
    nii_data = nii_img.get_fdata()

    # validate NIfTI data dimensions
    if nii_data.ndim != 3:
        raise ValueError(f"Error: expected a 3D NIfTI file. Got shape '{nii_data.shape}' instead.")

    # validate saving_mode and data_mode
    if saving_mode not in ["slice", "volume"]:
        raise ValueError("Error: saving_mode must be either 'slice' or 'volume'.")
    if data_mode not in ["center", "box", "radius"]:
        raise ValueError("Error: data_mode must be either 'center', 'box', or 'radius'.")

    # extract filename prefix
    prefix = strip_nifti_ext(os.path.basename(nii_path))
    # also strip trailing _bbox suffix if present
    if prefix.endswith("_bbox"):
        prefix = prefix[:-5]

    # calculate padding offsets if target_size is specified
    pad_offset_x = 0
    pad_offset_y = 0

    if target_size is not None:
        target_h, target_w = target_size

        # determine slice dimensions based on view
        # Note: When slices are extracted, they follow numpy convention [height, width]
        # Axial: slice is nii_data[:, :, z] → shape (X_3D, Y_3D) → [height=X_3D, width=Y_3D]
        # Coronal: slice is nii_data[:, y, :] → shape (X_3D, Z_3D) → [height=X_3D, width=Z_3D]
        # Sagittal: slice is nii_data[x, :, :] → shape (Y_3D, Z_3D) → [height=Y_3D, width=Z_3D]
        if view == "axial":
            current_h, current_w = nii_data.shape[0], nii_data.shape[1]  # X_3D, Y_3D
        elif view == "coronal":
            current_h, current_w = nii_data.shape[0], nii_data.shape[2]  # X_3D, Z_3D
        else:  # sagittal
            current_h, current_w = nii_data.shape[1], nii_data.shape[2]  # Y_3D, Z_3D

        # validate that target size is not smaller than current size
        if target_h < current_h or target_w < current_w:
            raise ValueError(
                f"Error: target_size {target_size} is smaller than slice dimensions "
                f"({current_h}, {current_w}). Target size must be >= slice dimensions."
            )

        # calculate padding offsets (matches extract_slices symmetric padding)
        pad_h = target_h - current_h
        pad_w = target_w - current_w
        pad_offset_y = pad_h // 2  # offset for Y dimension (height, rows)
        pad_offset_x = pad_w // 2  # offset for X dimension (width, columns)

    # identify all 3D bounding boxes
    unique_labels = np.unique(nii_data)
    unique_labels = unique_labels[unique_labels > 0]

    bounding_boxes = []

    # iterate over unique labels to extract bounding box information
    for label in tqdm(unique_labels, desc=f"Processing {prefix}", unit="box"):
        positions = np.argwhere(nii_data == label)

        # get min/max for bounding box in 3D coordinates (numpy indexing: [X_3D, Y_3D, Z_3D])
        min_x, min_y, min_z = np.min(positions, axis=0)
        max_x, max_y, max_z = np.max(positions, axis=0)

        # calculate 3D centers and radii
        center_x_3d = (min_x + max_x) / 2
        center_y_3d = (min_y + max_y) / 2
        center_z_3d = (min_z + max_z) / 2
        radius_x_3d = (max_x - min_x) / 2
        radius_y_3d = (max_y - min_y) / 2
        radius_z_3d = (max_z - min_z) / 2

        # Apply coordinate mapping based on view
        # Image coordinates (x, y) map to (column, row) = (width, height)
        # So 2D X (in CSV) corresponds to column/width position
        # And 2D Y (in CSV) corresponds to row/height position

        if view == "axial":
            # Axial: 2D X=Y_3D, 2D Y=X_3D, Slice=Z_3D
            coord_2d_x = center_y_3d + pad_offset_x  # 2D X = Y_3D
            coord_2d_y = center_x_3d + pad_offset_y  # 2D Y = X_3D
            slice_min = min_z
            slice_max = max_z

            if data_mode == "center":
                box_data = [coord_2d_x, coord_2d_y, slice_min, slice_max]
            elif data_mode == "radius":
                box_data = [center_x_3d, center_y_3d, center_z_3d, radius_x_3d, radius_y_3d, radius_z_3d]
            else:  # box
                box_data = [min_y + pad_offset_x, min_x + pad_offset_y, slice_min,
                            max_y + pad_offset_x, max_x + pad_offset_y, slice_max]

        elif view == "coronal":
            # Coronal: 2D X=Z_3D, 2D Y=X_3D, Slice=Y_3D
            # Slice shape is [X_3D, Z_3D] (height, width)
            # 2D X (column position) maps to Z_3D
            # 2D Y (row position) maps to X_3D
            coord_2d_x = center_z_3d + pad_offset_x
            coord_2d_y = center_x_3d + pad_offset_y
            slice_min = min_y
            slice_max = max_y

            if data_mode == "center":
                box_data = [coord_2d_x, coord_2d_y, slice_min, slice_max]
            elif data_mode == "radius":
                box_data = [center_x_3d, center_y_3d, center_z_3d, radius_x_3d, radius_y_3d, radius_z_3d]
            else:  # box
                box_data = [min_z + pad_offset_x, min_x + pad_offset_y, slice_min,
                            max_z + pad_offset_x, max_x + pad_offset_y, slice_max]

        else:  # sagittal
            # Sagittal: 2D X=Z_3D, 2D Y=Y_3D, Slice=X_3D
            # Slice shape is [Y_3D, Z_3D] (height, width)
            # 2D X (column position) maps to Z_3D
            # 2D Y (row position) maps to Y_3D
            coord_2d_x = center_z_3d + pad_offset_x
            coord_2d_y = center_y_3d + pad_offset_y
            slice_min = min_x
            slice_max = max_x

            if data_mode == "center":
                box_data = [coord_2d_x, coord_2d_y, slice_min, slice_max]
            elif data_mode == "radius":
                box_data = [center_x_3d, center_y_3d, center_z_3d, radius_x_3d, radius_y_3d, radius_z_3d]
            else:  # box
                box_data = [min_z + pad_offset_x, min_y + pad_offset_y, slice_min,
                            max_z + pad_offset_x, max_y + pad_offset_y, slice_max]

        bounding_boxes.append(box_data)

    # collect saved CSV file paths
    saved_paths = []

    # handle extraction mode
    if saving_mode == "volume":
        csv_file = os.path.join(output_path, f"{prefix}.csv")
        with open(csv_file, mode="w", newline="") as csvfile:
            writer = csv.writer(csvfile)

            # write appropriate headers based on data_mode
            if data_mode == "center":
                writer.writerow(["CENTER_X", "CENTER_Y", "CENTER_Z"])
                bounding_boxes_vol = [[box[0], box[1], box[2]] for box in bounding_boxes]
            elif data_mode == "radius":
                writer.writerow(["CENTER_X", "CENTER_Y", "CENTER_Z", "RADIUS_X", "RADIUS_Y", "RADIUS_Z"])
                bounding_boxes_vol = bounding_boxes
            else:  # box
                writer.writerow(["X_MIN", "Y_MIN", "Z_MIN", "X_MAX", "Y_MAX", "Z_MAX"])
                bounding_boxes_vol = bounding_boxes

            writer.writerows(bounding_boxes_vol)

        saved_paths.append(csv_file)

        # debug logging
        if debug:
            padding_info = f"\nPadding adjustment applied: target size {target_size}" if target_size else "\nNo padding adjustment applied"
            logger.info(f"\nInput file: '{nii_path}'\nOutput path: '{output_path}'"
                  f"{padding_info}\nTotal volume annotations extracted: {len(bounding_boxes)}")

    else:
        slice_annotations = {}

        # build a dict record for each slice with annotation
        for box in bounding_boxes:
            if data_mode == "radius":
                # For radius mode, use the appropriate 3D coordinate for slice number
                if view == "axial":
                    center_slice = box[2]  # center_z_3d
                    radius_slice = box[5]  # radius_z_3d
                elif view == "coronal":
                    center_slice = box[1]  # center_y_3d
                    radius_slice = box[4]  # radius_y_3d
                else:  # sagittal
                    center_slice = box[0]  # center_x_3d
                    radius_slice = box[3]  # radius_x_3d

                slice_min = int(center_slice - radius_slice)
                slice_max = int(center_slice + radius_slice)
            else:
                slice_min = int(box[2])
                slice_max = int(box[3] if data_mode == "center" else box[5])

            for slice_num in range(slice_min, slice_max + 1):
                if slice_num not in slice_annotations:
                    slice_annotations[slice_num] = []

                if data_mode == "center":
                    # store only 2D X and Y for center mode
                    # box format: [coord_2d_x, coord_2d_y, slice_min, slice_max]
                    slice_annotations[slice_num].append([box[0], box[1]])
                elif data_mode == "radius":
                    # For radius mode, box contains 3D data: [center_x_3d, center_y_3d, center_z_3d, radius_x_3d, radius_y_3d, radius_z_3d]
                    # We need to convert to 2D coordinates with correct mapping
                    if view == "axial":
                        # Axial: 2D X=Y_3D, 2D Y=X_3D
                        coord_2d_x = box[1] + pad_offset_x  # 2D X = Y_3D = center_y_3d
                        coord_2d_y = box[0] + pad_offset_y  # 2D Y = X_3D = center_x_3d
                        rx = box[4]  # radius for 2D X = radius_y_3d
                        ry = box[3]  # radius for 2D Y = radius_x_3d
                        slice_annotations[slice_num].append([coord_2d_x, coord_2d_y, rx, ry])
                    elif view == "coronal":
                        # Coronal: 2D X=Z_3D, 2D Y=X_3D
                        coord_2d_x = box[2] + pad_offset_x  # 2D X = Z_3D = center_z_3d
                        coord_2d_y = box[0] + pad_offset_y  # 2D Y = X_3D = center_x_3d
                        rx = box[5]  # radius for 2D X = radius_z_3d
                        ry = box[3]  # radius for 2D Y = radius_x_3d
                        slice_annotations[slice_num].append([coord_2d_x, coord_2d_y, rx, ry])
                    else:  # sagittal
                        # Sagittal: 2D X=Z_3D, 2D Y=Y_3D
                        coord_2d_x = box[2] + pad_offset_x  # 2D X = Z_3D = center_z_3d
                        coord_2d_y = box[1] + pad_offset_y  # 2D Y = Y_3D = center_y_3d
                        rx = box[5]  # radius for 2D X = radius_z_3d
                        ry = box[4]  # radius for 2D Y = radius_y_3d
                        slice_annotations[slice_num].append([coord_2d_x, coord_2d_y, rx, ry])
                else:  # box
                    # box format for center and box modes: already has 2D coords in [0] and [1]
                    # and bounds in [3] and [4]
                    slice_annotations[slice_num].append([box[0], box[1], box[3], box[4]])

        # process each slice
        num_slices = len(slice_annotations)
        for slice_num, boxes in tqdm(slice_annotations.items(), desc=f"Processing {prefix} (Slices)", unit="slice"):
            slice_filename = f"{prefix}_{view}_{str(slice_num).zfill(3)}.csv"
            slice_file = os.path.join(output_path, slice_filename)

            with open(slice_file, mode="w", newline="") as csvfile:
                writer = csv.writer(csvfile)
                if data_mode == "center":
                    writer.writerow(["CENTER_X", "CENTER_Y"])
                elif data_mode == "radius":
                    writer.writerow(["CENTER_X", "CENTER_Y", "RADIUS_X", "RADIUS_Y"])
                else:  # box
                    writer.writerow(["X_MIN", "Y_MIN", "X_MAX", "Y_MAX"])
                writer.writerows(boxes)

            saved_paths.append(slice_file)

        # debug logging
        if debug:
            padding_info = f"\nPadding adjustment applied: target size {target_size}" if target_size else "\nNo padding adjustment applied"
            logger.info(f"\nInput file: '{nii_path}'\nOutput path: '{output_path}'"
                  f"{padding_info}\nTotal slices with annotations extracted: {num_slices}")

    return saved_paths


def extract_annotations_dataset(nii_folder: str,
                                output_path: str,
                                view: str = "axial",
                                saving_mode: str = "case",
                                extraction_mode: str = "slice",
                                data_mode: str = "center",
                                target_size: Optional[Tuple[int, int]] = None,
                                save_stats: bool = False) -> int:
    """
    Extracts annotations from all NIfTI annotation files in a dataset folder and saves them as CSV, based on the selected view and named with:

        <NIFTI FILENAME>_<VIEW>_<PROGRESSIVE SLICE NUMBER>.csv

    or

        <NIFTI FILENAME>.csv

    :param nii_folder:
        Path to the folder containing all .nii.gz files with shape (X, Y, Z).

    :param output_path:
        Path where the extracted annotations will be saved.

    :param view:
        Anatomical view for annotation extraction:

        - ``"axial"`` → extracts along the Z-axis.
        - ``"coronal"`` → extracts along the Y-axis.
        - ``"sagittal"`` → extracts along the X-axis.

    :param saving_mode:
        - ``"case"`` → creates a folder for each case.
        - ``"view"`` → saves all CSVs inside a single folder.

    :param extraction_mode:
        - ``"slice"`` → generates a CSV per slice.
        - ``"volume"`` → generates a single CSV for the whole volume.

    :param data_mode:
        - ``"center"`` → saves the center (X, Y, Z) of the bounding box.
        - ``"box"`` → saves the bounding box coordinates.

    :param target_size:
        Optional target dimensions (height, width) for coordinate adjustment. If specified,
        coordinates will be adjusted to account for padding. This should match the
        ``target_size`` used in ``extract_slices_dataset`` to ensure alignment between
        images and annotations. If ``None``, coordinates are saved at their original values.

        Example: ``target_size=(512, 512)``

    :param save_stats:
        If ``True``, saves a CSV file with FILENAME and NUM_ANNOTATIONS information per
        case as ``<VIEW>_annotations_stats.csv``.

    :returns:
        Total annotation count across all NIfTI files.

    :raises FileNotFoundError:
        If the dataset folder does not exist or contains no .nii.gz files.

    :raises ValueError:
        If an invalid view, saving_mode or data_mode is provided.

    Example
    -------
    >>> from nidataset.slices import extract_annotations_dataset
    >>>
    >>> # define paths
    >>> nii_folder = "path/to/dataset"
    >>> output_path = "path/to/output_directory"
    >>>
    >>> # choose the anatomical view ('axial', 'coronal', or 'sagittal')
    >>> view = "axial"
    >>>
    >>> # extract annotations without padding adjustment
    >>> extract_annotations_dataset(nii_folder=nii_folder,
    ...                             output_path=output_path,
    ...                             view=view,
    ...                             saving_mode="view",
    ...                             extraction_mode="slice",
    ...                             data_mode="center",
    ...                             save_stats=True)
    >>>
    >>> # extract annotations with padding adjustment to match 512x512 images
    >>> extract_annotations_dataset(nii_folder=nii_folder,
    ...                             output_path=output_path,
    ...                             view=view,
    ...                             saving_mode="case",
    ...                             extraction_mode="slice",
    ...                             data_mode="box",
    ...                             target_size=(512, 512),
    ...                             save_stats=True)
    """

    # get all NIfTI files in the dataset folder (sorted; raises if missing)
    nii_files = list_nifti_files(nii_folder)

    # validate the view parameter
    validate_view(view)

    # validate modes
    if saving_mode not in ["case", "view"]:
        raise ValueError(f"Error: saving_mode must be either 'case' or 'view'. Got '{saving_mode}' instead.")
    if extraction_mode not in ["slice", "volume"]:
        raise ValueError(f"Error: extraction_mode must be either 'slice' or 'volume'. Got '{extraction_mode}' instead.")
    if data_mode not in ["center", "box", "radius"]:
        raise ValueError(f"Error: data_mode must be either 'center', 'box', 'radius'. Got '{data_mode}' instead.")

    # create output dir if it does not exist
    ensure_dir(output_path)

    # create a single folder for the chosen view if using "view" mode
    if saving_mode == "view":
        view_output_dir = os.path.join(output_path, view)
        ensure_dir(view_output_dir)

    # initialize statistics tracking
    stats = []
    total_annotations = 0
    stats_file = os.path.join(output_path, f"{view}_annotations_stats.csv") if save_stats else None

    # iterate over nii.gz files with tqdm progress bar
    for nii_file in tqdm(nii_files, desc="Processing NIfTI files", unit="file"):
        # nii.gz file path
        nii_path = os.path.join(nii_folder, nii_file)

        # extract the filename prefix (case ID)
        prefix = strip_nifti_ext(os.path.basename(nii_path))

        # update tqdm description with the current file prefix
        tqdm.write(f"Processing: {prefix}")

        # determine the number of annotations **before** calling extract_annotations
        try:
            nii_img = nib.load(nii_path)
            nii_data = nii_img.get_fdata()
            unique_labels = np.unique(nii_data)
            num_annotations = len(unique_labels[unique_labels > 0])  # Count non-zero annotations
        except Exception as e:
            tqdm.write(f"Error processing {nii_file} for statistical analysis: {e}")
            continue  # skip this file if an error occurs

        # keep track of the total number of annotations
        total_annotations += num_annotations
        if save_stats:
            stats.append([nii_file, num_annotations])

        # determine the appropriate output folder
        if saving_mode == "case":
            case_output_dir = os.path.join(output_path, prefix, view)
            ensure_dir(case_output_dir)
            extract_annotations(nii_path, case_output_dir, view, extraction_mode,
                              data_mode, target_size=target_size, debug=False)
        else:
            extract_annotations(nii_path, view_output_dir, view, extraction_mode,
                              data_mode, target_size=target_size, debug=False)

    # save statistics if enabled
    if save_stats:
        with open(stats_file, mode="w", newline="") as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(["FILENAME", "NUM_ANNOTATIONS"])
            writer.writerows(stats)
            writer.writerow(["TOTAL_ANNOTATIONS", total_annotations])

        padding_info = f" (with padding adjustment to {target_size})" if target_size else ""
        logger.info(f"\nAnnotation statistics saved in: '{stats_file}'{padding_info}")

    return total_annotations
