#!/usr/bin/env bash
# run_udp_producer.sh — Run a MARTe2 app that streams N sine channels at 1 Msps.
#
# A producer only: no StreamHub, no clients. Point whatever consumer you like at
# the UDP port (StreamHub, the Go hub, or Test/E2E tooling).
#
# Each channel is an N-element float32 array published every cycle by a
# real-time thread — 10000 samples x 100 Hz = 1 Msps per channel by default. A
# parallel uint64 time array gives every sample its own timestamp
# (TimeMode=FullArray), so consumers reconstruct the waveform at full rate
# rather than one point per cycle. When a cycle exceeds the UDP datagram limit
# it is split across several datagrams (fragmented) and reassembled by the
# receiver.
#
# One extra channel, HV, emulates a charged-capacitor pulse discharge on a
# high-voltage bus: an emulated EPICS "start" setpoint (SlowControlGAM) raises
# a rising edge, on which PulseGeneratorGAM ramps the output to -40 kV in 1 ms,
# holds it flat for the emulated EPICS "duration" setpoint, then discharges
# back to 0 over 100 ms. Always-on gaussian noise (3-sigma ~ +-1000 V) and
# random +-5000 V EMI spikes (on and off phase) ride on top.
#
# Usage:
#   ./run_udp_producer.sh [OPTIONS]
#
# Options:
#   -n <CHANNELS>   Number of 1 Msps sine channels (default 4, max 12 — see below)
#   -p <PORT>       UDP port to stream on (default 44501)
#   -t <MS>         HV trigger period (time between pulses, default 5000)
#   -d <MS>         HV plateau duration (EPICS "duration" setpoint, default 500)
#   -b <TARGET>     Build target (default: $TARGET or x86-linux)
#   -s              Skip the component rebuild
#   -k              Keep the generated .cfg on exit and print its path
#   -h              Show this help
#
# Channels are capped at 12 as a practical bound on the bytes the background
# thread copies and sends each cycle (the uint64 time array plus (CHANNELS+1)
# float32 arrays). Cycles larger than one UDP datagram are fragmented into
# 60000 B datagrams and reassembled by the receiver by sequence counter, so the
# datagram limit no longer caps the channel count — only the per-cycle send
# budget does.
#
# Environment:
#   MARTe2_DIR             must be set (or source env.sh first)
#   MARTe2_Components_DIR  must be set (or source env.sh first)

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BUILD_TARGET="${TARGET:-x86-linux}"
CHANNELS=4
PORT=44501
SKIP_BUILD=0
KEEP_CFG=0
HV_TRIG_PERIOD=5000
HV_PLATEAU_MS=500

MAX_CHANNELS=12

while getopts "n:p:t:d:b:skh" opt; do
  case "$opt" in
  n) CHANNELS="$OPTARG" ;;
  p) PORT="$OPTARG" ;;
  t) HV_TRIG_PERIOD="$OPTARG" ;;
  d) HV_PLATEAU_MS="$OPTARG" ;;
  b) BUILD_TARGET="$OPTARG" ;;
  s) SKIP_BUILD=1 ;;
  k) KEEP_CFG=1 ;;
  h)
    sed -n '2,43p' "$0" | sed 's/^# \?//'
    exit 0
    ;;
  *)
    echo "Unknown option: -$OPTARG" >&2
    exit 1
    ;;
  esac
done

# ── Validate ──────────────────────────────────────────────────────────────────
if ! [[ "$CHANNELS" =~ ^[0-9]+$ ]] || ((CHANNELS < 1 || CHANNELS > MAX_CHANNELS)); then
  echo "ERROR: -n must be 1..${MAX_CHANNELS} (got '${CHANNELS}')." >&2
  exit 1
fi
if [[ -z "${MARTe2_DIR:-}" || -z "${MARTe2_Components_DIR:-}" ]]; then
  echo "ERROR: MARTe2_DIR / MARTe2_Components_DIR not set." >&2
  echo "  source ${SCRIPT_DIR}/env.sh" >&2
  exit 1
fi

MARTE2_BIN="${MARTe2_DIR}/Build/${BUILD_TARGET}/App/MARTeApp.ex"
if [[ ! -x "$MARTE2_BIN" ]]; then
  echo "ERROR: MARTeApp.ex not found at ${MARTE2_BIN}" >&2
  exit 1
fi

# ── Build ─────────────────────────────────────────────────────────────────────
if [[ "$SKIP_BUILD" -eq 0 ]]; then
  echo "==> Building components (TARGET=${BUILD_TARGET})..."
  make -C "${SCRIPT_DIR}" -f Makefile.gcc TARGET="${BUILD_TARGET}" core 2>&1 | tail -5
fi

# ── Generate the config ───────────────────────────────────────────────────────
# Distinct amplitude/frequency/phase per channel so traces stay tellable apart
# (and so a shared-Y-axis view has a spread of magnitudes to cope with).
AMPS=(1.0 2.5 0.5 5.0 1.5 3.0 0.8 4.0 2.0 0.3 6.0 1.2 3.5)
FREQS=(1000 2000 5000 500 10000 3000 20000 1500 7000 50000 800 4000 15000)
PHASES=(0.0 0.7854 1.5708 2.3562 3.1416 3.9270 4.7124 5.4978 0.3927 1.1781 1.9635 2.7489 3.5343)

ELEMS=10000 # samples per cycle
RATE=100    # cycles per second -> 1 Msps
CYCLE_BYTES=$((8 * ELEMS + (CHANNELS + 1) * 4 * ELEMS + 8))
# UDP datagrams cap at 65507 B (IPv4) minus the 17 B UDPS header; keep a margin
# and use 60000 B of payload per fragment.
MAX_UDP_PAYLOAD=60000
if ((CYCLE_BYTES + 2000 <= MAX_UDP_PAYLOAD)); then
  PAYLOAD=$((CYCLE_BYTES + 2000)) # whole cycle in one datagram
else
  PAYLOAD=${MAX_UDP_PAYLOAD}      # fragment the cycle across datagrams
fi
FRAGMENTS=$(((CYCLE_BYTES + MAX_UDP_PAYLOAD - 1) / MAX_UDP_PAYLOAD))

sine_gams=""
iogam_in=""
iogam_out=""
stream_sigs=""
func_list="TimerGAM"

for ((i = 1; i <= CHANNELS; i++)); do
  k=$((i - 1))
  amp="${AMPS[$k]}"
  frq="${FREQS[$k]}"
  pha="${PHASES[$k]}"

  sine_gams+="
    +SineGAM${i} = {
      Class = SineArrayGAM
      Frequency = ${frq}.0
      Amplitude = ${amp}
      Phase = ${pha}
      Offset = 0.0
      SamplingRate = 1000000.0
      OutputSignals = {
        Ch${i} = {
          DataSource = DDB1
          Type = float32
          NumberOfDimensions = 1
          NumberOfElements = ${ELEMS}
        }
      }
    }
"
  iogam_in+="
        Ch${i} = {
          DataSource = DDB1
          Type = float32
          NumberOfDimensions = 1
          NumberOfElements = ${ELEMS}
        }"
  iogam_out+="
        Ch${i} = {
          DataSource = Streamer
          Type = float32
          NumberOfDimensions = 1
          NumberOfElements = ${ELEMS}
        }"
  stream_sigs+="
        Ch${i} = {
          Type = float32
          Unit = \"V\"
          NumberOfDimensions = 1
          NumberOfElements = ${ELEMS}
          RangeMin = -${amp}
          RangeMax = ${amp}
          TimeMode = \"FullArray\"
          TimeSignal = TimeArray
        }"
  func_list+=", SineGAM${i}"
done
hv_gams="
    +SlowControlGAM1 = {
      Class = SlowControlGAM
      TriggerPeriodMs = ${HV_TRIG_PERIOD}
      TriggerWidthMs = 10
      PlateauMs = ${HV_PLATEAU_MS}
      CycleFrequency = ${RATE}
      OutputSignals = {
        HVTrigger = {
          DataSource = DDB1
          Type = float32
        }
        HVPlateauMs = {
          DataSource = DDB1
          Type = float32
        }
      }
    }

    +PulseGeneratorGAM1 = {
      Class = PulseGeneratorGAM
      SamplingRate = 1000000.0
      RampUpMs = 1.0
      RampDownMs = 100.0
      HighLevel = -40000.0
      NoiseStdDev = 333.33
      EMIAmplitude = 5000.0
      EMIProbabilityPerSample = 0.00001
      EMISpikeSamples = 5
      AutoTriggerPeriodMs = 0.0
      InputSignals = {
        HVTrigger = {
          DataSource = DDB1
          Type = float32
        }
        HVPlateauMs = {
          DataSource = DDB1
          Type = float32
        }
      }
      OutputSignals = {
        HV = {
          DataSource = DDB1
          Type = float32
          NumberOfDimensions = 1
          NumberOfElements = ${ELEMS}
        }
      }
    }
"

iogam_in+="
        HV = {
          DataSource = DDB1
          Type = float32
          NumberOfDimensions = 1
          NumberOfElements = ${ELEMS}
        }
        HVTrigger = {
          DataSource = DDB1
          Type = float32
        }
        HVPlateauMs = {
          DataSource = DDB1
          Type = float32
        }"
iogam_out+="
        HV = {
          DataSource = Streamer
          Type = float32
          NumberOfDimensions = 1
          NumberOfElements = ${ELEMS}
        }
        HVTrigger = {
          DataSource = Streamer
          Type = float32
        }
        HVPlateauMs = {
          DataSource = Streamer
          Type = float32
        }"
stream_sigs+="
        HV = {
          Type = float32
          Unit = \"V\"
          NumberOfDimensions = 1
          NumberOfElements = ${ELEMS}
          RangeMin = -50000.0
          RangeMax = 10000.0
          TimeMode = \"FullArray\"
          TimeSignal = TimeArray
        }
        HVTrigger = {
          Type = float32
          Unit = \"1\"
        }
        HVPlateauMs = {
          Type = float32
          Unit = \"ms\"
        }"
func_list+=", TimeArrayGAM1, SlowControlGAM1, PulseGeneratorGAM1, StreamerGAM"

CFG="$(mktemp /tmp/udp_producer_XXXXXX.cfg)"

cat >"$CFG" <<EOF
/**
 * udp_producer — auto-generated by run_udp_producer.sh
 * ${CHANNELS} channel(s), ${ELEMS} elem x ${RATE} Hz = 1 Msps each, port ${PORT}.
 */
\$App = {
  Class = RealTimeApplication

  +Functions = {
    Class = ReferenceContainer

    +TimerGAM = {
      Class = IOGAM
      InputSignals = {
        Time = {
          DataSource = Timer
          Type = uint32
          Frequency = ${RATE}
        }
      }
      OutputSignals = {
        Time = {
          DataSource = DDB1
          Type = uint32
        }
      }
    }
${sine_gams}${hv_gams}
    // Expands the cycle's scalar timestamp into one timestamp per sample, so
    // consumers place all ${ELEMS} samples instead of collapsing them to a point.
    +TimeArrayGAM1 = {
      Class = TimeArrayGAM
      SamplingRate = 1000000.0
      Anchor = "Continuous"
      InputSignals = {
        Time = {
          DataSource = DDB1
          Type = uint32
        }
      }
      OutputSignals = {
        TimeArray = {
          DataSource = DDB1
          Type = uint64
          NumberOfDimensions = 1
          NumberOfElements = ${ELEMS}
        }
      }
    }

    +StreamerGAM = {
      Class = IOGAM
      InputSignals = {
        TimeArray = {
          DataSource = DDB1
          Type = uint64
          NumberOfDimensions = 1
          NumberOfElements = ${ELEMS}
        }${iogam_in}
      }
      OutputSignals = {
        TimeArray = {
          DataSource = Streamer
          Type = uint64
          NumberOfDimensions = 1
          NumberOfElements = ${ELEMS}
        }${iogam_out}
      }
    }
  }

  +Data = {
    Class = ReferenceContainer
    DefaultDataSource = DDB1

    +DDB1 = {
      Class = GAMDataSource
    }

    +Timer = {
      Class = LinuxTimer
      SleepNature = "Default"
      Signals = {
        Counter = {
          Type = uint32
        }
        Time = {
          Type = uint32
        }
      }
    }

    +Streamer = {
      Class = UDPStreamer
      Port = ${PORT}
      // One cycle is ${CYCLE_BYTES} B; it is sent as ${FRAGMENTS} datagram(s)
      // of up to ${PAYLOAD} B each (the receiver reassembles by counter).
      MaxPayloadSize = ${PAYLOAD}
      PublishingMode = "Strict"
      Signals = {
        TimeArray = {
          Type = uint64
          Unit = "ns"
          NumberOfDimensions = 1
          NumberOfElements = ${ELEMS}
        }${stream_sigs}
      }
    }

    +Timings = {
      Class = TimingDataSource
    }
  }

  +States = {
    Class = ReferenceContainer
    +Running = {
      Class = RealTimeState
      +Threads = {
        Class = ReferenceContainer
        +Thread1 = {
          Class = RealTimeThread
          CPUs = 0x2
          Functions = { ${func_list} }
        }
      }
    }
  }

  +Scheduler = {
    Class = GAMScheduler
    TimingDataSource = Timings
  }
}
EOF

# ── Run ───────────────────────────────────────────────────────────────────────
BUILD_DIR="${SCRIPT_DIR}/Build/${BUILD_TARGET}"
# UDPStream is not used directly here, but UDPStreamer.so carries a NEEDED entry
# on it, so dlopen of the DataSource fails without it on the path.
export LD_LIBRARY_PATH="\
${MARTe2_DIR}/Build/${BUILD_TARGET}/Core:\
${MARTe2_Components_DIR}/Build/${BUILD_TARGET}/Components/DataSources/LinuxTimer:\
${MARTe2_Components_DIR}/Build/${BUILD_TARGET}/Components/GAMs/IOGAM:\
${BUILD_DIR}/Components/DataSources/UDPStreamer:\
${BUILD_DIR}/Components/GAMs/SineArrayGAM:\
${BUILD_DIR}/Components/GAMs/TimeArrayGAM:\
${BUILD_DIR}/Components/GAMs/PulseGeneratorGAM:\
${BUILD_DIR}/Components/GAMs/SlowControlGAM:\
${BUILD_DIR}/Components/Interfaces/UDPStream:\
${LD_LIBRARY_PATH:-}"

cleanup() {
  if [[ "$KEEP_CFG" -eq 1 ]]; then
    echo ""
    echo "==> Config kept at ${CFG}"
  else
    rm -f "$CFG"
  fi
}
trap cleanup EXIT INT TERM

echo ""
echo "==> Streaming on udp/${PORT}"
echo "    Channels : ${CHANNELS} x 1 Msps (${ELEMS} elem @ ${RATE} Hz)"
for ((i = 1; i <= CHANNELS; i++)); do
  k=$((i - 1))
  printf '      Ch%-2d %8s Hz  %s V\n' "$i" "${FREQS[$k]}" "${AMPS[$k]}"
done
echo "      HV    pulse: -40000 V, ramp up 1 ms, flat ${HV_PLATEAU_MS} ms, discharge 100 ms"
echo "            noise +-1000 V (3 sigma), EMI +-5000 V, trigger every ${HV_TRIG_PERIOD} ms"
echo "    Cycle    : ${CYCLE_BYTES} B (${FRAGMENTS} x ${PAYLOAD} B datagram)"
echo "    Config   : ${CFG}"
echo ""
echo "    Consume with e.g.:"
echo "      Addr = \"127.0.0.1\"  Port = ${PORT}   (StreamHub source)"
echo ""
echo "    Press Ctrl-C to stop."
echo ""

exec "${MARTE2_BIN}" \
  -l RealTimeLoader \
  -f "${CFG}" \
  -s Running \
  -m StateMachine:START