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

Reports wall time, waveforms/s, CPU utilization (utime+stime over wall, in
cores), and verifies every variant against the exact reference image.

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, set_fft_backend  # 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 bench(sras, pads, backends):
    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
    print(f"{n_wf} waveforms x {spf} samples, {sras.n_angles} angle(s)")
    print(f"{'pad':>4} {'backend':>8} {'variant':>16} {'wall':>9} "
          f"{'wf/s':>10} {'util':>6}  match")

    for pad in pads:
        n_fft = spf * pad if pad > 1 else None
        for backend in backends:
            set_fft_backend(backend)

            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(exact=True))
            rows = [("exact(serial)", ref, wall, util, True)]
            for label, kw in (("zoom(serial)", dict(max_workers=1)),
                              ("zoom(pool)", {})):
                img, wall, util = _timed(lambda: run(**kw))
                rows.append((label, img, wall, util, bool(np.array_equal(img, ref))))
            for label, img, wall, util, ok in rows:
                print(f"{pad:>4} {backend:>8} {label:>16} {wall:>8.2f}s "
                      f"{n_wf / wall:>10.0f} {util:>5.1f}x  "
                      f"{'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("--backends", default=None,
                   help="comma-separated (default: scipy,pyfftw if available)")
    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.backends:
        backends = args.backends.split(",")
    else:
        backends = ["scipy"] + (["pyfftw"] if compute.PYFFTW_AVAILABLE else [])

    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, backends)
    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, backends)


if __name__ == "__main__":
    main()