Source code for at_py.readwrite.bdry3d

"""3D boundary ``.bty`` reader (port of Matlab ``readbdry3d.m``)."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from at_py.readwrite.bty import _quoted_inner
from at_py.readwrite.vector import format_readvector_lines, parse_readvector_lines


def _collect_float_tokens(lines: list[str], start: int) -> list[float]:
    """Collect floats from ``lines[start:]`` (skip ``!``, ``/``)."""
    out: list[float] = []
    for j in range(start, len(lines)):
        s = lines[j].strip()
        if not s:
            continue
        if "!" in s:
            s = s.split("!", 1)[0].strip()
        for tok in s.replace(",", " ").split():
            if tok == "/":
                continue
            try:
                out.append(float(tok))
            except ValueError:
                continue
    return out




@dataclass(frozen=True)
class Bdry3DRead:
    """3D grid: ``x_km``, ``y_km``, depths ``z_km`` shaped ``(n_y, n_x)`` (``readbdry3d``)."""

    interp_type: str
    """``R`` piecewise-linear or ``C`` curvilinear (first letter, padded like Matlab)."""
    x_km: np.ndarray
    y_km: np.ndarray
    z_km: np.ndarray





def parse_bdry3d(text: str) -> Bdry3DRead:
    """Parse a Bellhop3D 3D boundary file (Matlab ``readbdry3d``)."""
    raw = text.splitlines()
    lines = [ln.rstrip("\n") for ln in raw]
    i = 0
    while i < len(lines) and not lines[i].strip():
        i += 1
    if i >= len(lines):
        raise ValueError("empty 3D boundary file")
    inner = _quoted_inner(lines[i])
    bdry_type = (inner + "  ")[:2]
    kind = bdry_type[0:1].upper()
    if kind not in {"R", "C"}:
        raise ValueError(f"unknown boundary type {inner!r} (expected R or C)")

    x, nx, i2 = parse_readvector_lines(lines, i + 1)
    y, ny, i3 = parse_readvector_lines(lines, i2)
    flat = _collect_float_tokens(lines, i3)
    need = int(nx * ny)
    if len(flat) < need:
        raise ValueError(f"3D boundary: need {need} z values after x/y, got {len(flat)}")
    flat = flat[:need]
    arr = np.asarray(flat, dtype=np.float64)
    # Matlab: zBot = fscanf('%f',[NbdryPtsx,NbdryPtsy]); zBot = zBot'
    m = arr.reshape((int(nx), int(ny)), order="F")
    z_km = m.T
    return Bdry3DRead(interp_type=bdry_type, x_km=x, y_km=y, z_km=z_km)





def parse_bdry3d_bytes(data: bytes) -> Bdry3DRead:
    """Parse 3D boundary from UTF-8 text bytes."""
    text = data.decode("utf-8", errors="replace")
    return parse_bdry3d(text)





def format_bdry3d(b: Bdry3DRead) -> str:
    """Format a Bellhop3D 3D boundary file (inverse of :func:`parse_bdry3d`)."""
    ch = b.interp_type.strip().upper()[:1]
    if ch not in ("R", "C"):
        raise ValueError(f"interp_type must start with R or C; got {b.interp_type!r}")
    x_km = np.asarray(b.x_km, dtype=np.float64).reshape(-1)
    y_km = np.asarray(b.y_km, dtype=np.float64).reshape(-1)
    z_km = np.asarray(b.z_km, dtype=np.float64)
    if z_km.ndim != 2:
        raise ValueError("z_km must be 2D (ny, nx)")
    ny, nx = z_km.shape
    if x_km.size != nx or y_km.size != ny:
        raise ValueError("x_km, y_km lengths must match z_km axes (nx, ny)")
    m = z_km.T
    flat = m.reshape(-1, order="F")
    parts: list[str] = [
        f"'{ch}'",
        format_readvector_lines(int(nx), x_km),
        format_readvector_lines(int(ny), y_km),
    ]
    parts.append(" ".join(f"{v:.17g}" for v in flat))
    return "\n".join(parts) + "\n"





def format_bdry3d_bytes(b: Bdry3DRead, *, encoding: str = "utf-8") -> bytes:
    """Like :func:`format_bdry3d` encoded to bytes."""
    return format_bdry3d(b).encode(encoding)