# %% imports
import asyncio
import time

import numpy as np
from bleak import BleakClient, BleakScanner
from matplotlib import pyplot as plt

# NOTE: this must be run in Jupyter. I rely on the Jupyter run loop for async


# %%
devices = await BleakScanner.discover()
for d in devices:
    print(d)


# %%
address = "D7:FC:DA:4B:02:D7"
MODEL_NBR_UUID = "2A24"
SENSORS_UUID = "99b594745570464bbe3764465aa050fd"


timeseries = list()
async with BleakClient(address) as client:
    # model_number = await client.read_gatt_char(MODEL_NBR_UUID)
    # print("Model Number: {0}".format("".join(map(chr, model_number))))
    start_time = time.time()
    print("Starting capture...")
    while time.time() < start_time + 3:
        buffer = await client.read_gatt_char(SENSORS_UUID)
        sensor_values = np.frombuffer(bytes(buffer), dtype=np.double)
        timeseries.append(sensor_values)
    print("DONE")

timeseries = np.array(timeseries)

# %%
plt.plot(timeseries, marker=".")
plt.show()

# %%
a_x, a_y, a_z, g_x, g_y, g_z, m_x, m_y, m_z = timeseries.T
t = np.linspace(0, 2, len(timeseries))
a = np.sqrt(a_x**2 + a_y**2 + a_z**2)
g = 9.81  # acceleration due to gravity
h = np.array([0, *np.cumsum((a - g)[1:] * np.diff(t))])

# %%
plt.plot(t, a_x, label="$A_x$")
plt.plot(t, a_y, label="$A_y$")
plt.plot(t, a_z, label="$A_z$")
plt.plot(t, a, color="k", label="$A_{total}$")
plt.ylabel("Acceleration [m/s^2]")
plt.xlabel("Time [s]")
plt.xlim(t[0], t[-1])
plt.legend()
plt.grid(True)
plt.show()

# %%
plt.plot(t, h, label="$Z$")
plt.ylabel("Displacement [m]")
plt.title("Displacement assuming only vertical motion")
plt.xlabel("Time [s]")
plt.xlim(t[0], t[-1])
plt.legend()
plt.grid(True)
plt.show()

# %%