Source code for at_py.readwrite.ssp2d

"""2D SSP grid (``SSPFIL``) reader (port of Matlab ``readssp2d.m``)."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np


def _tokens(text: str) -> list[str]:
    """Whitespace tokens from SSPFIL text, stripping ``!`` comments and ``/``."""
    out: list[str] = []
    for line in text.splitlines():
        s = line.split("!", 1)[0].strip()
        if not s:
            continue
        for t in s.replace(",", " ").split():
            if t == "/":
                continue
            out.append(t)
    return out




@dataclass(frozen=True)
class SSP2DRead:
    """2D sound-speed field: ``cmat`` is ``(NSSP, NProf)``; ``r_prof_km`` has length ``NProf``."""

    n_prof: int
    nssp: int
    r_prof_km: np.ndarray
    cmat: np.ndarray





def parse_ssp2d(text: str) -> SSP2DRead:
    """Parse a 2D SSPFIL (Matlab ``readssp2d``).

    Layout: integer ``NProf``, ``NProf`` range values (km), then a flat block of
    ``NProf * NSSP`` sound speeds read into ``[NProf, NSSP]`` column-wise and
    transposed to ``[NSSP, NProf]``.
    """
    tok = _tokens(text)
    if len(tok) < 1:
        raise ValueError("empty SSPFIL (2D)")
    n_prof = int(tok[0])
    need = 1 + n_prof
    if len(tok) < need:
        raise ValueError(
            f"SSPFIL (2D): need at least {need} tokens (NProf + ranges), got {len(tok)}"
        )
    r_prof = np.array([float(tok[1 + i]) for i in range(n_prof)], dtype=np.float64)
    rest = [float(x) for x in tok[need:]]
    if not rest:
        raise ValueError("SSPFIL (2D): missing sound-speed block")
    if len(rest) % n_prof != 0:
        raise ValueError(
            f"SSPFIL (2D): sound-speed count {len(rest)} not divisible by NProf={n_prof}"
        )
    nssp = len(rest) // n_prof
    arr = np.asarray(rest, dtype=np.float64)
    raw = arr.reshape((n_prof, nssp), order="F")
    cmat = raw.T
    return SSP2DRead(n_prof=n_prof, nssp=nssp, r_prof_km=r_prof, cmat=cmat)





def parse_ssp2d_bytes(data: bytes) -> SSP2DRead:
    """Parse 2D SSPFIL from UTF-8 text bytes."""
    text = data.decode("utf-8", errors="replace")
    return parse_ssp2d(text)



def _fmt_ssp_float(x: float) -> str:
    return f"{x:.17g}"




def format_ssp2d(s: SSP2DRead) -> str:
    """Format a 2D SSPFIL text (inverse of :func:`parse_ssp2d`).

    Emits ``NProf``, profile ranges (km), then sound speeds in the same flat order
    the parser expects (column-major block then transpose to ``cmat``).
    """
    cmat = np.asarray(s.cmat, dtype=np.float64)
    r_prof = np.asarray(s.r_prof_km, dtype=np.float64).reshape(-1)
    if int(r_prof.size) != int(s.n_prof):
        raise ValueError("r_prof_km length must match n_prof")
    if cmat.shape != (int(s.nssp), int(s.n_prof)):
        exp = f"(nssp, n_prof)=({s.nssp}, {s.n_prof})"
        raise ValueError(f"cmat shape must be {exp}; got {cmat.shape}")
    raw = cmat.T
    flat = raw.reshape(-1, order="F")
    lines = [
        str(int(s.n_prof)),
        " ".join(_fmt_ssp_float(float(x)) for x in r_prof),
        " ".join(_fmt_ssp_float(float(x)) for x in flat.ravel()),
    ]
    return "\n".join(lines) + "\n"





def format_ssp2d_bytes(s: SSP2DRead, *, encoding: str = "utf-8") -> bytes:
    return format_ssp2d(s).encode(encoding)