"""Render side-by-side comparison plots for JS vs Python benchmark runs.

Two columns × three rows: each row is a channel (left antenna, right
antenna, head Y); left column is the JS run, right column is the Python
run. Same y range per row so the comparison is honest.
"""

import csv
import sys
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt


def load(path):
    rows = []
    with open(path) as f:
        for r in csv.DictReader(f):
            rows.append({k: float(v) if v else None for k, v in r.items()})
    return rows


def col(rows, key):
    return [r[key] for r in rows]


def plot_pair(ax, time, cmd, act, title, ylabel, ylim=None):
    ax.plot(time, cmd, color="C0", linestyle="--", linewidth=1.2, label="commanded")
    ax.plot(time, act, color="C3", linewidth=1.2, label="measured")
    ax.set_title(title, fontsize=10)
    ax.set_xlabel("t (s)", fontsize=9)
    ax.set_ylabel(ylabel, fontsize=9)
    ax.grid(True, alpha=0.3)
    if ylim is not None:
        ax.set_ylim(*ylim)
    ax.legend(fontsize=8, loc="upper right")
    ax.tick_params(labelsize=8)


def main(js_csv, py_csv, out_png):
    js = load(js_csv)
    py = load(py_csv)

    js_t = col(js, "t_s")
    py_t = col(py, "t_s")

    # Shared y ranges (per row), computed across BOTH runs.
    def span(*lists, pad=3):
        vals = [v for lst in lists for v in lst if v is not None]
        return min(vals) - pad, max(vals) + pad

    yl_l = span(col(js, "left_commanded_deg"),  col(js, "left_actual_deg"),
                col(py, "left_commanded_deg"),  col(py, "left_actual_deg"))
    yl_r = span(col(js, "right_commanded_deg"), col(js, "right_actual_deg"),
                col(py, "right_commanded_deg"), col(py, "right_actual_deg"))
    yl_h = span(col(js, "head_y_commanded_cm"), col(js, "head_y_actual_cm"),
                col(py, "head_y_commanded_cm"), col(py, "head_y_actual_cm"), pad=0.5)

    fig, axes = plt.subplots(3, 2, figsize=(13, 9))
    plot_pair(axes[0, 0], js_t, col(js, "left_commanded_deg"),  col(js, "left_actual_deg"),
              "Left antenna — JS",  "deg", yl_l)
    plot_pair(axes[0, 1], py_t, col(py, "left_commanded_deg"),  col(py, "left_actual_deg"),
              "Left antenna — Python", "deg", yl_l)
    plot_pair(axes[1, 0], js_t, col(js, "right_commanded_deg"), col(js, "right_actual_deg"),
              "Right antenna — JS", "deg", yl_r)
    plot_pair(axes[1, 1], py_t, col(py, "right_commanded_deg"), col(py, "right_actual_deg"),
              "Right antenna — Python", "deg", yl_r)
    plot_pair(axes[2, 0], js_t, col(js, "head_y_commanded_cm"), col(js, "head_y_actual_cm"),
              "Head Y — JS", "cm", yl_h)
    plot_pair(axes[2, 1], py_t, col(py, "head_y_commanded_cm"), col(py, "head_y_actual_cm"),
              "Head Y — Python", "cm", yl_h)

    fig.suptitle("Streaming benchmark: JS (browser) vs Python (script)",
                 fontsize=12, fontweight="bold")
    fig.tight_layout(rect=(0, 0, 1, 0.97))
    fig.savefig(out_png, dpi=140)
    print(f"Wrote {out_png}")


if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2], sys.argv[3])