#!/usr/bin/env python3
"""Real-time telemetry plotter for ESP32 Drone.

The drone sends one JSON packet per update over UDP:

    {"alt": 1.23, "roll": -0.5, "pitch": 0.2, "yaw": 178.0}

This script listens on UDP port 8888 and draws the four values in real time.

Usage
-----
    1. Connect your computer to the drone Wi-Fi hotspot (ESP-Drone).
    2. Run:  python plotter.py

Requirements
------------
    pip install matplotlib

Tutorial: https://drone.imagrove.com/intermediate/telemetry
"""

import json
import queue
import socket
import threading
from collections import deque

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

UDP_PORT = 8888
WINDOW = 300  # samples kept on screen

FIELDS = ("alt", "roll", "pitch", "yaw")
TITLES = {
    "alt": "Altitude (m)",
    "roll": "Roll (deg)",
    "pitch": "Pitch (deg)",
    "yaw": "Yaw (deg)",
}


def listen(q, stop):
    """Receive UDP datagrams and queue up every packet that parses as JSON."""
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind(("", UDP_PORT))
    sock.settimeout(0.5)
    print(f"Listening on UDP :{UDP_PORT} — waiting for the drone...")
    try:
        while not stop.is_set():
            try:
                raw, _addr = sock.recvfrom(2048)
            except socket.timeout:
                continue
            try:
                q.put(json.loads(raw.decode("utf-8", "ignore")))
            except (ValueError, UnicodeDecodeError):
                continue  # malformed packet, ignore it
    finally:
        sock.close()


def main():
    packets = queue.Queue()
    stop = threading.Event()

    receiver = threading.Thread(target=listen, args=(packets, stop), daemon=True)
    receiver.start()

    history = {field: deque([0.0] * WINDOW, maxlen=WINDOW) for field in FIELDS}

    fig, axes = plt.subplots(2, 2, figsize=(10, 6))
    fig.suptitle("ESP32 Drone Telemetry")
    try:
        fig.canvas.manager.set_window_title("ESP32 Drone Telemetry")
    except AttributeError:
        pass  # not every backend has a window manager (e.g. Agg)

    lines = {}
    for ax, field in zip(axes.flat, FIELDS):
        (line,) = ax.plot(range(WINDOW), history[field], linewidth=1.2)
        ax.set_title(TITLES[field])
        ax.set_xlim(0, WINDOW - 1)
        ax.grid(True, alpha=0.3)
        lines[field] = line
    fig.tight_layout()

    def update(_frame):
        # Drain everything received since the last frame.
        while True:
            try:
                packet = packets.get_nowait()
            except queue.Empty:
                break
            for field in FIELDS:
                if field in packet:
                    try:
                        history[field].append(float(packet[field]))
                    except (TypeError, ValueError):
                        pass

        for ax, field in zip(axes.flat, FIELDS):
            data = history[field]
            lines[field].set_ydata(data)

            # Keep the curve visible: rescale only when it leaves the view.
            low, high = min(data), max(data)
            bottom, top = ax.get_ylim()
            if low < bottom or high > top:
                pad = max((high - low) * 0.1, 1.0)
                ax.set_ylim(low - pad, high + pad)

        return list(lines.values())

    # Keep a reference: matplotlib does not hold the animation alive itself.
    animation = FuncAnimation(fig, update, interval=50, blit=False, cache_frame_data=False)

    try:
        plt.show()
    finally:
        stop.set()
        receiver.join(timeout=1.0)
        del animation


if __name__ == "__main__":
    main()