#!/usr/bin/env python3
# -*- coding: utf-8 -*-
r"""六层链 `energy_share` 步 —— **谱线能量占比**（最后补上的那一步脚本）。

## 它在链里的位置与口径

`src/windcms/pipeline.py::analyze` 的步序是 `oem_scan → **energy_share** → model_run → fusion → report`：
扫线（哪些特征频率有线）之后，本步回答"**这条线占了多少能量**"。此前这一步脚本没随包
（`chain_gap_check` 报"1 个步骤脚本未随包"），于是链上少一环、也没有任何件回答这个问题。

口径（写死在这里，供复核）：

    band  = 目标特征频率 ±tol_bins·dx（默认 ±2 bin，与 model_run 的取峰窗同宽）
    峰能量  = Σ y[i]² （i ∈ band）
    总能量  = Σ y[j]² （整条谱，去掉 DC 前 3 bin）
    energy_pct = 峰能量 ÷ 总能量 × 100

    ★只报**占比**，不定级：占比小不等于没问题（宽带抬升会摊薄单线占比），
      故消费端（model_run 的 `vib_verdict_and_writeback`）把 energy_pct<0.01% 的结论降为「参考」。

输入: `m5_cms_tcm/windows/<窗>/index.parquet` + `spectra/*`（振动摄入的产物）
      候选线来自 `reference/rudong/oem_scan_plan.json`（与 `rudong_tcm_oem_scan.py` 同一份配置）
输出: `m5_cms_tcm/line_energy_share.parquet`
      列: window, turbine, sensor, meas, component, line, hz, band_hz, energy_pct, peak, total, n_bins

用法: python scripts/rudong_line_energy_share.py [--window w0316] [--limit-turbines N]
"""
from __future__ import annotations

import argparse
import pathlib
import sys
import time

import numpy as np
import pandas as pd

ROOT = pathlib.Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(ROOT / 'scripts'))
from src import paths as P                                              # noqa: E402
import rudong_tcm_oem_scan as OS                                        # noqa: E402  候选线(分母)同一实现
import spectra_lookup as SL                                             # noqa: E402

LINES = (('BPFI', 'BPFI_hz'), ('BPFO', 'BPFO_hz'), ('BSF', 'BSF_hz'))


def windows_of(explicit=None) -> list[str]:
    m5 = P.m5()
    got = sorted(d.name for d in (m5 / 'windows').glob('w[0-9][0-9][0-9][0-9]')
                 if (d / 'index.parquet').is_file())
    return [explicit] if explicit else got


def main() -> int:
    ap = argparse.ArgumentParser(description='六层链 energy_share 步: 谱线能量占比')
    ap.add_argument('--window', default=None)
    ap.add_argument('--tol-bins', type=int, default=2)
    ap.add_argument('--limit-turbines', type=int, default=0)
    ap.add_argument('--out', default=None)
    a = ap.parse_args()
    _plan, comps = OS.load_plan()
    wins = windows_of(a.window)
    if not wins:
        print('[跳过] 没有窗索引（先跑振动摄入 scripts/vib_raw_build.py）')
        return 0
    rows = []
    t0 = time.time()
    for w in wins:
        mp = SL.meta_path(w)
        if not mp.is_file():
            continue
        meta = pd.read_parquet(mp, columns=['turbine', 'sensor', 'meas_name'])
        avail = set(map(tuple, meta.drop_duplicates().values))
        turbines = sorted({t for t, _s, _m in avail})
        if a.limit_turbines:
            turbines = turbines[:a.limit_turbines]
        for c in comps:
            sen, meas = c['sensor'], c['shard']
            for t in turbines:
                if (t, sen, meas) not in avail:
                    continue
                try:
                    rec = SL.load(w, t, sen, meas)
                except BaseException:
                    continue
                if not rec:
                    continue
                y = np.asarray(rec['v'], dtype=float)
                dx = float(rec.get('x_delta') or c.get('dx_hz') or 0)
                if not dx or y.size < 8:
                    continue
                y2 = y[3:] ** 2                       # 去前 3 bin(DC/趋势), 与"总能量"口径一致
                total = float(np.nansum(y2))
                if not np.isfinite(total) or total <= 0:
                    continue
                for label, col in LINES:
                    hz = c.get(col)
                    if not hz or not np.isfinite(hz) or hz <= 0:
                        continue
                    k = int(round(float(hz) / dx))
                    lo, hi = max(k - a.tol_bins - 3, 0), min(k + a.tol_bins - 3, y2.size)
                    if hi <= lo:
                        continue
                    band = y2[lo:hi]
                    peak = float(np.nansum(band))
                    rows.append(dict(window=w, turbine=t, sensor=sen, meas=meas,
                                     component=c['component'], line=label, hz=round(float(hz), 3),
                                     band_hz=round(2 * a.tol_bins * dx, 3),
                                     energy_pct=round(100.0 * peak / total, 6),
                                     peak=peak, total=total, n_bins=int(hi - lo)))
    if not rows:
        print(f'[X] 一条占比也没算出来（窗 {wins} 的谱库不可用?）—— 不写空文件')
        return 2
    df = pd.DataFrame(rows).sort_values(['window', 'turbine', 'component', 'line'])
    out = pathlib.Path(a.out) if a.out else (P.m5() / 'line_energy_share.parquet')
    out.parent.mkdir(parents=True, exist_ok=True)
    df.to_parquet(out, index=False)
    top = df.nlargest(3, 'energy_pct')[['turbine', 'component', 'line', 'energy_pct']].to_dict('records')
    print(f'已写 {P.rel(out)}: {len(df)} 行 × {len(df.columns)} 列（窗 {",".join(wins)}，{time.time()-t0:.0f}s）· '
          f'占比最高: {top}')
    try:
        from src import derived_manifest as DM
        DM.record(P.out_root(), {out.relative_to(P.out_root()).as_posix():
                  'scripts/rudong_line_energy_share.py (线能量占比 = ±2bin 带内能量 ÷ 全谱能量; 只报占比不定级)'},
                  by='rudong_line_energy_share')
        print('   已自登记 → _derived_manifest.json')
    except Exception as e:
        print(f'   [i] 自登记跳过: {type(e).__name__}: {e}')
    return 0


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