Source code for at_py.readwrite.flp

"""FIELD ``.flp`` parameter file reader (port of Matlab ``read_flp.m``)."""

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass

import numpy as np

from at_py.readwrite.vector import format_readvector_lines, parse_readvector_lines


def _title_from_line(line: str) -> str:
    """Match Matlab: extract between first two quotes if present, else strip ``!`` comments."""
    s = line.strip()
    parts = s.split("'")
    if len(parts) >= 3:
        return parts[1].strip()
    if "!" in s:
        s = s.split("!", 1)[0].strip()
    return s


def _opt_from_line(line: str) -> str:
    """Quoted ``Opt`` field from a ``.flp`` line."""
    parts = line.split("'")
    if len(parts) < 3:
        raise ValueError(f"expected quoted Opt line: {line!r}")
    return parts[1]


def _normalize_flp_opt(raw: str) -> str:
    """Apply ``read_flp.m`` defaults for 3rd and 4th option characters."""
    o = raw
    if len(o) <= 2:
        chars = list(o)
        while len(chars) < 3:
            chars.append(" ")
        chars[2] = "O"
        o = "".join(chars[:3])
    if len(o) <= 3:
        chars = list(o.ljust(4))
        chars[3] = "C"
        o = "".join(chars)
    return o


def _mlimit_from_line(line: str) -> int:
    """First numeric token on a line (``MLimit`` or similar)."""
    s = line.strip()
    if "!" in s:
        s = s.split("!", 1)[0].strip()
    first = s.replace(",", " ").split()[0]
    return int(float(first))




@dataclass(frozen=True)
class FieldFlpResult:
    """In-memory ``field`` driver parameters from a ``.flp`` file."""

    title: str
    opt: str
    """Normalized option string (length ≥ 4 after Matlab defaults)."""
    component: str
    """``Comp`` from ``Opt(3)`` (1-based), default ``P``."""
    m_limit: int
    r_prof_km: np.ndarray
    n_prof: int
    receiver_range_m: np.ndarray
    source_z_m: np.ndarray
    receiver_z_m: np.ndarray
    range_offset_m: np.ndarray
    n_range_offset: int





def parse_field_flp(text: str) -> FieldFlpResult:
    """Parse FIELD ``.flp`` text (Matlab ``read_flp.m``).

    Receiver ranges in the file are km; results use **meters** for ranges (``read_flp`` ×1000).
    Profile ranges ``r_prof_km`` stay in km as in the file.
    """
    lines = text.splitlines()
    i = 0
    if i >= len(lines):
        raise ValueError("empty .flp file")
    title = _title_from_line(lines[i])
    i += 1
    if i >= len(lines):
        raise ValueError(".flp: missing Opt line")
    opt_raw = _opt_from_line(lines[i])
    i += 1
    opt = _normalize_flp_opt(opt_raw)
    comp = opt[2] if len(opt) >= 3 else "P"

    if i >= len(lines):
        raise ValueError(".flp: missing MLimit line")
    m_limit = _mlimit_from_line(lines[i])
    i += 1

    r_prof_km, n_prof, i = parse_readvector_lines(lines, i)
    rr_km, _nrr, i = parse_readvector_lines(lines, i)
    receiver_range_m = rr_km * 1000.0

    sd, _nsd, i = parse_readvector_lines(lines, i)
    rd, _nrd, i = parse_readvector_lines(lines, i)

    ro, n_ro, i = parse_readvector_lines(lines, i)

    for j in range(i, len(lines)):
        if lines[j].strip():
            raise ValueError(f"unexpected line after .flp body: {lines[j]!r}")

    if np.max(np.abs(ro)) > 0.0:
        raise ValueError(
            "receiver range offsets (array tilt) are not supported (field.m limitation)"
        )

    return FieldFlpResult(
        title=title,
        opt=opt,
        component=comp,
        m_limit=m_limit,
        r_prof_km=r_prof_km,
        n_prof=n_prof,
        receiver_range_m=receiver_range_m,
        source_z_m=sd,
        receiver_z_m=rd,
        range_offset_m=ro,
        n_range_offset=n_ro,
    )





def parse_field_flp_bytes(data: bytes, *, encoding: str = "utf-8") -> FieldFlpResult:
    """Like :func:`parse_field_flp` after decoding ``data`` with ``encoding``."""
    return parse_field_flp(data.decode(encoding, errors="replace"))





def format_field_flp(r: FieldFlpResult) -> str:
    """Format a FIELD 2D ``.flp`` file (inverse of :func:`parse_field_flp`).

    Receiver range offsets must be zero (same restriction as the reader).
    """
    ro = np.asarray(r.range_offset_m, dtype=np.float64).reshape(-1)
    if ro.size != int(r.n_range_offset) or np.max(np.abs(ro)) > 0.0:
        raise ValueError("range_offset_m must be all zeros (field.m limitation)")
    ns = int(np.asarray(r.source_z_m).size)
    nr = int(np.asarray(r.receiver_z_m).size)
    nrr = int(np.asarray(r.receiver_range_m).size)
    parts: list[str] = [
        f"'{r.title}'",
        f"'{r.opt}'",
        str(int(r.m_limit)),
        format_readvector_lines(int(r.n_prof), r.r_prof_km),
        format_readvector_lines(nrr, r.receiver_range_m / 1000.0),
        format_readvector_lines(ns, r.source_z_m),
        format_readvector_lines(nr, r.receiver_z_m),
        format_readvector_lines(int(r.n_range_offset), r.range_offset_m),
    ]
    return "\n".join(parts) + "\n"





def format_field_flp_bytes(r: FieldFlpResult, *, encoding: str = "utf-8") -> bytes:
    """UTF-8 (or other) bytes for :func:`format_field_flp`."""
    return format_field_flp(r).encode(encoding)





def format_fields_flp(receiver_range_km: np.ndarray | Sequence[float]) -> str:
    """Minimal Scooter-style FIELD ``.flp`` text (Matlab ``write_fieldsflp.m``).

    Emits the ``'RP'`` option and one line with ``rMin``, ``rMax`` (km), and ``NRr`` from
    ``min(Pos.r.r)``, ``max(Pos.r.r)``, and ``length(Pos.r.r)``. This is **not** the full
    FIELD driver from ``read_flp`` / ``write_fieldflp.m``; use :func:`format_field_flp` for that.
    """
    rr = np.asarray(receiver_range_km, dtype=np.float64).reshape(-1)
    if rr.size < 1:
        raise ValueError("receiver_range_km must be non-empty")
    r_min = float(np.min(rr))
    r_max = float(np.max(rr))
    nrr = int(rr.size)
    line1 = "'RP' \t \t ! Option"
    line2 = f"{r_min:6.2f} {r_max:6.2f} {nrr:5d} \t ! rMin rMax (km) NRr"
    return f"{line1}\n{line2}\n"





def format_fields_flp_bytes(
    receiver_range_km: np.ndarray | Sequence[float],
    *,
    encoding: str = "utf-8",
) -> bytes:
    """UTF-8 (or other) bytes for :func:`format_fields_flp`."""
    return format_fields_flp(receiver_range_km).encode(encoding)