"""
api_example_v5_multi_detector.py
=================================

Example: simultaneous PK + QP + AV measurement using three parallel connections.

Each detector requires its own WebSocket connection to the server.  Using
``asyncio.gather`` the three connections run concurrently, so the total
measurement time equals the slowest individual sweep — not the sum of all three.

All three connections must share the same ``session_UUID``; the server rejects
connections that arrive with a different UUID while a session is already active
(WebSocket close code 4003).

Prerequisites
-------------
    pip install websockets matplotlib

Usage
-----
    1. Set INSTRUMENT_URL to the IP address of your EMScope instrument.
    2. Adjust the measurement parameters in the ``# Measurement parameters``
       section of main().
    3. Run:

           python api_example_v5_multi_detector.py

How it works
------------
The workflow is the same as api_example_v5.py, but repeated three times in
parallel — one connection per detector.  ``asyncio.gather`` runs all three
concurrently so the measurement time is no longer than a single detector run.

Each connection configures its own detector, collects sweeps for
``measure_duration`` seconds, and returns the last received spectrum.  The
results are then plotted together on the same axes.

See also
--------
    api_example_v5.py — single-detector example; read this first.
"""

import asyncio
import time

import matplotlib.pyplot as plt

from api_v5 import EmscopeConnection


# ── Instrument address ─────────────────────────────────────────────────────────
INSTRUMENT_URL = "ws://10.10.14.45:8010"

# Shared session UUID — all three connections must use the same value.
SESSION_UUID = "multi-detector-example"


async def _collect(
    url: str,
    detector: str,
    rbw: str,
    three_phase: bool,
    display_range,
    channel: str,
    amp_units: str,
    trace_type: str,
    average_count: int,
    reference_level: int,
    input_attenuator,
    sweep_time: float,
    measure_duration: float,
) -> tuple:
    """
    Open one connection, collect sweeps for ``measure_duration`` seconds, and
    return ``(detector, sn, values, active_range)`` where *values* is the last
    received sweep and *active_range* is the frequency window applied by set_rbw().
    """
    async with EmscopeConnection(url, session_uuid=SESSION_UUID) as em:
        # set_rbw() must come first — it may trigger a firmware reload and also
        # sets em.display_range to the natural frequency window for the band.
        await em.set_rbw(rbw, three_phase, display_range=display_range)

        # Measurement configuration — order does not matter after set_rbw().
        # trace_type must be sent to activate data streaming for this connection.
        await em.set_detector_type(detector)
        await em.set_trace_type(trace_type)
        await em.set_measure_channel(channel)
        await em.set_average(average_count)
        await em.set_amp_units(amp_units)
        await em.set_reference_level(reference_level)
        await em.set_input_attenuator(input_attenuator)
        await em.set_sweep_time(sweep_time)
        await em.set_visible(True)

        deadline = time.monotonic() + measure_duration
        sweeps = 0

        while time.monotonic() < deadline or sweeps == 0:
            msg = await em.parse_messages()

            if msg is None:
                print("[{}] connection closed by server.".format(detector))
                break

            if not em.has_sweep_data(msg):
                continue

            if em.overload:
                print("[{}] WARNING: ADC overload — reduce reference level or increase attenuator.".format(detector))

            sweeps += 1

        print("[{}] done ({} sweep{})".format(detector, sweeps, "s" if sweeps != 1 else ""))
        return detector, em.SN, list(em.values), em.display_range


async def main() -> None:

    # ── Measurement parameters ─────────────────────────────────────────────────
    #
    # rbw — Resolution bandwidth and frequency band selector.
    #
    #   Value      RBW       Frequency range      Standard
    #   -------    -------   ------------------   ----------
    #   "200"      200 Hz    9 kHz  – 150 kHz     CISPR 16-1-1
    #   "9"        9 kHz     150 kHz – 30 MHz     CISPR 16-1-1
    #   "120"      120 kHz   30 MHz  – 110 MHz    CISPR 16-1-1
    #   "1"        1 kHz     10 kHz  – 150 kHz    MIL-STD-461
    #   "10"       10 kHz    150 kHz – 30 MHz     MIL-STD-461
    #   "200_9"    dual      9 kHz  – 30 MHz      CISPR dual-band (200 Hz + 9 kHz)
    #   "1_10"     dual      10 kHz – 30 MHz      MIL  dual-band  (1 kHz + 10 kHz)
    #
    rbw = "9"

    # three_phase — True for RX4 (three-phase) models, False for RX2.
    three_phase = False

    # channel — input channel to measure.
    #   RX2 models: "lg" (Line), "ng" (Neutral), "cm" (Common Mode), "dm" (Differential)
    #   RX4 models: "l1", "l2", "l3", "n"
    channel = "lg"

    # amp_units — amplitude units for the received values.
    #   "dbuv"  → dBµV   (standard for CISPR / MIL conducted emissions)
    #   "dbmv"  → dBmV
    #   "dbm"   → dBm
    #   "volts" → V (linear)
    #   "watts" → W (linear)
    amp_units = "dbuv"

    # trace_type — how consecutive sweeps are combined on the displayed trace.
    #   "clearwrite" → each sweep replaces the previous one  (default)
    #   "maxhold"    → keeps the highest value seen at each frequency bin
    #   "minhold"    → keeps the lowest value seen
    #   "average"    → running average over `average_count` sweeps
    #   "freeze"     → trace is frozen; the instrument keeps sweeping internally
    #                  but does not update the displayed trace
    trace_type = "clearwrite"

    # average_count — number of sweeps to average when trace_type="average".
    # Typical range: 10–20. Ignored for other trace types.
    average_count = 10

    # reference_level — top of the amplitude scale, in the selected amp_units.
    # Must be consistent with amp_units. Typical values:
    #   100 dBuV, 80 dBmV, 20 dBm
    reference_level = 115

    # scale_div — amplitude per division on the plot (same units as amp_units).
    # num_divisions — number of vertical divisions shown below the reference level.
    scale_div     = 10
    num_divisions = 10

    # input_attenuator — input hardware attenuator in dB (0–78), or "auto".
    # "auto" lets the instrument choose the optimal value based on reference_level.
    input_attenuator = "auto"

    # sweep_time — duration of each sweep in seconds (1–15).
    # For Quasi-Peak measurements, a minimum of 1 s is required by CISPR 16-1-1.
    sweep_time = 1

    # display_range — optional (from_Hz, to_Hz) window override.
    # When None, set_rbw() applies the natural band for the selected RBW.
    display_range = None

    # measure_duration — how many seconds to collect sweeps per detector.
    measure_duration = 5

    # ── Run all three detectors in parallel ────────────────────────────────────
    #
    # Each detector gets its own WebSocket connection. All share the same
    # configuration except detector_type. asyncio.gather runs them concurrently
    # so the total time is no longer than a single-detector run.
    #
    print("Starting parallel PK / QP / AV measurement ...")
    t0 = time.monotonic()

    cfg = dict(
        url=INSTRUMENT_URL, rbw=rbw, three_phase=three_phase,
        display_range=display_range, channel=channel, amp_units=amp_units,
        trace_type=trace_type, average_count=average_count,
        reference_level=reference_level, input_attenuator=input_attenuator,
        sweep_time=sweep_time, measure_duration=measure_duration,
    )
    results = await asyncio.gather(
        _collect(detector="pk", **cfg),
        _collect(detector="qp", **cfg),
        _collect(detector="av", **cfg),
    )

    elapsed = time.monotonic() - t0
    print("Done in {:.1f} s".format(elapsed))

    # ── Plot ───────────────────────────────────────────────────────────────────
    sn = results[0][1]

    _, ax = plt.subplots(figsize=(12, 5))
    ax.set_xscale("log")

    colors = {"pk": "tab:red", "qp": "tab:blue", "av": "tab:green"}

    for detector, _, values, active_range in results:
        if not values:
            print("WARNING: no data received for detector '{}'".format(detector))
            continue

        if active_range is not None:
            visible = [v for v in values if active_range[0] <= v[0] <= active_range[1]]
        else:
            visible = values

        freqs = [v[0] for v in visible]
        amps  = [v[1] for v in visible]
        ax.plot(freqs, amps, linewidth=0.8, color=colors[detector],
                label=detector.upper())

    y_max = reference_level
    y_min = reference_level - scale_div * num_divisions
    ax.set_ylim(y_min, y_max)
    ax.set_yticks([y_max - scale_div * i for i in range(num_divisions + 1)])
    ax.set_ylabel("{} ({} {}/div)".format(amp_units, scale_div, amp_units))
    ax.set_xlabel("Frequency (Hz)")
    ax.set_title(
        "EMScope  SN={}  |  PK + QP + AV  |  {} · RBW {} kHz · RL {} {} · IA {}".format(
            sn, channel.upper(), rbw, reference_level, amp_units, input_attenuator)
    )
    ax.legend()
    ax.grid(True, alpha=0.4)
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    asyncio.run(main())
