#!/usr/bin/env python3
"""Load a periodic pulse into the AD9102 SRAM over the UART protocol.

Builds the BEGIN (0xCCCC) / DATA (0xDDDD) / COMMIT packet sequence for the
custom-waveform upload path. Pattern = <width> of full-scale top samples
followed by base samples; BEGIN param1 carries the repetition period in
MICROSECONDS (firmware converts it to PAT_PERIOD/PATTERN_PERIOD_BASE for
its 150 MHz DAC clock; 0 = legacy back-to-back playback).

Example (1 us pulse at 1 kHz):
    python3 ad9102_pulse.py --port /dev/ttyUSB0 --width-us 1 --period-us 1000
Use --dry-run to print the packets without opening a port.
"""

import argparse
import sys

HEADER_WAVE_CONTROL = 0xCCCC
HEADER_WAVE_DATA = 0xDDDD
OPCODE_BEGIN = 0x0001
OPCODE_COMMIT = 0x0002
OPCODE_CANCEL = 0x0003

DAC_CLOCK_HZ = 150_000_000  # board oscillator on AD9102 CLKIN
DATA_WORDS = 14             # count + 12 samples + checksum
MAX_CHUNK_SAMPLES = 12
MAX_SRAM_SAMPLES = 4096
MAX_PERIOD_US = 6553        # 65535 ticks x base 15 at 150 MHz
SAMPLE_MIN = -8192
SAMPLE_MAX = 8191


def xor_checksum(words):
    checksum = words[0]
    for word in words[1:]:
        checksum ^= word
    return checksum


def pack_packet(header, words):
    words = list(words) + [xor_checksum(words)]
    raw = bytearray()
    raw += header.to_bytes(2, "little")
    for word in words:
        raw += (word & 0xFFFF).to_bytes(2, "little")
    return bytes(raw)


def build_packets(samples, period_us):
    packets = [pack_packet(HEADER_WAVE_CONTROL,
                           [OPCODE_BEGIN, len(samples), period_us])]
    for offset in range(0, len(samples), MAX_CHUNK_SAMPLES):
        chunk = samples[offset:offset + MAX_CHUNK_SAMPLES]
        words = [len(chunk)] + [s & 0xFFFF for s in chunk]
        words += [0] * (DATA_WORDS - 1 - len(words))
        packets.append(pack_packet(HEADER_WAVE_DATA, words))
    packets.append(pack_packet(HEADER_WAVE_CONTROL, [OPCODE_COMMIT, 0, 0]))
    return packets


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--port", help="serial port, e.g. /dev/ttyUSB0 or COM5")
    parser.add_argument("--baud", type=int, default=115200)
    parser.add_argument("--width-us", type=float, required=True,
                        help="pulse width, microseconds")
    parser.add_argument("--period-us", type=int, required=True,
                        help="repetition period, integer microseconds")
    parser.add_argument("--top", type=int, default=SAMPLE_MAX,
                        help="pulse top code, -8192..8191 (default max)")
    parser.add_argument("--base", type=int, default=SAMPLE_MIN,
                        help="baseline code, -8192..8191 (default min)")
    parser.add_argument("--dry-run", action="store_true",
                        help="print packets, do not open the port")
    args = parser.parse_args()

    width_samples = max(1, round(DAC_CLOCK_HZ * args.width_us * 1e-6))
    tail_samples = max(1, min(8, MAX_SRAM_SAMPLES - width_samples))
    total_samples = width_samples + tail_samples

    if not (SAMPLE_MIN <= args.base <= SAMPLE_MAX
            and SAMPLE_MIN <= args.top <= SAMPLE_MAX):
        sys.exit("top/base out of DAC range -8192..8191")
    if total_samples > MAX_SRAM_SAMPLES:
        sys.exit(f"pulse too long: {width_samples} samples > SRAM {MAX_SRAM_SAMPLES}")
    if not 0 <= args.period_us <= MAX_PERIOD_US:
        sys.exit(f"period must be 0..{MAX_PERIOD_US} us")
    pattern_us = total_samples * 1e6 / DAC_CLOCK_HZ
    if args.period_us and args.period_us <= pattern_us:
        print(f"note: period <= pattern ({pattern_us:.2f} us) -> back-to-back playback")

    samples = [args.top] * width_samples + [args.base] * tail_samples
    packets = build_packets(samples, args.period_us)

    print(f"pulse:  {width_samples} samples = {width_samples / DAC_CLOCK_HZ * 1e6:.3f} us")
    if args.period_us:
        print(f"period: {args.period_us} us ({1e6 / args.period_us:.1f} Hz)")
    print(f"pattern {total_samples} samples, {len(packets)} packets")

    if args.dry_run or not args.port:
        for packet in packets:
            print(packet.hex(" "))
        return

    import serial  # pyserial
    with serial.Serial(args.port, args.baud, timeout=1) as link:
        for index, packet in enumerate(packets):
            link.write(packet)
            status = link.read(2)
            if len(status) < 2:
                sys.exit(f"packet {index}: no status response")
            if status[0] != 0:
                sys.exit(f"packet {index}: device error flags 0x{status[0]:02X}")
            print(f"packet {index}: ok, status {status.hex(' ')}")
        print("committed; RUN bit:", "on" if status[1] & 0x01 else "off")


if __name__ == "__main__":
    main()