dataCollection.py

import argparse
import serial
import os
import numpy as np
import matplotlib.pyplot as plt
import pickle
import signal
import sys
import time


def bytes2samples(byte_list, samp_per_segment, bytes_per_sample):
    """
    Parse a list of byte values into samples.

    Parameters
    ----------
    byte_list : list of int
        Raw bytes read from file (each element 0-255).
    samp_per_segment : int
        Number of samples per segment (e.g. 800).
    bytes_per_sample : int
        Number of bytes per sample (e.g. 2).

    Returns
    -------
    samples : numpy.ndarray
        Parsed sample array.
    """
    bytes_per_segment = samp_per_segment * bytes_per_sample

    # Compute number of full segments and trim
    n_segments = len(byte_list) // bytes_per_segment
    byte_list = byte_list[: n_segments * bytes_per_segment]

    total_bytes = bytes_per_segment * n_segments

    # Preallocate
    data = np.zeros(samp_per_segment * n_segments, dtype=np.float64)

    samp_ctr = 0
    i = 0
    while i < total_bytes - 1:
        byte1 = byte_list[i]
        byte2 = byte_list[i + 1]
        sample = byte1 * 256 + byte2

        if sample < 7000:
            data[samp_ctr] = sample
            samp_ctr += 1
            i += 2
        else:
            i += 1

    data = data[:samp_ctr]

    flags = np.where(data > 5000)[0]
    print(f"Found {len(flags)} pulses")

    return data


def collect_data(port, filename, baudrate, buffer_size, plot_results, output_format):
    """
    Open a serial port, collect data until the user presses Ctrl+C,
    then process and save the results.

    Parameters
    ----------
    port : str
        Serial port name (e.g. 'COM3' or '/dev/ttyUSB0').
    filename : str
        Base filename for saving data.
    baudrate : int
        Serial baud rate.
    buffer_size : int
        Number of bytes to read per chunk.
    plot_results : bool
        Whether to plot the processed samples after collection.
    output_format : str
        Output format: 'pkl' or 'mat'.
    """
    raw_filename = filename + ".raw"
    stop_flag = False

    def signal_handler(sig, frame):
        nonlocal stop_flag
        stop_flag = True
        print("\nStopping collection...")

    signal.signal(signal.SIGINT, signal_handler)

    # Open serial port
    ser = serial.Serial(
        port=port,
        baudrate=baudrate,
        bytesize=serial.EIGHTBITS,
        stopbits=serial.STOPBITS_ONE,
        parity=serial.PARITY_NONE,
    )

    # Windows only
    # ser.set_buffer_size(rx_size=buffer_size)

    fid = open(raw_filename, "wb")

    print(f"Collecting data on {port} at {baudrate} baud...")
    print("Press Ctrl+C to stop collection.\n")

    # Collection loop
    while not stop_flag:
        if ser.in_waiting >= buffer_size:
            data = ser.read(buffer_size)
            fid.write(data)
        else:
            time.sleep(0.001)

    ser.close()
    fid.close()
    print("Serial port closed.")

    # Read back raw file and process
    with open(raw_filename, "rb") as f:
        file_bytes = list(f.read())

    if len(file_bytes) == 0:
        print("No data collected.")
        os.remove(raw_filename)
        return

    samp_per_pulse = 800
    bytes_per_samp = 2
    samples = bytes2samples(file_bytes, samp_per_pulse, bytes_per_samp)

    # Remove temporary raw file
    os.remove(raw_filename)

    # Save processed data
    if output_format == "mat":
        from scipy.io import savemat
        save_path = filename + ".mat"
        savemat(save_path, {"samples": samples})
    else:
        save_path = filename + ".pkl"
        with open(save_path, "wb") as f:
            pickle.dump(samples, f)

    print(f"Saved {len(samples)} samples to {save_path}")

    # Plot if requested
    if plot_results:
        plt.figure()
        plt.plot(samples)
        plt.title("Collected Samples")
        plt.xlabel("Sample Index")
        plt.ylabel("Value")
        plt.show()


def main():
    parser = argparse.ArgumentParser(
        description="Serial data collection tool. Collects raw bytes from a serial "
                    "port, processes them into samples, and saves the result."
    )

    parser.add_argument(
        "-p", "--port",
        type=str,
        required=True,
        help="Serial port name (e.g. COM3 or /dev/ttyACM0)."
    )
    parser.add_argument(
        "-f", "--filename",
        type=str,
        default="Test",
        help="Base filename for saving data (default: Test)."
    )
    parser.add_argument(
        "-b", "--baudrate",
        type=int,
        default=921600,
        help="Serial baud rate (default: 921600)."
    )
    parser.add_argument(
        "-s", "--buffer-size",
        type=int,
        default=800 * 2,
        help="Read buffer size in bytes (default: 1600)."
    )
    parser.add_argument(
        "--plot",
        action="store_true",
        default=False,
        help="Enable plotting after collection."
    )
    parser.add_argument(
        "-o", "--output-format",
        type=str,
        choices=["pkl", "mat"],
        default="pkl",
        help="Output file format: 'pkl' (pickle) or 'mat' (MATLAB). Default: pkl."
    )

    args = parser.parse_args()

    collect_data(
        port=args.port,
        filename=args.filename,
        baudrate=args.baudrate,
        buffer_size=args.buffer_size,
        plot_results=args.plot,
        output_format=args.output_format,
    )


if __name__ == "__main__":
    main()