1"""
  2GRIB2 Codec for Zarr / numcodecs
  3=================================
  4
  5Provides :class:`Grib2Codec`, a ``numcodecs.abc.Codec`` implementation that
  6decodes raw GRIB2 section 7 bytes into NumPy arrays.  This codec is registered
  7with the numcodecs registry so that Zarr can transparently decode GRIB2 data
  8chunks when reading through a Kerchunk reference manifest or an Icechunk
  9virtual store.
 10
 11The codec is **decode-only**; encoding (writing GRIB2 data) is handled by
 12:func:`grib2io.open`.
 13"""
 14
 15from typing import List, Optional
 16
 17import numpy as np
 18
 19from . import g2clib
 20
 21
 22# ---------------------------------------------------------------------------
 23# Lazy import guard
 24# ---------------------------------------------------------------------------
 25
 26
 27def _ensure_numcodecs():
 28    """Raise ``ImportError`` if *numcodecs* is not available."""
 29    try:
 30        import numcodecs  # noqa: F401
 31    except ImportError:
 32        raise ImportError("numcodecs is required for the GRIB2 codec. Install with: pip install grib2io[kerchunk]")
 33
 34
 35# ---------------------------------------------------------------------------
 36# Codec implementation
 37# ---------------------------------------------------------------------------
 38
 39try:
 40    from numcodecs.abc import Codec
 41    from numcodecs.registry import register_codec
 42
 43    _HAS_NUMCODECS = True
 44except ImportError:
 45    _HAS_NUMCODECS = False
 46    Codec = object  # type: ignore[assignment,misc]
 47
 48    def register_codec(cls):  # type: ignore[misc]
 49        pass
 50
 51
 52class Grib2Codec(Codec):
 53    """Zarr codec for decoding raw GRIB2 section 7 bytes.
 54
 55    All constructor parameters are JSON-serializable integers or lists
 56    extracted from the GRIB2 message's section metadata at
 57    reference-generation time.  Together they carry everything needed to
 58    call ``g2clib.unpack7()`` — the same low-level routine used by
 59    ``grib2io._data()``.
 60    """
 61
 62    codec_id = "grib2io"
 63
 64    def __init__(
 65        self,
 66        drtn: int,
 67        drt: List[int],
 68        gdtn: int,
 69        gdt: List[int],
 70        gds: List[int],
 71        nx: int,
 72        ny: int,
 73        bitmap_flag: int,
 74        bitmap_offset: Optional[int] = None,
 75        bitmap_length: Optional[int] = None,
 76        scan_mode_flags: Optional[List[int]] = None,
 77        type_of_values: int = 0,
 78        number_of_data_points: int = 0,
 79        number_of_packed_values: int = 0,
 80    ):
 81        if not _HAS_NUMCODECS:
 82            raise ImportError("numcodecs is required for Grib2Codec. Install with: pip install grib2io[kerchunk]")
 83        self.drtn = drtn
 84        self.drt = list(drt)
 85        self.gdtn = gdtn
 86        self.gdt = list(gdt)
 87        self.gds = list(gds)
 88        self.nx = nx
 89        self.ny = ny
 90        self.bitmap_flag = bitmap_flag
 91        self.bitmap_offset = bitmap_offset
 92        self.bitmap_length = bitmap_length
 93        self.scan_mode_flags = list(scan_mode_flags) if scan_mode_flags is not None else None
 94        self.type_of_values = type_of_values
 95        self.number_of_data_points = number_of_data_points
 96        self.number_of_packed_values = number_of_packed_values
 97
 98    # ---- encode (not supported) ------------------------------------------
 99
100    def encode(self, buf):
101        """Not supported — GRIB2 encoding is handled by ``grib2io.open()``."""
102        raise NotImplementedError("Grib2Codec is decode-only; GRIB2 encoding is handled by grib2io.open()")
103
104    # ---- decode ----------------------------------------------------------
105
106    def decode(self, buf, out=None):
107        """Decode raw GRIB2 section 7 bytes into a NumPy array."""
108        if not isinstance(buf, bytes):
109            buf = bytes(buf)
110
111        fld = _decode_grib2_bytes(
112            buf,
113            drtn=self.drtn,
114            drt=list(self.drt) if not isinstance(self.drt, list) else self.drt,
115            gdtn=self.gdtn,
116            gdt=list(self.gdt) if not isinstance(self.gdt, list) else self.gdt,
117            nx=self.nx,
118            ny=self.ny,
119            bitmap_flag=self.bitmap_flag,
120            bitmap_length=self.bitmap_length,
121            scan_mode_flags=list(self.scan_mode_flags) if self.scan_mode_flags is not None else None,
122            type_of_values=self.type_of_values,
123            number_of_data_points=self.number_of_data_points,
124            number_of_packed_values=self.number_of_packed_values,
125        )
126
127        if out is not None:
128            np.copyto(out, fld)
129            return out
130
131        return fld
132
133    # ---- config serialization --------------------------------------------
134
135    def get_config(self):
136        """Return a JSON-serializable dict of all codec parameters."""
137        return {
138            "id": self.codec_id,
139            "drtn": self.drtn,
140            "drt": self.drt,
141            "gdtn": self.gdtn,
142            "gdt": self.gdt,
143            "gds": self.gds,
144            "nx": self.nx,
145            "ny": self.ny,
146            "bitmap_flag": self.bitmap_flag,
147            "bitmap_offset": self.bitmap_offset,
148            "bitmap_length": self.bitmap_length,
149            "scan_mode_flags": self.scan_mode_flags,
150            "type_of_values": self.type_of_values,
151            "number_of_data_points": self.number_of_data_points,
152            "number_of_packed_values": self.number_of_packed_values,
153        }
154
155    @classmethod
156    def from_config(cls, config):
157        """Reconstruct the codec from a configuration dict.
158
159        Parameters
160        ----------
161        config : dict
162            Dictionary as returned by :meth:`get_config`.
163
164        Returns
165        -------
166        Grib2Codec
167        """
168        # Remove the 'id' key which is not a constructor parameter
169        cfg = {k: v for k, v in config.items() if k != "id"}
170        return cls(**cfg)
171
172
173register_codec(Grib2Codec)
174
175
176# ---------------------------------------------------------------------------
177# Zarr v3 ArrayBytesCodec implementation
178# ---------------------------------------------------------------------------
179
180try:
181    from dataclasses import dataclass
182    from zarr.abc.codec import ArrayBytesCodec
183    from zarr.core.array_spec import ArraySpec
184    from zarr.core.buffer import Buffer, NDBuffer
185    import zarr.registry as _zarr_registry
186
187    _HAS_ZARR_V3 = True
188except ImportError:
189    _HAS_ZARR_V3 = False
190
191
192if _HAS_ZARR_V3:
193
194    @dataclass(frozen=True)
195    class Grib2SerializerCodec(ArrayBytesCodec):
196        """Zarr v3 ``ArrayBytesCodec`` (serializer) for GRIB2 virtual chunks.
197
198        Receives the raw GRIB2 section 7 bytes fetched from a virtual chunk
199        reference and decodes them into a NumPy array using
200        ``g2clib.unpack7()``.  This is the zarr v3 counterpart of
201        :class:`Grib2Codec` (which targets numcodecs / zarr v2).
202
203        All constructor parameters mirror those of :class:`Grib2Codec` and
204        are stored in ``zarr.json`` so they travel with the array metadata.
205        """
206
207        is_fixed_size: bool = False  # GRIB2 compressed size varies
208
209        # ---- GRIB2 decode parameters (mirror Grib2Codec) ----
210        drtn: int = 0
211        drt: tuple = ()
212        gdtn: int = 0
213        gdt: tuple = ()
214        gds: tuple = ()
215        nx: int = 0
216        ny: int = 0
217        bitmap_flag: int = 255
218        bitmap_offset: Optional[int] = None
219        bitmap_length: Optional[int] = None
220        scan_mode_flags: Optional[tuple] = None
221        type_of_values: int = 0
222        number_of_data_points: int = 0
223        number_of_packed_values: int = 0
224
225        # dataclass frozen=True uses __init__ auto-generated; we need a
226        # custom __init__ to accept list args and coerce to tuples.
227        def __new__(cls, **kwargs):
228            # coerce list args to tuples so the frozen dataclass is hashable
229            for f in ("drt", "gdt", "gds", "scan_mode_flags"):
230                if f in kwargs and isinstance(kwargs[f], list):
231                    kwargs[f] = tuple(kwargs[f])
232            obj = object.__new__(cls)
233            return obj
234
235        def __init__(
236            self,
237            *,
238            drtn: int = 0,
239            drt=(),
240            gdtn: int = 0,
241            gdt=(),
242            gds=(),
243            nx: int = 0,
244            ny: int = 0,
245            bitmap_flag: int = 255,
246            bitmap_offset: Optional[int] = None,
247            bitmap_length: Optional[int] = None,
248            scan_mode_flags=None,
249            type_of_values: int = 0,
250            number_of_data_points: int = 0,
251            number_of_packed_values: int = 0,
252        ):
253            object.__setattr__(self, "drtn", int(drtn))
254            object.__setattr__(self, "drt", tuple(drt) if drt is not None else ())
255            object.__setattr__(self, "gdtn", int(gdtn))
256            object.__setattr__(self, "gdt", tuple(gdt) if gdt is not None else ())
257            object.__setattr__(self, "gds", tuple(gds) if gds is not None else ())
258            object.__setattr__(self, "nx", int(nx))
259            object.__setattr__(self, "ny", int(ny))
260            object.__setattr__(self, "bitmap_flag", int(bitmap_flag))
261            object.__setattr__(self, "bitmap_offset", bitmap_offset)
262            object.__setattr__(self, "bitmap_length", bitmap_length)
263            object.__setattr__(self, "scan_mode_flags", tuple(scan_mode_flags) if scan_mode_flags is not None else None)
264            object.__setattr__(self, "type_of_values", int(type_of_values))
265            object.__setattr__(self, "number_of_data_points", int(number_of_data_points))
266            object.__setattr__(self, "number_of_packed_values", int(number_of_packed_values))
267
268        # ---- zarr v3 codec interface --------------------------------
269
270        @classmethod
271        def from_dict(cls, data: dict) -> "Grib2SerializerCodec":
272            cfg = data.get("configuration", {})
273            return cls(**cfg)
274
275        def to_dict(self) -> dict:
276            return {
277                "name": "grib2io",
278                "configuration": {
279                    "drtn": self.drtn,
280                    "drt": list(self.drt),
281                    "gdtn": self.gdtn,
282                    "gdt": list(self.gdt),
283                    "gds": list(self.gds),
284                    "nx": self.nx,
285                    "ny": self.ny,
286                    "bitmap_flag": self.bitmap_flag,
287                    "bitmap_offset": self.bitmap_offset,
288                    "bitmap_length": self.bitmap_length,
289                    "scan_mode_flags": list(self.scan_mode_flags) if self.scan_mode_flags else None,
290                    "type_of_values": self.type_of_values,
291                    "number_of_data_points": self.number_of_data_points,
292                    "number_of_packed_values": self.number_of_packed_values,
293                },
294            }
295
296        def compute_encoded_size(self, input_byte_length: int, chunk_spec: "ArraySpec") -> int:
297            raise NotImplementedError("Grib2SerializerCodec has variable encoded size")
298
299        async def _decode_single(
300            self,
301            chunk_bytes: "Buffer",
302            chunk_spec: "ArraySpec",
303        ) -> "NDBuffer":
304            raw = bytes(chunk_bytes.as_array_like())
305            fld = _decode_grib2_bytes(
306                raw,
307                drtn=self.drtn,
308                drt=list(self.drt),
309                gdtn=self.gdtn,
310                gdt=list(self.gdt),
311                nx=self.nx,
312                ny=self.ny,
313                bitmap_flag=self.bitmap_flag,
314                bitmap_length=self.bitmap_length,
315                scan_mode_flags=list(self.scan_mode_flags) if self.scan_mode_flags else None,
316                type_of_values=self.type_of_values,
317                number_of_data_points=self.number_of_data_points,
318                number_of_packed_values=self.number_of_packed_values,
319            )
320            # Reshape to the full chunk shape zarr expects (e.g. [1,1,1,1,ny,nx])
321            if fld.shape != chunk_spec.shape:
322                fld = fld.reshape(chunk_spec.shape)
323            return chunk_spec.prototype.nd_buffer.from_ndarray_like(fld)
324
325        async def _encode_single(
326            self,
327            chunk_array: "NDBuffer",
328            chunk_spec: "ArraySpec",
329        ) -> Optional["Buffer"]:
330            raise NotImplementedError("Grib2SerializerCodec is decode-only")
331
332    _zarr_registry.register_codec("grib2io", Grib2SerializerCodec)
333    # VirtualiZarr converts numcodecs {"id": "grib2io"} → zarr v3 {"name": "numcodecs.grib2io"}
334    # so we must also register under the prefixed name.
335    _zarr_registry.register_codec("numcodecs.grib2io", Grib2SerializerCodec)
336
337
338def _decode_grib2_bytes(
339    raw: bytes,
340    *,
341    drtn: int,
342    drt: List[int],
343    gdtn: int,
344    gdt: List[int],
345    nx: int,
346    ny: int,
347    bitmap_flag: int,
348    bitmap_length: Optional[int],
349    scan_mode_flags: Optional[List[int]],
350    type_of_values: int,
351    number_of_data_points: int,
352    number_of_packed_values: int,
353) -> np.ndarray:
354    """Shared GRIB2 decode logic used by both :class:`Grib2Codec` (numcodecs)
355    and :class:`Grib2SerializerCodec` (zarr v3).
356    """
357    gdt_arr = np.array(gdt, dtype=np.int64)
358    drt_arr = np.array(drt, dtype=np.int64)
359
360    storageorder = "C"
361    if scan_mode_flags is not None and len(scan_mode_flags) > 2:
362        storageorder = "F" if scan_mode_flags[2] else "C"
363
364    # --- Dynamic per-chunk parsing -------------------------------------------
365    # If the raw bytes start with GRIB2 Section 5 (byte[4] == 5), parse the
366    # data representation template and bitmap from the chunk bytes rather than
367    # using the static codec-config values.  This allows arrays whose messages
368    # have different DRT values (reference value, scale factors, etc.) to all
369    # decode correctly from a single shared .zarray codec config.
370    if len(raw) > 11 and raw[4] == 5:
371        # Parse section 5
372        sec5_len = int.from_bytes(raw[0:4], "big")
373        _drtn, drt_arr, _npts, _ipos = g2clib.unpack5(bytes(raw[:sec5_len]))
374        drtn = int(_drtn)
375        number_of_packed_values = int(_npts)
376        raw = raw[sec5_len:]
377
378    # Now raw starts with section 6 (bitmap indicator section).
379    # Always present (6 bytes minimum with indicator=255 when no bitmap).
380    if len(raw) > 5 and raw[4] == 6:
381        sec6_len = int.from_bytes(raw[0:4], "big")
382        bitmap_indicator = raw[5]
383        if bitmap_indicator in {0, 254}:
384            # Bitmap is present; unpack6 reads it from the sec6 bytes
385            bmap_bytes = bytes(raw[:sec6_len])
386            raw = raw[sec6_len:]
387            _bmapflag, bmap, _bpos = g2clib.unpack6(bmap_bytes, number_of_data_points)
388        else:
389            bmap = None
390            raw = raw[sec6_len:]
391
392        fld1, _ipos = g2clib.unpack7(
393            raw,
394            gdtn,
395            gdt_arr,
396            drtn,
397            drt_arr,
398            number_of_packed_values,
399            storageorder=storageorder,
400        )
401
402        if bitmap_indicator in {0, 254}:
403            if bmap is not None:
404                fld = np.full(number_of_data_points, np.nan, dtype=np.float32)
405                np.put(fld, np.nonzero(bmap), fld1)
406            else:
407                fld = fld1
408        else:
409            fld = fld1
410
411    else:
412        # Legacy path: raw already starts at sec6 or sec7 (old manifest format)
413        bmap = None
414        sec7_buf = raw
415        if bitmap_flag in {0, 254} and bitmap_length is not None and bitmap_length > 0:
416            bmap_bytes = raw[:bitmap_length]
417            sec7_buf = raw[bitmap_length:]
418            _bmapflag, bmap, _bpos = g2clib.unpack6(bmap_bytes, number_of_data_points)
419
420        fld1, _ipos = g2clib.unpack7(
421            sec7_buf,
422            gdtn,
423            gdt_arr,
424            drtn,
425            drt_arr,
426            number_of_packed_values,
427            storageorder=storageorder,
428        )
429
430        if bitmap_flag in {0, 254}:
431            if bmap is not None:
432                fld = np.full(number_of_data_points, np.nan, dtype=np.float32)
433                np.put(fld, np.nonzero(bmap), fld1)
434            else:
435                fld = fld1
436        else:
437            fld = fld1
438    # -------------------------------------------------------------------------
439
440    fld = np.reshape(fld, (ny, nx))
441
442    if scan_mode_flags is not None and len(scan_mode_flags) > 3:
443        if scan_mode_flags[3]:
444            fldsave = fld.astype(np.float32)
445            fld[1::2, :] = fldsave[1::2, ::-1]
446
447    if type_of_values == 1:
448        fld = fld.astype(np.int32)
449
450    return fld