#!/usr/bin/env python3
# -*- coding: utf-8 -*-
r"""谱库查表：**记录 → npz 分片**（2026-09-17 用户令"把映射建起来"）。

## 结论先说：这个映射**已经是产品的一部分**，不需要重建

`scripts/rudong_tcm_spectra.py` 落盘时同时写了 `spectra_meta.parquet`，每行一条谱：

    turbine · sensor · meas_name · trigger_time · shard · shard_row · x_offset · x_delta · …

取数配方（该脚本头部原文）：

    v = np.load(<store>/<shard> + '.npz')['values'][shard_row]
    x = x_offset + arange(len(v)) * x_delta

所以"某条记录对应哪个 npz 的哪一行"**直接查表就有**；本器把它封装成一条命令，
让"在特征频率附近取幅值"这类下一步工作只需几行代码。

## 用法

    python scripts/spectra_lookup.py --window w0316 --turbine WTG01 --sensor Gear_HS_generator_side \
        --meas FFT_10000_Tr [--at 0.223] [--list]
    python scripts/spectra_lookup.py --window w0316 --stats        # 各 meas 的谱件/记录数、频率轴分辨率

退出码: 0 找到 · 2 窗/记录/谱件缺（会指明缺哪一件）
"""
from __future__ import annotations

import argparse
import pathlib
import sys

import numpy as np
import pandas as pd

ROOT = pathlib.Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from src import paths as P                                              # noqa: E402


def meta_path(win: str) -> pathlib.Path:
    """window → spectra_meta.parquet（两份落点都认：窗根与 spectra/ 下）。"""
    m5 = P.m5()
    if win == 'w0127':                       # 首窗特例: 索引在 m5 根, 谱在 m5/spectra
        cands = [m5 / 'spectra' / 'spectra_meta.parquet', m5 / 'spectra_meta.parquet']
    else:
        w = m5 / 'windows' / win
        cands = [w / 'spectra_meta.parquet', w / 'spectra' / 'spectra_meta.parquet']
    for c in cands:
        if c.is_file():
            return c
    return cands[0]


def store_of(mp: pathlib.Path) -> pathlib.Path:
    """npz 分片所在根：meta 里 shard 是**相对该根**的（含 p<NN>/ 子目录）。

    ★ 两份 meta 落点对应同一个根：`<窗>/spectra_meta.parquet` 与 `<窗>/spectra/spectra_meta.parquet`
      的 shard 都是相对 **`<窗>/spectra/`** 的（实测：meta 记 `p00/FFT_4_Tr_n401_0000.npz`，
      文件在 `<窗>/spectra/p00/…`）。第一版按 "meta 的父目录" 解析，于是报"分片不在位"。
    """
    return mp.parent if mp.parent.name == 'spectra' else (mp.parent / 'spectra')


def load(win: str, turbine: str, sensor: str, meas: str, trigger: str | None = None):
    mp = meta_path(win)
    if not mp.is_file():
        raise SystemExit(f'[X] 没有谱库元数据: {P.rel(mp)} —— 该窗的谱还没摄入 '
                         f'(scripts/rudong_tcm_spectra.py)；注意 w0127(1 月窗) 的谱库当前不在盘上。')
    m = pd.read_parquet(mp)
    sel = m[(m['turbine'] == turbine) & (m['sensor'] == sensor) & (m['meas_name'] == meas)]
    if sel.empty:
        raise SystemExit(f'[X] 表里没有 {turbine}/{sensor}/{meas} 的谱')
    if trigger:
        t = pd.to_datetime(sel['trigger_time'], errors='coerce')
        sel = sel.loc[[(t - pd.to_datetime(trigger)).abs().idxmin()]]
    r = sel.iloc[0]
    npz = store_of(mp) / str(r['shard'])
    npz = npz.with_suffix('.npz') if npz.suffix != '.npz' else npz
    if not npz.is_file():
        raise SystemExit(f'[X] 分片不在位: {P.rel(npz)}（meta 记的 shard={r["shard"]}）')
    with np.load(npz, allow_pickle=False) as z:
        v = z['values'][int(r['shard_row'])]
    x = float(r['x_offset']) + np.arange(len(v)) * float(r['x_delta'])
    return dict(row=r, npz=npz, v=np.asarray(v, dtype=float), x=x,
                x_offset=float(r['x_offset']), x_delta=float(r['x_delta']))


def main() -> int:
    ap = argparse.ArgumentParser(description='谱库查表: 记录 → npz 分片 (取数配方见脚本头)')
    ap.add_argument('--window', default='w0316')
    ap.add_argument('--turbine', default=None)
    ap.add_argument('--sensor', default=None)
    ap.add_argument('--meas', default=None)
    ap.add_argument('--trigger', default=None, help='指定 trigger_time（不指定取该组合第一条）')
    ap.add_argument('--at', type=float, default=None, help='看这个频率(Hz)处的幅值（最近线）')
    ap.add_argument('--stats', action='store_true', help='只打印该窗谱库的概况')
    a = ap.parse_args()
    mp = meta_path(a.window)
    if not mp.is_file():
        print(f'[X] 谱库元数据不在位: {P.rel(mp)}')
        print(f'    该窗的谱还没摄入。重新摄入: python scripts/rudong_tcm_index.py --root <decode 目录> '
              f'--out <窗>/index.parquet 然后 python scripts/rudong_tcm_spectra.py --root <同> --out <窗>/spectra')
        return 2
    m = pd.read_parquet(mp)
    if a.stats or not a.turbine:
        print(f'== 谱库概况 · 窗 {a.window} · {P.rel(mp)} ==')
        print(f'   谱件记录 {len(m)} 条 · 分片 {m["shard"].nunique()} 个 · 测点 {m["sensor"].nunique()} 个'
              f' · 机组 {m["turbine"].nunique()} 台')
        g = m.groupby('meas_name').agg(条数=('shard_row', 'size'), 分片=('shard', 'nunique'),
                                       dx=('x_delta', 'first'), x0=('x_offset', 'first'))
        print(g.sort_values('条数', ascending=False).head(18).to_string())
        print(f'   取数配方: v = npz["values"][shard_row]; x = x_offset + arange(len(v))*x_delta')
        return 0
    got = load(a.window, a.turbine, a.sensor, a.meas, a.trigger)
    r = got['row']
    print(f'== 记录 → npz ==')
    print(f'   {r["turbine"]} / {r["sensor"]} / {r["meas_name"]} @ {r["trigger_time"]}')
    print(f'   分片 {P.rel(got["npz"])} · 第 {int(r["shard_row"])} 行 · {len(got["v"])} 点'
          f' · x: {got["x_offset"]} + i*{got["x_delta"]} Hz ⇒ {got["x"][0]:.3f} ~ {got["x"][-1]:.3f} Hz')
    print(f'   幅值: min={got["v"].min():.6g} max={got["v"].max():.6g} 均值={got["v"].mean():.6g}')
    if a.at is not None:
        i = int(np.argmin(np.abs(got['x'] - a.at)))
        print(f'   @{a.at} Hz（最近线 {got["x"][i]:.3f} Hz, 第 {i} 条）: {got["v"][i]:.6g}')
    return 0


if __name__ == '__main__':
    for _s in (sys.stdout, sys.stderr):
        try:
            _s.reconfigure(errors='replace')
        except Exception:
            pass
    sys.exit(main())