# 16ノード×3軸フルスペクトル(Base64量子化)を pyqtgraph で表示

import socket, json, base64, numpy as np	# UDP受信/JSON解析/Base64復号/数値計算
import pyqtgraph as pg	# 高速プロッタ
from pyqtgraph.Qt import QtCore	# Qt タイマー等

HOST, PORT = "0.0.0.0", 50000	# 受信待ちIPとポート（全IFで受ける）
NROWS, NCOLS = 4, 4	# グリッド（4×4 = 16ノード）
NODES = [f"N{i:02d}" for i in range(1, 16+1)]	# 期待するノードID一覧 N01..N16
X_MAX = 500	# 横軸の表示上限 [Hz]
Y_LIM = (-50, 0)	# 縦軸レンジ dB(µT/√Hz)
COLORS = {"X": (230, 80, 80), "Y": (80, 180, 80), "Z": (80, 140, 220)}	# 軸ごとの色
LINE_W = 2.0	# 線の太さ

def grid_pos(nid):	# ノードID→グリッド位置(行,列)
    i = int(nid[1:]) - 1	# "N07" → 6（0始まり）
    return i // NCOLS, i % NCOLS	# 行=商, 列=余り

def mains_lines(fs=1000.0):	# 50/60Hz高調波の補助線位置を作る
    out = []
    for base in (50, 60):	# 50Hzと60Hzの両方
        h = 1
        while base*h < fs/2:	# ナイキスト未満まで
            out.append(base*h); h += 1
    return out

# ---------- GUI ----------
pg.setConfigOptions(antialias=False, useOpenGL=False)	# 軽量設定（アンチエイリアス/GL無効）
pg.setConfigOptions(background='w', foreground='k', antialias=False)	# さらに白背景・黒前景
app = pg.mkQApp("Spectrum Viewer (pyqtgraph)")	# Qtアプリ生成

win = pg.GraphicsLayoutWidget(show=True, title="3-axis spectrum")	# グリッドウィンドウ
win.resize(1200, 800)	# ウィンドウサイズ

plots, curves = {}, {}	# 各ノードのPlotとCurveを保持
x_axis = {}	# 各ノードの周波数軸配列
fs_of  = {nid: None for nid in NODES}	# ノードごとのサンプリング周波数
n_of   = {nid: None for nid in NODES}	# ノードごとのNFFT

def safe_set_downsampling(plot):	# バージョン差吸収：間引き設定
    """pyqtgraph の版差を吸収"""
    try:
        plot.setDownsampling(auto=True, mode='peak')	# 新しい版の呼び方
    except TypeError:
        try:
            plot.setDownsampling(True, 'peak')	# 古い版の呼び方（位置引数）
        except Exception:
            pass	# どちらも無ければ諦める

for r in range(NROWS):	# グリッド行ループ
    for c in range(NCOLS):	# グリッド列ループ
        nid = NODES[r*NCOLS + c]	# このセルのノードID
        p = win.addPlot(row=r, col=c, title=nid)	# サブプロット追加（タイトルにID）
        p.setMenuEnabled(False)	# 右クリックメニュー無効
        p.setMouseEnabled(x=False, y=False)	# マウス操作無効（誤作動防止）
        p.setClipToView(True)	# ビュー外の描画を抑制
        safe_set_downsampling(p)	# 間引き設定
        try:
            p.enableAutoRange(x=False, y=False)	# 自動レンジ無効（固定レンジ）
        except Exception:
            pass
        p.setXRange(0, X_MAX)	# X軸[0..X_MAX]固定
        p.setYRange(*Y_LIM)	# Y軸固定

	# 50/60Hz 補助線を引く
        for f in mains_lines():
            v = pg.InfiniteLine(pos=f, angle=90, pen=pg.mkPen(120, 120, 120, 50, width=1))	# 薄い縦線
            p.addItem(v)

        plots[nid] = p	# Plot参照を保持
        curves[nid] = {}	# 軸ごとのCurveを入れる dict
        for ax in ("X","Y","Z"):	# X/Y/Zそれぞれ作成
            item = pg.PlotCurveItem(pen=pg.mkPen(COLORS[ax], width=LINE_W))	# 色・太さ
            p.addItem(item)	# サブプロットに追加
            curves[nid][ax] = item	# 参照を保存

def ensure_buffers(nid, fs, n):	# 周波数軸や曲線の初期化/再初期化
    if x_axis.get(nid) is None or fs_of[nid] != fs or n_of[nid] != n:
        bins = n // 2	# 片側スペクトルビン数
        x_axis[nid] = np.linspace(0, fs/2, bins, endpoint=False, dtype=np.float32)	# 周波数軸
	# 初期化：とりあえず空データをセット（描画リセット用）
        for ax in ("X","Y","Z"):
            curves[nid][ax].setData([], [])
        fs_of[nid], n_of[nid] = fs, n	# このノードのfs/NFFTを更新

# ---------- UDP ----------
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)	# UDPソケット生成
sock.bind((HOST, PORT))	# 受信バインド
sock.setblocking(False)	# ノンブロッキング受信

def handle_msg(msg):	# 1メッセージを取り込んで描画更新
    nid = msg.get("id")	# 送信元ノードID
    if nid not in NODES:	# 想定外のIDは無視
        return
    fs = float(msg.get("fs", 1000.0))	# サンプリング周波数
    n  = int(msg.get("n", 1024))	# NFFT
    ensure_buffers(nid, fs, n)	# 必要なら軸等を作り直す

    if "dby_b64" in msg and "axis" in msg:	# フルスペクトル(Base64)パケット？
        axis = msg["axis"]	# どの軸(X/Y/Z)
        if axis not in ("X","Y","Z"):
            return
        qstep = float(msg.get("qstep", 0.5))	# 量子化ステップ [dB/LSB]
        raw = base64.b64decode(msg["dby_b64"])	# Base64 → 512バイト
        arr = np.frombuffer(raw, dtype=np.uint8)	# uint8配列に見る
        if arr.size != (n//2):	# 破損/サイズ不一致は無視
            return

        qbias = float(msg.get("qbias",128.0))	# dBのオフセット（送信側規約）
        # print(nid,qbias,qstep)                          	# デバッグ用
        dB = arr.astype(np.float32) * qstep - qbias	# 復元：dB = q*step - bias

	# 表示上限を超える周波数は描画しない（軽量化）
        xmax = min(X_MAX, fs/2)	# 実際のナイキストとも比較
        xa = x_axis[nid]	# このノードの周波数軸
        valid = xa < xmax	# X_MAX未満の部分だけ
        curves[nid][axis].setData(xa[valid], dB[valid], connect='finite')	# 曲線更新

def poll_udp():	# 受信ループ（タイマーコールバックから呼ぶ）
    while True:
        try:
            data, _ = sock.recvfrom(65535)	# 受信（無ければ例外で抜ける）
        except BlockingIOError:
            break	# 受信バッファ空 → ループ終了
        except Exception:
            break	# 何かあれば念のため抜ける
        for line in data.decode("utf-8", "ignore").splitlines():	# 改行区切りで複数JSONに対応
            if not line:
                continue
            try:
                handle_msg(json.loads(line))	# JSONとして解釈→描画更新へ
            except Exception:
                pass	# 壊れた行は捨てる

timer = QtCore.QTimer()	# Qtタイマー生成
timer.timeout.connect(poll_udp)	# 周期ごとに受信処理を回す
timer.start(40)	# 40msごと（≈25fps相当）

if __name__ == "__main__":	# スクリプト実行直後のエントリ
    app.exec()	# Qtアプリのイベントループ開始
