import serial
import time
import struct
import matplotlib.pyplot as plt

# --- Serial Setup ---
ser = serial.Serial('/dev/ttyACM2', 115200, timeout=1) # Added timeout for safety

# Trigger the robot to send data
message = struct.pack(">bb", - ord('U'), 0)
ser.write(message)

reply = b""
expected_bytes = 1280 # 160 samples * 4 mics * 2 bytes
while len(reply) < expected_bytes:
    chunk = ser.read(expected_bytes - len(reply))
    if not chunk: # Break if timeout occurs to avoid infinite loop
        break
    reply += chunk

ser.close()

if len(reply) < expected_bytes:
    print(f"Error: Only received {len(reply)} bytes.")
else:
    # --- Data Unpacking ---
    # ">" = Big Endian (consistent with your command). Use "<" if the plot looks like noise.
    # "640h" = 640 signed shorts (16-bit)
    all_samples = struct.unpack(f"<640h", reply)

    # De-interleave: right, left, back, front
    right_mic = all_samples[0::4]
    left_mic  = all_samples[1::4]
    back_mic  = all_samples[2::4]
    front_mic = all_samples[3::4]

    # --- Plotting ---
    mics = [right_mic, left_mic, back_mic, front_mic]
    names = ["Right Mic", "Left Mic", "Back Mic", "Front Mic"]
    colors = ['r', 'g', 'b', 'm']

    fig, axs = plt.subplots(4, 1, figsize=(10, 8), sharex=True)
    fig.suptitle('Microphone Array Samples')

    for i in range(4):
        axs[i].plot(mics[i], color=colors[i])
        axs[i].set_ylabel('Amplitude')
        axs[i].set_title(names[i])
        axs[i].grid(True)

    plt.xlabel('Sample Index')
    plt.tight_layout(rect=[0, 0.03, 1, 0.95])
    plt.show()

    # --- CSV Fallback (Optional) ---
    import csv
    with open('mic_data.csv', 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(["Right", "Left", "Back", "Front"])
        writer.writerows(zip(right_mic, left_mic, back_mic, front_mic))