#!/usr/bin/env python3
"""
Teiknar step response-plott ur CSV-gognum sem Arduino-verkefnid prentar ut
(dalkar: time_ms, ref_rpm, actual_rpm, pwm_requested, pwm_actual).

pwm_requested = Kp * error * 2.55, oklippt (getur farid langt yfir +-255)
pwm_actual    = sama gildi en klippt vid +-255, þ.e. þad sem raunverulega
                var sent a motor.set() / vélbúnaðinn

Notkun:
    python plot_step_response.py [inntak.csv] [uttak.png] [titill]

Sjalfgefid: data/step_response.csv -> data/step_response.png
"""
import sys
import csv
import matplotlib.pyplot as plt


def load_csv(path):
    cols = {"time_ms": [], "ref_rpm": [], "actual_rpm": [],
            "pwm_requested": [], "pwm_actual": []}
    with open(path, newline="") as f:
        reader = csv.DictReader(f)
        for row in reader:
            try:
                for key in cols:
                    cols[key].append(float(row[key]))
            except (ValueError, KeyError):
                # sleppir linum sem eru ekki gild data (t.d. rusl fyrir/eftir)
                for key in cols:
                    if len(cols[key]) > len(cols["time_ms"]):
                        cols[key].pop()
                continue
    return cols


def plot(cols, out_path, title_suffix=""):
    time_s = [t / 1000.0 for t in cols["time_ms"]]
    ref_rpm = cols["ref_rpm"]
    actual_rpm = cols["actual_rpm"]
    pwm_requested = cols["pwm_requested"]
    pwm_actual = cols["pwm_actual"]

    fig, ax1 = plt.subplots(figsize=(10, 5))

    ax1.plot(time_s, ref_rpm, label="Reference (ref_rpm)", color="#666666",
              linestyle="--", linewidth=1.5)
    ax1.plot(time_s, actual_rpm, label="Measured (actual_rpm)", color="#1f77b4",
              linewidth=1.8)
    ax1.set_xlabel("Time [s]")
    ax1.set_ylabel("Angular velocity [RPM]")
    ax1.grid(True, alpha=0.3)

    ax2 = ax1.twinx()
    ax2.plot(time_s, pwm_requested, label="PWM requested (unclamped)",
              color="#d62728", linewidth=1.0, alpha=0.5, linestyle=":")
    ax2.plot(time_s, pwm_actual, label="PWM actual (clamped)",
              color="#d62728", linewidth=1.4)
    ax2.axhline(255, color="#d62728", linewidth=0.6, alpha=0.3)
    ax2.axhline(-255, color="#d62728", linewidth=0.6, alpha=0.3)
    ax2.set_ylabel("PWM [-255, 255]")

    # Merkjum stokkpunktinn (fyrsta stadin thar sem ref hoppar fra 0)
    step_idx = next((i for i, r in enumerate(ref_rpm) if r > 0), None)
    if step_idx is not None:
        ax1.axvline(time_s[step_idx], color="black", linewidth=0.8,
                     linestyle=":")
        ax1.text(time_s[step_idx], max(actual_rpm) * 1.02, " step",
                  fontsize=9, va="bottom")

    lines1, labels1 = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax1.legend(lines1 + lines2, labels1 + labels2, loc="lower right",
               fontsize=8)

    title = "Double step response: RPM and PWM over time"
    if title_suffix:
        title += f" ({title_suffix})"
    fig.suptitle(title)
    fig.tight_layout()
    fig.savefig(out_path, dpi=150)
    print(f"Vistad: {out_path}")


if __name__ == "__main__":
    in_path = sys.argv[1] if len(sys.argv) > 1 else "data/twostep.csv"
    out_path = sys.argv[2] if len(sys.argv) > 2 else "data/twostep.png"
    title_suffix = sys.argv[3] if len(sys.argv) > 3 else ""
    data = load_csv(in_path)
    plot(data, out_path, title_suffix)