90class open:
 91    """
 92    Open class for tdlpackio.
 93
 94    Parameters
 95    ----------
 96    path : str
 97        File name.
 98
 99    mode : {'r', 'w'}, default 'r'
100        File handle mode.
101
102    format : {'sequential', 'random-access'}, optional
103        File type when creating a new file.
104
105    ra_template : {'small', 'large'}, optional
106        Template used when creating random-access files.
107    """
108
109    _filetype_map = {
110        "random-access": 1,
111        "sequential": 2,
112    }
113
114    def __init__(self, path, mode="r", format=None, ra_template=None):
115        if mode not in {"r", "w", "a"}:
116            raise ValueError(f"Invalid mode: {mode}")
117        if mode in {"r", "w"}:
118            mode = mode + "b"
119        elif mode == "a":
120            mode = "wb"
121        self.path = path
122        self.mode = mode
123        self.format = format
124        self.ra_template = ra_template
125        self._hasindex = False
126        self._index = {}
127        self._lun = -1
128        self.mode = mode
129        self.name = os.path.abspath(path)
130        self.records = 0
131
132        # Perform indexing on read
133        if "r" in self.mode:
134            self._filehandle = builtins.open(path, mode=mode)
135            self.filetype = self._get_tdlpack_file_type()
136            self._build_index()
137
138        elif "w" in self.mode:
139            self.bytes_written = 0
140            self.records_written = 0
141            self.filetype = format if format is not None else "sequential"
142            if self.filetype == "random-access":
143                ra_template = "small" if ra_template is None else ra_template
144                iret, self._lun = tdlpacklib.open_tdlpack_file(self.name, self.mode, self._ifiletype, ra_template=ra_template)
145            elif self.filetype == "sequential":
146                self._filehandle = builtins.open(path, mode=mode)
147                iret, self._lun = tdlpacklib.open_tdlpack_file(self.name, self.mode, self._ifiletype, ra_template=ra_template)
148            self._station_id_record_hashes = []
149
150        # Add self to file data store
151        _open_file_store[self.name] = self
152
153    @property
154    def _ifiletype(self):
155        """Return numeric filetype"""
156        return self._filetype_map[self.filetype]
157
158    @property
159    def size(self):
160        """Return the file size."""
161        return os.path.getsize(self.name)
162
163    def __enter__(self):
164        """"""
165        return self
166
167    def __exit__(self, atype, value, traceback):
168        """"""
169        self.close()
170
171    def __iter__(self):
172        """"""
173        yield from self._index["record"]
174
175    def __repr__(self):
176        """"""
177        strings = []
178        keys = list(self.__dict__.keys())
179        for k in keys:
180            if not k.startswith("_"):
181                strings.append(f"{k} = {self.__dict__[k]}\n")
182        # Attach size property.
183        strings.append(f"size = {self.size}\n")
184        return "".join(strings)
185
186    def __getitem__(self, key):
187        """"""
188        if isinstance(key, slice):
189            return self._index["record"][key]
190        elif isinstance(key, int):
191            return self._index["record"][key]
192        else:
193            raise KeyError("Key must be an integer record number or a slice")
194
195    def _get_tdlpack_file_type(self):
196        """Determine the type of TDLPACK file"""
197        self._filehandle.seek(0)
198        a = struct.unpack(">i", self._filehandle.read(4))[0]
199        b = struct.unpack(">i", self._filehandle.read(4))[0]
200        self._filehandle.seek(0)
201        return "random-access" if [a, b] == [0, 4] else "sequential"
202
203    def _build_index(self):
204        """Record Indexer"""
205        # Initialize index dictionary
206        self._index["offset"] = []
207        self._index["size"] = []
208        self._index["type"] = []
209        self._index["record"] = []
210
211        if self.filetype == "random-access":
212            self._randomaccess_file_indexer()
213        elif self.filetype == "sequential":
214            self._sequential_file_indexer()
215        self._hasindex = True
216
217    def _randomaccess_file_indexer(self):
218        """Indexer for random-access TDLPACK files"""
219        # Read master key
220        version, nids, nwords, nkyrec, maxent, lastky = struct.unpack(">iiiiii", self._filehandle.read(24))
221        nbytes = nwords * NBYPWD
222        last_key_check = [99999999] if lastky > 9999 else [9999, 99999999]
223        self.master_key = dict(
224            version=version,
225            nids=nids,
226            nwords=nwords,
227            nkyrec=nkyrec,
228            maxent=maxent,
229            lastky=lastky,
230        )
231        self.key_records = []
232        # Set file position to first key record
233        self._filehandle.seek(nbytes)
234
235        last_station_id_rec = -1
236        last_station_lat_rec = -1
237        last_station_lon_rec = -1
238
239        # Iterate over all key records
240        while True:
241            # Read key record "header" data
242            nkeys, prec_this_key, prec_next_key = struct.unpack(">iii", self._filehandle.read(12))
243            self.key_records.append(
244                dict(
245                    nkeys=nkeys,
246                    prec_this_key=prec_this_key,
247                    prec_next_key=prec_next_key,
248                )
249            )
250
251            ids = list()
252            nsize = list()
253            prec_begin = list()
254
255            # Read key record information
256            for i in range(nkeys):
257                id1, id2, id3, id4, nd, bprec = struct.unpack(">iiiiii", self._filehandle.read(24))
258                ids.append([id1, id2, id3, id4])
259                nsize.append(nd)
260                prec_begin.append(bprec)
261
262            # Using key record info, move around file to inventory TDLPACK data
263            for m, n, b in zip(ids, nsize, prec_begin):
264                # Disect prec_begin
265                prec1 = int(b / 1000.0)
266
267                # Offset to the data record
268                offset = (prec1 - 1) * nbytes
269                self._filehandle.seek(offset)
270                self._index["offset"].append(offset)
271                self._index["size"].append(n * NBYPWD)
272
273                # Determine record type.  Since the 4-word ID is stored in the
274                # key record, we can use it to determine station call letter
275                # record or TDLPACK data record.
276                if m[0] == 400001000:
277                    # Station ID record...not in TDLPACK format
278                    rec = TdlpackStationRecord()
279                    last_station_id_rec = self.records
280                    rec._recnum = self.records
281                    rec._source = self.name
282                    rec._nsta_expected = int(n / 2)
283                    self._index["record"].append(rec)
284                    self._index["type"].append("station")
285                else:
286                    if m[0] == 400006000:
287                        last_station_lat_rec = self.records
288                    elif m[0] == 400007000:
289                        last_station_lon_rec = self.records
290                    # TDLPACK data record
291                    ipack = np.frombuffer(self._filehandle.read(132), dtype=">i4").astype(np.int32)
292                    iret, is0, is1, is2, is4 = tdlpacklib.unpack_meta(ipack)
293                    if np.all(is2 == 0):
294                        is2 = None
295                    rec = TdlpackRecord(is0, is1, is2, is4)
296                    rec._recnum = self.records
297                    rec._linked_station_id_record = last_station_id_rec
298                    rec._linked_station_lat_record = last_station_lat_rec
299                    rec._linked_station_lon_record = last_station_lon_rec
300                    rec._source = self.name
301                    shape = (rec.ny, rec.nx) if rec.type == "grid" else (rec.numberOfPackedValues,)
302                    ndim = len(shape)
303                    dtype = "float32"
304                    rec._data = TdlpackRecordOnDiskArray(
305                        shape,
306                        ndim,
307                        dtype,
308                        self.filetype,
309                        self._filehandle,
310                        rec,
311                        self._index["offset"][-1],
312                        self._index["size"][-1],
313                    )
314                    self._index["record"].append(rec)
315                    self._index["type"].append("data")
316                self.records += 1
317
318            # Hold the record number of the last station ID record
319            if self._index["type"][-1] == "station":
320                _last_station_id_record = self.records  # This should be OK.
321
322            # Break loop here, at last key record
323            if prec_next_key in last_key_check:
324                break
325            offset = (prec_next_key - 1) * nbytes
326            self._filehandle.seek(offset)
327
328        # Make key_records immutable
329        self.key_records = tuple(self.key_records)
330
331    def _sequential_file_indexer(self):
332        """Indexer for sequential TDLPACK files"""
333        last_station_id_rec = -1
334        last_station_lat_rec = -1
335        last_station_lon_rec = -1
336
337        # Iterate
338        while True:
339            try:
340                # First read 4-byte Fortran record header
341                pos = self._filehandle.tell()
342                fortran_header = struct.unpack(">i", self._filehandle.read(4))[0]
343                if fortran_header >= 132:
344                    bytes_to_read = 132
345                else:
346                    bytes_to_read = fortran_header
347
348                pos = self._filehandle.tell()
349                ioctet = np.frombuffer(self._filehandle.read(8), dtype=">i8").astype(np.int64)[0]
350                ipack = np.frombuffer(self._filehandle.read(bytes_to_read - 8), dtype=">i4").astype(np.int32)
351                _header = struct.unpack(">4s", ipack[0])[0].decode()
352
353                # Check to first 4 bytes of the data record to determine the data
354                # record type.
355                if _header == "PLDT":
356                    if ipack[5] == 400006000:
357                        last_station_lat_rec = self.records
358                    elif ipack[5] == 400007000:
359                        last_station_lon_rec = self.records
360                    # TDLPACK data record
361                    iret, is0, is1, is2, is4 = tdlpacklib.unpack_meta(ipack)
362                    self._index["offset"].append(pos)
363                    self._index["size"].append(fortran_header)  # Size given by Fortran header
364                    if np.all(is2 == 0):
365                        is2 = None
366                    rec = TdlpackRecord(is0, is1, is2, is4)
367                    rec._recnum = self.records
368                    rec._linked_station_id_record = last_station_id_rec
369                    rec._linked_station_lat_record = last_station_lat_rec
370                    rec._linked_station_lon_record = last_station_lon_rec
371                    rec._source = self.name
372                    shape = (rec.ny, rec.nx) if rec.type == "grid" else (rec.numberOfPackedValues,)
373                    ndim = len(shape)
374                    dtype = "float32"
375                    rec._data = TdlpackRecordOnDiskArray(
376                        shape,
377                        ndim,
378                        dtype,
379                        self.filetype,
380                        self._filehandle,
381                        rec,
382                        self._index["offset"][-1],
383                        self._index["size"][-1],
384                    )
385                    self._index["record"].append(rec)
386                    self._index["type"].append("data")
387                else:
388                    if ioctet == 24 and ipack[4] == 9999:
389                        # Trailer record
390                        rec = TdlpackTrailerRecord()
391                        rec._recnum = self.records
392                        rec._source = self.name
393                        self._index["offset"].append(pos)
394                        self._index["size"].append(fortran_header)
395                        self._index["type"].append("trailer")
396                        self._index["record"].append(rec)
397                    else:
398                        # Station ID record
399                        rec = TdlpackStationRecord()
400                        last_station_id_rec = self.records
401                        rec._recnum = self.records
402                        rec._source = self.name
403                        rec._nsta_expected = int(ioctet / 8)
404                        self._index["offset"].append(pos)
405                        self._index["size"].append(fortran_header)
406                        self._index["type"].append("station")
407                        self._index["record"].append(rec)
408
409                # At this point we have successfully identified a TDLPACK record from
410                # the file. Increment self.records and position the file pointer to
411                # now read the Fortran trailer.
412                self.records += 1  # Includes trailer records
413                self._filehandle.seek(fortran_header - bytes_to_read, 1)
414                fortran_trailer = struct.unpack(">i", self._filehandle.read(4))[0]
415
416                # Check Fortran header and trailer for the record.
417                if fortran_header != fortran_trailer:
418                    raise IOError("Bad Fortran record.")
419
420                # Hold the record number of the last station ID record
421                if self._index["type"][-1] == "station":
422                    _last_station_id_record = self.records  # This should be OK.
423
424            except struct.error:
425                self._filehandle.seek(0)
426                break
427
428    def read(self, n):
429        """
430        Read record from file.
431
432        Parameters
433        ----------
434        n : int
435            Record number.
436
437        Returns
438        -------
439        numpy.ndarray
440            Record data as a NumPy array. The returned dtype depends on the
441            record type:
442
443            - ``data`` or ``trailer`` : ``int32`` array.
444            - ``station`` : fixed-width byte string array (``S8``).
445
446        Notes
447        -----
448        The file pointer is positioned using the internal index before reading.
449        Record size is determined from the file header (sequential) or index
450        (random-access).
451        """
452        if "w" in self.mode:
453            pass  # Remove this at some point....
454        # Position file pointer to the beginning of the TDLPACK record.
455        self._filehandle.seek(self._index["offset"][n])
456        if self.filetype == "sequential":
457            size = np.frombuffer(self._filehandle.read(8), dtype=">i8").astype(np.int64)[0]
458        elif self.filetype == "random-access":
459            size = self._index["size"][n]
460
461        if self._index["type"][n] in {"data", "trailer"}:
462            return np.frombuffer(self._filehandle.read(size), dtype=">i4").astype(np.int32)
463        elif self._index["type"][n] == "station":
464            return np.frombuffer(self._filehandle.read(size), dtype="S8")
465
466    def write(self, record):
467        """
468        Write record(s) to file.
469
470        Parameters
471        ----------
472        record : TdlpackStationRecord, _TdlpackRecord, TdlpackTrailerRecord, or list
473            Record or list of records to write.
474
475            - ``TdlpackStationRecord`` : Station record. Station identifiers
476              are padded to ``NCHAR``.
477            - ``_TdlpackRecord`` : Packed data record.
478            - ``TdlpackTrailerRecord`` : Trailer record.
479            - ``list`` : List of supported record types.
480
481        Returns
482        -------
483        None
484
485        Notes
486        -----
487        Updates ``bytes_written``, ``records_written``, ``records``, and
488        ``_type_lastrecord_written``.
489        """
490        if isinstance(record, list):
491            for rec in record:
492                self.write(rec)
493            return
494
495        nreplace, ncheck = 0, 0
496
497        if isinstance(record, TdlpackStationRecord):
498            # Adjust string length of each station to NCHAR.
499            stns = [s.ljust(NCHAR) for s in record.stations]
500            iret, self.bytes_written, self.records_written = tdlpacklib.write_station_record(
501                self.name,
502                self._lun,
503                self._ifiletype,
504                stns,
505                self.bytes_written,
506                self.records_written,
507                nreplace,
508                ncheck,
509            )
510            self._station_id_record_hashes.append(hash(record))
511
512        elif issubclass(record.__class__, _TdlpackRecord):
513            # When writing a vector record to a sequential file, we need to make sure a station call letter
514            # record exists and is valid for the data record.
515            if self.filetype == "sequential" and record.type == "vector":
516                if record._source is not None:
517                    if hash(_open_file_store[record._source][record._linked_station_id_record]) not in self._station_id_record_hashes:
518                        if self.records > 0:
519                            self.write(TdlpackTrailerRecord())
520                        # Write the appropriate station call letter record
521                        self.write(_open_file_store[record._source][record._linked_station_id_record])
522
523            if not hasattr(record, "_ipack"):
524                record.pack()
525
526            iret, self.bytes_written, self.records_written = tdlpacklib.write_tdlpack_record(
527                self.name,
528                self._lun,
529                self._ifiletype,
530                record._ipack,
531                self.bytes_written,
532                self.records_written,
533                nreplace,
534                ncheck,
535            )
536
537        elif isinstance(record, TdlpackTrailerRecord):
538            iret, self.bytes_written, self.records_written = tdlpacklib.write_trailer_record(
539                self._lun, self._ifiletype, self.bytes_written, self.records_written
540            )
541
542        self._type_lastrecord_written = record.type
543        self.records += 1
544
545    def close(self):
546        """
547        Close the file.
548
549        Returns
550        -------
551        None
552
553        Notes
554        -----
555        For write mode, a trailer record may be written for sequential files
556        if the last record type requires it. The underlying TDLpack file handle
557        is then closed. The file is removed from the internal open file store.
558        """
559        if "r" in self.mode:
560            self._filehandle.close()
561        if "w" in self.mode:
562            if self.filetype == "sequential":
563                if self._type_lastrecord_written == "vector":
564                    iret = tdlpacklib.write_trailer_record(
565                        self._lun,
566                        self._ifiletype,
567                        self.bytes_written,
568                        self.records_written,
569                    )
570            iret = tdlpacklib.close_tdlpack_file(self._lun, self._ifiletype)
571            if iret != 0:
572                raise ValueError("return from tdlpacklib.close_tdlpack_file is non-zero")
573        if self.name in _open_file_store.keys():
574            del _open_file_store[self.name]
575
576    def select(self, **kwargs):
577        """
578        Select records by attribute.
579
580        Parameters
581        ----------
582        **kwargs : dict
583            Keyword arguments specifying ``TdlpackRecord`` attributes and
584            values to match.
585
586        Returns
587        -------
588        list
589            List of records matching all provided attribute/value pairs.
590
591        Notes
592        -----
593        Selection is performed against records in the internal index. Records
594        must match all specified attributes to be included in the result.
595        """
596        _id_keys = [
597            "word1",
598            "word2",
599            "word3",
600            "word4",
601            "ccc",
602            "fff",
603            "b",
604            "dd",
605            "v",
606            "llll",
607            "uuuu",
608            "t",
609            "rr",
610            "o",
611            "hh",
612            "tau",
613            "thresh",
614            "i",
615            "s",
616            "g",
617        ]
618
619        # TODO: Add ability to process multiple values for each keyword (attribute)
620        idxs = []
621        nkeys = len(kwargs.keys())
622        for k, v in kwargs.items():
623            if k in _id_keys:
624                for rec in self._index["record"]:
625                    # if hasattr(rec, k) and getattr(rec, k) == v:
626                    if getattr(rec.id, k) == v:
627                        idxs.append(rec._recnum)
628            else:
629                for rec in self._index["record"]:
630                    if hasattr(rec, k) and getattr(rec, k) == v:
631                        idxs.append(rec._recnum)
632        idxs = np.array(idxs, dtype=">i4")
633        return [self._index["record"][i] for i in [ii[0] for ii in collections.Counter(idxs).most_common() if ii[1] == nkeys]]

636class TdlpackRecord:
637    """
638    Creation class for TDLPACK record objects.
639
640    This class dynamically constructs and returns an instance of a
641    ``_TdlpackRecord`` subclass based on the provided section arrays and
642    optional keyword arguments. Record classes are cached by type to avoid
643    repeated class construction.
644
645    Parameters
646    ----------
647    is0 : numpy.ndarray, optional
648        Section 0 array. Default is zero-initialized ``int32`` array of size ``ND7``.
649    is1 : numpy.ndarray, optional
650        Section 1 array. Default is zero-initialized ``int32`` array of size ``ND7``.
651    is2 : numpy.ndarray, optional
652        Section 2 array. Default is zero-initialized ``int32`` array of size ``ND7``.
653    is4 : numpy.ndarray, optional
654        Section 4 array. Default is zero-initialized ``int32`` array of size ``ND7``.
655    *args : tuple
656        Additional positional arguments passed to the constructed record class.
657    **kwargs : dict
658        Optional keyword arguments. The following key is recognized:
659
660        - ``type`` : str, optional
661          Record type. Default is ``"vector"``. If ``is2`` contains any
662          non-zero values, the type is automatically set to ``"grid"``.
663        - ``grid`` : str, optional
664          Predefined grid. Valid values can be found in ``tdlpackio.grids.GRIDS``.
665          This will force the type to be ``"grid"`` and override values in the ``is2``
666          argument.
667
668    Returns
669    -------
670    _TdlpackRecord
671        Instance of a dynamically generated subclass of ``_TdlpackRecord``.
672
673    Notes
674    -----
675    - Record subclasses are created dynamically and cached in
676      ``_record_class_store`` using the record type as the key.
677    - For ``"grid"`` records, ``templates.GridDefinitionSection`` is added
678      as a base class and ``is1[1]`` is set to indicate the presence of a
679      grid definition section.
680    """
681
682    def __new__(
683        self,
684        is0: np.array = np.zeros((ND7), dtype=np.int32),
685        is1: np.array = np.zeros((ND7), dtype=np.int32),
686        is2: np.array = np.zeros((ND7), dtype=np.int32),
687        is4: np.array = np.zeros((ND7), dtype=np.int32),
688        *args,
689        **kwargs,
690    ):
691
692        rectype = "vector"
693        if "type" in kwargs.keys():
694            rectype = kwargs["type"]
695            if rectype not in {"grid", "vector"}:
696                raise ValueError('Invalid "type" argument. Expected "grid" or "vector".')
697
698        if bool(np.any(is2)):
699            rectype = "grid"
700
701        if "grid" in kwargs.keys():
702            rectype = "grid"
703            if kwargs["grid"] not in grids.GRIDS.keys():
704                raise ValueError('Invalid "grid" argument.')
705            is2 = grids.get_grid(kwargs["grid"]).to_is2()
706
707        bases = list()
708        if rectype == "grid":
709            bases.append(templates.GridDefinitionSection)
710            is1[1] = 1  # Flag in is1 to state that a grid definition section exists
711
712        try:
713            Record = _record_class_store[rectype]
714        except KeyError:
715
716            @dataclass(init=False, repr=False)
717            class Record(_TdlpackRecord, *bases):
718                __hash__ = _TdlpackRecord.__hash__
719                pass
720
721            _record_class_store[rectype] = Record
722
723        return Record(is0, is1, is2, is4, *args)

1149@dataclass
1150class TdlpackStationRecord:
1151    """
1152    TDLPACK Station Record class
1153    """
1154
1155    stations_in: InitVar[Optional[Iterable[str]]] = None
1156
1157    type: str = field(init=False, repr=False, default="station")
1158    stations: ClassVar[templates.Stations] = templates.Stations()
1159
1160    # Private class variable holding the list
1161    _stations: Optional[list[str]] = field(init=False, repr=False, default=None)
1162
1163    def __post_init__(self, stations_in):
1164        self._nsta_expected = 0
1165        self._recnum = -1
1166        self._source = None
1167        self._stations = None
1168        self._type = "station"
1169        self.id = TdlpackID([400001000, 0, 0, 0], self)
1170
1171        if stations_in is not None:
1172            self.stations = stations_in
1173
1174    def __hash__(self):
1175        return hash(tuple(self.stations))
1176
1177    def __str__(self):
1178        return f"{self._recnum}:d=0000000000:STATION CALL LETTER RECORD:{self.numberOfStations}"
1179
1180    @property
1181    def numberOfStations(self):
1182        if self._source is not None and (isinstance(self._stations, list) or self._stations is None):
1183            return self._nsta_expected
1184        return 0 if self.stations is None else len(self.stations)
1185
1186    @property
1187    def data(self):
1188        pass
1189
1190    def pack(self):
1191        pass

1194@dataclass
1195class TdlpackTrailerRecord:
1196    """
1197    TDLPACK Trailer Record class
1198    """
1199
1200    type: str = field(init=False, repr=False, default="trailer")
1201
1202    def __post_init__(self):
1203        """"""
1204        self._recnum = -1
1205        self._type = "trailer"
1206        self.id = TdlpackID([0, 0, 0, 0], self)
1207
1208    def __hash__(self):
1209        return hash((""))
1210
1211    def __str__(self):
1212        """"""
1213        return f"{self._recnum}:d=0000000000:TRAILER RECORD"
1214
1215    @property
1216    def data(self):
1217        """"""
1218        pass
1219
1220    def pack(self):
1221        """"""
1222        pass