'''
ESPNow visualizer

- Receives data from the serial line
- Display graphics dynamically
- saves data when closing the window

made with the help of Perplexity

2026-01-31 mchris
'''

import serial
import csv
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from datetime import datetime, date
import numpy as np
import sys
from collections import deque
import matplotlib.dates as mdates

# ------------------------------------------------------------
# Configuration
# ------------------------------------------------------------
PORT = '/dev/ttyUSB0'
BAUD = 115200
FILENAME = 'sensor_data.csv'
MAX_POINTS = 10000

# ------------------------------------------------------------
# Data buffers
# ------------------------------------------------------------
data_buffers = {
    'temp': deque(maxlen=MAX_POINTS),
    'pressure': deque(maxlen=MAX_POINTS),
    'humidity': deque(maxlen=MAX_POINTS),
    'vbat': deque(maxlen=MAX_POINTS),
    'counter': deque(maxlen=MAX_POINTS),
    'time': deque(maxlen=MAX_POINTS)  # datetime objects for x‑axis
}

# ------------------------------------------------------------
# CSV initialization
# ------------------------------------------------------------
with open(FILENAME, 'w', newline='') as csvfile:
    writer = csv.writer(csvfile)
    writer.writerow([
        'timestamp',
        'temperature',
        'pressure',
        'humidity',
        'vbat',
        'counter'
    ])

# ------------------------------------------------------------
# Serial connection
# ------------------------------------------------------------
try:
    ser = serial.Serial(PORT, BAUD, timeout=1)
    ser.flushInput()
    print(f"Connected to {PORT} at {BAUD} baud")
except serial.SerialException as e:
    print(f"Serial error: {e}")
    sys.exit(1)

# ------------------------------------------------------------
# Plot setup - FIXED X-AXIS DISPLAY
# ------------------------------------------------------------
fig, axes = plt.subplots(3, 2, figsize=(12, 10), sharex='col')  # Share x-axis by column
fig.suptitle('Real-time Sensor Data', fontsize=16)

plots = [
    ('temp',     axes[0, 0], 'Temperature (°C)', 'tab:blue'),
    ('pressure', axes[0, 1], 'Pressure (hPa)',   'tab:green'),
    ('humidity', axes[1, 0], 'Humidity (%)',     'tab:red'),
    ('vbat',     axes[2, 0], 'VBAT (V)',         'orange'),
    ('counter',  axes[1, 1], 'Counter',          'tab:purple')
]

#axes[1, 1].axis('off')  # counter position

lines = {}

for key, ax, ylabel, color in plots:
    line, = ax.plot([], [], color=color, linewidth=2, label=key)
    lines[key] = line
    ax.set_ylabel(ylabel)
    ax.legend()
    ax.grid(True, alpha=0.3)

    # Format x‑axis as HH:MM:SS on EVERY subplot
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%H:%M:%S'))

    # Show x-labels on ALL plots (not just bottom row)
    ax.tick_params(axis='x', labelbottom=True)

# Hide the unused subplot
axes[2, 1].axis('off')

# Rotate all x-labels for readability
fig.autofmt_xdate()
plt.tight_layout()

# ------------------------------------------------------------
# Animation callbacks
# ------------------------------------------------------------
def init():
    for line in lines.values():
        line.set_data([], [])
    return list(lines.values())

def update(frame):
    try:
        raw = ser.readline().decode('utf-8').strip()
        if raw:
            parts = raw.split(',')
            if len(parts) == 6:
                ts_str = parts[0].strip()   # e.g. "07:26:50"
                temp   = float(parts[1])
                pressure = float(parts[2])
                humidity = float(parts[3])
                vbat   = float(parts[4])
                counter = float(parts[5])

                # Parse "HH:MM:SS" to a datetime object with today's date
                t_only = datetime.strptime(ts_str, '%H:%M:%S').time()
                ts = datetime.combine(date.today(), t_only)

                data_buffers['temp'].append(temp)
                data_buffers['pressure'].append(pressure)
                data_buffers['humidity'].append(humidity)
                data_buffers['vbat'].append(vbat)
                data_buffers['counter'].append(counter)
                data_buffers['time'].append(ts)

                # Log to CSV using the same timestamp
                with open(FILENAME, 'a', newline='') as csvfile:
                    writer = csv.writer(csvfile)
                    writer.writerow([
                        ts.isoformat(),
                        temp, pressure, humidity, vbat, counter
                    ])

                print(
                    f"{ts_str}: "
                    f"T={temp:.2f}°C "
                    f"P={pressure:.1f} "
                    f"RH={humidity:.1f}% "
                    f"V={vbat:.3f} "
                    f"C={counter}"
                )

    except (ValueError, IndexError):
        # Ignore malformed lines
        pass

    # Convert datetime objects to matplotlib's float format
    xdata = mdates.date2num(list(data_buffers['time']))

    for key in lines:
        ydata = np.asarray(data_buffers[key])
        lines[key].set_data(xdata, ydata)

    # Rescale axes
    for ax in (axes[0,0], axes[0,1], axes[1,0], axes[1,1], axes[2,0]):
        ax.relim()
        ax.autoscale_view()

    return list(lines.values())

# ------------------------------------------------------------
# Animation
# ------------------------------------------------------------
ani = animation.FuncAnimation(
    fig,
    update,
    init_func=init,
    interval=100,
    blit=False,
    cache_frame_data=False
)

try:
    plt.show()
finally:
    ser.close()
    print(f"Data saved to {FILENAME}")