"""
Base snapshot interface for AMUN simulation outputs.

This module defines the Amun base class, which provides the common
interface and lifecycle for reading AMUN simulation snapshots,
independent of the underlying storage format.

Responsibilities include:
- snapshot initialization and validation
- access to global attributes and datasets
- reconstruction of derived quantities
- consistency checks shared by all snapshot formats

Concrete snapshot readers (e.g. AmunXML, AmunH5) implement the
format-specific parsing and data access logic on top of this base class.
"""
import os, sys
import numpy as np

from .exceptions import (
    AmunPyError,
    AmunPyAttributeError,
    AmunPyChunkError,
    AmunPyVariableError,
    AmunPyDatasetError,
)
from .interpolation import interpolate

class Amun:
    """
    Base class for AMUN snapshot readers.

    The Amun class defines the common snapshot lifecycle and user-facing
    interface shared by all snapshot formats (e.g. AmunXML, AmunH5).
    Concrete subclasses implement the format-specific hooks:
      - __init_snapshot__(): validate the format and set self.dataformat
      - __fill_attributes__(): populate self.attributes
      - __fill_chunks__(): populate self.chunks (per-chunk metadata)
      - __fill_variables__(): populate self.variables (available datasets)
      - __read_binary_data__(): load a dataset array from a given chunk

    The base class provides:
      - mandatory metadata checks and completion of derived attributes
      - derived dataset reconstruction (velocity magnitude, energies, etc.)
      - resampling to a uniform mesh (dataset())
    """
    def __init__(self, path, version=-1):
        """
        Construct a snapshot reader and fully initialize snapshot metadata.

        Parameters
        ----------
        path : str
            Path to the snapshot. May be a directory (XML snapshots) or a file
            (HDF5 snapshots), depending on the concrete reader implementation.
        version : int, optional
            Snapshot format version used for compatibility handling (e.g. axis
            ordering). If not provided, subclasses may infer it from metadata.

        Notes
        -----
        Initialization sequence:
        1) __init_snapshot__()
        2) __fill_attributes__(), __fill_chunks__(), __fill_variables__()
        3) __check_*() and __complete_*() validation/completion steps
        """
        if os.path.exists(path):
            if os.path.isdir(path):
                self.path_is_file = False
                self.dirname      = path
                self.filename     = None
            else:
                self.path_is_file = True
                self.dirname      = os.path.dirname(path)
                self.filename     = os.path.basename(path)
        else:
            raise AmunPyError("Path '{}' does not exist!".format(path))

        self.path        = path
        self.dataformat  = None
        self.version     = version
        self.chunkname   = '{:d}'
        self.attributes  = dict()
        self.chunks      = dict()
        self.variables   = dict()

        self.__init_snapshot__()
        self.__fill_attributes__()
        self.__fill_chunks__()
        self.__fill_variables__()

        self.__check_attributes__()
        self.__check_chunks__()
        self.__check_variables__()

        self.__complete_attributes__()
        self.__complete_chunks__()
        self.__complete_variables__()


    def __init_snapshot__(self):
        """
        Format-specific initialization hook.

        Subclasses must validate the snapshot format and set:
        - self.dataformat (string identifier)
        - self.version (if available from metadata)

        This method is called during Amun.__init__().
        """
        pass


    def __fill_attributes__(self):
        """
        Format-specific attribute parsing hook.

        Subclasses must populate `self.attributes` with the snapshot's global
        metadata (e.g. isnap, nchunks, ndims, bounds, block layout, etc.).
        """
        pass


    def __fill_chunks__(self):
        """
        Format-specific chunk metadata hook.

        Subclasses must populate `self.chunks`, indexed by chunk number.
        Each chunk must include at least:
        - filename : str
        - dblocks  : int
        - levels   : array-like, shape (dblocks,)
        - coords   : array-like, shape (dblocks, ndims)
        - bounds   : array-like, shape (dblocks, 2, ndims)
        """
        pass


    def __fill_variables__(self):
        """
        Format-specific variable discovery hook.

        Subclasses must populate `self.variables` with available dataset names
        (including the base primitive fields needed for derived quantities).
        """
        pass


    def __check_attributes__(self):
        """
        Validate that all mandatory global attributes are present and consistent.

        Raises
        ------
        AmunPyAttributeError
            If a required attribute is missing or invalid in `self.attributes`.
        """
        mandatory_attributes = [
            'eqsys', 'eos', 'isnap', 'nchunks', 'ndims',
            'maxlev', 'nleafs', 'ncells', 'nghosts',
            'xmin', 'ymin', 'zmin', 'xmax', 'ymax', 'zmax',
            'xblocks', 'yblocks', 'zblocks',
        ]

        for attr in mandatory_attributes:
            if attr not in self.attributes:
                raise AmunPyAttributeError(
                    "Mandatory attribute '{}' not set!".format(attr)
                )

        # ---- basic dimensionality and domain sanity checks (dataset() prerequisites)

        ndims = self.attributes['ndims']
        if ndims not in (2, 3):
            raise AmunPyAttributeError(
                "Unsupported 'ndims' value {} (expected 2 or 3).".format(ndims)
            )

        xmin, xmax = self.attributes['xmin'], self.attributes['xmax']
        ymin, ymax = self.attributes['ymin'], self.attributes['ymax']
        zmin, zmax = self.attributes['zmin'], self.attributes['zmax']

        if not (xmin < xmax):
            raise AmunPyAttributeError(
                "Invalid x-domain bounds: xmin={} must be < xmax={}.".format(xmin, xmax)
            )
        if not (ymin < ymax):
            raise AmunPyAttributeError(
                "Invalid y-domain bounds: ymin={} must be < ymax={}.".format(ymin, ymax)
            )
        if ndims == 3 and not (zmin < zmax):
            raise AmunPyAttributeError(
                "Invalid z-domain bounds: zmin={} must be < zmax={}.".format(zmin, zmax)
            )

        # ---- EOS/eqsys presence (already mandatory) + EOS-specific requirements

        eos = self.attributes['eos']
        if eos == 'adi':
            if 'adiabatic_index' not in self.attributes:
                raise AmunPyAttributeError(
                    "Mandatory attribute 'adiabatic_index' not set for eos='adi'!"
                )
            gamma = self.attributes['adiabatic_index']
            try:
                gamma = float(gamma)
            except (TypeError, ValueError):
                raise AmunPyAttributeError(
                    "Invalid 'adiabatic_index' value {!r} (expected a real number > 1).".format(gamma)
                )
            if not (gamma > 1.0):
                raise AmunPyAttributeError(
                    "Invalid 'adiabatic_index' value {} (expected > 1).".format(gamma)
                )

        elif eos == 'iso':
            # Only required if AmunPy supports pressure synthesis for isothermal EOS
            # (used e.g. in __get_dataset__ for 'pressure').
            if 'sound_speed' not in self.attributes:
                raise AmunPyAttributeError(
                    "Mandatory attribute 'sound_speed' not set for eos='iso'!"
                )
            cs = self.attributes['sound_speed']
            try:
                cs = float(cs)
            except (TypeError, ValueError):
                raise AmunPyAttributeError(
                    "Invalid 'sound_speed' value {!r} (expected a real number > 0).".format(cs)
                )
            if not (cs > 0.0):
                raise AmunPyAttributeError(
                    "Invalid 'sound_speed' value {} (expected > 0).".format(cs)
                )

        else:
            raise AmunPyAttributeError(
                "Unsupported 'eos' value {!r} (expected 'adi' or 'iso').".format(eos)
            )


    def __check_chunks__(self):
        """
        Validate that chunk metadata is present and complete.

        Ensures:
        - at least one chunk exists
        - chunk indexing matches `nchunks` (0..nchunks-1)
        - each chunk contains required fields:
            filename, dblocks, levels, bounds, coords
        - required arrays have consistent shapes:
            levels.shape == (dblocks,)
            coords.shape  == (dblocks, ndims)
            bounds.shape  == (dblocks, 2, ndims)
        - refinement levels are within the expected range
        """
        if len(self.chunks) < 1:
            raise AmunPyChunkError("The snapshot {} has no chunks!".format(self.path))

        nchunks = self.attributes.get('nchunks')
        if nchunks is None:
            raise AmunPyChunkError("Mandatory attribute 'nchunks' not set before chunk validation!")

        # Enforce contiguous integer chunk indexing (0..nchunks-1)
        expected_keys = set(range(int(nchunks)))
        got_keys = set(self.chunks.keys())
        if got_keys != expected_keys:
            raise AmunPyChunkError(
                "Chunk index set mismatch for snapshot {}: expected keys {} but got {}.".format(
                    self.path, sorted(expected_keys), sorted(got_keys)
                )
            )

        ndims = self.attributes.get('ndims')
        if ndims is None:
            raise AmunPyChunkError("Mandatory attribute 'ndims' not set before chunk validation!")

        maxlev = self.attributes.get('maxlev')
        if maxlev is None:
            raise AmunPyChunkError("Mandatory attribute 'maxlev' not set before chunk validation!")

        xblocks = self.attributes.get('xblocks')
        yblocks = self.attributes.get('yblocks')
        if xblocks is None or yblocks is None:
            raise AmunPyChunkError("Mandatory attributes 'xblocks'/'yblocks' not set before chunk validation!")

        if int(ndims) == 3:
            zblocks = self.attributes.get('zblocks')
            if zblocks is None:
                raise AmunPyChunkError("Mandatory attribute 'zblocks' not set before chunk validation for ndims=3!")
        else:
            zblocks = None

        mandatory_fields = ['filename', 'dblocks', 'levels', 'bounds', 'coords']

        for nc in range(int(nchunks)):
            chunk = self.chunks[nc]

            for field in mandatory_fields:
                if field not in chunk:
                    raise AmunPyChunkError(
                        "Mandatory field '{}' not set in chunk no. {}!".format(field, nc)
                    )

            dblocks = chunk['dblocks']
            try:
                dblocks = int(dblocks)
            except (TypeError, ValueError):
                raise AmunPyChunkError(
                    "Invalid 'dblocks' value {!r} in chunk no. {} (expected integer).".format(dblocks, nc)
                )
            if dblocks < 0:
                raise AmunPyChunkError(
                    "Invalid 'dblocks' value {} in chunk no. {} (expected >= 0).".format(dblocks, nc)
                )

            # Shape checks
            levels = np.asarray(chunk['levels'])
            coords = np.asarray(chunk['coords'])
            bounds = np.asarray(chunk['bounds'])

            if chunk['levels'] is None or chunk['coords'] is None or chunk['bounds'] is None:
                raise AmunPyChunkError(f"Chunk no. {nc}: levels/coords/bounds must be arrays (empty arrays allowed), not None.")

            if levels.shape != (dblocks,):
                raise AmunPyChunkError(
                    "Invalid 'levels' shape {} in chunk no. {} (expected ({},)).".format(
                        levels.shape, nc, dblocks
                    )
                )

            if coords.ndim != 2 or coords.shape[0] != dblocks or coords.shape[1] < int(ndims):
                raise AmunPyChunkError(
                    "Invalid 'coords' shape {} in chunk no. {} (expected ({}, >= {})).".format(
                        coords.shape, nc, dblocks, int(ndims)
                    )
                )

            if bounds.ndim != 3 or bounds.shape[0] != dblocks or bounds.shape[1] != 2 or bounds.shape[2] < int(ndims):
                raise AmunPyChunkError(
                    "Invalid 'bounds' shape {} in chunk no. {} (expected ({}, 2, >= {})).".format(
                        bounds.shape, nc, dblocks, int(ndims)
                    )
                )

            # ------------------------------------------------------------------
            # Empty chunks (metadata-only ranks) are valid: skip min/max semantics
            # checks that would fail on empty arrays, but keep shape/finite checks.
            # ------------------------------------------------------------------
            if dblocks == 0:
                # Nothing further to validate semantically (no coordinates/levels).
                # Shapes already enforced above.
                continue

            # Numeric sanity (finite)
            if not np.isfinite(levels.astype(float)).all():
                raise AmunPyChunkError(
                    "Non-finite values found in 'levels' for chunk no. {}.".format(nc)
                )
            if not np.isfinite(coords.astype(float)).all():
                raise AmunPyChunkError(
                    "Non-finite values found in 'coords' for chunk no. {}.".format(nc)
                )
            if not np.isfinite(bounds.astype(float)).all():
                raise AmunPyChunkError(
                    "Non-finite values found in 'bounds' for chunk no. {}.".format(nc)
                )

            # ------------------------------------------------------------------
            # coords semantics: integer index-space block coordinates
            # ------------------------------------------------------------------
            coords_view = coords[:, :int(ndims)]  # active dimensions only

            # Require integer dtype OR strictly integer-valued floats (avoid silent truncation)
            if np.issubdtype(coords_view.dtype, np.integer):
                coords_int = coords_view.astype(np.int64, copy=False)
            else:
                if not np.all(np.equal(coords_view, np.round(coords_view))):
                    raise AmunPyChunkError(
                        "Chunk no. {}: 'coords' must be integer-valued index-space block coordinates; "
                        "got dtype {} with non-integer values.".format(nc, coords_view.dtype)
                    )
                coords_int = np.round(coords_view).astype(np.int64)

            # Non-negative
            if np.min(coords_int) < 0:
                raise AmunPyChunkError(
                    "Chunk no. {}: invalid 'coords' contains negative entries (min={}).".format(
                        nc, int(np.min(coords_int))
                    )
                )

            # AMR level range check (assumes 1-based levels as used by cell_size[l+1])
            lmin = int(np.min(levels))
            lmax = int(np.max(levels))
            if lmin < 1 or lmax > int(maxlev):
                raise AmunPyChunkError(
                    "Refinement level out of range in chunk no. {}: levels in [{}, {}], expected within [1, {}].".format(
                        nc, lmin, lmax, int(maxlev)
                    )
                )


    def __check_variables__(self):
        """
        Validate that the snapshot exposes at least one dataset.

        Raises
        ------
        Exception
            If `self.variables` is empty.
        """
        if len(self.variables) < 1:
            raise AmunPyVariableError("The snapshot {} has no variables stored!".format(self.path))


    def __complete_attributes__(self):
        """
        Fill in derived/global convenience attributes.

        Adds defaults and derived values when missing, including:
        - toplev (top refinement level present in the snapshot)
        - domain lengths (xlen, ylen, zlen)
        - bcells (block cells including ghost zones)
        - resistivity / viscosity (default 0)
        - cell_size[l] mapping for refinement levels
        """
        if not 'toplev' in self.attributes:
            self.attributes['toplev'] = max([max(chunk['levels']) if chunk['dblocks'] > 0 else 1 for chunk in self.chunks.values()])

        if not 'xlen' in self.attributes:
            self.attributes['xlen'] = self.attributes['xmax'] - self.attributes['xmin']
        if not 'ylen' in self.attributes:
            self.attributes['ylen'] = self.attributes['ymax'] - self.attributes['ymin']
        if not 'zlen' in self.attributes:
            self.attributes['zlen'] = self.attributes['zmax'] - self.attributes['zmin']

        if not 'bcells' in self.attributes:
            self.attributes['bcells'] = self.attributes['ncells'] + 2 * self.attributes['nghosts']

        if not 'resistivity' in self.attributes:
            self.attributes['resistivity'] = 0
        if not 'viscosity' in self.attributes:
            self.attributes['viscosity'] = 0

        self.cell_size = dict()
        for l in range(self.attributes['maxlev']):
            self.cell_size[l+1] = self.attributes['xlen'] / (self.attributes['xblocks'] * self.attributes['ncells'] * 2**l)


    def __complete_chunks__(self):
        """
        Fill in derived chunk attributes.

        Ensures each chunk has a `dims` field describing the shaped dataset
        dimensions:
        dims = [dblocks] + [bcells] * ndims
        """
        for n in range(len(self.chunks)):
            if not 'dims' in self.chunks[n]:
                self.chunks[n]['dims'] = [ self.attributes['bcells'] ]*self.attributes['ndims']
                self.chunks[n]['dims'][:0] = [ self.chunks[n]['dblocks'] ]


    def __complete_variables__(self):
        """
        Extend the available dataset catalog with derived quantities.

        Based on the primitive fields present in `self.variables` and the
        equation system / EOS in `self.attributes`, this method adds derived
        datasets such as:
        - vector fields and magnitudes (velocity, magnetic field, vorticity)
        - divergence and curl-like diagnostics
        - energies and thermodynamic quantities (eint, ekin, emag, temp, etot)
        - relativistic Lorentz factor (lore)

        Primitive dataset keys are removed from `self.variables` after the
        derived catalog is assembled, leaving user-facing names.
        """
        denflag = 'dens' in self.variables
        preflag = 'pres' in self.variables
        velflag = all(v in self.variables for v in ['velx','vely','velz'])
        magflag = all(v in self.variables for v in ['magx','magy','magz'])

        # Preserve raw snapshot variables (4-char IDs) exactly as loaded by AmunXML/AmunH5
        self.raw_variables = dict(self.variables) if self.variables is not None else {}

        # Build alias mapping from user-friendly names to raw dataset IDs
        self.variables = {}
        self.variables['refinement level']                  = 'mlev'
        self.variables['total energy']                      = 'etot'
        if denflag:
            self.variables['density']                       = 'dens'
        if preflag or self.attributes['eos'] == 'iso':
            self.variables['pressure']                      = 'pres'
        if velflag:
            self.variables['velocity']                      = 'vvec'
            self.variables['velocity magnitude']            = 'velo'
            self.variables['x-velocity']                    = 'velx'
            self.variables['y-velocity']                    = 'vely'
            self.variables['z-velocity']                    = 'velz'
            self.variables['divergence of velocity']        = 'divv'
            self.variables['x-vorticity']                   = 'vorx'
            self.variables['y-vorticity']                   = 'vory'
            self.variables['z-vorticity']                   = 'vorz'
            self.variables['vorticity']                     = 'wvec'
            self.variables['vorticity magnitude']           = 'vort'
        if magflag:
            self.variables['magnetic field']                = 'bvec'
            self.variables['magnetic field magnitude']      = 'magn'
            self.variables['x-magnetic field']              = 'magx'
            self.variables['y-magnetic field']              = 'magy'
            self.variables['z-magnetic field']              = 'magz'
            self.variables['divergence of magnetic field']  = 'divb'
            self.variables['x-current density']             = 'curx'
            self.variables['y-current density']             = 'cury'
            self.variables['z-current density']             = 'curz'
            self.variables['current density']               = 'jvec'
            self.variables['current density magnitude']     = 'curr'
        if preflag and 'adiabatic_index' in self.attributes:
            self.variables['internal energy']               = 'eint'
        if denflag and preflag:
            self.variables['temperature']                   = 'temp'
        if self.attributes['eqsys'] in ['hd','mhd'] and denflag and velflag:
            self.variables['kinetic energy']                = 'ekin'
        if self.attributes['eqsys'] in ['mhd','srmhd'] and magflag:
            self.variables['magnetic energy']               = 'emag'
            self.variables['magnetic pressure']             = 'emag'
        if self.attributes['eqsys'] in ['srhd','srmhd'] and velflag:
            self.variables['Lorentz factor']                = 'lore'
        if 'bpsi' in self.variables:
            self.variables['magnetic divergence potential'] = 'bpsi'

        # Sort variables by name for consistency
        self.variables = dict(sorted(self.variables.items()))


    def __read_binary_data__(self, dataset_id, chunk_number):
        """
        Format-specific dataset loader hook.

        Subclasses must return a NumPy array for `dataset_id` from the
        specified `chunk_number`, shaped as:
        [dblocks] + [bcells] * ndims

        The returned array must use consistent axis ordering; the base class
        provides __swap__ for compatibility with legacy versions.
        """
        pass


    def __swap__(self, arr):
        """
        Apply version-dependent axis permutation for legacy snapshots.

        For snapshots with version < 1, AMUN used a different axis ordering.
        This helper normalizes arrays to the current axis convention.

        Parameters
        ----------
        arr : np.ndarray
            Array to be permuted.

        Returns
        -------
        np.ndarray
            View (or copy) of the array with standardized axis ordering.
        """
        if self.version < 1:
            return np.transpose(arr, np.roll(np.arange(arr.ndim), 1))
        else:
            return arr

    def __etot_newtonian__(self, chunk_number):
        """
        Total energy density for Newtonian HD/MHD.

        For eqsys='hd' or 'mhd':
        etot = 1/2 rho v^2 + etherm + 1/2 B^2

        EOS handling:
        - eos='adi': etherm = p/(gamma-1)
        - eos='iso': etherm is omitted (0.0), since adiabatic internal energy
        is not defined for isothermal closures.

        Returns
        -------
        np.ndarray
            Energy density with the same chunk/block layout as primitives.
        """
        eqsys = str(self.attributes.get("eqsys", "mhd")).lower()
        eos   = str(self.attributes.get("eos", "adi")).lower()

        if eqsys not in ("hd", "mhd"):
            raise ValueError(f"__etot_newtonian__ called for eqsys={eqsys!r}")

        # Primitives
        rho = self.__read_binary_data__("dens", chunk_number)

        vx = self.__read_binary_data__("velx", chunk_number)
        vy = self.__read_binary_data__("vely", chunk_number)
        vz = self.__read_binary_data__("velz", chunk_number)

        # Kinetic
        v2 = vx*vx + vy*vy + vz*vz
        ekin = 0.5 * rho * v2

        # Thermal/internal (only for adiabatic EOS)
        if eos == "adi":
            p     = self.__read_binary_data__("pres", chunk_number)
            gamma = float(self.attributes.get("adiabatic_index", 0.0))
            if gamma <= 1.0:
                raise ValueError(
                    "Invalid adiabatic_index for eos='adi': "
                    f"gamma={gamma} (must be > 1)."
                )
            etherm = p / (gamma - 1.0)
        elif eos == "iso":
            etherm = 0.0
        else:
            raise ValueError(f"Unknown eos={eos!r} (expected 'adi' or 'iso').")

        # Magnetic (only for MHD when magnetic field exists)
        emag = 0.0
        has_B = (eqsys == "mhd") and ("magn" in self.variables)
        if has_B:
            bx = self.__read_binary_data__("magx", chunk_number)
            by = self.__read_binary_data__("magy", chunk_number)
            bz = self.__read_binary_data__("magz", chunk_number)
            emag = 0.5 * (bx*bx + by*by + bz*bz)

        return ekin + etherm + emag


    def __etot_special_relativity__(self, chunk_number):
        """
        Total energy density for special relativistic HD/SRMHD (lab frame).

        For eqsys='srhd' or 'srmhd', returns the lab-frame total energy density
        including rest-mass:
        E = rho h W^2 - p + (magnetic terms, if srmhd)

        where:
        W = 1/sqrt(1-v^2)
        h = 1 + gamma/(gamma-1) * p/rho   (ideal-gas EOS, eos='adi')

        EOS handling:
        - eos='adi': uses ideal-gas enthalpy above
        - eos='iso': not inferred here (requires explicit barotropic closure);
        raises ValueError to avoid silent misuse.

        Magnetic contribution (ideal SRMHD, common lab-frame form):
        + 1/2 (B^2 + (v·B)^2)

        Returns
        -------
        np.ndarray
            Lab-frame energy density including rest mass.
        """
        eqsys = str(self.attributes.get("eqsys", "srmhd")).lower()
        eos   = str(self.attributes.get("eos", "adi")).lower()

        if eqsys not in ("srhd", "srmhd"):
            raise ValueError(f"__etot_special_relativity__ called for eqsys={eqsys!r}")

        # Primitives
        rho = self.__read_binary_data__("dens", chunk_number)

        vx = self.__read_binary_data__("velx", chunk_number)
        vy = self.__read_binary_data__("vely", chunk_number)
        vz = self.__read_binary_data__("velz", chunk_number)

        v2 = vx*vx + vy*vy + vz*vz

        # Lorentz factor: prefer stored 'lore' if available
        if "lore" in self.variables:
            W = self.__read_binary_data__("lore", chunk_number)
        else:
            # Guard against tiny overshoots due to numerical noise
            v2c = np.minimum(v2, 1.0 - 1e-14)
            W = 1.0 / np.sqrt(1.0 - v2c)

        # Enthalpy closure
        if eos == "adi":
            p     = self.__read_binary_data__("pres", chunk_number)
            gamma = float(self.attributes.get("adiabatic_index", 0.0))
            if gamma <= 1.0:
                raise ValueError(
                    "Invalid adiabatic_index for eos='adi': "
                    f"gamma={gamma} (must be > 1)."
                )
            h = 1.0 + (gamma / (gamma - 1.0)) * (p / rho)
        elif eos == "iso":
            raise ValueError(
                "SR etot for eos='iso' requires an explicit barotropic closure "
                "(cannot infer enthalpy from p, rho, gamma)."
            )
        else:
            raise ValueError(f"Unknown eos={eos!r} (expected 'adi' or 'iso').")

        # Hydro contribution (including rest mass)
        E = rho * h * (W*W) - p

        # Magnetic contribution for SRMHD (only if B exists)
        has_B = (eqsys == "srmhd") and ("magn" in self.variables)
        if has_B:
            bx = self.__read_binary_data__("magx", chunk_number)
            by = self.__read_binary_data__("magy", chunk_number)
            bz = self.__read_binary_data__("magz", chunk_number)

            B2 = bx*bx + by*by + bz*bz
            vB = vx*bx + vy*by + vz*bz

            E = E + 0.5 * (B2 + vB*vB)

        return E


    def __read_vec3__(self, base, chunk_number):
        """
        Read a 3-component vector field (x,y,z) from a chunk.
        """
        return (
            self.__read_binary_data__(f"{base}x", chunk_number),
            self.__read_binary_data__(f"{base}y", chunk_number),
            self.__read_binary_data__(f"{base}z", chunk_number),
        )


    def __chunk_cell_size_broadcast__(self, chunk_number, ndim):
        """
        Return per-block cell size broadcastable to a chunk field.

        Parameters
        ----------
        chunk_number : int
            Chunk index.
        ndim : int
            Number of dimensions of the field array (e.g. 3 for [blk,ny,nx],
            4 for [blk,nz,ny,nx]).

        Returns
        -------
        np.ndarray
            Array shaped [dblocks, 1, 1] (2D) or [dblocks, 1, 1, 1] (3D)
            suitable for broadcasting divisions.
        """
        levels = self.chunks[chunk_number]["levels"]
        h = np.array([self.cell_size[l] for l in levels], dtype=float)
        return h.reshape((-1,) + (1,) * (ndim - 1))


    def __cdiff2_index__(self, arr, axis, nghosts=None):
        """
        Second-order central difference in index space (no 1/dx scaling).

        Returns (f[i+1]-f[i-1]) / 2 computed on the stencil-valid interior.
        Boundaries/ghost zones are set to 0.

        Array layout:
        3D: [dblocks, nz, ny, nx]
        2D: [dblocks, ny, nx]
        """
        if nghosts is None:
            nghosts = int(self.attributes.get("nghosts", 0))

        out = np.zeros_like(arr)

        axis = axis if axis >= 0 else arr.ndim + axis
        n = arr.shape[axis]
        i0 = nghosts
        i1 = n - nghosts

        if i1 - i0 < 3:
            return out

        slc_c = [slice(None)] * arr.ndim
        slc_p = [slice(None)] * arr.ndim
        slc_m = [slice(None)] * arr.ndim

        slc_c[axis] = slice(i0, i1)
        slc_p[axis] = slice(i0 + 1, i1 + 1)
        slc_m[axis] = slice(i0 - 1, i1 - 1)

        out[tuple(slc_c)] = 0.5 * (arr[tuple(slc_p)] - arr[tuple(slc_m)])
        return out


    def __div_chunk__(self, chunk_number, ux, uy, uz=None, nghosts=None):
        """
        Divergence of a vector field for an entire chunk (all blocks at once).

        Array layout conventions:
        - 2D: components [dblocks, ny, nx]
        - 3D: components [dblocks, nz, ny, nx]

        The derivative stencil is second-order central difference in index
        space; scaling by physical cell size is applied per block using
        self.cell_size[level].

        Parameters
        ----------
        chunk_number : int
            Chunk index (used to access block levels and cell sizes).
        ux, uy, uz : np.ndarray
            Vector components. For 2D, pass uz=None.
        nghosts : int, optional
            Ghost zone count used by the stencil. Defaults to attributes['nghosts'].

        Returns
        -------
        np.ndarray
            Divergence field, same shape as ux.
        """
        if nghosts is None:
            nghosts = int(self.attributes.get("nghosts", 0))

        # index-space divergence (no 1/h scaling yet)
        div = self.__cdiff2_index__(ux, axis=-1, nghosts=nghosts) \
            + self.__cdiff2_index__(uy, axis=-2, nghosts=nghosts)

        if uz is not None:
            div = div + self.__cdiff2_index__(uz, axis=-3, nghosts=nghosts)

        # physical scaling (assumes isotropic spacing per block)
        h = self.__chunk_cell_size_broadcast__(chunk_number, ux.ndim)
        return div / h


    def __curl_chunk__(self, chunk_number, vx, vy, vz, nghosts=None):
        """
        Curl of a 3-component vector field for an entire chunk (all blocks at once).

        Array layout:
        - 2D snapshots: [dblocks, ny, nx]   (no z-derivatives)
        - 3D snapshots: [dblocks, nz, ny, nx]

        Returns
        -------
        tuple[np.ndarray, np.ndarray, np.ndarray]
            (curl_x, curl_y, curl_z) in physical units, each matching input shape.
        """
        if nghosts is None:
            nghosts = int(self.attributes.get("nghosts", 0))

        # per-block isotropic spacing broadcast
        h = self.__chunk_cell_size_broadcast__(chunk_number, vx.ndim)

        if vx.ndim == 3:
            # 2D: [blk, ny, nx], with ∂/∂z = 0
            # axes: d/dx -> -1, d/dy -> -2
            dvz_dy = self.__cdiff2_index__(vz, axis=-2, nghosts=nghosts)
            dvz_dx = self.__cdiff2_index__(vz, axis=-1, nghosts=nghosts)
            dvy_dx = self.__cdiff2_index__(vy, axis=-1, nghosts=nghosts)
            dvx_dy = self.__cdiff2_index__(vx, axis=-2, nghosts=nghosts)

            curl_x = dvz_dy
            curl_y = -dvz_dx
            curl_z = dvy_dx - dvx_dy

            return curl_x / h, curl_y / h, curl_z / h

        if vx.ndim == 4:
            # 3D: [blk, nz, ny, nx]
            # axes: d/dx -> -1, d/dy -> -2, d/dz -> -3
            dvz_dy = self.__cdiff2_index__(vz, axis=-2, nghosts=nghosts)
            dvy_dz = self.__cdiff2_index__(vy, axis=-3, nghosts=nghosts)

            dvx_dz = self.__cdiff2_index__(vx, axis=-3, nghosts=nghosts)
            dvz_dx = self.__cdiff2_index__(vz, axis=-1, nghosts=nghosts)

            dvy_dx = self.__cdiff2_index__(vy, axis=-1, nghosts=nghosts)
            dvx_dy = self.__cdiff2_index__(vx, axis=-2, nghosts=nghosts)

            curl_x = dvz_dy - dvy_dz
            curl_y = dvx_dz - dvz_dx
            curl_z = dvy_dx - dvx_dy

            return curl_x / h, curl_y / h, curl_z / h

        raise ValueError(
            f"__curl_chunk__: unsupported array rank {vx.ndim}; expected 3 (2D) or 4 (3D)."
        )


    def __grad_chunk__(self, chunk_number, s, nghosts=None):
        """
        Gradient of a scalar field for an entire chunk (all blocks at once).

        Array layout conventions:
        - 2D snapshots: [dblocks, ny, nx]
        - 3D snapshots: [dblocks, nz, ny, nx]

        The derivative stencil is second-order central difference in index
        space; scaling by physical cell size is applied per block using
        self.cell_size[level]. Ghost zones are not differentiated (as per
        __cdiff2_index__ behavior).

        Parameters
        ----------
        chunk_number : int
            Chunk index (used to access block levels and cell sizes).
        s : np.ndarray
            Scalar field.
        nghosts : int, optional
            Ghost zone count used by the stencil. Defaults to attributes['nghosts'].

        Returns
        -------
        tuple[np.ndarray, np.ndarray, np.ndarray]
            (ds_dx, ds_dy, ds_dz) in physical units.
            For 2D inputs, ds_dz is returned as a zeros_like(s) array.
        """
        if nghosts is None:
            nghosts = int(self.attributes.get("nghosts", 0))

        # per-block isotropic spacing broadcast
        h = self.__chunk_cell_size_broadcast__(chunk_number, s.ndim)

        if s.ndim == 3:
            # 2D: [blk, ny, nx] with ∂/∂z = 0
            ds_dx = self.__cdiff2_index__(s, axis=-1, nghosts=nghosts) / h
            ds_dy = self.__cdiff2_index__(s, axis=-2, nghosts=nghosts) / h
            ds_dz = np.zeros_like(s)
            return ds_dx, ds_dy, ds_dz

        if s.ndim == 4:
            # 3D: [blk, nz, ny, nx]
            ds_dx = self.__cdiff2_index__(s, axis=-1, nghosts=nghosts) / h
            ds_dy = self.__cdiff2_index__(s, axis=-2, nghosts=nghosts) / h
            ds_dz = self.__cdiff2_index__(s, axis=-3, nghosts=nghosts) / h
            return ds_dx, ds_dy, ds_dz

        raise ValueError(
            f"__grad_chunk__: unsupported array rank {s.ndim}; expected 3 (2D) or 4 (3D)."
        )


    def __build_dataset_resolvers__(self):
        """
        Build dispatch table mapping internal dataset IDs -> resolver functions.

        A resolver is a callable: resolver(chunk_number) -> ndarray or list[ndarray]
        """
        # NOTE: store bound callables (they capture `self`)
        R = {}

        # --- derived datasets (do NOT exist as raw 4-char IDs)
        R["mlev"] = lambda n: self.__dataset_mlev__(n)

        # velocity family
        R["vvec"] = lambda n: self.__read_vec3__("vel", n)
        R["velo"] = lambda n: self.__dataset_velo__(n)
        R["ekin"] = lambda n: self.__dataset_ekin__(n)
        R["lore"] = lambda n: self.__dataset_lore__(n)
        R["divv"] = lambda n: self.__dataset_divv__(n)
        R["wvec"] = lambda n: self.__dataset_vorticity_vec__(n)
        R["vorx"] = lambda n: self.__dataset_vorticity_comp__(n, "x")
        R["vory"] = lambda n: self.__dataset_vorticity_comp__(n, "y")
        R["vorz"] = lambda n: self.__dataset_vorticity_comp__(n, "z")
        R["vort"] = lambda n: self.__dataset_vorticity_mag__(n)

        # magnetic family
        R["bvec"] = lambda n: self.__read_vec3__("mag", n)
        R["magn"] = lambda n: self.__dataset_magn__(n)
        R["emag"] = lambda n: self.__dataset_emag__(n)
        R["divb"] = lambda n: self.__dataset_divb__(n)
        R["jvec"] = lambda n: self.__dataset_current_vec__(n)
        R["curx"] = lambda n: self.__dataset_current_comp__(n, "x")
        R["cury"] = lambda n: self.__dataset_current_comp__(n, "y")
        R["curz"] = lambda n: self.__dataset_current_comp__(n, "z")
        R["curr"] = lambda n: self.__dataset_current_mag__(n)

        # thermo/energetics
        R["eint"] = lambda n: self.__dataset_eint__(n)
        R["pres"] = lambda n: self.__dataset_pres__(n)
        R["temp"] = lambda n: self.__dataset_temp__(n)
        R["etot"] = lambda n: self.__dataset_etot__(n)

        self._dataset_resolvers = R


    def __dataset_mlev__(self, chunk_number):
        dset = np.zeros(self.chunks[chunk_number]["dims"])
        for p in range(self.chunks[chunk_number]["dblocks"]):
            dset[p, ...] = self.chunks[chunk_number]["levels"][p]
        return dset


    def __dataset_velo__(self, chunk_number):
        vx, vy, vz = self.__read_vec3__("vel", chunk_number)
        return np.sqrt(vx*vx + vy*vy + vz*vz)


    def __dataset_ekin__(self, chunk_number):
        dn = self.__read_binary_data__("dens", chunk_number)
        vx, vy, vz = self.__read_vec3__("vel", chunk_number)
        return 0.5 * dn * (vx*vx + vy*vy + vz*vz)


    def __dataset_lore__(self, chunk_number):
        vx, vy, vz = self.__read_vec3__("vel", chunk_number)
        vv = np.maximum(1e-14, vx*vx + vy*vy + vz*vz)
        return 1.0 / np.sqrt(1.0 - vv)


    def __dataset_divv__(self, chunk_number):
        vx = self.__read_binary_data__("velx", chunk_number)
        vy = self.__read_binary_data__("vely", chunk_number)
        vz = self.__read_binary_data__("velz", chunk_number) if int(self.attributes["ndims"]) == 3 else None
        return self.__div_chunk__(chunk_number, vx, vy, vz)


    def __dataset_vorticity_vec__(self, chunk_number):
        vx, vy, vz = self.__read_vec3__("vel", chunk_number)
        wx, wy, wz = self.__curl_chunk__(chunk_number, vx, vy, vz)
        return [wx, wy, wz]


    def __dataset_vorticity_comp__(self, chunk_number, comp):
        wx, wy, wz = self.__dataset_vorticity_vec__(chunk_number)
        return {"x": wx, "y": wy, "z": wz}[comp]


    def __dataset_vorticity_mag__(self, chunk_number):
        wx, wy, wz = self.__dataset_vorticity_vec__(chunk_number)
        return np.sqrt(wx*wx + wy*wy + wz*wz)


    def __dataset_magn__(self, chunk_number):
        bx, by, bz = self.__read_vec3__("mag", chunk_number)
        return np.sqrt(bx*bx + by*by + bz*bz)


    def __dataset_emag__(self, chunk_number):
        bx, by, bz = self.__read_vec3__("mag", chunk_number)
        return 0.5 * (bx*bx + by*by + bz*bz)


    def __dataset_divb__(self, chunk_number):
        bx = self.__read_binary_data__("magx", chunk_number)
        by = self.__read_binary_data__("magy", chunk_number)
        bz = self.__read_binary_data__("magz", chunk_number) if int(self.attributes["ndims"]) == 3 else None
        return self.__div_chunk__(chunk_number, bx, by, bz)


    def __dataset_current_vec__(self, chunk_number):
        bx, by, bz = self.__read_vec3__("mag", chunk_number)
        jx, jy, jz = self.__curl_chunk__(chunk_number, bx, by, bz)
        return [jx, jy, jz]


    def __dataset_current_comp__(self, chunk_number, comp):
        jx, jy, jz = self.__dataset_current_vec__(chunk_number)
        return {"x": jx, "y": jy, "z": jz}[comp]


    def __dataset_current_mag__(self, chunk_number):
        jx, jy, jz = self.__dataset_current_vec__(chunk_number)
        return np.sqrt(jx*jx + jy*jy + jz*jz)


    def __dataset_eint__(self, chunk_number):
        gamma = float(self.attributes["adiabatic_index"])
        pr = self.__read_binary_data__("pres", chunk_number).copy()
        return pr / (gamma - 1.0)


    def __dataset_pres__(self, chunk_number):
        if self.attributes["eos"] == "iso":
            dn = self.__read_binary_data__("dens", chunk_number).copy()
            return dn * (float(self.attributes["sound_speed"]) ** 2)
        return self.__read_binary_data__("pres", chunk_number).copy()


    def __dataset_temp__(self, chunk_number):
        dn = self.__read_binary_data__("dens", chunk_number)
        pr = self.__read_binary_data__("pres", chunk_number)
        return pr / np.maximum(1e-14, dn)


    def __dataset_etot__(self, chunk_number):
        eqsys = str(self.attributes.get("eqsys", "mhd")).lower()
        if eqsys in ("hd", "mhd"):
            return self.__etot_newtonian__(chunk_number)
        if eqsys in ("srhd", "srmhd"):
            return self.__etot_sr__(chunk_number)
        raise ValueError(f"Unknown eqsys={eqsys!r}")


    def __get_dataset__(self, dataset_id, chunk_number):
        """
        Return a primitive or derived dataset for a given chunk.

        Parameters
        ----------
        dataset_id : str
            Internal dataset 4-character ID.
        chunk_number : int
            Chunk index.

        Returns
        -------
        np.ndarray or list[np.ndarray]
            Chunk/block-shaped dataset.
        """
        # Build dispatch table once (lazy init)
        if not hasattr(self, "_dataset_resolvers"):
            self.__build_dataset_resolvers__()

        # 1) derived/registered dataset
        resolver = self._dataset_resolvers.get(dataset_id)
        if resolver is not None:
            return resolver(chunk_number)

        # 2) raw dataset fallback (read directly)
        #    (works for dens/velx/... that exist in snapshots)
        if hasattr(self, "raw_variables") and dataset_id in self.raw_variables:
            return self.__read_binary_data__(dataset_id, chunk_number)

        # 3) not resolvable
        raise AmunPyVariableError(
            "Dataset id {!r} is not resolvable (not derived and not present in raw_variables).".format(dataset_id)
        )


    def __resolve_dataset_id__(self, dataset_name: str) -> str:
        if dataset_name in self.variables:
            return self.variables[dataset_name]
        if dataset_name in self.raw_variables:
            return dataset_name
        raise AmunPyVariableError(
            "Dataset '{}' is not available!\n"
            "Available datasets (aliases): {}\n"
            "Available raw variables: {}\n".format(
                dataset_name,
                sorted(self.variables.keys()),
                sorted(self.raw_variables.keys()),
            )
        )


    def __resolve_target_level__(self, target_level, maxlev, shrink):
        """
        Resolve target refinement level to use for uniform-grid output.

        target_level is the preferred knob (1..maxlev), defaulting to toplev.
        maxlev/shrink are legacy knobs.
        """
        toplev = int(self.attributes["toplev"])
        max_level = int(self.attributes["maxlev"])

        # Preferred
        if target_level is not None:
            if not isinstance(target_level, (int, np.integer)):
                raise AmunPyDatasetError(
                    "Argument 'target_level' must be an integer between 1 and {}.".format(max_level)
                )
            if not (1 <= int(target_level) <= max_level):
                raise AmunPyDatasetError(
                    "Argument 'target_level' must be between 1 and {}.".format(max_level)
                )
            # If legacy args are also provided, enforce deterministic behavior
            if maxlev is not None or shrink is not None:
                raise AmunPyDatasetError(
                    "Use either 'target_level' or legacy ('maxlev'/'shrink'), not both."
                )
            return int(target_level)

        # Legacy compatibility: maxlev means "target level"
        if maxlev is not None:
            if not isinstance(maxlev, (int, np.integer)):
                raise AmunPyDatasetError(
                    "Argument 'maxlev' must be an integer between 1 and {}.\n".format(max_level)
                )
            if not (1 <= int(maxlev) <= max_level):
                raise AmunPyDatasetError(
                    "Argument 'maxlev' should be between 1 and {}.\n".format(max_level)
                )
            if shrink is not None:
                # You can choose: error or deterministic override.
                # I recommend error (clear semantics).
                raise AmunPyDatasetError("Do not pass both 'maxlev' and 'shrink'. Use 'target_level'.")
            return int(maxlev)

        # Legacy shrink: interpret as downsample factor relative to toplev
        if shrink is not None:
            if not isinstance(shrink, (int, np.integer)) or int(shrink) < 1:
                raise AmunPyDatasetError("Argument 'shrink' must be a positive integer.")
            sh = int(shrink)
            # shrink = 2**(maxlev - target_level) previously; you used maxlev sometimes, toplev elsewhere.
            # Here we define it relative to toplev (more intuitive, and deterministic).
            # Enforce power-of-two shrink to match AMR level steps.
            if sh & (sh - 1) != 0:
                raise AmunPyDatasetError("Argument 'shrink' must be a power of two.")
            # Compute target_level such that dm ~ 2**(toplev-1)/sh
            # sh = 2**(toplev - target_level)
            delta = int(np.log2(sh))
            tl = toplev - delta
            if tl < 1:
                raise AmunPyDatasetError(
                    "Argument 'shrink'={} implies target_level={} < 1.".format(sh, tl)
                )
            return tl

        # Default
        return toplev


    def __xyz_to_arr_tuple__(self, v_xyz: np.ndarray, nd: int) -> tuple:
        """
        Convert an xyz-ordered integer vector into array-order tuple.

        For nd=2: (x,y) -> (y,x)
        For nd=3: (x,y,z) -> (z,y,x)
        """
        if nd == 3:
            return (int(v_xyz[2]), int(v_xyz[1]), int(v_xyz[0]))
        return (int(v_xyz[1]), int(v_xyz[0]))


    def __slice_tuple__(self, beg: tuple, end: tuple) -> tuple:
        """Build a tuple of slices from begin/end index tuples."""
        return tuple(slice(b, e) for b, e in zip(beg, end))


    def __normalize_components__(self, dset):
        """
        Normalize dataset return value to a tuple of components.

        - scalar ndarray -> (ndarray,)
        - vector list/tuple -> tuple(list/tuple)
        """
        return tuple(dset) if isinstance(dset, (list, tuple)) else (dset,)


    def __probe_dataset_components__(self, dataset_id: str):
        """
        Probe the first non-empty chunk to determine dataset component count
        and reuse the probed dataset to avoid a redundant read.

        Returns
        -------
        probe_chunk : int | None
        probe_comps : tuple[np.ndarray, ...] | None
        """
        nchunks = int(self.attributes["nchunks"])
        for n in range(nchunks):
            if self.chunks[n]["dblocks"] > 0:
                dset = self.__get_dataset__(dataset_id, n)
                return n, self.__normalize_components__(dset)
        return None, None


    def dataset(
        self,
        dataset_name,
        extent=None,
        target_level=None,
        maxlev=None,
        shrink=None,
        interpolation="rebin",
        progress=False,
    ):
        """
        Resample a dataset to a uniform mesh over the requested domain region.

        Parameters
        ----------
        dataset_name : str
            User-facing dataset name (key in `self.variables`) or raw 4-char ID.
        extent : sequence[float], optional
            Physical bounds of the returned region:
            [xmin, xmax, ymin, ymax] (2D) or
            [xmin, xmax, ymin, ymax, zmin, zmax] (3D).
            If None, the full snapshot domain is used.
        target_level : int, optional
            Target AMR refinement level of the output uniform grid (1..maxlev).
            If None, defaults to `toplev`.
        maxlev : int, optional
            Legacy: interpreted as `target_level`. Mutually exclusive with
            `target_level` and `shrink`.
        shrink : int, optional
            Legacy: downsampling factor relative to `toplev`. Must be a power of two.
            Mutually exclusive with `target_level` and `maxlev`.
        interpolation : str, optional
            Interpolation method for prolongation (passed to interpolate()).
            'rebin' uses averaging/cell replication.
        progress : bool, optional
            If True, print progress while scanning blocks.

        Returns
        -------
        np.ndarray or list[np.ndarray]
            Uniform-grid scalar field, or three component arrays for vector fields.
        """
        if self.dataformat is None:
            raise AmunPyError("Snapshot object has not been properly initialized!")

        # Resolve dataset name to internal 4-character ID
        dataset_id = self.__resolve_dataset_id__(dataset_name)

        # Resolve output target level (single resolution knob) with legacy compatibility
        target_level = self.__resolve_target_level__(target_level, maxlev, shrink)

        # Cache frequently used scalars
        nd = int(self.attributes["ndims"])
        nc = int(self.attributes["ncells"])
        ng = int(self.attributes["nghosts"])
        toplev = int(self.attributes["toplev"])
        nchunks = int(self.attributes["nchunks"])
        nleafs = int(self.attributes.get("nleafs", 0))

        # ---------- Region setup in xyz order (slice to nd)
        dlo = np.array([self.attributes["xmin"], self.attributes["ymin"], self.attributes["zmin"]], dtype=float)[:nd]
        dup = np.array([self.attributes["xmax"], self.attributes["ymax"], self.attributes["zmax"]], dtype=float)[:nd]
        dln = dup - dlo

        slo = dlo.copy()
        sup = dup.copy()
        if extent is not None:
            if len(extent) != 2 * nd:
                raise AmunPyDatasetError("Wrong dimensions of the argument 'extent'!")
            slo = np.array(extent[0::2], dtype=float)
            sup = np.array(extent[1::2], dtype=float)
            if np.any(slo >= sup):
                raise AmunPyDatasetError("Wrong order of the dimensions in the argument 'extent'!")
            if np.any(slo > dup) or np.any(sup < dlo):
                raise AmunPyDatasetError("Wrong order of the dimensions in the argument 'extent'!")

        # ---------- Full-domain cell counts at requested target level (xyz)
        blocks_xyz = np.array(
            [self.attributes["xblocks"], self.attributes["yblocks"], self.attributes["zblocks"]],
            dtype=int,
        )[:nd]
        domain_cells_xyz = blocks_xyz * nc * (2 ** (target_level - 1))

        # Crop indices for requested extent (xyz)
        ll_xyz = np.floor(domain_cells_xyz * (slo - dlo) / dln).astype(int)
        uu_xyz = np.ceil(domain_cells_xyz * (sup - dlo) / dln).astype(int)
        ll_xyz = np.maximum(ll_xyz, 0)
        uu_xyz = np.minimum(uu_xyz, domain_cells_xyz)
        out_cells_xyz = uu_xyz - ll_xyz

        # Output shape in array order (2D: y,x ; 3D: z,y,x)
        if nd == 3:
            dm = (int(out_cells_xyz[2]), int(out_cells_xyz[1]), int(out_cells_xyz[0]))
        else:
            dm = (int(out_cells_xyz[1]), int(out_cells_xyz[0]))

        # Base cells per block in xyz (used later with nl)
        bm_xyz = np.array([nc] * nd, dtype=int)

        # Deterministic equivalent of legacy shrink from target_level
        shrink_equiv = 2 ** (toplev - target_level)

        # ---------- Probe for scalar/vector nature and reuse first chunk dataset
        probe_chunk, probe_comps = self.__probe_dataset_components__(dataset_id)
        if probe_comps is None:
            return np.zeros(dm)

        ncomp = len(probe_comps)
        arr = [np.zeros(dm, dtype=float) for _ in range(ncomp)]

        if progress:
            sys.stdout.write("Snapshot's path:\n  '{}'\n".format(self.path))
            m = 0

        for n in range(nchunks):
            if self.chunks[n]["dblocks"] <= 0:
                continue

            # Fetch datasets once per chunk; reuse probe for first chunk
            if n == probe_chunk:
                comps = probe_comps
            else:
                comps = self.__normalize_components__(self.__get_dataset__(dataset_id, n))

            if len(comps) != ncomp:
                raise AmunPyDatasetError(
                    "Inconsistent dataset component count for '{}': expected {} components, got {}.".format(
                        dataset_name, ncomp, len(comps)
                    )
                )

            dblocks = int(self.chunks[n]["dblocks"])

            for p in range(dblocks):
                nl = 2 ** (toplev - int(self.chunks[n]["levels"][p]))
                method = "rebin" if nl <= shrink_equiv else interpolation

                # Block cell size at requested output resolution (xyz)
                cm_xyz = bm_xyz * nl // shrink_equiv

                # Block indices: coords are xyz
                il_xyz = self.chunks[n]["coords"][p, :nd] * cm_xyz
                iu_xyz = il_xyz + cm_xyz

                # Overlap test in xyz
                if not (np.all(iu_xyz > ll_xyz) and np.all(il_xyz < uu_xyz)):
                    if progress:
                        m += 1
                    continue

                # Overlap bounds in xyz (global output grid indices)
                nb_xyz = il_xyz - ll_xyz
                ne_xyz = iu_xyz - ll_xyz
                ib_xyz = np.maximum(nb_xyz, 0)
                ie_xyz = np.minimum(ne_xyz, uu_xyz - ll_xyz)

                # Corresponding source bounds inside interpolated block (xyz)
                jb_xyz = ib_xyz - nb_xyz
                je_xyz = ie_xyz - ne_xyz + cm_xyz

                # Convert to array-order tuples and slice tuples
                ib = self.__xyz_to_arr_tuple__(ib_xyz, nd)
                ie = self.__xyz_to_arr_tuple__(ie_xyz, nd)
                jb = self.__xyz_to_arr_tuple__(jb_xyz, nd)
                je = self.__xyz_to_arr_tuple__(je_xyz, nd)
                cm_arr = self.__xyz_to_arr_tuple__(cm_xyz, nd)

                dst = self.__slice_tuple__(ib, ie)
                src = self.__slice_tuple__(jb, je)

                # Insert components (scalar: ncomp=1, vector: ncomp=3)
                cm_arr_np = np.array(cm_arr, dtype=int)
                for di in range(ncomp):
                    blk = interpolate(
                        comps[di][p, ...],
                        cm_arr_np,
                        ng,
                        method=method,
                    )
                    arr[di][dst] = blk[src]

                if progress:
                    cfile = self.chunkname.format(n)
                    m += 1
                    sys.stdout.write("\r")
                    sys.stdout.write(
                        "Reading '{}' from '{}': block {} of {}".format(dataset_name, cfile, m, nleafs)
                    )
                    sys.stdout.flush()

        if progress:
            sys.stdout.write("\n")
            sys.stdout.flush()

        return arr[0] if ncomp == 1 else arr