"""Utilities for loading and converting AUTC DAT files."""

import shutil
import warnings
from pathlib import Path
from typing import TYPE_CHECKING, Any, Optional, TypeAlias, TypedDict

import numpy as np
import numpy.typing as npt
import zarr

if TYPE_CHECKING:
    from rich.progress import Progress

from confusius.io.utils import check_path


class AUTCDATFileHeader(TypedDict):
    """TypedDict representing the header structure of an AUTC DAT file.

    Attributes
    ----------
    acquisition_block_index : int
        Index of the acquisition block.
    n_data_items : int
        Number of data items in the block.
    n_z : int
        Number of depth samples.
    n_x : int
        Number of lateral samples.
    n_frames : int
        Number of frames in the block.
    """

    acquisition_block_index: int
    n_data_items: int
    n_z: int
    n_x: int
    n_frames: int


def _get_dat_dtype(
    block_shape: tuple[int, int, int],
    header_dtype: npt.DTypeLike = np.int32,
    data_dtype: npt.DTypeLike = np.complex64,
) -> np.dtype:
    """Get the Numpy dtype for an AUTC DAT file with a specific block shape.

    Parameters
    ----------
    block_shape : tuple[int, int, int]
        Block shape with ``(n_x, n_z, n_frames)``.
    header_dtype : dtype_like, default: numpy.int32
        Data type used for the header fields.
    data_dtype : dtype_like, default: numpy.complex64
        Data type used for the data field.

    Returns
    -------
    numpy.dtype
        Numpy dtype representing the structure of the DAT file.
    """
    return np.dtype(
        [
            ("acquisition_block_index", header_dtype),
            ("n_data_items", header_dtype),
            ("n_z", header_dtype),
            ("n_x", header_dtype),
            ("n_frames", header_dtype),
            ("field5", header_dtype),
            ("field6", header_dtype),
            ("field7", header_dtype),
            ("field8", header_dtype),
            ("field9", header_dtype),
            ("data", data_dtype, block_shape),
        ]
    )


NumpyIndex: TypeAlias = (
    bool  # Single boolean
    | int  # Single integer
    | slice  # Slice object
    | list[int]  # List of integers
    | npt.NDArray[np.integer]  # Integer array
    | npt.NDArray[np.bool_]  # Boolean array
    | tuple[int | slice | list[int] | npt.NDArray, ...]  # Multi-dimensional
)
"""Type alias for numpy-style indexing objects.

This type represents valid indexing objects that can be used with numpy arrays,
including single values, slices, lists, arrays, and multi-dimensional combinations.
"""


class AUTCDAT:
    """Beamformed IQ data stored in AUTC DAT format.

    Parameters
    ----------
    path : str or pathlib.Path
        Path to the AUTC DAT file.
    n_header_items : int, default: 10
        Number of header items to read from each block.
    header_dtype : dtype_like, default: numpy.int32
        Data type used for header fields.
    data_dtype : dtype_like, default: numpy.complex64
        Data type used for data fields.

    Attributes
    ----------
    path : pathlib.Path
        Resolved path to the AUTC DAT file.
    shape : tuple[int, int, int, int]
        Shape of the data in the AUTC DAT file in the format ``(blocks, x, z, frames)``.
    n_blocks : int
        Number of frame blocks in the AUTC DAT file.
    n_frames_per_block : int
        Number of frames per block in the AUTC DAT file.
    acquisition_block_indices : numpy.ndarray
        Array of acquisition block indices contained in the AUTC DAT file.
    """

    def __init__(
        self,
        path: str | Path,
        n_header_items: int = 10,
        header_dtype: npt.DTypeLike = np.int32,
        data_dtype: npt.DTypeLike = np.complex64,
    ):
        self.path = check_path(path, type="file")

        self._header_dtype = header_dtype
        first_header = self._read_first_header(
            dtype=header_dtype, n_items=n_header_items
        )

        n_x = first_header["n_x"]
        n_z = first_header["n_z"]
        n_frames = first_header["n_frames"]

        self._memmap = np.memmap(
            self.path,
            dtype=_get_dat_dtype(
                block_shape=(n_x, n_z, n_frames),
                header_dtype=header_dtype,
                data_dtype=data_dtype,
            ),
            mode="r",
        )

        self._acquisition_block_indices = self._read_acquisition_block_indices()

    def _read_first_header(
        self, dtype: npt.DTypeLike, n_items: int
    ) -> AUTCDATFileHeader:
        """Internal method to read the first block's header.

        Parameters
        ----------
        dtype : dtype_like
            Data type used for header fields.
        n_items : int
            Number of header items to read.

        Returns
        -------
        AUTCDATFileHeader
            Dictionary containing the first block's header information.
        """
        header = np.fromfile(self.path, dtype=dtype, count=n_items)
        return {
            # Frame index is 1-based in the AUTC DAT format.
            "acquisition_block_index": header[0] - 1,
            "n_data_items": header[1],
            "n_z": header[2],
            "n_x": header[3],
            "n_frames": header[4],
        }

    def _read_acquisition_block_indices(self) -> npt.NDArray:
        """Read acquisition block indices from the AUTC DAT file.

        We read the acquisition block indices directly from the file to avoid loading
        the entire memmap into memory.

        Returns
        -------
        numpy.ndarray
            Array of acquisition block indices.
        """
        acquisition_block_indices = np.empty(self.n_blocks, dtype=np.int32)
        with open(self.path, "r") as file:
            for block_index in range(self.n_blocks):
                file.seek(block_index * self._memmap.dtype.itemsize)

                header = np.fromfile(file, dtype=self._header_dtype, count=1)
                acquisition_block_indices[block_index] = header[0]
        return acquisition_block_indices

    @property
    def shape(self) -> tuple[int, int, int, int]:
        """Shape of the data in this AUTC DAT file.

        Returns
        -------
        tuple[int, int, int, int]
            Tuple of the form ``(blocks, x, z, frames)``.
        """
        return self._memmap["data"].shape

    @property
    def n_blocks(self) -> int:
        """Number of frame blocks in this AUTC DAT file.

        Returns
        -------
        int
            Number of blocks.
        """
        return self.shape[0]

    @property
    def n_frames_per_block(self) -> int:
        """Number of frames per block in this AUTC DAT file.

        Returns
        -------
        int
            Number of frames per block.
        """
        return self.shape[3]

    @property
    def acquisition_block_indices(self) -> npt.NDArray:
        """Acquisition-wide frame indices contained in this AUTC DAT file.

        Returns
        -------
        numpy.ndarray
            Array of acquisition block indices.
        """
        return self._acquisition_block_indices

    def __len__(self) -> int:
        """Number of blocks in this AUTC DAT file.

        Returns
        -------
        int
            Number of blocks.
        """
        return self.shape[0]

    def __getitem__(self, slice_object: NumpyIndex) -> npt.NDArray:
        """Index the data contained in this AUTC DAT file.

        Parameters
        ----------
        slice_object : NumpyIndex
            Indexing object for the data.

        Returns
        -------
        numpy.ndarray
            The requested data from the AUTC DAT file.
        """
        return self._memmap["data"][slice_object]


class AUTCDATsLoader:
    """Load a beamformed IQ acquisition from a set of AUTC DAT files.

    Parameters
    ----------
    root : str or pathlib.Path
        Path to the root directory containing the AUTC DAT files.
    pattern : str, default: "bf*part*.dat"
        Glob pattern used to search for AUTC DAT files inside `root`.
    show_progress : bool, default: True
        Whether to show progress during file indexing. If ``False``, no progress bars
        are displayed.
    progress : rich.progress.Progress, optional
        External `rich.progress.Progress` instance to add tasks to. If provided and
        `show_progress` is ``True``, a task will be added to this
        `rich.progress.Progress` instance instead of creating a new progress bar with
        `rich.progress.track`.
    track_kwargs : dict, optional
        Additional keyword arguments to pass to `rich.progress.track` if using internal
        progress tracking (only used if `show_progress` is ``True`` and `progress` is
        ``None``).

    Attributes
    ----------
    root : pathlib.Path
        Resolved path to the root directory.
    pattern : str
        Glob pattern used for file searching.
    dats : dict[pathlib.Path, AUTCDAT]
        Dictionary mapping AUTC DAT file paths to `AUTCDAT` instances.
    block_index_to_file : dict[int, pathlib.Path]
        Mapping from block indices to their corresponding file paths.
    """

    def __init__(
        self,
        root: str | Path,
        pattern: str = "bf*part*.dat",
        show_progress: bool = True,
        progress: Optional["Progress"] = None,
        track_kwargs: dict[str, Any] | None = None,
    ):
        self.root = check_path(root, label="root", type="dir")
        self.pattern = pattern
        self.dats: dict[Path, AUTCDAT] = {}
        self.block_index_to_file: dict[int, Path] = {}
        self._index_dat_files(
            show_progress=show_progress, progress=progress, track_kwargs=track_kwargs
        )

    def _index_dat_files(
        self,
        show_progress: bool = True,
        progress: Optional["Progress"] = None,
        track_kwargs: dict[str, Any] | None = None,
    ):
        """Locate AUTCDAT files in the root folder and index their frames."""
        from rich.progress import track

        dat_paths = list(self.root.glob(self.pattern))
        if len(dat_paths) == 0:
            raise FileNotFoundError(
                f"No AUTC DAT files found in {self.root} with pattern {self.pattern}."
            )

        task_id = None

        if not show_progress:
            iterable = enumerate(dat_paths)
        elif progress is not None:
            task_id = progress.add_task(
                "Indexing AUTC DAT files...", total=len(dat_paths)
            )
            iterable = enumerate(dat_paths)
        else:
            kwargs = track_kwargs or {}
            kwargs.setdefault("description", "Indexing AUTC DAT files...")
            kwargs.setdefault("total", len(dat_paths))
            iterable = track(enumerate(dat_paths), **kwargs)

        for dat_index, dat_path in iterable:
            if dat_path not in self.dats:
                n_total_bytes = dat_path.stat().st_size
                if n_total_bytes == 0:
                    warnings.warn(f"Skipping empty AUTC DAT file: {dat_path}")
                    continue

                dat = AUTCDAT(dat_path)
                self.dats[dat_path] = dat
                for frame_label in dat.acquisition_block_indices:
                    self.block_index_to_file[frame_label] = dat_path

                if task_id is not None:
                    progress.update(task_id, advance=1)  # type: ignore[union-attr]

    def __getitem__(
        self, slice_object: int | slice | tuple[int | slice, ...]
    ) -> npt.NDArray:
        """Index the data across all AUTC DAT files.

        Parameters
        ----------
        slice_object : int or slice or tuple[int or slice, ...]
            Indexing object for the data. Can be:
            - ``int``: Single block index;
            - ``slice``: Range of block indices;
            - ``tuple``: Multi-dimensional indexing ``(block, x, z, frame)``.

        Returns
        -------
        numpy.ndarray
            The requested data from the AUTC DAT files.
        """
        if isinstance(slice_object, tuple):
            block_slice = slice_object[0]
            remaining_slices = slice_object[1:]
        elif isinstance(slice_object, (int, slice)):
            block_slice = slice_object
            remaining_slices = ()
        else:
            raise TypeError(
                "slice_object must be an int, slice, or tuple of (int or slice)."
            )

        if isinstance(block_slice, int):
            block_slice = slice(block_slice, block_slice + 1)

        block_indices = self.blocks[block_slice]

        # Pre-allocate to avoid repeated allocations during loading.
        loader_shape = self.shape
        output_shape = [len(block_indices)]

        for i, slice_obj in enumerate(remaining_slices):
            if i + 1 < len(loader_shape):
                dim_size = loader_shape[i + 1]
                if isinstance(slice_obj, slice):
                    start, stop, step = slice_obj.indices(dim_size)
                    output_shape.append(len(range(start, stop, step)))
                elif isinstance(slice_obj, int):
                    output_shape.append(1)

        for i in range(len(remaining_slices) + 1, len(loader_shape)):
            output_shape.append(loader_shape[i])

        result = np.empty(tuple(output_shape), dtype=self.dtype)

        # We group blocks by AUTC DAT file for efficient access.
        blocks_by_file: dict[Path, list] = {}
        for i, acquisition_block_index in enumerate(block_indices):
            dat_path = self.block_index_to_file[acquisition_block_index]
            if dat_path not in blocks_by_file:
                blocks_by_file[dat_path] = []
            blocks_by_file[dat_path].append((i, acquisition_block_index))

        for dat_path, file_blocks in blocks_by_file.items():
            dat = self.dats[dat_path]

            for original_index, acquisition_block_index in file_blocks:
                block_position = np.where(
                    dat.acquisition_block_indices == acquisition_block_index
                )[0][0]

                full_slice = (block_position,) + remaining_slices
                result[original_index] = dat[full_slice]

        return result

    @property
    def blocks(self) -> npt.NDArray:
        """List of all blocks available in the acquisition.

        Returns
        -------
        numpy.ndarray
            Sorted array of all available block indices.
        """
        return np.asarray(sorted(self.block_index_to_file.keys()))

    @property
    def shape(self) -> tuple[int, int, int, int]:
        """Shape of the complete dataset across all AUTC DAT files.

        Returns
        -------
        tuple[int, int, int, int]
            Shape in the format ``(n_blocks, n_x, n_z, n_frames)``.
        """
        first_dat_shape = list(self.dats.values())[0].shape
        return (
            len(self.blocks),
            first_dat_shape[1],
            first_dat_shape[2],
            first_dat_shape[3],
        )

    @property
    def dtype(self) -> np.dtype:
        """Data type of the arrays in the AUTC DAT files.

        Returns
        -------
        numpy.dtype
            Data type of the stored arrays.
        """
        return list(self.dats.values())[0]._memmap["data"].dtype


def convert_autc_dats_to_zarr(
    dats_root: str | Path,
    output_path: str | Path,
    dats_pattern: str = "bf*part*.dat",
    frames_per_chunk: int | None = None,
    frames_per_shard: int | None = None,
    batch_size: int = 100,
    overwrite: bool = False,
    zarr_kwargs: dict[str, Any] | None = None,
    lateral_coords: npt.ArrayLike | None = None,
    axial_coords: npt.ArrayLike | None = None,
    transmit_frequency: float | None = None,
    probe_n_elements: int | None = None,
    probe_pitch: float | None = None,
    speed_of_sound: float | None = None,
    plane_wave_angles: npt.ArrayLike | None = None,
    n_transmissions: int | None = None,
    compound_sampling_frequency: float | None = None,
    beamforming_method: str | None = None,
    skip_first_blocks: int = 0,
    skip_last_blocks: int = 0,
    block_times: npt.ArrayLike | None = None,
    show_progress: bool = True,
    progress: Optional["Progress"] = None,
    track_kwargs: dict[str, Any] | None = None,
) -> "zarr.Group":
    """Convert AUTC DAT files to Zarr format compatible with Xarray.

    Beamformed IQ data is converted to a Zarr group with an ``iq`` array of shape
    ``(time, z, y, x)`` chunked along the first dimension. Coordinates are stored as
    separate Zarr arrays following Xarray conventions, allowing the data to be opened
    directly with ``xarray.open_zarr()``.

    Parameters
    ----------
    dats_root : str or pathlib.Path
        Path to the directory containing AUTC DAT files.
    output_path : str or pathlib.Path
        Path where the Zarr group will be saved.
    dats_pattern : str, default: "bf*part*.dat"
        Glob pattern used to search for AUTC DAT files inside `dats_root`.
    frames_per_chunk : int, optional
        Number of frames to include in each Zarr chunk. If not provided, defaults to
        the number of frames per block in the AUTC DAT files.
    frames_per_shard : int, optional
        Number of frames to include in each shard. If provided, enables Zarr sharding
        to reduce the number of files on disk. Must be a multiple of
        `frames_per_chunk`. If not provided, sharding is disabled.
    batch_size : int, default: 100
        Number of blocks to process in each batch.
    overwrite : bool, default: False
        Whether to overwrite existing Zarr group at the output path.
    zarr_kwargs : dict, optional
        Additional keyword arguments to pass to `zarr.create_array` for the main data
        array.
    lateral_coords : array_like, optional
        Lateral coordinates in millimeters. These define the x-axis positions.
    axial_coords : array_like, optional
        Axial (depth) coordinates in millimeters. These define the y-axis positions.
    transmit_frequency : float, optional
        Central frequency of the ultrasound probe in hertz.
    probe_n_elements : int, optional
        Number of probe transducers.
    probe_pitch : float, optional
        Inter-element pitch of the probe in millimeters.
    speed_of_sound : float, optional
        Speed of sound in meters per second.
    plane_wave_angles : array_like, optional
        Angles at which tilted plane waves are emitted in degrees.
    n_transmissions : int, optional
        Number of plane wave transmissions.
    compound_sampling_frequency : float, optional
        Sampling frequency of the compounded frames in hertz.
    beamforming_method : str, optional
        Beamforming method used (e.g. ``"DAS"``).
    skip_first_blocks : int, default: 0
        Number of blocks to skip from the beginning of the acquisition. This is useful
        when the first blocks are known to be corrupted or unusable.
    skip_last_blocks : int, default: 0
        Number of blocks to skip from the end of the acquisition. This is useful
        when the last blocks are known to be corrupted or unusable.
    block_times : (n_blocks_after_skip,) array_like, optional
        Start time of each IQ block in seconds, for the retained blocks only. If
        provided, individual volume times will be computed and stored as a time
        coordinate. Requires `compound_sampling_frequency` to be provided. If not
        provided, time coordinate will be computed based on
        `compound_sampling_frequency` or set to frame indices.
    show_progress : bool, default: True
        Whether to show progress during conversion. If ``False``, no progress bars are
        displayed.
    progress : rich.progress.Progress, optional
        External `rich.progress.Progress` instance to add tasks to. If provided and
        `show_progress` is ``True``, a task will be added to this
        `rich.progress.Progress` instance instead of creating a new progress bar with
        `rich.progress.track`.
    track_kwargs : dict, optional
        Additional keyword arguments to pass to `rich.progress.track` if using internal
        progress tracking (only used if `show_progress` is ``True`` and `progress` is
        ``None``).

    Returns
    -------
    zarr.Group
        The created Zarr group containing the ``iq`` data array and coordinate arrays.
        Can be opened directly with ``xarray.open_zarr()``.

    Notes
    -----
    The output Zarr group follows Xarray conventions and can be opened with::

        import xarray as xr
        ds = xr.open_zarr("output.zarr")
        iq = ds["iq"]

    Metadata attributes (e.g., ``transmit_frequency``, ``sound_velocity``) are stored
    on the ``iq`` DataArray (accessible via ``iq.attrs``), consistent with how
    reduction functions return DataArrays with attributes.

    The group contains:

    - ``iq``: The main data array with dimensions ``(time, z, y, x)``.
    - ``time``: Time coordinate array.
    - ``z``: Elevation coordinate array (always ``[0]`` for 2D data).
    - ``y``: Axial (depth) coordinate array (if ``axial_coords`` is provided, otherwise indices).
    - ``x``: Lateral coordinate array (if ``lateral_coords`` is provided, otherwise indices).
    """
    from rich.progress import track

    loader = AUTCDATsLoader(
        root=dats_root,
        pattern=dats_pattern,
        show_progress=show_progress,
        progress=progress,
        track_kwargs=track_kwargs,
    )

    output_path = Path(output_path)
    n_blocks, n_x, n_z, n_frames = loader.shape

    total_skip = skip_first_blocks + skip_last_blocks
    if total_skip >= n_blocks:
        raise ValueError(
            f"Cannot skip {total_skip} blocks (skip_first_blocks={skip_first_blocks}, "
            f"skip_last_blocks={skip_last_blocks}) from {n_blocks} total blocks."
        )

    n_blocks_after_skip = n_blocks - total_skip
    n_total_frames = n_blocks_after_skip * n_frames

    if frames_per_chunk is None:
        frames_per_chunk = n_frames

    output_shape = (n_total_frames, 1, n_z, n_x)
    chunks = (frames_per_chunk, 1, n_z, n_x)

    if frames_per_shard is not None:
        if frames_per_shard % frames_per_chunk != 0:
            raise ValueError(
                f"frames_per_shard ({frames_per_shard}) must be a multiple of "
                f"frames_per_chunk ({frames_per_chunk})."
            )
        shards = (frames_per_shard, 1, n_z, n_x)
    else:
        shards = None

    zarr_group = zarr.open_group(output_path, mode="w" if overwrite else "w-")

    # Dimension names required for Zarr v3 / Xarray compatibility.
    dim_names = ["time", "z", "y", "x"]

    create_array_kwargs = zarr_kwargs.copy() if zarr_kwargs else {}
    handled_keys = {"shape", "chunks", "shards", "dtype", "dimension_names"}
    overridden_keys = handled_keys & create_array_kwargs.keys()
    if overridden_keys:
        warnings.warn(
            f"zarr_kwargs contains keys that are handled by function parameters and "
            f"will be overridden: {overridden_keys}.",
            stacklevel=2,
        )
    create_array_kwargs["shape"] = output_shape
    create_array_kwargs["chunks"] = chunks
    create_array_kwargs["dtype"] = loader.dtype
    create_array_kwargs["dimension_names"] = dim_names

    if shards is not None:
        create_array_kwargs["shards"] = shards

    zarr_iq = zarr_group.create_array("iq", **create_array_kwargs)

    if block_times is not None:
        if compound_sampling_frequency is None:
            raise ValueError(
                "compound_sampling_frequency is required when block_times is provided "
                "to compute individual volume times."
            )
        block_times_array = np.asarray(block_times)
        if block_times_array.size != n_blocks_after_skip:
            raise ValueError(
                f"block_times length ({block_times_array.size}) does not match "
                f"number of blocks after skipping ({n_blocks_after_skip})."
            )
        time_values = np.concatenate(
            [
                block_start + np.arange(n_frames) / compound_sampling_frequency
                for block_start in block_times_array
            ]
        )
    elif compound_sampling_frequency is not None:
        skipped_first_frames = skip_first_blocks * n_frames
        time_values = (
            np.arange(n_total_frames, dtype=np.float64) / compound_sampling_frequency
            + skipped_first_frames / compound_sampling_frequency
        )
    else:
        time_values = np.arange(n_total_frames, dtype=np.float64)

    zarr_group.create_array("time", data=time_values, dimension_names=["time"])
    zarr_group["time"].attrs["units"] = "s"
    zarr_group["time"].attrs["long_name"] = "Time"

    # z coordinate is the stacking dimension (size 1 for 2D data).
    z_values = np.array([0.0])
    zarr_group.create_array("z", data=z_values, dimension_names=["z"])
    zarr_group["z"].attrs["units"] = "mm"
    zarr_group["z"].attrs["long_name"] = "Elevation"

    if axial_coords is not None:
        y_values = np.asarray(axial_coords)
    else:
        y_values = np.arange(n_z, dtype=np.float64)
    zarr_group.create_array("y", data=y_values, dimension_names=["y"])
    zarr_group["y"].attrs["units"] = "mm"
    zarr_group["y"].attrs["long_name"] = "Depth"

    if lateral_coords is not None:
        x_values = np.asarray(lateral_coords)
    else:
        x_values = np.arange(n_x, dtype=np.float64)
    zarr_group.create_array("x", data=x_values, dimension_names=["x"])
    zarr_group["x"].attrs["units"] = "mm"
    zarr_group["x"].attrs["long_name"] = "Lateral"

    lateral_coords_dim = (
        float(np.diff(np.asarray(lateral_coords)).mean())
        if lateral_coords is not None
        else 0.0
    )
    axial_coords_dim = (
        float(np.diff(np.asarray(axial_coords)).mean())
        if axial_coords is not None
        else 0.0
    )
    # TODO: we should compute the actual y-axis voxdim from the elevation beam width,
    # but we're currently missing some information for that, such as the elevation
    # aperture and elevation focus.
    # voxdim is ordered as [z, y, x].
    zarr_iq.attrs["voxdim"] = [0.4, axial_coords_dim, lateral_coords_dim]

    if transmit_frequency is not None:
        zarr_iq.attrs["transmit_frequency"] = transmit_frequency

    if probe_n_elements is not None:
        zarr_iq.attrs["probe_n_elements"] = probe_n_elements

    if probe_pitch is not None:
        zarr_iq.attrs["probe_pitch"] = probe_pitch

    if speed_of_sound is not None:
        zarr_iq.attrs["sound_velocity"] = speed_of_sound

    if plane_wave_angles is not None:
        zarr_iq.attrs["plane_wave_angles"] = np.asarray(plane_wave_angles).tolist()

    if compound_sampling_frequency is not None:
        zarr_iq.attrs["compound_sampling_frequency"] = compound_sampling_frequency

        if n_transmissions is not None:
            zarr_iq.attrs["pulse_repetition_frequency"] = (
                compound_sampling_frequency * n_transmissions
            )

    if beamforming_method is not None:
        zarr_iq.attrs["beamforming_method"] = beamforming_method

    first_block = skip_first_blocks
    last_block = n_blocks - skip_last_blocks

    n_batches = (n_blocks_after_skip + batch_size - 1) // batch_size
    batches = range(first_block, last_block, batch_size)
    task_id = None

    if not show_progress:
        iterable = batches
    elif progress is not None:
        task_id = progress.add_task(
            "Converting AUTC DAT files to Zarr...", total=n_batches
        )
        iterable = batches
    else:
        kwargs = track_kwargs or {}
        kwargs.setdefault("description", "Converting AUTC DAT files to Zarr...")
        iterable = track(batches, **kwargs)

    try:
        for idx, start_block in enumerate(iterable):
            end_block = min(start_block + batch_size, last_block)

            batch_data = np.transpose(
                loader[start_block:end_block], axes=(0, 3, 2, 1)
            ).reshape(-1, 1, n_z, n_x)

            output_start = (start_block - skip_first_blocks) * n_frames
            output_end = (end_block - skip_first_blocks) * n_frames
            zarr_iq[output_start:output_end] = batch_data

            if task_id is not None:
                progress.update(task_id, advance=1)  # type: ignore[union-attr]

        # Consolidate metadata for faster opening with Xarray.
        # Suppress warning about consolidated metadata not being in Zarr v3 spec yet.
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", message="Consolidated metadata")
            zarr.consolidate_metadata(output_path)

    except Exception:
        # Clean up incomplete zarr store to avoid leaving unconsolidated data.
        if output_path.exists():
            shutil.rmtree(output_path)
        raise

    return zarr_group
