scanengine-3/core/sras_format.py

"""SRAS binary scan-file format (v7/v11, reading v6/v10 too) — the single
implementation.

Full byte-level spec: scan_format.md.  Summary:

    header  >4sBHfffffffIdBB   magic ver n_angles xs_nom ys_nom xd_nom yd_nom
                               row_spacing velocity laser_freq spf sample_rate
                               bytes_per_sample n_channels
    angle table                n_angles × >f
    geometry table             n_angles × >ffIH  (x_start x_delta n_frames n_rows)
    row tables (ragged)        per angle: n_rows × >f  (y positions)
    preambles                  n_channels × (>H length + utf-8 WFMOutpre string)
    data block                 per angle: [background][rows]
                               background = >I length + raw int8 CH1 average
                               rows       = for each row, for each channel,
                               n_frames × samples_per_frame × bytes_per_sample

Every angle carries its own background: the operator switches the Genesis
laser off before each angle and the engine averages a fresh CH1 record, so
the reference a reader subtracts was taken minutes — not hours — from the
data it is subtracted from.  That is the whole difference between v7 and v6,
which held a single background for the entire file, ahead of the data block.

Incomplete files are valid: the data block is one contiguous append-only
stream, so the readable prefix defines a single frontier past which nothing
has been written yet (see ``SrasFile.angle_status``).  Because the background
blocks are length-prefixed, the per-angle offsets come from a walk of that
stream at parse time rather than from arithmetic on the geometry table.

Version 11 is the SAW quality check (core.saw_check): byte layout identical
to v7, but every angle declares exactly one row — the row-wise middle of the
ROI.  The version byte is the whole difference, and it exists so a reader can
tell a one-row-per-angle check from a full scan that was aborted after its
first row.  ``create_scan_file`` enforces the one-row rule at write time.

v6 and v10 are the pre-per-angle-background versions of the same two files.
They are still read (every scan taken before this change is one); nothing
writes them any more, and a v6 file cannot be resumed into, since appending
v7 blocks to it would shift its data.
"""
from __future__ import annotations

import mmap
import struct
from dataclasses import dataclass, field
from pathlib import Path
from typing import BinaryIO

import numpy as np

from core.scan_geometry import AngleGeometry, ScanPlan

MAGIC = b"SRAS"
VERSION = 7
# One row per angle, taken from the middle of the ROI — see core.saw_check.
VERSION_SAW_CHECK = 11
# v6/v10: the same two files with one background for the whole scan, written
# ahead of the data block instead of once per angle.  Read-only.
LEGACY_VERSIONS = (6, 10)
WRITABLE_VERSIONS = (VERSION, VERSION_SAW_CHECK)
SUPPORTED_VERSIONS = tuple(sorted(WRITABLE_VERSIONS + LEGACY_VERSIONS))
SAW_CHECK_VERSIONS = (10, VERSION_SAW_CHECK)
HDR_FMT = ">4sBHfffffffIdBB"
HDR_SIZE = struct.calcsize(HDR_FMT)   # 49 bytes
GEOM_FMT = ">ffIH"
GEOM_SIZE = struct.calcsize(GEOM_FMT)  # 14 bytes
BG_LEN_FMT = ">I"
BG_LEN_SIZE = struct.calcsize(BG_LEN_FMT)   # 4 bytes

# Oscilloscope channels recorded, in on-disk order.
SCAN_CHANNELS = [1, 3, 4]

STATUS_OK = "OK"
STATUS_TRUNCATED = "TRUNCATED"
STATUS_MISSING = "MISSING"


@dataclass
class ScanHeader:
    """The fixed v6 global header (everything but magic/version)."""
    n_angles: int
    x_start_nominal: float
    y_start_nominal: float
    x_delta_nominal: float
    y_delta_nominal: float
    row_spacing: float
    velocity: float
    laser_freq: float
    samples_per_frame: int
    sample_rate: float
    bytes_per_sample: int
    n_channels: int


@dataclass
class AngleStatus:
    """How much of one angle's declared data is actually on disk."""
    index: int
    angle_deg: float
    n_rows: int              # declared
    row_bytes: int
    bg_offset: int           # start of this angle's background block
    data_offset: int         # start of its rows, i.e. just past that block
    n_rows_available: int
    status: str              # STATUS_OK / STATUS_TRUNCATED / STATUS_MISSING

    @property
    def complete(self) -> bool:
        return self.status == STATUS_OK

    @property
    def bg_bytes(self) -> int:
        """Size of the background block ahead of the rows (0 on v6/v10)."""
        return self.data_offset - self.bg_offset


def create_scan_file(path: Path, plan: ScanPlan, samples_per_frame: int,
                     sample_rate: float, preambles: list[str],
                     version: int = VERSION) -> BinaryIO:
    """Create a new .sras file and write the header + tables.

    ``version`` selects which kind of file this is — VERSION for a full scan,
    VERSION_SAW_CHECK for a middle-row quality check.  The layout is the same
    either way; the one-row-per-angle rule that gives v11 its meaning is
    checked here, since nothing downstream can recover from a v11 file that
    breaks it.

    Returns an open binary file positioned at the start of the data block.
    The caller writes each angle as ``write_background_block()`` followed by
    that angle's rows, and must close the file (try/finally).
    """
    if version not in WRITABLE_VERSIONS:
        raise ValueError(
            f"Cannot write SRAS format version {version} "
            f"(writable: {', '.join(str(v) for v in WRITABLE_VERSIONS)})"
        )
    if version == VERSION_SAW_CHECK:
        bad = [f"{pa.angle_deg:.1f}° has {pa.n_rows}"
               for pa in plan.per_angle if pa.n_rows != 1]
        if bad:
            raise ValueError(
                "A v11 SAW-check file holds exactly one row per angle, but "
                + ", ".join(bad) + " — build the plan with "
                "core.saw_check.middle_row_plan()."
            )

    path.parent.mkdir(parents=True, exist_ok=True)
    f = open(path, "wb")
    f.write(struct.pack(
        HDR_FMT, MAGIC, version,
        plan.n_angles,
        plan.x_start_nominal, plan.y_start_nominal,
        plan.x_delta_nominal, plan.y_delta_nominal,
        plan.row_spacing,
        plan.velocity_mm_s, plan.laser_freq_hz,
        samples_per_frame,
        sample_rate,
        1,                      # bytes_per_sample: int8 from scope default
        len(SCAN_CHANNELS),
    ))
    f.write(struct.pack(f">{plan.n_angles}f", *plan.angles))
    for pa in plan.per_angle:
        f.write(struct.pack(GEOM_FMT, pa.x_start, pa.x_delta, pa.n_frames, pa.n_rows))
    for pa in plan.per_angle:
        f.write(struct.pack(f">{pa.n_rows}f", *pa.y_positions))
    for p in preambles:
        enc = p.encode("utf-8")
        f.write(struct.pack(">H", len(enc)))
        f.write(enc)
    return f


def write_background_block(f: BinaryIO, waveform: bytes) -> int:
    """Write one angle's background block; returns the bytes written.

    Every angle's rows are preceded by one of these, so a reader walking the
    data block knows where that angle's frames start.
    """
    f.write(struct.pack(BG_LEN_FMT, len(waveform)))
    f.write(waveform)
    return BG_LEN_SIZE + len(waveform)


@dataclass
class SrasFile:
    """Parsed .sras file (v7/v11, or legacy v6/v10): header, tables, and lazy
    (memmap) access.

    Parsing reads only the header/tables and the per-angle background blocks
    — never the waveform block — so opening a multi-GB file is cheap.
    ``load_angle``/``load_row`` return read-only numpy views backed by a
    shared mmap; no data is copied until the caller computes on it.

    ``backgrounds[i]`` is angle *i*'s own background (v7/v11) or the file's
    single background repeated for every angle (v6/v10), so a reader never
    has to branch on the version to subtract the right one.
    """
    path: Path
    version: int = field(init=False)
    header: ScanHeader = field(init=False)
    per_angle: list[AngleGeometry] = field(init=False)
    preambles: list[str] = field(init=False)
    preambles_raw: list[bytes] = field(init=False)
    backgrounds: list[bytes] = field(init=False)
    data_start_offset: int = field(init=False)
    file_size: int = field(init=False)

    def __post_init__(self):
        self.path = Path(self.path)
        self._mmap: mmap.mmap | None = None
        self._parse()

    def _parse(self):
        self.file_size = self.path.stat().st_size
        with open(self.path, "rb") as f:
            raw = f.read(HDR_SIZE)
            if len(raw) < HDR_SIZE:
                raise ValueError(f"{self.path.name}: file too short to contain a valid header")
            (magic, version, n_angles, x_start_nominal, y_start_nominal,
             x_delta_nominal, y_delta_nominal, row_spacing, velocity, laser_freq,
             samples_per_frame, sample_rate, bytes_per_sample,
             n_channels) = struct.unpack(HDR_FMT, raw)
            if magic != MAGIC:
                raise ValueError(f"{self.path.name}: not a valid SRAS file (bad magic)")
            if version not in SUPPORTED_VERSIONS:
                raise ValueError(
                    f"{self.path.name}: unsupported SRAS format version {version} "
                    f"(supported: {', '.join(str(v) for v in SUPPORTED_VERSIONS)})"
                )
            self.version = version
            self.header = ScanHeader(
                n_angles=n_angles,
                x_start_nominal=x_start_nominal, y_start_nominal=y_start_nominal,
                x_delta_nominal=x_delta_nominal, y_delta_nominal=y_delta_nominal,
                row_spacing=row_spacing, velocity=velocity, laser_freq=laser_freq,
                samples_per_frame=samples_per_frame, sample_rate=sample_rate,
                bytes_per_sample=bytes_per_sample, n_channels=n_channels,
            )

            angles = struct.unpack(f">{n_angles}f", f.read(4 * n_angles))

            self.per_angle = []
            for a in angles:
                x_start, x_delta, n_frames, n_rows = struct.unpack(GEOM_FMT, f.read(GEOM_SIZE))
                self.per_angle.append(AngleGeometry(
                    angle_deg=a, x_start=x_start, x_delta=x_delta,
                    n_frames=n_frames, n_rows=n_rows,
                ))

            for pa in self.per_angle:
                pa.y_positions = list(struct.unpack(f">{pa.n_rows}f", f.read(4 * pa.n_rows)))

            self.preambles_raw = []
            for _ in range(n_channels):
                (plen,) = struct.unpack(">H", f.read(2))
                self.preambles_raw.append(f.read(plen))
            self.preambles = [p.decode("utf-8", errors="replace") for p in self.preambles_raw]

            shared_bg = None
            if self.is_legacy_layout:
                # v6/v10: one background for the whole file, ahead of the data.
                (n_bg,) = struct.unpack(BG_LEN_FMT, f.read(BG_LEN_SIZE))
                shared_bg = f.read(n_bg)

            self.data_start_offset = f.tell()
            self._walk_data_block(f, shared_bg)

    def _walk_data_block(self, f, shared_bg: bytes | None):
        """Locate every angle's background block and the rows behind it.

        v7 interleaves a length-prefixed background ahead of each angle's
        rows, so the offsets are no longer pure arithmetic over the geometry
        table — the walk reads each prefix as it goes.  Past the frontier of
        a partial file there is nothing to read, so the remaining offsets are
        predicted from the block a writer would have produced (a full record),
        which is exactly where a resumed scan writes.
        """
        expected_bg = BG_LEN_SIZE + self.header.samples_per_frame
        self.backgrounds, self._bg_present = [], []
        self._bg_offsets, self._data_offsets = [], []
        cursor = self.data_start_offset
        for ai, pa in enumerate(self.per_angle):
            if shared_bg is not None:
                bg, bg_bytes, present = shared_bg, 0, True
            else:
                bg, bg_bytes, present = self._read_background(f, cursor, expected_bg)
            self.backgrounds.append(bg)
            self._bg_present.append(present)
            self._bg_offsets.append(cursor)
            cursor += bg_bytes
            self._data_offsets.append(cursor)
            cursor += self.row_bytes(ai) * pa.n_rows

    def _read_background(self, f, offset: int, expected_bytes: int):
        """One angle's background block as (waveform, block_bytes, present).

        A block that runs past the end of the file was never written: the
        walk keeps going with the size a writer would have used, and the
        angle is reported MISSING.
        """
        if offset + BG_LEN_SIZE > self.file_size:
            return b"", expected_bytes, False
        f.seek(offset)
        (n_bg,) = struct.unpack(BG_LEN_FMT, f.read(BG_LEN_SIZE))
        if offset + BG_LEN_SIZE + n_bg > self.file_size:
            return b"", expected_bytes, False
        return f.read(n_bg), BG_LEN_SIZE + n_bg, True

    @property
    def is_saw_check(self) -> bool:
        """True for a middle-row SAW quality check rather than a scan."""
        return self.version in SAW_CHECK_VERSIONS

    @property
    def is_legacy_layout(self) -> bool:
        """True for v6/v10: one background for the file, not one per angle."""
        return self.version in LEGACY_VERSIONS

    # ── Frontier / truncation analysis ───────────────────────────────────────

    def row_bytes(self, angle_idx: int) -> int:
        pa = self.per_angle[angle_idx]
        return (self.header.n_channels * pa.n_frames
                * self.header.samples_per_frame * self.header.bytes_per_sample)

    def angle_status(self) -> list[AngleStatus]:
        """Walk declared per-row byte counts against the actual file size.

        Because the data is one contiguous append-only stream, once an angle
        is found short every later angle is necessarily absent too — there is
        a single frontier past which nothing has been written yet.  An angle
        whose background block never made it to disk is short by definition,
        even though no row of it was due yet.
        """
        statuses = []
        frontier_seen = False
        for ai, pa in enumerate(self.per_angle):
            row_bytes = self.row_bytes(ai)
            data_offset = self._data_offsets[ai]
            if frontier_seen or not self._bg_present[ai]:
                n_rows_available = 0
                status = STATUS_MISSING
                frontier_seen = True
            else:
                declared_bytes = row_bytes * pa.n_rows
                if row_bytes > 0 and data_offset + declared_bytes <= self.file_size:
                    n_rows_available = pa.n_rows
                    status = STATUS_OK
                else:
                    remaining = max(0, self.file_size - data_offset)
                    n_rows_available = remaining // row_bytes if row_bytes > 0 else 0
                    status = STATUS_MISSING if n_rows_available == 0 else STATUS_TRUNCATED
                    frontier_seen = True
            statuses.append(AngleStatus(
                index=ai, angle_deg=pa.angle_deg, n_rows=pa.n_rows,
                row_bytes=row_bytes, bg_offset=self._bg_offsets[ai],
                data_offset=data_offset,
                n_rows_available=n_rows_available, status=status,
            ))
        return statuses

    def angle_data_offset(self, angle_idx: int) -> int:
        """Where angle ``angle_idx``'s rows start (past its background)."""
        return self._data_offsets[angle_idx]

    def angle_block_offset(self, angle_idx: int) -> int:
        """Where angle ``angle_idx``'s block starts, background included.

        Equal to ``angle_data_offset`` on v6/v10, which have no per-angle
        background block.
        """
        return self._bg_offsets[angle_idx]

    # ── Lazy data access ─────────────────────────────────────────────────────

    def _ensure_mmap(self) -> mmap.mmap:
        if self._mmap is None:
            # The mapping stays valid after the file object is closed, so
            # don't hold the descriptor open for the (long) life of a viewer
            # session.
            with open(self.path, "rb") as f:
                self._mmap = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
        return self._mmap

    def _dtype(self) -> np.dtype:
        return np.dtype(np.int16 if self.header.bytes_per_sample == 2 else np.int8)

    def background_array(self, angle_idx: int) -> np.ndarray | None:
        """One angle's background as float32 ADC counts, or None if absent."""
        bg = np.frombuffer(self.backgrounds[angle_idx], dtype=self._dtype())
        return bg.astype(np.float32) if bg.size else None

    def load_angle(self, angle_idx: int, n_rows: int | None = None) -> np.ndarray:
        """Read-only view of one angle's data block, shape
        (n_rows, n_channels, n_frames, samples_per_frame).

        ``n_rows`` limits the view to the rows actually on disk (pass
        ``AngleStatus.n_rows_available`` for truncated files); default is the
        declared row count.
        """
        pa = self.per_angle[angle_idx]
        h = self.header
        if n_rows is None:
            n_rows = pa.n_rows
        start = self.angle_data_offset(angle_idx)
        count = n_rows * h.n_channels * pa.n_frames * h.samples_per_frame
        arr = np.frombuffer(self._ensure_mmap(), dtype=self._dtype(),
                            count=count, offset=start)
        arr = arr.reshape(n_rows, h.n_channels, pa.n_frames, h.samples_per_frame)
        arr.flags.writeable = False
        return arr

    def load_row(self, angle_idx: int, row: int, channel_idx: int) -> np.ndarray:
        """Read-only view of one row/channel, shape (n_frames, samples_per_frame)."""
        pa = self.per_angle[angle_idx]
        h = self.header
        ch_bytes = pa.n_frames * h.samples_per_frame * h.bytes_per_sample
        start = (self.angle_data_offset(angle_idx) + row * self.row_bytes(angle_idx)
                 + channel_idx * ch_bytes)
        arr = np.frombuffer(self._ensure_mmap(), dtype=self._dtype(),
                            count=pa.n_frames * h.samples_per_frame, offset=start)
        arr = arr.reshape(pa.n_frames, h.samples_per_frame)
        arr.flags.writeable = False
        return arr

    def close(self):
        """Release this file's hold on the mapping.

        Views handed out earlier stay valid — they keep the mapping alive
        until they are garbage-collected, at which point the OS frees it.
        """
        if self._mmap is not None:
            try:
                self._mmap.close()
            except BufferError:
                pass  # live numpy views still reference the buffer
            self._mmap = None

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    # ── Axes helpers (viewer conveniences, derived from header fields) ───────

    def pixel_pitch_x_mm(self) -> float:
        """Distance between adjacent frames along X."""
        return self.header.velocity / self.header.laser_freq

    def x_axis_mm(self, angle_idx: int) -> np.ndarray:
        pa = self.per_angle[angle_idx]
        return pa.x_start + np.arange(pa.n_frames) * self.pixel_pitch_x_mm()

    def time_axis_ns(self) -> np.ndarray:
        h = self.header
        return np.arange(h.samples_per_frame) / h.sample_rate * 1e9

    def freq_axis_mhz(self, nfft: int) -> np.ndarray:
        return np.fft.rfftfreq(nfft, d=1.0 / self.header.sample_rate) / 1e6


def plan_from_header(sras: SrasFile) -> ScanPlan:
    """Reconstruct the ScanPlan a file was written with (for resume)."""
    h = sras.header
    return ScanPlan(
        x_start_nominal=h.x_start_nominal, y_start_nominal=h.y_start_nominal,
        x_delta_nominal=h.x_delta_nominal, y_delta_nominal=h.y_delta_nominal,
        row_spacing=h.row_spacing,
        velocity_mm_s=h.velocity, laser_freq_hz=h.laser_freq,
        per_angle=list(sras.per_angle),
    )