"""
api_example_v5_modal.py
=======================

Example: user-defined modal channels on a three-phase (RX4) instrument.

Measures the common and differential mode of each phase against the neutral:

    V_CM_Lx = (V_Lx + V_N) / 2
    V_DM_Lx = (V_Lx - V_N) / 2      (add 6.02 dB for V_Lx - V_N)

The instrument has four modal channels (cm, dm1, dm2, dm3), so the six modes
are measured with two coefficient configurations, one after the other.

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

    The instrument needs the ``advanced_modal`` license.

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

           python api_example_v5_modal.py

How it works
------------
A control connection selects the band, the neutral, the modal mode and the
coefficients.  Then one connection per modal channel collects sweeps in
parallel, as in api_example_v5_multi_detector.py.  The coefficients are only
changed between measurements, never while the traces are running.

The instrument measures two modal channels per sweep: cm with dm1, and dm2 with
dm3.  Each phase uses one of those pairs, so its common and differential modes
are measured at the same instant.

See also
--------
    api_example_v5_multi_detector.py — parallel connections; read it 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 — every connection must use the same value.
SESSION_UUID = "modal-example"

# Coefficients c1..c4 weight the inputs L1, L2, L3, N.
CONFIGURATIONS = [
    {
        # Measurement A: L1 on the pair cm/dm1, L2 on the pair dm2/dm3
        "cm":  ("V_CM_L1", {"c1": 0.5, "c2": 0,   "c3": 0,   "c4": 0.5}),
        "dm1": ("V_DM_L1", {"c1": 0.5, "c2": 0,   "c3": 0,   "c4": -0.5}),
        "dm2": ("V_CM_L2", {"c1": 0,   "c2": 0.5, "c3": 0,   "c4": 0.5}),
        "dm3": ("V_DM_L2", {"c1": 0,   "c2": 0.5, "c3": 0,   "c4": -0.5}),
    },
    {
        # Measurement B: L3 on the pair cm/dm1
        "cm":  ("V_CM_L3", {"c1": 0,   "c2": 0,   "c3": 0.5, "c4": 0.5}),
        "dm1": ("V_DM_L3", {"c1": 0,   "c2": 0,   "c3": 0.5, "c4": -0.5}),
    },
]


async def _configure(url: str, rbw: str, coefficients: dict) -> None:
    """Select the band, neutral, modal mode and coefficients."""
    async with EmscopeConnection(url, session_uuid=SESSION_UUID) as em:
        await em.set_rbw(rbw, threephase=True)
        await em.request_licenses()
        while True:
            msg = await em.parse_messages()
            if msg is None or "licenses" in msg:
                break
        if "advanced_modal" not in em.licenses:
            raise RuntimeError("The instrument has no 'advanced_modal' license.")
        await em.set_neutral(True)
        await em.set_mode("modal")
        await em.set_modal_coefficients({"emi_3p_neutral": coefficients})


async def _collect(url: str, rbw: str, channel: str, detector: str,
                   sweep_time: float, measure_duration: float) -> list:
    """Collect sweeps of one modal channel and return the last one."""
    async with EmscopeConnection(url, session_uuid=SESSION_UUID) as em:
        await em.set_rbw(rbw, threephase=True)
        await em.set_detector_type(detector)
        await em.set_measure_channel(channel)
        await em.set_sweep_time(sweep_time)
        # Only visible traces choose the inputs measured: without it dm2/dm3 get no data
        await em.set_visible(True)
        # trace_type activates data streaming for this connection
        await em.set_trace_type("clearwrite")

        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(channel))
                break
            if not em.has_sweep_data(msg):
                continue
            if em.overload:
                print("[{}] WARNING: ADC overload.".format(channel))
            sweeps += 1
        return list(em.values)


async def main() -> None:

    # ── Measurement parameters ─────────────────────────────────────────────────
    # rbw — "200", "9", "1", "10", "200_9" or "1_10" (the modes are fixed for "120")
    rbw = "9"
    detector = "pk"
    sweep_time = 1
    # With four channels each pair is measured every other sweep: collect long
    # enough for several sweeps of each pair.
    measure_duration = 10

    results = {}
    for configuration in CONFIGURATIONS:
        await _configure(INSTRUMENT_URL, rbw,
                         {mode: coefs for mode, (_, coefs) in configuration.items()})
        channels = list(configuration)
        spectra = await asyncio.gather(*[
            _collect(INSTRUMENT_URL, rbw, channel, detector, sweep_time, measure_duration)
            for channel in channels
        ])
        for channel, values in zip(channels, spectra):
            results[configuration[channel][0]] = values

    # ── Plot ───────────────────────────────────────────────────────────────────
    _, ax = plt.subplots(figsize=(12, 5))
    ax.set_xscale("log")
    for name, values in results.items():
        if not values:
            print("WARNING: no data received for {}".format(name))
            continue
        ax.plot([v[0] for v in values], [v[1] for v in values], linewidth=0.8, label=name)
    ax.set_xlabel("Frequency (Hz)")
    ax.set_ylabel("dBuV")
    ax.set_title("EMScope modal channels  |  {} · RBW {} kHz".format(detector.upper(), rbw))
    ax.legend()
    ax.grid(True, alpha=0.4)
    plt.tight_layout()
    plt.show()


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