#!/usr/bin/env python3
"""Benchmark the FFT peak-search path: serial vs pooled.

Reports wall time, waveforms/s, CPU utilization (utime+stime over wall, in
cores), the estimated per-thread resident pyFFTW plan footprint at each pad
factor (see sras_compute._fft_block_for), and verifies pooled output against
the serial reference.

Usage:
    python tools/bench_fft.py                       # synthetic, pads 1/8/40
    python tools/bench_fft.py --pads 40 --spf 2500 --rows 8 --frames 1024
    python tools/bench_fft.py --real /path/big.sras --real-rows 32 --pads 40
"""

import argparse
import resource
import sys
import tempfile
import time
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import sras_compute as compute                              # noqa: E402
from sras_compute import compute_rf_image                    # noqa: E402
from sras_format import SrasFile                             # noqa: E402
import tools.make_test_sras as gen                           # noqa: E402
from tools.check_equivalence import row_slice                # noqa: E402


def _timed(fn):
    r0 = resource.getrusage(resource.RUSAGE_SELF)
    t0 = time.perf_counter()
    out = fn()
    wall = time.perf_counter() - t0
    r1 = resource.getrusage(resource.RUSAGE_SELF)
    cpu = (r1.ru_utime - r0.ru_utime) + (r1.ru_stime - r0.ru_stime)
    return out, wall, cpu / max(wall, 1e-9)


def _plan_mb(spf: int, n_len: int, n_workers: int) -> float:
    """Estimated resident pyFFTW plan-buffer footprint across the whole
    pool at this transform length (see sras_compute._fft_block_for)."""
    block = compute._fft_block_for(spf, n_len)
    bytes_per_wf = 4 * spf + 8 * (n_len // 2 + 1)
    return block * bytes_per_wf * n_workers / (1024 * 1024)


def bench(sras, pads):
    n_wf = sum(int(sras.n_rows[a]) * int(sras.n_frames[a])
               for a in range(sras.n_angles))
    spf = sras.samples_per_frame
    n_workers = compute._MAX_WORKERS
    print(f"{n_wf} waveforms x {spf} samples, {sras.n_angles} angle(s), "
          f"{n_workers} workers")
    print(f"{'pad':>4} {'variant':>8} {'wall':>9} {'wf/s':>10} {'util':>6} "
          f"{'plan MB':>9}  match")

    for pad in pads:
        n_fft = spf * pad if pad > 1 else None
        n_len = n_fft if n_fft is not None else spf
        plan_mb = _plan_mb(spf, n_len, n_workers)

        def run(**kw):
            imgs = [compute_rf_image(sras, a, dc_threshold_mv=None,
                                     apply_bg_sub=True, n_fft=n_fft, **kw)
                    for a in range(sras.n_angles)]
            return np.concatenate([i.ravel() for i in imgs])

        ref, wall, util = _timed(lambda: run(max_workers=1))
        rows = [("serial", ref, wall, util, True)]
        img, wall, util = _timed(lambda: run())
        rows.append(("pooled", img, wall, util, bool(np.array_equal(img, ref))))
        for label, img, wall, util, ok in rows:
            print(f"{pad:>4} {label:>8} {wall:>8.2f}s {n_wf / wall:>10.0f} "
                  f"{util:>5.1f}x {plan_mb:>8.1f}  {'OK' if ok else 'MISMATCH'}")


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--pads", default="1,8,40",
                   help="comma-separated pad factors (default 1,8,40)")
    p.add_argument("--spf", type=int, default=2500)
    p.add_argument("--rows", type=int, default=8)
    p.add_argument("--frames", type=int, default=1024)
    p.add_argument("--real", help="path to a real .sras file")
    p.add_argument("--real-rows", type=int, default=32,
                   help="rows of angle 0 to use from the real file")
    args = p.parse_args()

    pads = [int(x) for x in args.pads.split(",")]

    if args.real:
        sras = row_slice(SrasFile(args.real), 0, args.real_rows)
        sras.data = [sras.data[0]]
        sras.n_angles = 1
        bench(sras, pads)
    else:
        with tempfile.TemporaryDirectory(prefix="sras_bench_") as tmp:
            path = Path(tmp) / "bench.sras"
            gen.write(path, n_angles=1, seed=0, samples_per_frame=args.spf,
                      geometry=[(args.rows, args.frames)])
            bench(SrasFile(str(path)), pads)


if __name__ == "__main__":
    main()