#!/usr/bin/env python3
"""
Continuous hardware-timed acquisition with a rolling CSV writer.

Designed for unattended logging: it keeps the acquisition buffer small, writes
in blocks, flushes on a timer, and stops cleanly on Ctrl+C or after a fixed
duration.

    pip install nidaqmx
    python python-continuous-acquisition.py --device Dev1 --channels ai0:7 \
        --rate 50000 --duration 60 --out run01.csv

Works with every model on the site. Set --rate to something the device can
actually sustain: 20 / 50 / 100 kS/s aggregate for the USB-6001 / 6002 / 6003,
250 kS/s for the USB-6229, and 102.4 kS/s per channel for the USB-4431 series.
"""

from __future__ import annotations

import argparse
import csv
import sys
import time
from datetime import datetime

try:
    import numpy as np
    import nidaqmx
    from nidaqmx.constants import AcquisitionType, TerminalConfiguration
    from nidaqmx.stream_readers import AnalogMultiChannelReader
except ImportError:  # pragma: no cover
    sys.exit("Requires nidaqmx and numpy. Run: pip install nidaqmx numpy")


def channel_names(channels: str, device: str) -> list[str]:
    """Expand 'ai0:3' or 'ai0,ai2,ai5' into physical channel names."""
    names: list[str] = []
    for part in channels.split(","):
        part = part.strip()
        if ":" in part:
            head, tail = part.split(":", 1)
            prefix = "".join(c for c in head if not c.isdigit())
            start = int("".join(c for c in head if c.isdigit()))
            end = int("".join(c for c in tail if c.isdigit()))
            names += [f"{prefix}{i}" for i in range(start, end + 1)]
        elif part:
            names.append(part)
    return [f"{device}/{n}" for n in names]


def main() -> int:
    parser = argparse.ArgumentParser(description="Continuous acquisition to CSV")
    parser.add_argument("--device", default="Dev1")
    parser.add_argument("--channels", default="ai0:7")
    parser.add_argument("--rate", type=float, default=50_000.0, help="aggregate sample rate, S/s")
    parser.add_argument("--block", type=int, default=5_000, help="samples per channel per read")
    parser.add_argument("--duration", type=float, default=30.0, help="seconds; 0 = until Ctrl+C")
    parser.add_argument("--out", default=None, help="output CSV path")
    parser.add_argument("--min", dest="min_val", type=float, default=-10.0)
    parser.add_argument("--max", dest="max_val", type=float, default=10.0)
    args = parser.parse_args()

    names = channel_names(args.channels, args.device)
    out_path = args.out or f"daq-{datetime.now():%Y%m%d-%H%M%S}.csv"

    print(f"Channels : {len(names)} ({names[0]} … {names[-1]})")
    print(f"Rate     : {args.rate:,.0f} S/s aggregate")
    print(f"Output   : {out_path}")
    print(f"Duration : {'until Ctrl+C' if args.duration <= 0 else f'{args.duration:g} s'}")

    buffer = np.zeros((len(names), args.block), dtype=np.float64)
    total = 0
    deadline = time.monotonic() + args.duration if args.duration > 0 else None

    with nidaqmx.Task() as task, open(out_path, "w", newline="", encoding="utf-8") as fh:
        for name in names:
            task.ai_channels.add_ai_voltage_chan(
                name,
                terminal_config=TerminalConfiguration.RSE,
                min_val=args.min_val,
                max_val=args.max_val,
            )
        task.timing.cfg_samp_clk_timing(
            rate=args.rate,
            sample_mode=AcquisitionType.CONTINUOUS,
            samps_per_chan=args.block * 4,
        )

        writer = csv.writer(fh)
        writer.writerow(["timestamp_s"] + [n.split("/", 1)[1] for n in names])

        reader = AnalogMultiChannelReader(task.in_stream)
        started = time.monotonic()

        try:
            for block_index in range(1_000_000):
                reader.read_many_sample(
                    buffer,
                    number_of_samples_per_channel=args.block,
                    timeout=10.0,
                )

                first = block_index * args.block / args.rate
                step = 1.0 / args.rate
                for i in range(args.block):
                    writer.writerow(
                        [f"{first + i * step:.6f}"] + [f"{buffer[c, i]:.6f}" for c in range(len(names))]
                    )

                total += args.block
                fh.flush()

                elapsed = time.monotonic() - started
                print(
                    f"\r  {total:,} samples  ({elapsed:6.1f} s, "
                    f"{total / max(elapsed, 1e-9):,.0f} S/s effective)",
                    end="",
                    flush=True,
                )

                if deadline and time.monotonic() >= deadline:
                    break
        except KeyboardInterrupt:
            print("\n  stopped by user")

    print(f"\nWrote {total:,} samples to {out_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())