"""Burst-mode FastFrame acquisition policy.

Per-row acquisition pays a full arm/stop/transfer round trip for every row,
and the transfer alone is one IEEE-488.2 block read per frame (~16k frames a
row).  A burst instead runs one FastFrame acquisition across as many complete
rows as the scope's frame memory holds, then pulls the whole thing in a single
transaction — amortising the round trip over `rows_per_burst` rows.

The scope reports its capacity with ``HORizontal:FASTframe:MAXFRames?`` once
the horizontal settings are fixed; ``rows_per_burst`` turns that into a row
count.  Everything here that computes rather than talks to hardware is a free
function, so the row-splitting logic is testable without a rig.

The catch is that the burst contains no row markers: the scope hands back one
flat run of frames.  Boundaries come from polling ``ACQuire:NUMFRAMESACQuired?``
after each row's acquiring pass, while the stage gate is already low — see
``split_row_counts``.
"""
from __future__ import annotations

import logging
import time

import numpy as np

logger = logging.getLogger(__name__)

# Peak transfer buffer, per channel.  The writer holds one channel at a time
# (see ScanEngine._scan_rows_burst), so this is the real high-water mark.
BURST_MEMORY_BUDGET_BYTES = 512 * 1024 * 1024

# Extra frames budgeted per row on top of n_frames.  0 gives the plain
# floor(max_frames / n_frames) row count; raise it if the acquiring pass
# routinely over-triggers (watch the pad/truncate warnings).
BURST_FRAME_HEADROOM = 0

BURST_ARM_SETTLE_S = 0.05   # after ACQuire:STATE RUN, before the first move
BURST_ROW_SETTLE_S = 0.05   # after the gate drops, before reading the counter


# ── Pure helpers ─────────────────────────────────────────────────────────────

def rows_per_burst(max_frames: int, n_frames: int, samples_per_frame: int,
                   rows_remaining: int,
                   memory_budget: int = BURST_MEMORY_BUDGET_BYTES,
                   headroom: int = BURST_FRAME_HEADROOM) -> int:
    """How many complete rows fit in one acquisition.

    Rounds down — a partial row is worthless, since a row must be transferred
    whole to be written.  Clamped by the transfer buffer budget and by the rows
    actually left in the angle, and never below 1 (a single row always goes,
    even if it exceeds the budget, so the scan can still make progress).
    """
    if n_frames < 1 or samples_per_frame < 1:
        raise ValueError(f"n_frames={n_frames} samples_per_frame={samples_per_frame}")

    by_scope = max_frames // (n_frames + headroom)
    by_memory = memory_budget // (n_frames * samples_per_frame)
    return max(1, min(by_scope, by_memory, rows_remaining))


def split_row_counts(cumulative: list[int]) -> list[int]:
    """Per-row frame counts from the cumulative counter sampled after each row.

    ``cumulative`` is ``ACQuire:NUMFRAMESACQuired?`` read once per row, already
    rebased on the value at burst start.
    """
    counts = []
    prev = 0
    for i, c in enumerate(cumulative):
        if c < prev:
            raise RuntimeError(
                f"FastFrame counter went backwards at row {i} ({prev} → {c}) — "
                "the acquisition was restarted mid-burst"
            )
        counts.append(c - prev)
        prev = c
    return counts


def normalize_row(buf, offset: int, count: int, n_frames: int,
                  samples_per_frame: int):
    """Coerce one row's frames to exactly ``n_frames``.

    The v6 format commits to n_frames per row in the header and has no per-row
    length field, so a row that over- or under-triggers must be squared up or
    every later row in the file shifts.  Short rows are zero-padded, long rows
    lose their trailing frames.  Returns something writable directly.
    """
    want = n_frames * samples_per_frame
    end = min(offset + count * samples_per_frame, offset + want, len(buf))
    chunk = memoryview(buf)[offset:end]
    if len(chunk) == want:
        return chunk
    return bytes(chunk) + bytes(want - len(chunk))


def frame_means_block(buf, offset: int, n_frames: int,
                      samples_per_frame: int) -> list[float]:
    """Per-frame DC mean over one row's slice of a burst buffer."""
    n = n_frames * samples_per_frame
    block = np.frombuffer(buf, dtype=np.int8, count=n, offset=offset)
    return block.reshape(n_frames, samples_per_frame).mean(
        axis=1, dtype=np.float32).tolist()


# ── Instrument control ───────────────────────────────────────────────────────

def max_frames(scope) -> int:
    """Frames the scope can hold under the current horizontal settings."""
    try:
        m = scope.get_fastframe_max_frames()
    except Exception as exc:
        raise RuntimeError(
            "Scope did not answer HORizontal:FASTframe:MAXFRames? — burst mode "
            "cannot size a burst without it.  Use per-row acquisition on this "
            f"firmware. ({exc})"
        ) from exc
    if m < 1:
        raise RuntimeError(f"Scope reports a FastFrame capacity of {m} frames")
    return m


def start_burst(scope, frame_count: int) -> int:
    """Arm one burst; returns the counter baseline to subtract from later reads.

    Reading the baseline back beats assuming the counter resets to 0 on RUN —
    any residual is simply subtracted out instead of being misattributed to the
    first row.
    """
    scope.set_fastframe_count(frame_count)
    scope.write("ACQuire:STATE RUN")
    time.sleep(BURST_ARM_SETTLE_S)
    return frames_acquired(scope)


def stop_burst(scope) -> None:
    time.sleep(BURST_ROW_SETTLE_S)
    scope.write("ACQuire:STATE STOP")


def frames_acquired(scope) -> int:
    return int(scope.query("ACQuire:NUMFRAMESACQuired?"))


def transfer_burst(scope, ch: int, frame_count: int, samples_per_frame: int):
    """Pull a whole burst for one channel in a single CURVe? transaction."""
    scope.set_data_source(ch)
    return scope.transfer_fastframe_bulk(frame_count, samples_per_frame)