#!/usr/bin/env python3
"""Generate small synthetic .sras files for testing.

Writes v6 files (per-angle geometry, ragged waveform blocks) matching
scan_format.md, with deterministic pseudo-random waveform content so a test
can compute expected DC/FFT images independently of the reader under test.

Usage:
    python tools/make_test_sras.py out.sras [--angles 3] [--seed 0]
"""

import argparse
import struct
from pathlib import Path

import numpy as np

HDR_FMT_V6 = ">4sBHfffffffIdBB"
GEO_FMT_V6 = ">ffIH"

# Per-angle (n_rows, n_frames) — deliberately different per angle so ragged
# geometry handling is actually exercised.
_GEOMETRY = [(5, 7), (4, 11), (6, 9), (3, 13), (7, 6)]

_SAMPLE_RATE_HZ = 6.25e9
_VELOCITY_MM_S = 20.0
_LASER_FREQ_HZ = 1000.0
_ROW_SPACING_MM = 0.05


def _preamble(ymult_v: float, yoff_adc: float, yzero_v: float) -> bytes:
    """A Tektronix WFMOutpre string in verbose (keyword) form — the reader
    pulls YMULT/YOFF/YZERO out of it by name, so the keywords must be
    present literally. YMULT/YZERO are in volts, as the scope reports them."""
    return (
        ":WFMOUTPRE:BYT_NR 1;BIT_NR 8;ENCDG BIN;BN_FMT RI;BYT_OR MSB;"
        'WFID "Ch1, DC coupling";NR_PT 2500;PT_FMT Y;'
        "XINCR 1.6000E-10;XZERO 0.0E0;XUNIT \"s\";"
        f"YMULT {ymult_v:.6E};YOFF {yoff_adc:.6E};YZERO {yzero_v:.6E};"
        'YUNIT "V"'
    ).encode("utf-8")


def build(n_angles: int, seed: int, samples_per_frame: int) -> tuple[bytes, dict]:
    rng = np.random.default_rng(seed)
    geom = [_GEOMETRY[a % len(_GEOMETRY)] for a in range(n_angles)]
    n_ch = 3
    bps = 1

    angles_deg = np.linspace(0.0, 60.0, n_angles, dtype=np.float32)
    # Distinct calibration per channel so a swapped-channel bug is visible.
    cal = [
        (1.5625e-3, -87.04, 0.0),
        (2.0000e-3, -60.00, 1.0e-3),
        (2.5000e-3, -40.00, -2.0e-3),
    ]

    out = bytearray()
    out += struct.pack(
        HDR_FMT_V6, b"SRAS", 6, n_angles,
        0.0, 0.0, 1.0, 1.0, _ROW_SPACING_MM,
        _VELOCITY_MM_S, _LASER_FREQ_HZ,
        samples_per_frame, _SAMPLE_RATE_HZ, bps, n_ch,
    )
    out += angles_deg.astype(">f4").tobytes()

    x_starts = []
    for a, (n_rows, n_frames) in enumerate(geom):
        x_start = -0.5 + 0.1 * a
        x_starts.append(x_start)
        out += struct.pack(GEO_FMT_V6, x_start, 1.0, n_frames, n_rows)

    y_positions = []
    for a, (n_rows, _) in enumerate(geom):
        y = (0.2 * a + np.arange(n_rows) * _ROW_SPACING_MM).astype(np.float32)
        y_positions.append(y)
        out += y.astype(">f4").tobytes()

    for ymult_v, yoff, yzero_v in cal:
        p = _preamble(ymult_v, yoff, yzero_v)
        out += struct.pack(">H", len(p)) + p

    background = rng.integers(-8, 9, size=samples_per_frame, dtype=np.int8)
    out += struct.pack(">I", samples_per_frame) + background.tobytes()

    # Waveform data. CH1 gets a sinusoid at a per-pixel frequency so the FFT
    # peak is predictable; CH3/CH4 get per-pixel DC levels so the mean is too.
    t = np.arange(samples_per_frame)
    waveforms = []
    for a, (n_rows, n_frames) in enumerate(geom):
        block = np.empty((n_rows, n_ch, n_frames, samples_per_frame), dtype=np.int8)
        for r in range(n_rows):
            for f in range(n_frames):
                bin_idx = 3 + ((a + r + f) % 17)
                phase = 2 * np.pi * bin_idx * t / samples_per_frame
                block[r, 0, f] = np.clip(
                    np.round(60 * np.sin(phase)), -128, 127).astype(np.int8)
                block[r, 1, f] = np.int8((a * 7 + r * 3 + f) % 100 - 50)
                block[r, 2, f] = np.int8((a * 5 + r * 11 + f * 2) % 120 - 60)
        waveforms.append(block)
        out += block.tobytes()

    meta = {
        "n_angles": n_angles,
        "geometry": geom,
        "angles_deg": angles_deg,
        "x_starts": x_starts,
        "y_positions": y_positions,
        "cal": cal,
        "background": background,
        "waveforms": waveforms,
        "samples_per_frame": samples_per_frame,
        "sample_rate_hz": _SAMPLE_RATE_HZ,
    }
    return bytes(out), meta


def write(path: Path, n_angles: int = 3, seed: int = 0,
          samples_per_frame: int = 64) -> dict:
    payload, meta = build(n_angles, seed, samples_per_frame)
    path.write_bytes(payload)
    return meta


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("output")
    p.add_argument("--angles", type=int, default=3)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--spf", type=int, default=64, help="samples per frame")
    args = p.parse_args()

    out = Path(args.output)
    meta = write(out, args.angles, args.seed, args.spf)
    print(f"Wrote {out} ({out.stat().st_size:,} bytes)")
    print(f"  angles   : {meta['n_angles']}")
    print(f"  geometry : {meta['geometry']}")
    print(f"  spf      : {meta['samples_per_frame']}")


if __name__ == "__main__":
    main()