#!/usr/bin/env python3
"""fir_probe.py — численная реплика FIR-цепи (Этап A3) против захватов.

Структура по дизасму (BLOCKMAP 24mm6/24mm8 + wrap/worker декод этого раунда):
  copy:    FIR[2j]=scr[j], FIR[2j+1]=0                (18004d900, 2049 пар)
  opA:     th2180 = INVERSE real-FFT (план buf548, N=4096, scale 1/4096)
  scale:   float[1..2047] *= 2.0 ; float[2049..4095] = 0   (52d920/52db50)
  opB:     th1a90 = FORWARD real-FFT
  EXP:     expf in-place по первым 2049 ФЛОАТАМ          (140b30, 52b708-716)
  opC:     th2180 = INVERSE
  window:  float[0..2047] *= WINfreq[2048..4095]         (52d990, падающий Hann)
           float[2048..4095] = 0                         (52db50)
  opD:     th1a90 = FORWARD
  fix:     FIR[0]=1.0, FIR[1]=0                          (52b7cd-e1)
  df0:     track_i := track_i ⊗ FIR (комплексное умножение, 18000b3c0)
Цель: воспроизвести cur@678 из trk@688 без свободных параметров.
"""
import numpy as np
import glob
import os
import sys

NFLOAT = 4098          # 2049 пар
NBINS = 2049           # n/2+1, n=[ctx+0x540534]=4096


def load_frame(npz):
    d = np.load(npz)
    S = {}
    for k in d.keys():
        if k.startswith('0x'):
            S[int(k[2:], 16)] = d[k]
    return S


def win_periodic_hann(N):
    return 0.5 * (1.0 - np.cos(2.0 * np.pi * np.arange(N) / N)).astype(np.float64)


def fir_chain(scr, winfall, variant='flat'):
    """scr: 2049 float (log-домен). Возвращает halfcomplex-спектр ядра F[2049]."""
    # copy/pack: пары (re=scr, im=0) -> inverse rfft вход (numpy: complex[2049])
    H = scr.astype(np.float64).astype(np.complex128)
    # opA: inverse real FFT, нормировка 1/N (план scale=2^-12 при активном флаге)
    y = np.fft.irfft(H, n=4096)              # уже содержит деление на 4096
    # scale/zero по asm: float[1..2047]*=2, float[2049..]=0 (f[2048] не трогаем)
    y[1:2048] *= 2.0
    y[2049:] = 0.0
    # opB: forward
    Y = np.fft.rfft(y, n=4096)               # complex[2049]
    # EXP по первым 2049 флоатам плоского массива
    flat = np.empty(NFLOAT)
    flat[0::2] = Y.real
    flat[1::2] = Y.imag
    if variant == 'flat':
        flat[:2049] = np.exp(flat[:2049])
    elif variant == 'cplx':
        Y = np.exp(Y.astype(np.complex128))
        flat[0::2] = Y.real
        flat[1::2] = Y.imag
    Y2 = flat[0::2] + 1j * flat[1::2]
    # opC: inverse
    w = np.fft.irfft(Y2, n=4096)
    # window: float[0..2047] *= падающая половина; хвост = 0
    w[:2048] *= winfall
    w[2048:] = 0.0
    # opD: forward
    F = np.fft.rfft(w, n=4096)
    # fix: FIR[0]=1.0, FIR[1]=0
    F[0] = 1.0 + 0.0j
    return F


def evaluate(ds, ph_file, verbose=True):
    S = load_frame(os.path.join(ds, ph_file))
    scr = S[0x540628][:NBINS].astype(np.float64)
    trk = S[0x540688][:NBINS].astype(np.float64)
    cur = S[0x540678][:NBINS].astype(np.float64)
    # проверка trk == exp(scr)
    m_ok = trk > 1e-30
    err_trk = np.abs(np.log(trk[m_ok]) - scr[m_ok]).max()
    # фит gamma
    sel = np.abs(scr) > 0.05
    g = float(np.sum(scr[sel] * np.log(cur[sel])) / np.sum(scr[sel] ** 2))
    rms_fit = float(np.sqrt(np.mean((np.log(cur[sel]) - g * scr[sel]) ** 2)))
    winfall = win_periodic_hann(4096)[2048:]
    out = []
    for variant in ('flat', 'cplx'):
        F = fir_chain(scr, winfall, variant)
        # маска = track ⊗ F (df0), берём реальную часть как применённую маску
        mask_sim = np.abs(trk * F) if variant == 'cplx' else trk * F.real
        mm = (cur > 1e-6) & np.isfinite(mask_sim)
        e_db = 20.0 / np.log(10) * np.log(np.abs(mask_sim[mm])) - \
               20.0 / np.log(10) * np.log(cur[mm])
        rms_db = float(np.sqrt(np.mean(e_db ** 2)))
        out.append((variant, rms_db, int(mm.sum())))
        if verbose:
            print(f'{ph_file} [{variant}] gamma_fit={g:.6f} (rms {rms_fit:.1e}) '
                  f'trk_err={err_trk:.2e} MASK rms={rms_db:.4f} дБ / {mm.sum()} бинов')
    return out


if __name__ == '__main__':
    jobs = [
        ('/tmp/opencode/sc_multi4b', 'ph073.npz'),
        ('/tmp/opencode/sc_multi6', 'ph037.npz'),
        ('/tmp/opencode/sc_multi6', 'ph034.npz'),
    ]
    if len(sys.argv) > 1:
        jobs = [(os.path.dirname(sys.argv[1]), os.path.basename(sys.argv[1]))]
    for ds, ph in jobs:
        evaluate(ds, ph)