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

pwm_actual = stýrimerkid klippt vid +-255, þ.e. þad sem raunverulega
             var sent a motor.set() / vélbúnaðinn

Hausinn er prentadur i hvert skipti sem farid er inn i Operational, svo ein
skra getur innihaldid nokkrar keyrslur. Adeins sidasta keyrslan er teiknud.

Skriptan metur lika sveiflutima (Tu) og utslag ef kerfid sveiflast, sem
nytist til ad finna Ku/Tu fyrir Ziegler-Nichols.

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

--no-pwm teiknar bara RPM (PWM-dalkurinn ma lika vanta i CSV-skranni)

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


COLUMNS = ["time_ms", "ref_rpm", "actual_rpm", "pwm_actual"]
REQUIRED = ["time_ms", "ref_rpm", "actual_rpm"]

# Onnur nofn sem eldri utgafur firmware-sins hafa notad fyrir sama dalk
ALIASES = {"pwm": "pwm_actual"}


def empty_cols():
    return {key: [] for key in COLUMNS}


def load_csv(path):
    cols = empty_cols()
    header = None
    with open(path, newline="") as f:
        for row in csv.reader(f):
            if row and row[0].strip() == "time_ms":
                header = [name.strip() for name in row]
                header = [ALIASES.get(name, name) for name in header]
                missing = [key for key in REQUIRED if key not in header]
                if missing:
                    sys.exit(f"Dalka vantar i haus {row}: {missing}")
                continue
            if header is None:
                # rusl a undan fyrsta haus (t.d. "State: Initialization")
                continue
            try:
                values = {key: float(row[header.index(key)])
                          for key in COLUMNS if key in header}
            except (ValueError, IndexError):
                # sleppir linum sem eru ekki gild data (t.d. "State: ...")
                continue

            # Timinn byrjar aftur a 0 -> ny keyrsla, hendum theirri gomlu
            if cols["time_ms"] and values["time_ms"] <= cols["time_ms"][-1]:
                cols = empty_cols()

            for key in values:
                cols[key].append(values[key])
    return cols


def estimate_oscillation(time_s, actual_rpm):
    """
    Metur sveiflutima ut fra uppleid i gegnum medaltal seinni helmings
    gagnanna (fyrri helmingur er transient). Hysteresis kemur i veg fyrir
    ad sudid i rpm-maelingunni gefi falskar skurdpunkta.
    Skilar (Tu [s], utslag [rpm]) eda None ef engin skyr sveifla.
    """
    half = len(actual_rpm) // 2
    t = time_s[half:]
    y = actual_rpm[half:]
    if len(y) < 20:
        return None

    mean = statistics.fmean(y)
    band = 0.25 * statistics.pstdev(y)
    if band == 0:
        return None

    crossings = []
    below = y[0] < mean - band
    for ti, yi in zip(t, y):
        if below and yi > mean + band:
            crossings.append(ti)
            below = False
        elif not below and yi < mean - band:
            below = True

    if len(crossings) < 3:
        return None

    periods = [b - a for a, b in zip(crossings, crossings[1:])]
    amplitude = (max(y) - min(y)) / 2.0
    return statistics.fmean(periods), amplitude


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

    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)

    lines, labels = ax1.get_legend_handles_labels()
    if show_pwm:
        ax2 = ax1.twinx()
        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]")
        lines2, labels2 = ax2.get_legend_handles_labels()
        lines += lines2
        labels += labels2

    # 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")

    osc = estimate_oscillation(time_s, actual_rpm)
    # if osc is not None:
    #     tu, amplitude = osc
    #     print(f"Sveifla: Tu = {tu * 1000:.1f} ms, utslag = +-{amplitude:.2f} RPM")
    #     ax1.text(0.99, 0.98, f"Tu ≈ {tu * 1000:.1f} ms\n"
    #                           f"amplitude ≈ ±{amplitude:.2f} RPM",
    #               transform=ax1.transAxes, fontsize=9, va="top", ha="right",
    #               bbox=dict(facecolor="white", alpha=0.8, edgecolor="#cccccc"))
    # else:
    #     print("Engin skyr sveifla fannst i seinni helmingi gagnanna")

    ax1.legend(lines, labels, loc="lower right", fontsize=8)

    if show_pwm:
        title = "Step response: RPM and PWM over time"
    else:
        title = "Step response: RPM 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__":
    show_pwm = "--no-pwm" not in sys.argv
    args = [a for a in sys.argv[1:] if a != "--no-pwm"]
    in_path = args[0] if len(args) > 0 else "data/step_response.csv"
    out_path = args[1] if len(args) > 1 else "data/step_response.png"
    title_suffix = args[2] if len(args) > 2 else ""
    data = load_csv(in_path)
    if not data["time_ms"]:
        sys.exit(f"Engin gogn fundust i {in_path}")
    plot(data, out_path, title_suffix, show_pwm)