#!/usr/bin/env python3
#datasheetPartOneStart
"""Capture the fixed Appendix B waveform and plot the result."""

import argparse
import time
from types import SimpleNamespace

import matplotlib.pyplot as plt
from pymodbus.client import ModbusSerialClient
from pymodbus.exceptions import ConnectionException
from pymodbus.exceptions import ModbusIOException

UNIT_ID = 1
SOURCE = 2                  # ADXL380
CHANNEL_MASK = 0x0007       # X, Y, Z
OUTPUT_RATE_HZ = 1000       # samples/s per axis
SAMPLES_PER_CHANNEL = 1024
FILTER_PROFILE = 3          # FIR decimation
DOWNLOAD_POLICY = 1         # enforce transfer budget
MAX_DOWNLOAD_MS = 30000
FULL_SCALE_G = 8
LSB_PER_G = 3750.0
CONFIGURED, ACQUIRING, PROCESSING, READY, FAILED = range(1, 6)
REQUEST = [SOURCE, CHANNEL_MASK, OUTPUT_RATE_HZ, 0,
           SAMPLES_PER_CHANNEL, FILTER_PROFILE, DOWNLOAD_POLICY,
           MAX_DOWNLOAD_MS, FULL_SCALE_G]
TRANSPORT_ERRORS = (ConnectionException, ModbusIOException)

def u32(high, low):
    return (high << 16) | low

def i16(word):
    return word - 0x10000 if word & 0x8000 else word

def open_modvibe(port, unit_id):
    global UNIT_ID
    UNIT_ID = unit_id
    mb = ModbusSerialClient(
        port, baudrate=115200, timeout=1.0, retries=0,
    )
    if not mb.connect():
        raise RuntimeError(f"cannot open {port}")
    return mb

def check(reply):
    if isinstance(reply, ModbusIOException):
        raise reply
    if reply.isError():
        raise RuntimeError(f"Modbus exception: {reply}")
    return reply

def read_regs(mb, address, count):
    reply = mb.read_holding_registers(
        address, count=count, device_id=UNIT_ID,
    )
    words = check(reply).registers
    if words is None or len(words) != count:
        actual = 0 if words is None else len(words)
        raise RuntimeError(
            f"short register read at {address}: expected {count}, got {actual}"
        )
    return words

def read_state(mb):
    words = read_regs(mb, 202, 23)
    return SimpleNamespace(
        status=words[0], error=words[1], request=words[3:12],
        output_words=u32(*words[12:14]),
        page_count=u32(*words[14:16]),
        waveform_id=words[18], maximum_words=u32(*words[19:21]),
        payload_words=words[21], header_words=words[22])

def write_regs(mb, address, values):
    reply = mb.write_registers(address, values, device_id=UNIT_ID)
    check(reply)

def verify_device(mb):
    identity = read_regs(mb, 0, 14)
    if identity[:3] != [0x4649, 0x5642, 0x02]:
        raise RuntimeError("incompatible identity/map")
    release_contract = (
        u32(identity[9], identity[10]), identity[3],
        *read_regs(mb, 200, 2),
        read_regs(mb, 1200, 1)[0],
    )
    if release_contract == (0x00000102, 2, 4, 3, 2):
        return False
    if release_contract == (0x00000103, 3, 5, 3, 3):
        return True
    raise RuntimeError("incompatible firmware/contract tuple")

#datasheetPartOneEnd
#datasheetPartTwoStart
def configure(mb, require_cleared_id=True):
    write_regs(mb, 205, REQUEST)
    state = read_state(mb)
    if ((state.status, state.error) != (CONFIGURED, 0)
            or state.request != REQUEST):
        raise RuntimeError("configuration rejected")
    if require_cleared_id and state.waveform_id != 0:
        raise RuntimeError("configuration retained stale waveform ID")
    return state.waveform_id

def start_once(mb, previous_id, invalidates_id):
    for attempt in range(2):
        try:
            check(mb.write_register(204, 1, device_id=UNIT_ID))
            return
        except TRANSPORT_ERRORS:
            state = read_state(mb)
            status, error = state.status, state.error
            waveform_id = state.waveform_id
            if error == 0 and status in (ACQUIRING, PROCESSING):
                return
            ready_id_matches = (
                waveform_id != 0 if invalidates_id
                else waveform_id == (1 if previous_id == 0xffff
                                     else previous_id + 1)
            )
            if (error == 0 and status == READY and state.request == REQUEST
                    and ready_id_matches):
                return
            not_started = status == CONFIGURED and error == 0
            same_request = state.request == REQUEST
            unchanged = waveform_id == previous_id and same_request
            if attempt == 0 and not_started and unchanged:
                continue
            raise RuntimeError("ambiguous START outcome")

def wait_ready(mb, deadline_s, previous_id, invalidates_id):
    time.sleep(SAMPLES_PER_CHANNEL / OUTPUT_RATE_HZ)
    deadline = time.monotonic() + deadline_s
    while time.monotonic() < deadline:
        state = read_state(mb)
        status, error = state.status, state.error
        ready_id_matches = (
            state.waveform_id != 0 if invalidates_id
            else state.waveform_id == (1 if previous_id == 0xffff
                                       else previous_id + 1)
        )
        if (status == READY and error == 0 and state.request == REQUEST
                and ready_id_matches):
            return state
        if status == FAILED:
            raise RuntimeError(f"capture error {error}")
        if status not in (ACQUIRING, PROCESSING) or error != 0:
            raise RuntimeError(f"unexpected state {status}/{error}")
        time.sleep(0.1)
    raise TimeoutError("capture deadline exceeded")

def download(mb, state):
    if read_state(mb).request != REQUEST:
        raise RuntimeError("capture changed before download")
    if state.output_words != 3072 or state.maximum_words < 3072:
        raise RuntimeError("unexpected output size")
    if (state.payload_words, state.header_words) != (122, 3):
        raise RuntimeError("unsupported page format")
    if state.page_count != (state.output_words + 121) // 122:
        raise RuntimeError("inconsistent page count")

    words = []
    while len(words) < state.output_words:
        offset = len(words)
        for _ in range(3):
            try:
                seek = [(offset >> 16) & 0xFFFF, offset & 0xFFFF]
                write_regs(mb, 228, seek)
                page = read_regs(mb, 1000, 125)
            except TRANSPORT_ERRORS:
                continue
            if page[0] != state.waveform_id:
                raise RuntimeError("waveform ID changed")
            if u32(page[1], page[2]) == offset:
                break
        else:
            raise RuntimeError("page retry exhausted")
        count = min(122, state.output_words - offset)
        payload = page[3:3 + count]
        words.extend(i16(word) for word in payload)
    return state.waveform_id, words

def plot_waveform(filename, waveform_id, words):
    count = SAMPLES_PER_CHANNEL
    time_s = [sample / OUTPUT_RATE_HZ for sample in range(count)]
    figure, axis = plt.subplots(layout="constrained")
    for number, label in enumerate("XYZ"):
        data = words[number * count:(number + 1) * count]
        axis.plot(time_s, [value / LSB_PER_G for value in data],
                  label=label)
    axis.set(xlabel="Time (s)", ylabel="Acceleration (g)",
             title=f"ModVibe waveform {waveform_id}")
    axis.grid(alpha=0.25)
    axis.legend()
    figure.savefig(filename, dpi=150)
    plt.close(figure)

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("port")
    parser.add_argument("unit_id", type=int)
    args = parser.parse_args()
    mb = open_modvibe(args.port, args.unit_id)
    invalidates_id = verify_device(mb)
    previous_id = configure(mb, invalidates_id)
    start_once(mb, previous_id, invalidates_id)
    state = wait_ready(mb, 10.0, previous_id, invalidates_id)
    waveform_id, words = download(mb, state)
    plot_waveform("modvibe_waveform.png", waveform_id, words)
    mb.close()
if __name__ == "__main__": main()
#datasheetPartTwoEnd