"""Behavioural tests for the compute/format layer.

Covers what the golden-hash harness can't: the v6->v7 cache round-trip
(including block carry-forward), parallel-vs-serial identity, the no-mask
fast path, and the ROI bounding-box mask optimisation.
"""

import subprocess
import sys
from pathlib import Path

import numpy as np
import pytest

import sras_compute as compute
from sras_compute import (
    cache_file, compute_dc_image, compute_rf_image, dc_image_mv,
)
from sras_format import CH3_IDX, CH4_IDX, SrasFile, adc_to_mv
import tools.make_test_sras as gen

REPO = Path(__file__).resolve().parent.parent


def test_cache_roundtrip(tmp_path):
    """v6 -> v7 for DC, then FFT, asserting the first block survives the
    second write (the carry-forward path in write_v7_cache)."""
    path = tmp_path / "roundtrip.sras"
    gen.write(path, n_angles=3, seed=1, samples_per_frame=64)

    src = SrasFile(str(path))
    assert src.version == 6, f"got v{src.version}"
    expect_dc3 = [dc_image_mv(src, a, CH3_IDX) for a in range(src.n_angles)]
    expect_dc4 = [dc_image_mv(src, a, CH4_IDX) for a in range(src.n_angles)]
    expect_fft = [compute_rf_image(src, a, dc_threshold_mv=None, apply_bg_sub=True)
                  for a in range(src.n_angles)]

    err = cache_file(str(path), "dc", True)
    assert err == "", err

    after_dc = SrasFile(str(path))
    assert after_dc.version == 7, f"got v{after_dc.version}"
    assert all(x is not None for x in after_dc.precomputed_dc3_mv)
    assert all(np.allclose(after_dc.precomputed_dc3_mv[a], expect_dc3[a], atol=1e-4)
               for a in range(after_dc.n_angles))
    assert all(np.allclose(after_dc.precomputed_dc4_mv[a], expect_dc4[a], atol=1e-4)
               for a in range(after_dc.n_angles))
    assert all(x is None for x in after_dc.precomputed_freq_mhz), "no fft block yet"
    assert (after_dc.precomputed_dc3_mv[0].dtype == np.float32
            and after_dc.precomputed_dc3_mv[0].dtype.byteorder in ("=", "|")), \
        "cached images are native float32"
    assert after_dc.precomputed_dc3_mv[0].flags.writeable

    err = cache_file(str(path), "fft", True)
    assert err == "", err

    both = SrasFile(str(path))
    assert all(x is not None for x in both.precomputed_freq_mhz)
    assert all(np.allclose(both.precomputed_freq_mhz[a], expect_fft[a], atol=1e-3)
               for a in range(both.n_angles))
    assert all(np.allclose(both.precomputed_dc3_mv[a], expect_dc3[a], atol=1e-4)
               for a in range(both.n_angles)), \
        "DC block carried forward through the FFT write"
    assert both.precomputed_bg_sub is True

    # The fast path must reproduce a fresh compute, and masking must still
    # apply on top of a cached (unmasked) image.
    fresh = SrasFile(str(path))
    fresh.precomputed_freq_mhz = [None] * fresh.n_angles
    dc4 = dc_image_mv(both, 0, CH4_IDX)
    thr = float(np.median(dc4))
    assert np.allclose(
        compute_rf_image(both, 0, dc_threshold_mv=None, apply_bg_sub=True),
        compute_rf_image(fresh, 0, dc_threshold_mv=None, apply_bg_sub=True),
        atol=1e-3), "cached fast path == fresh compute (unmasked)"
    assert np.allclose(
        compute_rf_image(both, 0, dc_threshold_mv=thr, apply_bg_sub=True),
        compute_rf_image(fresh, 0, dc_threshold_mv=thr, apply_bg_sub=True),
        atol=1e-3), "cached fast path == fresh compute (masked)"

    # Waveform data must be byte-identical to the pre-cache file.
    orig = tmp_path / "roundtrip_orig.sras"
    gen.write(orig, n_angles=3, seed=1, samples_per_frame=64)
    o, n = SrasFile(str(orig)), SrasFile(str(path))
    assert all(np.array_equal(np.asarray(o.data[a]), np.asarray(n.data[a]))
               for a in range(o.n_angles)), \
        "waveform data untouched by the cache write"


def test_partial_v7_cache(tmp_path):
    """Only some angles cached: uncached angles must compute, not read zeros.
    This is the v5 bug the ragged normalisation fixed, checked via v7."""
    path = tmp_path / "partial.sras"
    gen.write(path, n_angles=3, seed=2, samples_per_frame=64)

    src = SrasFile(str(path))
    expected = [compute_rf_image(src, a, dc_threshold_mv=None, apply_bg_sub=True)
                for a in range(src.n_angles)]
    partial = [expected[0], None, expected[2]]     # angle 1 deliberately absent
    src.write_v7_cache(new_freq_mhz=partial, new_bg_sub=True)

    reread = SrasFile(str(path))
    assert reread.precomputed_freq_mhz[1] is None
    assert (reread.precomputed_freq_mhz[0] is not None
            and reread.precomputed_freq_mhz[2] is not None)
    img1 = compute_rf_image(reread, 1, dc_threshold_mv=None, apply_bg_sub=True)
    assert np.any(img1 != 0) and np.allclose(img1, expected[1], atol=1e-3), \
        "uncached angle computes rather than returning zeros"


def test_parallel_identity(tmp_path, monkeypatch):
    """Forcing 1 worker vs many must give identical output — catches
    chunk-boundary and race bugs."""
    path = tmp_path / "parallel.sras"
    # Many rows, so the row loop actually splits into several chunks.
    n_rows, n_frames, spf = 48, 9, 256
    gen.write(path, n_angles=1, seed=3, samples_per_frame=spf,
              geometry=[(n_rows, n_frames)])
    sras = SrasFile(str(path))

    # Shrink the budget so the outer row loop splits into many chunks, and
    # the block size so every chunk splits into many FFT tasks — the worst
    # case for boundary bugs.
    monkeypatch.setattr(compute, "_TOTAL_BYTES_BUDGET", 8 * n_frames * spf * 4)
    monkeypatch.setattr(compute, "_FFT_BLOCK_MAX", 4)
    fft_rows = compute._plan_fft_rows(n_frames, spf, compute._TOTAL_BYTES_BUDGET)
    assert fft_rows < n_rows, \
        f"FFT work actually splits into multiple chunks ({fft_rows} of {n_rows})"
    dc_rows = compute._chunk_rows_for(n_frames, spf, compute._TOTAL_BYTES_BUDGET)
    assert dc_rows < n_rows, \
        f"DC work actually splits into multiple chunks ({dc_rows} of {n_rows})"

    monkeypatch.setattr(compute, "_MAX_WORKERS", 1)
    dc_serial = compute_dc_image(sras, 0, CH4_IDX)
    rf_serial = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
    dc4 = adc_to_mv(dc_serial, *sras.cal(CH4_IDX))
    thr = float(np.median(dc4))
    rf_masked_serial = compute_rf_image(sras, 0, dc_threshold_mv=thr,
                                        apply_bg_sub=True)
    rf_pad_serial = compute_rf_image(sras, 0, dc_threshold_mv=thr,
                                     apply_bg_sub=True, n_fft=spf * 8)

    monkeypatch.setattr(compute, "_MAX_WORKERS", 8)
    dc_par = compute_dc_image(sras, 0, CH4_IDX)
    rf_par = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=True)
    rf_masked_par = compute_rf_image(sras, 0, dc_threshold_mv=thr, apply_bg_sub=True)
    rf_pad_par = compute_rf_image(sras, 0, dc_threshold_mv=thr,
                                  apply_bg_sub=True, n_fft=spf * 8)

    assert np.array_equal(dc_serial, dc_par), "dc image identical"
    assert np.array_equal(rf_serial, rf_par), "rf image identical (unmasked)"
    assert np.array_equal(rf_masked_serial, rf_masked_par), \
        "rf image identical (masked)"
    assert np.array_equal(rf_pad_serial, rf_pad_par), \
        "rf image identical (masked, padded/zoom)"


def test_peak_bins_fuzz():
    """Hammer _peak_bins directly with adversarial spectra: noise,
    un-subtracted DC offsets, on-bin and off-bin tones, near-tie tone pairs,
    and all-zero rows — against an independent scipy.fft reference."""
    import scipy.fft as scipy_fft

    rng = np.random.default_rng(42)
    for _ in range(25):
        spf = int(rng.integers(16, 220))
        pad = int(rng.choice([4, 5, 8, 16, 40]))
        n_fft = spf * pad
        n_wf = 24
        w = rng.normal(scale=20.0, size=(n_wf, spf))
        t = np.arange(spf)
        # rows 0-5: pure/noisy tones (some off-bin), row 6-7: near-tie pair,
        # row 8: big DC offset, row 9: all zeros, rest: plain noise.
        for r in range(6):
            f = rng.uniform(1.0, spf / 2 - 1)
            w[r] = 60 * np.sin(2 * np.pi * f * t / spf) + w[r] * (r % 2)
        f1, f2 = rng.uniform(2.0, spf / 2 - 2, size=2)
        w[6] = 50 * np.sin(2 * np.pi * f1 * t / spf) \
            + 49.9 * np.sin(2 * np.pi * f2 * t / spf)
        w[7] = 50 * np.sin(2 * np.pi * f1 * t / spf) \
            + 50 * np.cos(2 * np.pi * f2 * t / spf)
        w[8] = 90 + rng.normal(scale=5.0, size=spf)
        w[9] = 0.0
        w = w.astype(np.float32)

        S = scipy_fft.rfft(w, n=n_fft, axis=-1, workers=1)
        P = S.real ** 2
        P += S.imag ** 2
        P[:, 0] = 0.0
        ref = np.argmax(P, axis=1)

        got = compute._peak_bins(w, n_fft)
        bad = np.nonzero(ref != got)[0]
        assert not len(bad), \
            (f"spf={spf} pad={pad}: rows {bad.tolist()} picked "
             f"{got[bad].tolist()} instead of {ref[bad].tolist()}")


def test_peak_bins_fuzz_min_freq():
    """Same adversarial-spectra fuzz as test_peak_bins_fuzz, but with a swept
    min_bin floor: bins below the floor must be excluded from the argmax
    exactly as the independent scipy.fft reference is, when zeroed the same
    way before argmax."""
    import scipy.fft as scipy_fft

    rng = np.random.default_rng(43)
    for _ in range(25):
        spf = int(rng.integers(16, 220))
        pad = int(rng.choice([4, 5, 8, 16, 40]))
        n_fft = spf * pad
        n_wf = 24
        w = rng.normal(scale=20.0, size=(n_wf, spf))
        t = np.arange(spf)
        for r in range(6):
            f = rng.uniform(1.0, spf / 2 - 1)
            w[r] = 60 * np.sin(2 * np.pi * f * t / spf) + w[r] * (r % 2)
        f1, f2 = rng.uniform(2.0, spf / 2 - 2, size=2)
        w[6] = 50 * np.sin(2 * np.pi * f1 * t / spf) \
            + 49.9 * np.sin(2 * np.pi * f2 * t / spf)
        w[7] = 50 * np.sin(2 * np.pi * f1 * t / spf) \
            + 50 * np.cos(2 * np.pi * f2 * t / spf)
        w[8] = 90 + rng.normal(scale=5.0, size=spf)
        w[9] = 0.0
        w = w.astype(np.float32)

        S = scipy_fft.rfft(w, n=n_fft, axis=-1, workers=1)
        P = S.real ** 2
        P += S.imag ** 2
        n_bins_fine = n_fft // 2 + 1
        min_bin = int(rng.integers(1, max(2, n_bins_fine // 3)))
        P[:, :min_bin] = 0.0
        ref = np.argmax(P, axis=1)

        got = compute._peak_bins(w, n_fft, min_bin)
        bad = np.nonzero(ref != got)[0]
        assert not len(bad), \
            (f"spf={spf} pad={pad} min_bin={min_bin}: rows {bad.tolist()} picked "
             f"{got[bad].tolist()} instead of {ref[bad].tolist()}")


@pytest.mark.parametrize("spf", [64, 500, 2500])
def test_fft_block_for(spf):
    """Block size == _FFT_BLOCK_MAX at natural resolution, shrinks and stays
    >= _FFT_BLOCK_MIN as n_len grows, and the implied per-thread byte
    estimate respects _FFT_PLAN_BYTES_BUDGET except when the floor is
    engaged."""
    block_natural = compute._fft_block_for(spf, spf)
    assert block_natural == compute._FFT_BLOCK_MAX

    prev = compute._FFT_BLOCK_MAX
    for pad in (2, 4, 8, 40, 500):
        n_len = spf * pad
        block = compute._fft_block_for(spf, n_len)
        assert compute._FFT_BLOCK_MIN <= block <= prev
        bytes_per_wf = 4 * spf + 8 * (n_len // 2 + 1)
        if block > compute._FFT_BLOCK_MIN:
            assert block * bytes_per_wf <= compute._FFT_PLAN_BYTES_BUDGET
        prev = block


def test_compute_rf_image_min_freq_mhz(tmp_path):
    """min_freq_mhz threads through compute_rf_image end-to-end, for both
    the natural-resolution and padded paths: 0.0 (default) must reproduce
    the pre-existing image exactly, and a floor above every real peak must
    collapse the image to bin 0 (0 MHz) — the same fallback the low-level
    search uses when nothing survives the floor."""
    path = tmp_path / "floor_e2e.sras"
    gen.write(path, n_angles=1, seed=12, samples_per_frame=64)
    sras = SrasFile(str(path))
    huge_floor = float(sras.freq_axis_mhz(None)[-1]) + 1.0   # above Nyquist

    for n_fft in (None, 64 * 8):
        unfiltered = compute_rf_image(sras, 0, dc_threshold_mv=None,
                                      apply_bg_sub=False, n_fft=n_fft)
        same = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=False,
                                n_fft=n_fft, min_freq_mhz=0.0)
        assert np.array_equal(unfiltered, same), \
            f"n_fft={n_fft}: min_freq_mhz=0.0 changed the output"
        assert unfiltered.any(), \
            f"n_fft={n_fft}: fixture should have real signal"

        collapsed = compute_rf_image(sras, 0, dc_threshold_mv=None, apply_bg_sub=False,
                                     n_fft=n_fft, min_freq_mhz=huge_floor)
        assert not collapsed.any(), \
            f"n_fft={n_fft}: floor above Nyquist should collapse to 0 MHz"


def test_nomask_equals_low_threshold(tmp_path):
    """dc_threshold_mv=None must equal a threshold below every pixel, while
    skipping the CH4 read."""
    path = tmp_path / "nomask.sras"
    gen.write(path, n_angles=2, seed=4, samples_per_frame=128)
    sras = SrasFile(str(path))
    for a in range(sras.n_angles):
        none_img = compute_rf_image(sras, a, dc_threshold_mv=None, apply_bg_sub=True)
        low_img = compute_rf_image(sras, a, dc_threshold_mv=-1e9, apply_bg_sub=True)
        assert np.array_equal(none_img, low_img), \
            f"angle {a}: None == -1e9 threshold"
        assert len(np.unique(none_img)) > 1, \
            f"angle {a}: image is degenerate ({len(np.unique(none_img))} unique)"


def test_roi_mask():
    """The bbox-restricted mask must equal a full-grid point-in-polygon test."""
    from matplotlib.path import Path as MplPath
    from sras_viewer import RoiQuad

    rng = np.random.default_rng(0)
    x = np.linspace(-2.0, 3.0, 137)
    y = np.linspace(1.0, 4.0, 91)

    cases = {
        "axis-aligned rect": np.array([[0.0, 1.5], [1.0, 1.5], [1.0, 3.0], [0.0, 3.0]]),
        "skewed quad": np.array([[-0.5, 1.2], [1.7, 1.9], [1.2, 3.4], [-1.0, 2.6]]),
        "entirely outside": np.array([[8.0, 8.0], [9.0, 8.0], [9.0, 9.0], [8.0, 9.0]]),
        "covers whole grid": np.array([[-9.0, -9.0], [9.0, -9.0], [9.0, 9.0], [-9.0, 9.0]]),
        "straddles left edge": np.array([[-4.0, 2.0], [0.5, 2.0], [0.5, 3.0], [-4.0, 3.0]]),
    }
    for _ in range(5):
        cases[f"random {_}"] = rng.uniform([-2.5, 0.5], [3.5, 4.5], size=(4, 2))

    for name, pts in cases.items():
        roi = RoiQuad(pts)
        fast = roi.mask_for_grid(x, y)
        X, Y = np.meshgrid(x.astype(np.float64), y.astype(np.float64))
        slow = MplPath(pts).contains_points(
            np.column_stack([X.ravel(), Y.ravel()])).reshape(X.shape)
        assert np.array_equal(fast, slow), f"{name} ({int(slow.sum())} px inside)"

    # Descending y axis (images are stored top-down in some scans).
    roi = RoiQuad(cases["skewed quad"])
    y_desc = y[::-1]
    fast = roi.mask_for_grid(x, y_desc)
    X, Y = np.meshgrid(x.astype(np.float64), y_desc.astype(np.float64))
    slow = MplPath(cases["skewed quad"]).contains_points(
        np.column_stack([X.ravel(), Y.ravel()])).reshape(X.shape)
    assert np.array_equal(fast, slow), "descending y axis"


def test_legacy_parse(tmp_path):
    """v2-v4 parsing against known written data."""
    for version in (2, 3, 4):
        path = tmp_path / f"legacy_v{version}.sras"
        meta = gen.write_legacy(path, version=version, n_angles=2, n_rows=4,
                                n_frames=12, samples_per_frame=32, seed=version)
        s = SrasFile(str(path))
        assert s.version == version, f"got v{s.version}"
        assert list(s.n_rows) == [4, 4] and list(s.n_frames) == [12, 12], \
            f"rows={list(s.n_rows)} frames={list(s.n_frames)}"
        assert all(np.array_equal(np.asarray(s.data[a]), meta["data"][a])
                   for a in range(s.n_angles)), \
            f"v{version} waveform data matches what was written"
        assert (s.background is not None) == (version >= 4), \
            f"v{version} background {'present' if version >= 4 else 'absent'}"
        assert (isinstance(s.precomputed_freq_mhz, list)
                and len(s.precomputed_freq_mhz) == s.n_angles), \
            f"v{version} precomputed stores are ragged lists"
        # DC image must equal a direct mean of the known input.
        expect = meta["data"][0][:, CH3_IDX, :, :].astype(np.float64).mean(axis=-1)
        assert np.allclose(compute_dc_image(s, 0, CH3_IDX), expect, atol=1e-3), \
            f"v{version} DC image equals a direct mean"


def test_sras_average(tmp_path):
    """The sras_average.py CLI: frame averaging with remainder handling."""
    src = tmp_path / "legacy_v4.sras"
    meta = gen.write_legacy(src, version=4, n_angles=2, n_rows=4, n_frames=12,
                            samples_per_frame=32, seed=4)
    dst = tmp_path / "legacy_v4_avg.sras"
    proc = subprocess.run(
        [sys.executable, str(REPO / "sras_average.py"), str(src), str(dst), "--n", "4"],
        capture_output=True, text=True, cwd=REPO)
    assert proc.returncode == 0, (proc.stderr or proc.stdout).strip()[-200:]

    avg = SrasFile(str(dst))
    assert avg.version == 4
    assert list(avg.n_frames) == [3, 3], f"{list(avg.n_frames)}"
    assert (avg.n_angles == 2 and list(avg.n_rows) == [4, 4]
            and avg.n_channels == meta["n_channels"])
    assert np.allclose(avg.ch_ymult_mv, SrasFile(str(src)).ch_ymult_mv), \
        "calibration preserved"
    assert np.array_equal(avg.background, SrasFile(str(src)).background), \
        "background preserved"
    src_data = meta["data"]
    expect0 = src_data[0][:, :, 0:4, :].astype(np.float32).mean(axis=2).astype(np.int16)
    assert np.array_equal(np.asarray(avg.data[0])[:, :, 0, :], expect0), \
        "first averaged group equals the mean of its 4 source frames"

    # Remainder handling: 12 frames / 5 -> 2 full groups + 1 partial.
    dst2 = tmp_path / "legacy_v4_avg5.sras"
    subprocess.run([sys.executable, str(REPO / "sras_average.py"),
                    str(src), str(dst2), "--n", "5"],
                   capture_output=True, text=True, cwd=REPO)
    assert list(SrasFile(str(dst2)).n_frames) == [3, 3], \
        "partial trailing group kept by default"
    dst3 = tmp_path / "legacy_v4_avg5d.sras"
    subprocess.run([sys.executable, str(REPO / "sras_average.py"),
                    str(src), str(dst3), "--n", "5", "--discard-remainder"],
                   capture_output=True, text=True, cwd=REPO)
    assert list(SrasFile(str(dst3)).n_frames) == [2, 2], \
        "--discard-remainder drops the partial group"


def test_unsupported_version_reported(tmp_path):
    """cache_file must report, not raise, for a file it can't handle."""
    bogus = tmp_path / "bogus.sras"
    bogus.write_bytes(b"SRAS" + bytes([99]) + b"\x00" * 200)
    err = cache_file(str(bogus), "dc", True)
    assert err, "bad version returns an error string"
    missing = cache_file(str(tmp_path / "does_not_exist.sras"), "dc", True)
    assert missing, "missing file returns an error string"