import asyncio
import math
import time
import threading

import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from bleak import BleakClient

# ====== HARD-CODED DEVICE ======
ADDRESS = "10:51:DB:1B:E7:1E"

# ====== ADXL345 CHARACTERISTIC UUID ======
CHAR_UUID = "7e2a1002-1111-2222-3333-444455556666"

# Plot refresh interval (ms)
PLOT_INTERVAL_MS = 50

latest = {
    "ax": 0.0,
    "ay": 0.0,
    "az": 1.0,
    "ts": 0.0,
    "raw": "",
    "count": 0,
}

def parse_csv_triplet(s: str):
    s = s.strip().strip('"')
    parts = s.split(",")
    if len(parts) != 3:
        return None
    try:
        return float(parts[0]), float(parts[1]), float(parts[2])
    except ValueError:
        return None

def roll_pitch_from_accel(ax, ay, az):
    roll = math.degrees(math.atan2(ay, az))
    pitch = math.degrees(math.atan2(-ax, math.sqrt(ay * ay + az * az)))
    return roll, pitch

def notify_handler(_sender: int, data: bytearray):
    s = data.decode(errors="ignore").strip()
    latest["raw"] = s

    v = parse_csv_triplet(s)
    if v is None:
        return

    ax, ay, az = v
    latest["ax"], latest["ay"], latest["az"] = ax, ay, az
    latest["ts"] = time.time()
    latest["count"] += 1

    # Debug print every 5 packets (avoid flooding)
    if latest["count"] % 5 == 0:
        print(f"RX[{latest['count']}]: {ax:+.4f}, {ay:+.4f}, {az:+.4f}")

async def ble_loop():
    while True:
        try:
            print(f"Connecting to {ADDRESS} ...")
            async with BleakClient(ADDRESS, timeout=15.0) as client:
                if not client.is_connected:
                    raise RuntimeError("Connect failed (not connected).")

                print("Connected.")
                print("Subscribing to notifications...")
                await client.start_notify(CHAR_UUID, notify_handler)
                print("Notify enabled. Streaming...\n")

                while True:
                    await asyncio.sleep(1.0)

        except Exception as e:
            print(f"[BLE] Error/disconnect: {e}")
            print("Reconnecting in 2 seconds...\n")
            await asyncio.sleep(2.0)

def start_ble_thread():
    t = threading.Thread(target=lambda: asyncio.run(ble_loop()), daemon=True)
    t.start()

def main():
    print("IMPORTANT: Disconnect nRF Connect (or any other BLE client) first.")
    start_ble_thread()

    # --- Matplotlib setup ---
    fig = plt.figure()
    ax3 = fig.add_subplot(111, projection="3d")
    ax3.set_title("ADXL345 live acceleration vector (g)")

    lim = 2.0
    ax3.set_xlim(-lim, lim)
    ax3.set_ylim(-lim, lim)
    ax3.set_zlim(-lim, lim)
    ax3.set_xlabel("X (g)")
    ax3.set_ylabel("Y (g)")
    ax3.set_zlabel("Z (g)")

    # Axes lines
    ax3.plot([-lim, lim], [0, 0], [0, 0])
    ax3.plot([0, 0], [-lim, lim], [0, 0])
    ax3.plot([0, 0], [0, 0], [-lim, lim])

    # Initial arrow + overlay
    q = ax3.quiver(0, 0, 0, 0, 0, 1, length=1.0, normalize=False)
    overlay = ax3.text2D(0.02, 0.95, "Waiting for data...", transform=ax3.transAxes)

    def update(_frame):
        nonlocal q

        # Remove old arrow
        try:
            q.remove()
        except Exception:
            pass

        vx, vy, vz = latest["ax"], latest["ay"], latest["az"]
        roll, pitch = roll_pitch_from_accel(vx, vy, vz)
        age_ms = (time.time() - latest["ts"]) * 1000.0 if latest["ts"] else float("inf")

        # Draw new arrow
        q = ax3.quiver(0, 0, 0, vx, vy, vz, length=1.0, normalize=False)

        overlay.set_text(
            f"packets: {latest['count']}\n"
            f"ax={vx:+.4f}  ay={vy:+.4f}  az={vz:+.4f}  (g)\n"
            f"roll={roll:+.1f}°  pitch={pitch:+.1f}°   age={age_ms:.0f} ms\n"
            f"raw: {latest['raw']}"
        )

        return q, overlay

    # KEY FIX: keep this object alive
    ani = FuncAnimation(fig, update, interval=PLOT_INTERVAL_MS, blit=False)

    plt.show()

if __name__ == "__main__":
    main()