Source code for at_py.readwrite.vector

"""Fortran-style vector shorthand (port of Matlab ``readvector.m``)."""

from __future__ import annotations

from collections.abc import Sequence

import numpy as np


def _fmt_fortran_float(x: float) -> str:
    """Emit a float that ``read_fortran_vector`` can parse back reliably."""
    if not np.isfinite(x):
        raise ValueError(f"non-finite value cannot be formatted for Fortran vector: {x!r}")
    return f"{x:.17g}"




def format_fortran_vector_data_line(nx: int, x: np.ndarray | Sequence[float]) -> str:
    """Emit one data line for a readvector block (inverse of :func:`read_fortran_vector`)."""
    arr = np.asarray(x, dtype=np.float64).reshape(-1)
    if arr.size != nx:
        raise ValueError(f"expected {nx} value(s), got {arr.size}")
    if nx < 1:
        raise ValueError("nx must be >= 1")

    if nx == 1:
        return _fmt_fortran_float(float(arr[0]))
    if nx == 2:
        return f"{_fmt_fortran_float(float(arr[0]))} {_fmt_fortran_float(float(arr[1]))}"

    if np.allclose(arr, arr[0]):
        return f"{_fmt_fortran_float(float(arr[0]))} /"

    a0, a1 = float(arr[0]), float(arr[-1])
    lin = np.linspace(a0, a1, nx, dtype=np.float64)
    if np.allclose(arr, lin):
        return f"{_fmt_fortran_float(a0)} {_fmt_fortran_float(a1)} /"

    return " ".join(_fmt_fortran_float(float(v)) for v in arr)





def format_readvector_block(nx: int, x: np.ndarray | Sequence[float]) -> tuple[str, str]:
    """Emit count line and data line for a Matlab ``readvector`` block."""
    arr = np.asarray(x, dtype=np.float64).reshape(-1)
    if arr.size != nx:
        raise ValueError(f"expected {nx} value(s), got {arr.size}")
    return str(nx), format_fortran_vector_data_line(nx, arr)





def format_readvector_lines(nx: int, x: np.ndarray | Sequence[float]) -> str:
    """Two-line ``readvector`` section (count, then data), newline-separated."""
    a, b = format_readvector_block(nx, x)
    return f"{a}\n{b}"





def read_fortran_vector(nx: int, spec_line: str) -> np.ndarray:
    """Parse one line after the count (``Nx``), emulating ``readvector.m``.

    Supports ``/`` termination:

    - ``Nx > 2`` with two values ``a b /`` → ``linspace(a, b, Nx)``
    - ``Nx > 2`` with one value ``a /`` → ``full(Nx, a)``
    - ``Nx <= 2`` or explicit list → ``Nx`` floats from the line
    """
    line = spec_line.strip()
    if "!" in line:
        line = line.split("!", 1)[0].strip()
    if "/" in line:
        tokens = [float(x) for x in line.replace("/", " ").split() if x]
        if nx == 1 and len(tokens) >= 1:
            return np.array([tokens[0]], dtype=np.float64)
        if nx > 2:
            if len(tokens) > 1:
                return np.linspace(tokens[0], tokens[1], nx, dtype=np.float64)
            return np.full(nx, tokens[0], dtype=np.float64)
        return np.array(tokens[:nx], dtype=np.float64)

    tokens = [float(x) for x in line.split()]
    # Shorthand used in many AT files: two endpoints without ``/`` (Matlab ``sscanf`` path
    # cannot fill ``Nx`` from one line; AT examples use ``a b /`` or rely on linspace).
    if nx > 2 and len(tokens) == 2:
        return np.linspace(tokens[0], tokens[1], nx, dtype=np.float64)
    return np.array(tokens[:nx], dtype=np.float64)





def parse_readvector_lines(lines: list[str], i: int) -> tuple[np.ndarray, int, int]:
    """Port of Matlab ``readvector.m`` using pre-split lines.

    The count is taken from ``lines[i]`` (first integer), data from ``lines[i+1]``.

    Returns:
        ``x``: length-``Nx`` vector
        ``nx``: requested count (matches ``x.size``)
        ``i_next``: index of first line after the readvector block
    """
    if i >= len(lines):
        raise ValueError("readvector: missing count line")
    parts = lines[i].split()
    if not parts:
        raise ValueError("readvector: empty count line")
    nx = int(parts[0])
    if i + 1 >= len(lines):
        raise ValueError("readvector: missing data line")
    x = read_fortran_vector(nx, lines[i + 1])
    if x.size != nx:
        raise ValueError(f"readvector: got {x.size} value(s), expected {nx}")
    return x, nx, i + 2