"""
api_example_v5.py
=================

Minimal working example of the EMScope Python API v5.

This script connects to an EMScope instrument, configures it for a single-band
EMC measurement, collects sweeps for a fixed duration, and plots the last
received spectrum with matplotlib.

Requirements
------------
    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.py

How it works
------------
The EMScope server streams measurement data continuously over WebSocket.
There is no explicit "start" command — the instrument starts sweeping as soon
as a client connects and has been configured.

Each call to ``parse_messages()`` blocks until one message arrives from the
server.  Two kinds of messages carry a completed spectrum sweep:

- **Binary amplitude frames** (EMI mode): the server sends one binary frame
  per sweep; the API decodes it internally and updates ``emscope.values``.
  ``parse_messages()`` returns ``{"amplitudes": True}`` for these frames.
- **JSON values** (impedance mode): ``{"values": [[freq, amp], ...], ...}``.

Use ``emscope.has_sweep_data(msg)`` to detect both cases uniformly, without
needing to know the underlying wire format.

Other messages carry configuration echoes, keepalive pings, or responses to
explicit requests (temperatures, standards, etc.) — ``has_sweep_data`` returns
``False`` for all of these.

The workflow is always:
    1. Connect  →  server sends initial device info (SN, MAC, SFP_SN).
    2. Configure  →  call set_rbw() first, then any other setters.
    3. Loop parse_messages()  →  read sweeps until done.
    4. Disconnect.
"""

import asyncio
import time

import matplotlib.pyplot as plt

from api_v5 import EmscopeConnection


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


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 = "200_9"

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

    # detector_type — detector algorithm applied to each frequency bin.
    #   "pk"  → Peak          (fastest, conservative)
    #   "qp"  → Quasi-Peak    (IEC standard; slow — increase sweep_time accordingly)
    #   "av"  → Average
    detector_type = "pk"

    # 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 before plotting.
    measure_duration = 5

    # ── Connect and configure ──────────────────────────────────────────────────
    #
    # The context manager calls connect() on entry and disconnect() on exit,
    # even if an exception occurs.
    #
    # connect() sends session_UUID to activate the client, then waits for the
    # device info reply (SN, MAC, SFP_SN).
    #
    async with EmscopeConnection(INSTRUMENT_URL, session_uuid="example-session-001") as emscope:

        print("Connected to EMScope  SN={}  MAC={}".format(emscope.SN, emscope.MAC))

        # Set the frequency band and RBW first — changing RBW may trigger a
        # firmware reload (a few seconds). set_rbw() blocks until the firmware
        # is ready and also applies the correct display_range for the band.
        await emscope.set_rbw(rbw, three_phase, display_range=display_range)

        # Measurement configuration — order does not matter after set_rbw().
        await emscope.set_detector_type(detector_type)
        await emscope.set_measure_channel(channel)
        await emscope.set_trace_type(trace_type)
        await emscope.set_average(average_count)
        await emscope.set_amp_units(amp_units)
        await emscope.set_reference_level(reference_level)
        await emscope.set_input_attenuator(input_attenuator)
        await emscope.set_sweep_time(sweep_time)
        await emscope.set_visible(True)

        # ── Measurement loop ───────────────────────────────────────────────────
        #
        # parse_messages() blocks until one message arrives.
        # has_sweep_data(msg) returns True for:
        #   - binary amplitude frames  (EMI mode — most common)
        #   - JSON "values" messages   (impedance mode)
        #
        # The server also sends periodic keepalive pings; the API replies
        # with pong automatically — no special handling needed.
        #
        print("Collecting sweeps for {} s ...".format(measure_duration))
        deadline = time.monotonic() + measure_duration
        sweeps_received = 0

        while time.monotonic() < deadline or sweeps_received == 0:
            msg = await emscope.parse_messages()

            if msg is None:
                print("Connection closed by server.")
                break

            if not emscope.has_sweep_data(msg):
                # Configuration echo, temperature response, ping, etc.
                continue

            if emscope.overload:
                print("WARNING: ADC overload — reduce reference level or increase attenuator.")

            sweeps_received += 1
            if emscope.values:
                peak = max(v[1] for v in emscope.values)
                print("  Sweep {}: {} points  peak={:.1f} {}".format(
                    sweeps_received, len(emscope.values), peak, amp_units))

        # ── Collect plot data before disconnecting ─────────────────────────────
        if not emscope.values:
            print("No measurement data received.")
            return

        # Apply display_range client-side (emscope.display_range was set by
        # set_rbw() — either the default for the band or the custom override).
        active_range = emscope.display_range
        if active_range is not None:
            visible = [(v[0], v[1]) for v in emscope.values
                       if active_range[0] <= v[0] <= active_range[1]]
        else:
            visible = emscope.values

        plot_data = (
            [v[0] for v in visible],
            [v[1] for v in visible],
            emscope.SN,
        )

    # ── async with block closed — WebSocket is disconnected ───────────────────
    # plt.show() is safe to call now.

    freqs, amps, sn = plot_data

    y_max = reference_level
    y_min = reference_level - scale_div * num_divisions

    _, ax = plt.subplots(figsize=(12, 5))
    ax.set_xscale("log")
    ax.plot(freqs, amps, linewidth=0.8)

    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={}  |  {} · {} · {}  |  RBW {} kHz · RL {} {} · IA {}".format(
            sn, detector_type.upper(), channel.upper(), trace_type,
            rbw, reference_level, amp_units, input_attenuator)
    )
    ax.grid(True, alpha=0.4)
    plt.tight_layout()
    plt.show()


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