MARTe-Integrated-Components/Test/E2E/suite/plots.py

#!/usr/bin/env python3
"""
plots.py — Report figures for the streaming-chain E2E suite.

For a scenario it reads ``received_<id>.bin`` (RCV1), ``metrics_<id>.json`` and
``checks_<id>.json`` from the artifact dir and writes:

  * ``wave_<id>.png`` — received waveform(s) vs analytic sinusoid fit (truth)
  * ``trig_<id>.png`` — trigger signal with threshold + fired-trigger markers
  * ``zoom_<id>.png`` — full received signal with the requested zoom spans shaded

Missing inputs are handled gracefully (a placeholder note is drawn) so the
orchestrator never aborts on a partial scenario.
"""
import argparse
import json
import os
import sys

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import scenarios as S  # noqa: E402
import gen_data as G    # noqa: E402
import validate_waveform as V  # noqa: E402


def _placeholder(path, msg):
    fig, ax = plt.subplots(figsize=(10, 3))
    ax.axis("off")
    ax.text(0.5, 0.5, msg, ha="center", va="center", fontsize=12)
    fig.savefig(path, dpi=110, bbox_inches="tight")
    plt.close(fig)


def _load(scenario, d):
    sid = scenario["id"]
    recv_p = os.path.join(d, f"received_{sid}.bin")
    checks_p = os.path.join(d, f"checks_{sid}.json")
    recv = V.read_received(recv_p) if os.path.exists(recv_p) else {}
    checks = json.load(open(checks_p)) if os.path.exists(checks_p) else {}
    return recv, checks


def _data_keys(recv, gt):
    """Non-time signal keys present in received, base-matched to ground truth."""
    out = []
    for k in sorted(recv):
        base = k.split("[")[0]
        g = gt.get(base)
        if g is None or g.get("is_time"):
            continue
        out.append((k, base, g))
    return out


def plot_wave(scenario, recv, gt, path):
    keys = _data_keys(recv, gt)
    if not keys:
        _placeholder(path, f"{scenario['id']}: no received signal data")
        return path
    keys = keys[:4]
    fig, axes = plt.subplots(len(keys), 1, figsize=(12, 3.0 * len(keys)),
                             squeeze=False)
    for ax, (k, base, g) in zip(axes[:, 0], keys):
        t, v = recv[k]
        if t.size == 0:
            ax.text(0.5, 0.5, f"{k}: empty", ha="center"); continue
        t0 = t[0]
        ax.plot(t - t0, v, ".", ms=3, label="received", color="tab:blue")
        if g["formula"] == "sine" and g["freq"] and t.size >= 8:
            corr, nrmse, amp = V.sine_shape(t, v, g["freq"])
            w = 2 * np.pi * g["freq"]
            A = np.column_stack([np.sin(w * t), np.cos(w * t), np.ones_like(t)])
            coef, *_ = np.linalg.lstsq(A, v, rcond=None)
            ts = np.linspace(t.min(), t.max(), 2000)
            As = np.column_stack([np.sin(w * ts), np.cos(w * ts), np.ones_like(ts)])
            ax.plot(ts - t0, As @ coef, "-", lw=1, color="tab:orange",
                    label=f"fit f={g['freq']}Hz corr={corr:.4f}")
        ax.set_title(f"{k}  ({g['type']}, quant={g['quant']}, {g['formula']})")
        ax.set_xlabel("t − t0 (s)")
        ax.legend(loc="upper right", fontsize=8)
        ax.grid(alpha=0.3)
    fig.suptitle(f"{scenario['id']} — received waveform vs truth")
    fig.tight_layout()
    fig.savefig(path, dpi=120, bbox_inches="tight")
    plt.close(fig)
    return path


def plot_trigger(scenario, recv, checks, gt, path):
    trig = checks.get("trigger", [])
    if not trig:
        _placeholder(path, f"{scenario['id']}: no trigger checks")
        return path
    key = trig[0].get("key", "")
    if key not in recv or recv[key][0].size == 0:
        _placeholder(path, f"{scenario['id']}: trigger signal {key} not recorded")
        return path
    t, v = recv[key]
    t0 = t[0]
    thr = float(np.mean(v))
    fig, ax = plt.subplots(figsize=(12, 4))
    ax.plot(t - t0, v, "-", lw=0.7, color="tab:blue", label=key)
    ax.axhline(thr, color="gray", ls="--", lw=1, label=f"~threshold {thr:.3g}")
    seen = set()
    for tr in trig:
        if not tr.get("fired"):
            continue
        tt = tr["trigTime"] - t0
        lbl = f"{tr['edge']}/{tr['mode']}"
        ax.axvline(tt, color="tab:red", lw=1, alpha=0.6,
                   label=("fired" if "fired" not in seen else None))
        seen.add("fired")
    ax.set_title(f"{scenario['id']} — trigger captures ({len(trig)} combos)")
    ax.set_xlabel("t − t0 (s)")
    ax.legend(loc="upper right", fontsize=8)
    ax.grid(alpha=0.3)
    fig.tight_layout()
    fig.savefig(path, dpi=120, bbox_inches="tight")
    plt.close(fig)
    return path


def plot_zoom(scenario, recv, checks, gt, path):
    zooms = checks.get("zoom", [])
    keys = _data_keys(recv, gt)
    if not keys:
        _placeholder(path, f"{scenario['id']}: no received data for zoom")
        return path
    k, base, g = keys[0]
    t, v = recv[k]
    t0 = t[0]
    fig, ax = plt.subplots(figsize=(12, 4))
    ax.plot(t - t0, v, "-", lw=0.6, color="tab:blue", label=k)
    for zc in zooms:
        rg = zc.get("range", [0, 0])
        ax.axvspan(rg[0] - t0, rg[1] - t0, alpha=0.15, color="tab:green")
    ax.set_title(f"{scenario['id']} — zoom spans (n={len(zooms)})")
    ax.set_xlabel("t − t0 (s)")
    ax.legend(loc="upper right", fontsize=8)
    ax.grid(alpha=0.3)
    fig.tight_layout()
    fig.savefig(path, dpi=120, bbox_inches="tight")
    plt.close(fig)
    return path


def main():
    p = argparse.ArgumentParser(description="Generate E2E chain report plots")
    p.add_argument("--scenario", required=True)
    p.add_argument("--dir", required=True)
    args = p.parse_args()
    sc = next((s for s in S.SCENARIOS if s["id"] == args.scenario), None)
    if sc is None:
        print(f"unknown scenario {args.scenario}", file=sys.stderr)
        sys.exit(2)
    gt = G.build_ground_truth(sc)
    recv, checks = _load(sc, args.dir)
    w = plot_wave(sc, recv, gt, os.path.join(args.dir, f"wave_{sc['id']}.png"))
    tr = plot_trigger(sc, recv, checks, gt,
                      os.path.join(args.dir, f"trig_{sc['id']}.png"))
    z = plot_zoom(sc, recv, checks, gt,
                  os.path.join(args.dir, f"zoom_{sc['id']}.png"))
    print(w)
    print(tr)
    print(z)


if __name__ == "__main__":
    main()