from machine import Pin, UART
import struct
import time

# Wiring:
# IMU TX -> ESP32 GPIO16
# IMU RX -> ESP32 GPIO17
# IMU GND -> ESP32 GND
IMU_RX_PIN = 16
IMU_TX_PIN = 17

# CP2102 USB serial. Binary output is used because JSON is too large at 200 Hz.
USB_BAUD = 230400

DIN_PINS = (25, 26, 27, 32)
digital_inputs = [Pin(pin, Pin.IN, Pin.PULL_UP) for pin in DIN_PINS]

# Two controllable digital outputs. Both default to LOW at boot.
DOUT_PINS = (18, 19)
digital_outputs = [Pin(pin, Pin.OUT, value=0) for pin in DOUT_PINS]

FRAME_LEN = 11
HEADER = 0x55

# WIT normal-protocol frame types
TYPE_ACC = 0x51
TYPE_GYRO = 0x52
TYPE_ANGLE = 0x53

# Latest raw values: acc XYZ, gyro XYZ, angle XYZ, temperature
raw_values = [0] * 10
rx_buffer = bytearray()
command_buffer = bytearray()


def write_register(uart, address, value):
    uart.write(bytes((
        0xFF,
        0xAA,
        address & 0xFF,
        value & 0xFF,
        (value >> 8) & 0xFF,
    )))


def has_valid_frame(uart, timeout_ms=400):
    test_buffer = bytearray()
    deadline = time.ticks_add(time.ticks_ms(), timeout_ms)

    while time.ticks_diff(deadline, time.ticks_ms()) > 0:
        count = uart.any()
        if count:
            chunk = uart.read(count)
            if chunk:
                test_buffer.extend(chunk)

        while len(test_buffer) >= FRAME_LEN:
            if test_buffer[0] != HEADER:
                del test_buffer[0]
                continue

            frame = test_buffer[:FRAME_LEN]
            if (sum(frame[:10]) & 0xFF) == frame[10]:
                return True

            del test_buffer[0]

        time.sleep_ms(1)

    return False


def open_imu_uart(baudrate):
    return UART(
        2,
        baudrate=baudrate,
        bits=8,
        parity=None,
        stop=1,
        rx=IMU_RX_PIN,
        tx=IMU_TX_PIN,
        timeout=0,
        rxbuf=2048,
    )


def start_imu_at_200hz():
    # If data already arrives at 115200, enforce 200 Hz for this session.
    # No SAVE is needed here, avoiding repeated writes to sensor flash.
    uart = open_imu_uart(115200)
    if has_valid_frame(uart):
        write_register(uart, 0x69, 0xB588)
        time.sleep_ms(200)
        write_register(uart, 0x02, 0x000E)  # ACC + GYRO + ANGLE only
        time.sleep_ms(50)
        write_register(uart, 0x03, 0x000B)  # RRATE = 200 Hz
        time.sleep_ms(100)
        return uart

    # Factory/default examples use 9600 baud.
    uart.deinit()
    uart = open_imu_uart(9600)

    if not has_valid_frame(uart):
        raise RuntimeError("IMU not detected at 9600 or 115200 baud")

    # Official protocol sequence:
    # unlock -> change baud -> switch local UART -> unlock -> set 200 Hz -> save
    write_register(uart, 0x69, 0xB588)
    time.sleep_ms(200)
    write_register(uart, 0x04, 0x0006)  # BAUD = 115200
    time.sleep_ms(100)

    uart.deinit()
    uart = open_imu_uart(115200)
    time.sleep_ms(200)

    write_register(uart, 0x69, 0xB588)
    time.sleep_ms(200)
    write_register(uart, 0x02, 0x000E)  # ACC + GYRO + ANGLE only
    time.sleep_ms(50)
    write_register(uart, 0x03, 0x000B)  # RRATE = 200 Hz
    time.sleep_ms(100)
    write_register(uart, 0x00, 0x0000)  # SAVE
    time.sleep_ms(500)

    return uart


def update_raw_values(frame):
    frame_type = frame[1]
    x, y, z, extra = struct.unpack_from("<hhhh", frame, 2)

    if frame_type == TYPE_ACC:
        raw_values[0:3] = (x, y, z)
        raw_values[9] = extra
    elif frame_type == TYPE_GYRO:
        raw_values[3:6] = (x, y, z)
    elif frame_type == TYPE_ANGLE:
        raw_values[6:9] = (x, y, z)
        return True

    return False


def poll_imu(uart):
    count = uart.any()
    if count:
        chunk = uart.read(count)
        if chunk:
            rx_buffer.extend(chunk)

    angle_updated = False

    while len(rx_buffer) >= FRAME_LEN:
        if rx_buffer[0] != HEADER:
            del rx_buffer[0]
            continue

        frame = rx_buffer[:FRAME_LEN]
        if (sum(frame[:10]) & 0xFF) != frame[10]:
            del rx_buffer[0]
            continue

        del rx_buffer[:FRAME_LEN]
        if update_raw_values(frame):
            angle_updated = True

    return angle_updated


def make_packet():
    din_mask = 0
    for index, pin in enumerate(digital_inputs):
        din_mask |= pin.value() << index

    dout_mask = 0
    for index, pin in enumerate(digital_outputs):
        dout_mask |= pin.value() << index

    # Packet layout:
    # A5 5A | version | DIN | DOUT | ticks_ms | 10 signed int16 | checksum
    packet = struct.pack(
        "<2sBBBI10h",
        b"\xA5\x5A",
        2,
        din_mask,
        dout_mask,
        time.ticks_ms() & 0xFFFFFFFF,
        *raw_values
    )
    return packet + bytes((sum(packet) & 0xFF,))


def set_gripper(mode):
    # GPIO HIGH means the external transistor/optocoupler is conducting.
    # Always release both inputs briefly before changing direction.
    digital_outputs[0].value(0)
    digital_outputs[1].value(0)
    time.sleep_ms(20)

    if mode == 1:       # OPEN: line 2 = 12 V, line 3 = 0 V
        digital_outputs[1].value(1)
    elif mode == 2:     # CLOSE: line 2 = 0 V, line 3 = 12 V
        digital_outputs[0].value(1)


def poll_usb_commands(uart):
    count = uart.any()
    if count:
        chunk = uart.read(count)
        if chunk:
            command_buffer.extend(chunk)

    # Command: C3 3C | mode | reserved | checksum
    # mode: 0=STOP, 1=OPEN, 2=CLOSE
    while len(command_buffer) >= 5:
        if command_buffer[0] != 0xC3 or command_buffer[1] != 0x3C:
            del command_buffer[0]
            continue

        command = command_buffer[:5]
        if (sum(command[:4]) & 0xFF) != command[4]:
            del command_buffer[0]
            continue

        del command_buffer[:5]
        mode = command[2]
        if mode <= 2:
            set_gripper(mode)


imu_uart = start_imu_at_200hz()
usb_uart = UART(0, baudrate=USB_BAUD, bits=8, parity=None, stop=1, tx=1, rx=3)

while True:
    poll_usb_commands(usb_uart)

    # The angle frame is used as the 200 Hz snapshot trigger.
    if poll_imu(imu_uart):
        usb_uart.write(make_packet())

    time.sleep_ms(0)