Source code for pycsamt.jones.heads

# Author: LKouadio <etanoyau@gmail.com>
# License: LGPL-3.0

from __future__ import annotations

from collections.abc import Sequence
from datetime import datetime
from pathlib import Path
from typing import Any

from ..exceptions import JParseError
from .base import JComponentBase
from .config import (
    ENCODING_DEFAULT,
    RE_BANNER,
    RE_BLANK,
    RE_COMMENT,
    RE_DATATYPE_UNITS,
    RE_INFO,
    RE_NPOINTS,
    RE_STATION,
    UNITS_CANONICAL,
)
from .property import JSiteProperty
from .utils import DataType, iter_lines, parse_datatype_units

__all__ = ["Head", "Info", "Heads", "HeadMixin", "InfoMixin", "Banner"]








[docs]
class Info(JComponentBase):
    r"""
    Parse and serialize the J-format information block.

    This component collects ``>KEY=VALUE`` records along with
    leading comment lines. It exposes a small set of convenience
    properties derived from :class:`JSiteProperty`.

    Parameters
    ----------
    j_info_list : sequence of str, optional
        A sequence containing the comment and information
        records. The parser stops at the first non-info,
        non-comment, non-blank line.
    verbose : int, default=0
        Verbosity for warnings during parsing.
    **kwargs
        Accepted for API forwards-compatibility; ignored here.

    Attributes
    ----------
    items : dict of (str -> str)
        Mapping of upper-cased keys to unmodified string values.
    comments : list of str
        Preserved leading comment lines (starting with ``#``).
    site : JSiteProperty
        Lazily parsed view providing normalized latitude,
        longitude, azimuth and elevation (see properties below).
    latitude : float or None
        Decimal degrees; hemisphere and DMS tolerated.
    longitude : float or None
        Decimal degrees in ``[-180, 180)``.
    azimuth : float or None
        Site X-axis azimuth (degrees, true north).
    elevation : float or None
        Elevation in metres.
    path : pathlib.Path or None
        Source path when constructed via :meth:`from_file`.
    encoding : str
        Text encoding used by :meth:`from_file`.
    verbose : int
        Verbosity level inherited from :class:`JComponentBase`.

    Methods
    -------
    from_file(j_fn, *, verbose=0)
        Read only the comment/info header from a file.
    from_lines(j_info_list, *, verbose=0)
        Build from an in-memory sequence.
    read(j_info_list)
        Parse the header and mark the instance as read.
    write(j_info_list=None)
        Render comments and ``>KEY = VALUE`` lines.

    Notes
    -----
    Unknown keys are preserved verbatim. Coordinate and azimuth
    values are normalized via :class:`JSiteProperty`, which
    handles ranges, hemispheres and DMS.

    Examples
    --------
    >>> info = Info.from_file('data/j/kb0-s001.txt')
    >>> info.latitude, info.longitude
    (41.9782, 140.8958)
    >>> lines = info.write()
    >>> Info.from_lines(lines).azimuth == info.azimuth
    True

    See Also
    --------
    Head : One data-block header (station / dtype / count).
    Heads : Combined view with convenience properties.
    JSiteProperty : Robust parsing used by this class.

    References
    ----------
    .. [1] A. G. Jones (1994). Magnetotelluric data file
       J-format, version 2.0.
    """

    def __init__(
        self,
        j_info_list: Sequence[str] | None = None,
        *,
        verbose: int = 0,
        **kwargs: Any,
    ) -> None:
        super().__init__(verbose=verbose)
        self.comments: list[str] = []
        self.items: dict[str, str] = {}
        self._site_cache: JSiteProperty | None = None
        if j_info_list is not None:
            self.read(j_info_list)



[docs]
    @classmethod
    def from_file(cls, j_fn: str | Path, *, verbose: int = 0) -> Info:
        lines = list(iter_lines(j_fn, encoding=ENCODING_DEFAULT))
        info_list = _extract_info_header_list(lines)
        return cls(info_list, verbose=verbose)





[docs]
    @classmethod
    def from_lines(
        cls,
        j_info_list: Sequence[str] | None = None,
        *,
        verbose: int = 0,
    ) -> Info:
        if j_info_list is None:
            raise ValueError("j_info_list is required")
        seq = list(j_info_list)
        info_list = _extract_info_header_list(seq)

        return cls(info_list, verbose=verbose)





[docs]
    def write(self, j_info_list: Sequence[str] | None = None) -> list[str]:
        if j_info_list is not None:
            self.read(j_info_list)
        # out: List[str] = []
        # for c in self.comments:
        #     out.append(c.rstrip("\n"))
        # for k, v in self.items.items():
        #     out.append(f">{k} = {v}")
        # return out
        out: list[str] = []
        for c in self.comments:
            # skip any original #WRITTEN BY … banner lines
            if RE_BANNER.match(c):
                continue
            out.append(c.rstrip("\n"))
        for k, v in self.items.items():
            out.append(f">{k:<10} = {v:>13}")
        return out





[docs]
    def read(self, j_info_list: Sequence[str] | None = None) -> Info:
        if j_info_list is None:
            raise ValueError("j_info_list is required")
        self.comments = []
        self.items = {}
        self._site_cache = None
        for ln in j_info_list:
            if _is_comment(ln):
                self.comments.append(ln)
                continue
            if _is_blank(ln):
                continue
            if _is_info(ln):
                k, v = ln.split("=", 1)
                key = k.strip().lstrip(">").strip().upper()
                self.items[key] = v.strip()
        self._mark_read(True)
        return self





[docs]
    @property
    def site(self) -> JSiteProperty:
        if self._site_cache is None:
            self._site_cache = JSiteProperty.from_lines(
                self.write(), verbose=self.verbose
            )
        return self._site_cache





[docs]
    @property
    def latitude(self) -> float | None:
        return self.site.latitude





[docs]
    @property
    def longitude(self) -> float | None:
        return self.site.longitude





[docs]
    @property
    def azimuth(self) -> float | None:
        return self.site.azimuth





[docs]
    @property
    def elevation(self) -> float | None:
        return self.site.elevation





[docs]
    def get(self, key: str, default: str | None = None) -> str | None:
        return self.items.get(key.upper(), default)





[docs]
    def keys(self) -> tuple[str, ...]:
        return tuple(self.items.keys())





[docs]
    def values(self) -> tuple[str, ...]:
        return tuple(self.items.values())





[docs]
    def items_map(self) -> dict[str, str]:
        # explicit copy to avoid accidental mutation
        return dict(self.items)





[docs]
    @property
    def lat(self) -> float | None:
        return self.latitude





[docs]
    @property
    def lon(self) -> float | None:
        return self.longitude



    def __contains__(self, key: str) -> bool:
        return key.upper() in self.items

    def __getitem__(self, key: str) -> str:
        return self.items[key.upper()]

    def __str__(self) -> str:
        n = len(self.items)
        return f"Info(items={n})"

    def __repr__(self) -> str:
        return f"Info(items={len(self.items)}, cmts={{len(self.comments)}})"










[docs]
class Heads(JComponentBase):
    r"""
    Minimal container for one :class:`Head` and one
    :class:`Info`.

    The class provides convenient accessors for station and
    site-level properties, and a small banner recorder for the
    top provenance line (``#WRITTEN BY ...``). It is intended as
    the lightest useful representation of a single J header
    section.

    Parameters
    ----------
    head : Head, optional
        Existing head to attach. If omitted, an empty
        :class:`Head` is created.
    info : Info, optional
        Existing info to attach. If omitted, an empty
        :class:`Info` is created.
    verbose : int, default=0
        Verbosity for warnings during parsing.

    Attributes
    ----------
    head : Head
        The parsed header triple.
    info : Info
        The parsed site information block.
    banner : Banner
        Parsed top comment provenance. The writer defaults the
        ``software`` field to ``PYSCAMT`` if missing.
    n : int
        ``0`` or ``1`` depending on whether a head has been
        parsed.
    station : str or None
        Shortcut to ``head.station``.
    latitude, longitude, elevation, azimuth : float or None
        Shortcuts to values provided by ``info``. ``azimuth``
        falls back to ``head.az_hint`` when the info block does
        not contain ``AZIMUTH``.

    Methods
    -------
    from_file(j_fn, *, verbose=0)
        Read the file then delegate to :meth:`read`.
    from_lines(lines, *, verbose=0)
        Build from an in-memory line sequence.
    read(text_or_lines)
        Extract info + head lists and parse both.
    write()
        Serialize banner, head and info back-to-back.

    Notes
    -----
    This class does *not* parse the subsequent data rows. It is
    focused on fast, dependency-free header discovery to support
    scanning tasks and metadata extraction.

    Examples
    --------
    >>> h = Heads.from_file('data/j/kb0-s001.txt')
    >>> h.station, h.latitude, h.software
    ('KB0001', 41.9782, 'GEOTOOLS')

    See Also
    --------
    Head : Single data-block header.
    Info : Site-level header.
    Banner : Top provenance line handler.

    References
    ----------
    .. [1] A. G. Jones (1994). Magnetotelluric data file
       J-format, version 2.0.
    """

    _repr_keys = ["n"]

    def __init__(
        self,
        head: Head | None = None,
        info: Info | None = None,
        *,
        verbose: int = 0,
    ) -> None:
        super().__init__(verbose=verbose)
        self.banner = Banner(verbose=verbose)
        self.head: Head = head if head is not None else Head(verbose=verbose)
        self.info: Info = info if info is not None else Info(verbose=verbose)



[docs]
    def read(self, text_or_lines: str | Sequence[str]) -> Heads:
        if isinstance(text_or_lines, str):
            lines = text_or_lines.splitlines()
        else:
            lines = list(text_or_lines)
        self.banner.read(lines)
        info_list = _extract_info_header_list(lines)
        head_list = _extract_first_head_list(lines)
        self.info.read(info_list)
        self.head.read(head_list)
        self._mark_read(True)
        return self





[docs]
    def write(self, include_origin=False) -> list[str]:
        out: list[str] = []
        out.extend(self.banner.write(new=True, include_origin=include_origin))
        out.extend(self.head.write())
        out.extend(self.info.write())
        return out





[docs]
    @classmethod
    def from_lines(cls, lines: Sequence[str], *, verbose: int = 0) -> Heads:
        inst = cls(verbose=verbose)
        inst.read(lines)
        return inst





[docs]
    @classmethod
    def from_file(cls, j_fn: str | Path, *, verbose: int = 0) -> Heads:
        lines = list(iter_lines(j_fn, encoding=ENCODING_DEFAULT))
        return cls.from_lines(lines, verbose=verbose)





[docs]
    @property
    def n(self) -> int:
        return 0 if self.head.n is None else 1





[docs]
    @property
    def station(self) -> str | None:
        return self.head.station





[docs]
    @property
    def latitude(self) -> float | None:
        return self.info.latitude





[docs]
    @property
    def longitude(self) -> float | None:
        return self.info.longitude





[docs]
    @property
    def elevation(self) -> float | None:
        return self.info.elevation





[docs]
    @property
    def azimuth(self) -> float | None:
        return (
            self.info.azimuth
            if self.info.azimuth is not None
            else self.head.az_hint
        )





[docs]
    @property
    def software(self) -> str | None:
        return self.banner.software



    def __str__(self) -> str:
        return f"Heads(n={self.n})"

    def __repr__(self) -> str:
        return self.__str__()






[docs]
class HeadMixin:
    r"""
    Mixin that provides a :class:`Head` on host classes.

    The mixin offers class-level and instance-level helpers that
    delegate to :class:`Head` while keeping a single copy of the
    header on the host.

    Attributes
    ----------
    head : Head
        Lazily created and cached on first use.

    Methods
    -------
    from_file(edi_fn)
        Class method that returns ``Head.from_file(edi_fn)``.
    read(j_header_list=None)
        Ensure a ``Head`` exists on the host, parse into it, and
        return it.
    write(head_list_infos=None)
        Serialize the host's ``Head`` to three header lines.

    Examples
    --------
    >>> class Host(HeadMixin):  # doctest: +SKIP
    ...     pass
    >>> Host.from_file('file.j')  # doctest: +SKIP
    Head(...)
    """

    head: Head



[docs]
    @classmethod
    def from_file(cls, edi_fn: str | Path) -> Head:
        return Head.from_file(edi_fn)





[docs]
    def read(self, j_header_list: Sequence[str] | None = None) -> Head:
        if not hasattr(self, "head") or self.head is None:
            self.head = Head()
        return self.head.read(j_header_list)





[docs]
    def write(
        self, head_list_infos: Sequence[str] | None = None
    ) -> list[str]:
        if not hasattr(self, "head") or self.head is None:
            self.head = Head()
        return self.head.write(head_list_infos)








[docs]
class InfoMixin:
    r"""
    Mixin that provides an :class:`Info` on host classes.

    The mixin mirrors :class:`HeadMixin` but for the information
    block. It centralizes parsing and writing of ``>KEY=VALUE``
    lines.

    Attributes
    ----------
    info : Info
        Lazily created and cached on first use.

    Methods
    -------
    from_file(edi_fn)
        Class method that returns ``Info.from_file(edi_fn)``.
    read(j_info_list=None)
        Ensure an ``Info`` exists on the host, parse into it, and
        return it.
    write(j_info_list=None)
        Serialize the host's ``Info`` to comment and info lines.

    Examples
    --------
    >>> class Host(InfoMixin):  # doctest: +SKIP
    ...     pass
    >>> Host().read(['>LATITUDE=10'])  # doctest: +SKIP
    Info(items=1)
    """

    info: Info



[docs]
    @classmethod
    def from_file(cls, edi_fn: str | Path) -> Info:
        return Info.from_file(edi_fn)





[docs]
    def read(self, j_info_list: Sequence[str] | None = None) -> Info:
        if not hasattr(self, "info") or self.info is None:
            self.info = Info()
        return self.info.read(j_info_list)





[docs]
    def write(self, j_info_list: Sequence[str] | None = None) -> list[str]:
        if not hasattr(self, "info") or self.info is None:
            self.info = Info()
        return self.info.write(j_info_list)






def _extract_info_header_list(lines: Sequence[str]) -> list[str]:
    header: list[str] = []
    for ln in lines:
        if _is_comment(ln) or _is_info(ln) or _is_blank(ln):
            header.append(ln)
            continue
        break
    return header


def _extract_first_head_list(lines: Sequence[str]) -> list[str]:
    i, n = 0, len(lines)
    while i < n and (
        _is_comment(lines[i]) or _is_info(lines[i]) or _is_blank(lines[i])
    ):
        i += 1
    while i < n and not RE_STATION.match(lines[i]):
        i += 1
    if i >= n:
        raise ValueError("no head triple found")
    s = i

    # first non-blank after station
    j = i + 1
    while j < n and _is_blank(lines[j]):
        j += 1
    if j >= n:
        raise ValueError("missing header after station")

    # case A: standard order (dtype then count)
    try:
        parse_datatype_units(lines[j])
        k = j + 1
        while k < n and _is_blank(lines[k]):
            k += 1
        if k >= n or not RE_NPOINTS.match(lines[k]):
            raise ValueError("missing count after data-type")
        return [lines[s], lines[j], lines[k]]
    except JParseError:
        pass  # not a dtype here

    # case B: count first, then dtype later
    if RE_NPOINTS.match(lines[j]):
        k = j
        # scan forward for first dtype before next station/EOF
        j = k + 1
        while j < n:
            if RE_STATION.match(lines[j]):
                break
            if _is_blank(lines[j]):
                j += 1
                continue
            try:
                parse_datatype_units(lines[j])
                return [lines[s], lines[j], lines[k]]
            except JParseError:
                # not a dtype; likely data row — keep scanning a bit
                j += 1
                continue
        raise ValueError("no data-type found after count")
    else:
        raise ValueError("bad header after station")


def _is_comment(s: str) -> bool:
    return bool(RE_COMMENT.match(s))


def _is_blank(s: str) -> bool:
    return bool(RE_BLANK.match(s))


def _is_info(s: str) -> bool:
    return bool(RE_INFO.match(s))


def _fmt_dtype(dt: DataType) -> str:
    base = f"{dt.kind}{dt.comp}"
    if dt.units:
        for k, v in UNITS_CANONICAL.items():
            if v == dt.units:
                return f"{base} {k}"
    return base