"""
HDF5 snapshot reader for AMUN simulations.

This module implements the AmunH5 class, a concrete Amun snapshot reader
for simulation outputs stored in HDF5 format. It exposes snapshot
attributes, coordinates, and datasets through the common Amun interface
and is intended to be semantically consistent with the AmunXML reader.
"""
import h5py, os
import numpy as np

from .amun import Amun

class AmunH5(Amun):
    """
    Snapshot reader for AMUN HDF5-format outputs.

    AmunH5 reads snapshot metadata and datasets from an HDF5 container and
    per-chunk HDF5 files. It implements the format-specific hooks required
    by the Amun base class:
    - snapshot identification and version handling
    - parsing of global attributes and available variables
    - chunk discovery and per-block coordinates/bounds
    - retrieval of per-chunk dataset arrays
    """
    def __init_snapshot__(self):
        """
        Validate that the provided path refers to an AMUN HDF5 snapshot and
        initialize format metadata.

        This method:
        - requires `self.path` to be a file (not a directory)
        - opens the HDF5 snapshot and checks for AMUN-specific markers:
        * a top-level 'code'/'codes' attribute identifying AMUN, or
        * required groups: 'attributes', 'coordinates', 'variables'
        - sets `self.dataformat = 'AmunH5'` upon successful validation
        - reads and stores the snapshot version (if present) for compatibility
        handling (e.g. axis ordering via __swap__).
        """
        if not self.path_is_file:
            raise Exception("AmunH5 requires a file not directory as the argument!")

        def _as_str(x):
            # robust conversion for bytes / numpy scalars / arrays
            if isinstance(x, (bytes, np.bytes_)):
                return x.decode("utf-8", errors="replace")
            if isinstance(x, np.ndarray):
                if x.size == 1:
                    return _as_str(x.reshape(()).item())
                return str(x.astype(str))
            return str(x)

        with h5py.File(self.path, 'r') as h5:
            code = None
            if 'code' in h5.attrs:
                code = _as_str(h5.attrs['code'])
            elif 'codes' in h5.attrs:
                code = _as_str(h5.attrs['codes'])

            if code is not None:
                if code.strip() == "AMUN":
                    self.dataformat = 'AmunH5'
                else:
                    raise Exception(f"'{self.path}' has code={code!r}, expected 'AMUN'.")
            elif 'attributes' in h5 and 'coordinates' in h5 and 'variables' in h5:
                self.dataformat = 'AmunH5'
            else:
                raise Exception("{} misses one of these groups: 'attributes', 'coordinates' or 'variables'!".format(self.path))
            if 'version' in h5.attrs:
                try:
                    self.version = int(np.array(h5.attrs['version']).reshape(()))
                except Exception:
                    # keep default; version-dependent swap only matters for legacy files anyway
                    pass


    def __fill_attributes__(self):
        """
        Load global snapshot attributes into `self.attributes`.

        Reads attributes from the HDF5 group 'attributes' and converts values
        into Python scalars/lists when appropriate:
        - numeric scalars are stored as numbers
        - numeric arrays are squeezed and converted to Python lists
        - non-numeric attributes are converted to strings

        Some snapshot-internal bookkeeping attributes are excluded.

        Normalizes legacy naming to the canonical Amun interface:
        - maps 'nprocs' -> 'nchunks' if 'nchunks' is not present
        - maps 'rdims'  -> 'xblocks', 'yblocks', ('zblocks' for 3D)
        """
        exclude_list = ['nseeds', 'seeds', 'dblocks', 'nproc', 'dims', 'dtn', 'last_id', 'mblocks']

        with h5py.File(self.path, 'r') as h5:
            for aname in h5['attributes'].attrs:
                if not aname in exclude_list:
                    attr = h5['attributes'].attrs[aname]
                    if attr.dtype in [ 'float64', 'float32', 'int64', 'int32' ]:
                        if isinstance(attr, np.ndarray):
                            self.attributes[aname] = np.squeeze(attr).tolist()
                        else:
                            self.attributes[aname] = attr
                    else:
                        self.attributes[aname] = np.squeeze(attr).astype(str)

        if not 'nchunks' in self.attributes and 'nprocs' in self.attributes:
            self.attributes['nchunks'] = self.attributes['nprocs']
            del self.attributes['nprocs']
        if 'rdims' in self.attributes:
            self.attributes['xblocks'] = self.attributes['rdims'][0]
            self.attributes['yblocks'] = self.attributes['rdims'][1]
            if self.attributes['ndims'] == 3:
                self.attributes['zblocks'] = self.attributes['rdims'][2]
            else:
                # keep schema uniform for 2D
                self.attributes['zblocks'] = 1
            del self.attributes['rdims']


    def __fill_variables__(self):
        """
        Populate `self.variables` with dataset names available in this snapshot.

        Dataset names are taken from the HDF5 group 'variables'. Each key is
        exposed through the common Amun dataset interface.
        """
        with h5py.File(self.path, 'r') as h5:
            for variable in h5['variables']:
                v = variable.strip()
                self.variables[v] = v


    def __fill_chunks__(self):
        """
        Discover chunk files and load per-block metadata for each chunk.

        For each chunk index n in [0, nchunks):
        - constructs the per-chunk filename based on snapshot number (isnap)
        - opens the chunk HDF5 file and reads:
        * number of datablocks ('dblocks')
        * per-block refinement levels
        * per-block coordinates and physical bounds

        Coordinates and bounds are passed through __swap__ to handle legacy
        axis ordering differences.
        """
        self.chunkname = 'p{:06d}'.format(int(self.attributes['isnap'])) + '_{:05d}.h5'
        nd = int(self.attributes['ndims'])
        for n in range(int(self.attributes['nchunks'])):
            self.chunks[n] = dict()
            self.chunks[n]['filename'] = self.chunkname.format(n)
            cname = os.path.join(self.dirname, self.chunks[n]['filename'])
            if os.path.exists(cname):
                with h5py.File(cname, 'r') as h5:
                    dblocks = int(np.array(h5['attributes'].attrs['dblocks']).reshape(()))
                    self.chunks[n]['dblocks'] = dblocks
                    if dblocks > 0:
                        self.chunks[n]['levels'] = np.array(h5['coordinates']['levels'])
                        self.chunks[n]['bounds'] = self.__swap__(np.array(h5['coordinates']['bounds']))
                        self.chunks[n]['coords'] = self.__swap__(np.array(h5['coordinates']['coords']))
                    else:
                        # represent empty chunks with well-typed empty arrays (not None)
                        self.chunks[n]['levels'] = np.zeros((0,), dtype=int)
                        self.chunks[n]['coords']  = np.zeros((0, nd), dtype=int)
                        self.chunks[n]['bounds']  = np.zeros((0, 2, nd), dtype=float)
            else:
                raise Exception("Snapshot's chunk '{}' not present!".format(cname))


    def __read_binary_data__(self, dataset_name, chunk_number):
        """
        Return a dataset array from a given chunk.

        Loads `dataset_name` from the chunk file referenced by `chunk_number`,
        squeezes singleton dimensions, and applies version-dependent axis
        swapping for legacy snapshots.

        Parameters
        ----------
        dataset_name : str
            Name of the dataset within the 'variables' group.
        chunk_number : int
            Chunk index.

        Returns
        -------
        np.ndarray
            Dataset array with consistent axis ordering.
        """
        cname = os.path.join(self.dirname, self.chunks[chunk_number]['filename'])
        dblocks = int(self.chunks[chunk_number]["dblocks"])
        ndims   = int(self.attributes["ndims"])
        with h5py.File(cname, "r") as h5:
            arr = np.array(h5["variables"][dataset_name])

        # Keep legacy axis ordering fix, but do not squeeze away axes.
        arr = self.__swap__(arr)

        # If payload is single-block and stored without a block axis, add it.
        if dblocks > 0 and arr.ndim == ndims:
            arr = arr[None, ...]

        return arr