#!/usr/bin/env python3
"""
gen_data.py — Deterministic typed/shaped data generator for the streaming-chain
E2E suite.

For a scenario (see scenarios.py) it writes a MARTe2 FileReader-compatible binary
``input_<id>.bin`` and returns a *ground-truth* dict the waveform validator uses
to reconstruct the expected stream without re-deriving the layout.

MARTe2 binary format
--------------------
    [u32 numSigs]
    per signal: [u16 TypeDescriptor.all][32B name (null-padded)][u32 numElements]
    then NUM_ROWS rows, each row = all signals' elements concatenated, every value
    little-endian at its native width.

Ground-truth dict schema (keyed "<src_id>:<sig_name>")
------------------------------------------------------
    {
      "t":         np.ndarray[float64]  # intended sample time (s), flattened
      "v":         np.ndarray           # native-dtype values, flattened
      "dt":        float                # per-sample spacing (s)
      "formula":   str
      "freq":      float | None
      "elements":  int
      "rows":      int
      "type":      str
      "quant":     str
      "range_min": float | None
      "range_max": float | None
      "is_time":   bool
    }

Values are the *raw native* values fed to the FileReader. Wire-side quantisation
is performed by the UDPStreamer, not here — the validator applies the quant
tolerance using the recorded quant/range fields.
"""
import argparse
import os
import struct
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import scenarios as S  # noqa: E402  (TYPE_CODES / NP_DTYPE / SCENARIOS)

# Buffer geometry lives in scenarios.py so the seamless-loop constraint
# (validate_scenario) and the data layout cannot drift apart.
NUM_ROWS = S.NUM_ROWS   # producer cycles written to the FileReader input
ROW_DT = S.ROW_DT       # seconds per producer cycle (row); 1 kHz producer


def _sample_dt(sig):
    """Per-element time spacing (s) for a signal."""
    if sig["sampling_rate"]:
        return 1.0 / sig["sampling_rate"]
    e = sig["elements"]
    return ROW_DT / e if e > 1 else ROW_DT


def _sample_times(sig):
    """Flattened intended sample times (s): NUM_ROWS*elements values."""
    e = sig["elements"]
    sdt = _sample_dt(sig)
    rows = np.arange(NUM_ROWS, dtype=np.float64).reshape(-1, 1) * ROW_DT
    cols = np.arange(e, dtype=np.float64).reshape(1, -1) * sdt
    return (rows + cols).reshape(-1)


def _values(sig, t):
    """Compute float64 values for flattened sample times ``t``."""
    f = sig["formula"]
    e = sig["elements"]
    idx = np.arange(t.size, dtype=np.float64)
    if f == "sine":
        freq = sig["freq"] if sig["freq"] else 1.0
        return np.sin(2.0 * np.pi * freq * t)
    if f == "ramp":
        # linear in the global element index, normalised to a modest range
        return (idx % 1000.0)
    if f == "counter":
        return idx
    if f == "time_ns":
        return np.round(t * 1.0e9)
    if f == "time_us":
        return np.round(t * 1.0e6)
    raise ValueError(f"unknown formula {f!r} for {sig['name']}")


def _native(sig, vals):
    """Cast float64 values to the signal's native numpy dtype."""
    dt = S.NP_DTYPE[sig["type"]]
    if sig["type"] in S.FLOAT_TYPES:
        return vals.astype(dt)
    # integer: round then cast (deterministic, no banker's rounding surprises)
    return np.rint(vals).astype(dt)


def build_ground_truth(scenario):
    """Return {"<src>:<sig>": gt_dict} for every signal in the scenario."""
    gt = {}
    for src in scenario["sources"]:
        for sig in src["signals"]:
            t = _sample_times(sig)
            v = _native(sig, _values(sig, t))
            gt[f"{src['id']}:{sig['name']}"] = {
                "t": t, "v": v, "dt": _sample_dt(sig),
                "formula": sig["formula"], "freq": sig["freq"],
                "elements": sig["elements"], "rows": NUM_ROWS,
                "type": sig["type"], "quant": sig["quant"],
                "range_min": sig["range_min"], "range_max": sig["range_max"],
                "is_time": sig["is_time"],
            }
    return gt


def write_input(scenario, path):
    """Write the MARTe binary for *one* source and return its ground-truth dict.

    The MARTe FileReader reads a single flat row layout, so one input file maps
    to one source. Multi-source scenarios call this once per source with distinct
    paths (the orchestrator handles the per-source filename); here we write the
    first source by default but accept an explicit ``src`` via the scenario when
    a single source is present.
    """
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    gt = build_ground_truth(scenario)

    # write each source to its own file: <path> for src[0], <path>.<srcid> else.
    written = {}
    for i, src in enumerate(scenario["sources"]):
        p = path if i == 0 else f"{path}.{src['id']}"
        _write_source_bin(src, p)
        written[src["id"]] = p
    gt["_files"] = written
    return gt


def _write_source_bin(src, path):
    sigs = src["signals"]
    # per-signal native 2D arrays [NUM_ROWS, elements]
    cols = {}
    for sig in sigs:
        t = _sample_times(sig)
        v = _native(sig, _values(sig, t))
        cols[sig["name"]] = v.reshape(NUM_ROWS, sig["elements"])

    with open(path, "wb") as f:
        f.write(struct.pack("<I", len(sigs)))
        for sig in sigs:
            f.write(struct.pack("<H", S.TYPE_CODES[sig["type"]]))
            name = (sig["name"] + "\0").encode()
            f.write((name + b"\0" * 32)[:32])
            f.write(struct.pack("<I", sig["elements"]))
        for r in range(NUM_ROWS):
            for sig in sigs:
                f.write(cols[sig["name"]][r].tobytes())


def main():
    p = argparse.ArgumentParser(description="Generate E2E chain input data")
    p.add_argument("--scenario", required=True, help="scenario id")
    p.add_argument("--out", required=True, help="output input_<id>.bin path")
    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 = write_input(sc, args.out)
    for sid, fp in gt["_files"].items():
        print(f"{sid}: {fp} ({os.path.getsize(fp)} bytes)")


if __name__ == "__main__":
    main()