#!/usr/bin/env python3
"""Plot packed 4-byte-integer channel data from a serial_logger.py CSV, dark mode.

Expects rows where `data_hex` is N 4-byte little-endian integers packed back
to back (default N=18, matching an 18-channel ADC_READ_ALL-style packet).
X-axis is the device-clock `tick_ms` column by default (or host_time).
"""

import argparse
import csv
import struct
import sys
from datetime import datetime

import matplotlib
import matplotlib.pyplot as plt
import matplotlib.dates as mdates


def build_parser():
	p = argparse.ArgumentParser(
		description="Plot packed 4-byte-integer channels from a serial_logger.py CSV (dark mode).",
		formatter_class=argparse.RawDescriptionHelpFormatter,
		epilog="""\
examples:
  %(prog)s -i adc_log.csv
  %(prog)s -i adc_log.csv --command ADC_READ_ALL -o plot.png
  %(prog)s -i adc_log.csv --channels 18 --unsigned --x host
""",
	)

	p.add_argument("-i", "--input", required=True, help="input CSV file (from serial_logger.py)")
	p.add_argument("-o", "--output", help="save the plot to this file instead of showing it interactively")

	p.add_argument("-c", "--command", help="only use rows with this command name (default: use whatever "
	                                        "PACKET rows have the right byte length)")
	p.add_argument("-n", "--channels", type=int, default=18,
	               help="number of 4-byte integer channels packed in data_hex (default: %(default)s)")
	p.add_argument("--signed", dest="signed", action="store_true", default=True,
	               help="interpret channels as signed int32 (default: on)")
	p.add_argument("--unsigned", dest="signed", action="store_false",
	               help="interpret channels as unsigned uint32")

	p.add_argument("--x", choices=["tick", "host"], default="tick",
	               help="x-axis source: device tick_ms, or host_time (default: %(default)s)")
	p.add_argument("--labels", help="comma-separated channel labels (default: ch0..chN-1)")
	p.add_argument("--title", default="Channel data", help="plot title (default: %(default)s)")

	return p


def load_rows(path, command, channel_bytes):
	"""Read the CSV and return rows that look like packed-channel PACKET rows."""
	rows = []
	with open(path, newline="", encoding="utf-8") as fh:
		reader = csv.DictReader(fh)
		for row in reader:
			if row.get("type") != "PACKET":
				continue
			if command and row.get("command") != command:
				continue
			data_hex = row.get("data_hex", "")
			if len(data_hex) != channel_bytes * 2:  # 2 hex chars per byte
				continue
			rows.append(row)
	return rows


def main():
	parser = build_parser()
	args = parser.parse_args()

	channel_bytes = args.channels * 4

	rows = load_rows(args.input, args.command, channel_bytes)
	if not rows:
		hint = f" with command={args.command!r}" if args.command else ""
		print(f"ERROR: no PACKET rows{hint} with {channel_bytes}-byte payloads "
		      f"({args.channels} x 4-byte ints) found in {args.input}", file=sys.stderr)
		return 1

	fmt = f"<{args.channels}{'i' if args.signed else 'I'}"

	xs = []
	channels = [[] for _ in range(args.channels)]

	for row in rows:
		if args.x == "tick":
			xs.append(int(row["tick_ms"]))
		else:
			xs.append(datetime.fromisoformat(row["host_time"]))

		values = struct.unpack(fmt, bytes.fromhex(row["data_hex"]))
		for i, v in enumerate(values):
			channels[i].append(v)

	# sort chronologically, just in case the file wasn't strictly ordered
	order = sorted(range(len(xs)), key=lambda i: xs[i])
	xs = [xs[i] for i in order]
	channels = [[ch[i] for i in order] for ch in channels]

	if args.labels:
		labels = [s.strip() for s in args.labels.split(",")]
		if len(labels) != args.channels:
			parser.error(f"--labels has {len(labels)} entries, expected {args.channels}")
	else:
		labels = [f"ch{i}" for i in range(args.channels)]

	# --- dark mode plot ---
	plt.style.use("dark_background")
	fig, ax = plt.subplots(figsize=(13, 7))

	cmap = matplotlib.colormaps["tab20"].resampled(args.channels)
	for i in range(args.channels):
		ax.plot(xs, channels[i], label=labels[i], color=cmap(i), linewidth=1.2)

	ax.set_title(args.title, fontsize=14, color="white")
	ax.set_xlabel("tick (ms)" if args.x == "tick" else "host time")
	ax.set_ylabel("value" + (" (int32)" if args.signed else " (uint32)"))
	ax.grid(True, alpha=0.25)

	if args.x == "host":
		fig.autofmt_xdate()
		ax.xaxis.set_major_formatter(mdates.DateFormatter("%H:%M:%S"))

	ax.legend(ncol=min(6, args.channels), fontsize=8, loc="upper center",
	          bbox_to_anchor=(0.5, -0.12), frameon=False)
	fig.tight_layout()

	if args.output:
		fig.savefig(args.output, dpi=150, facecolor=fig.get_facecolor())
		print(f"Saved plot to {args.output} ({len(xs)} samples, {args.channels} channels)")
	else:
		plt.show()

	return 0


if __name__ == "__main__":
	sys.exit(main())