"""Tests for GPU helper module."""

import numpy as np
import pytest


def test_has_gpu_attribute():
    """HAS_GPU must be a boolean."""
    from lidar_pipeline.gpu import HAS_GPU
    assert isinstance(HAS_GPU, bool)


def test_to_gpu_returns_array():
    """to_gpu returns a float32 array with correct values."""
    from lidar_pipeline.gpu import to_gpu, to_cpu, HAS_GPU
    arr = np.array([1.0, 2.0, 3.0])
    result = to_gpu(arr)
    # to_gpu converts to float32 to reduce GPU memory usage
    assert result.dtype == np.float32
    # Always bring back to CPU for comparison
    np.testing.assert_array_equal(to_cpu(result), [1.0, 2.0, 3.0])


def test_to_cpu_noop_numpy():
    """to_cpu on a numpy array is a no-op."""
    from lidar_pipeline.gpu import to_cpu
    arr = np.array([1.0, 2.0])
    result = to_cpu(arr)
    assert result is arr


def test_xp_gaussian_filter():
    """xp_gaussian_filter blurs a point source correctly."""
    from lidar_pipeline.gpu import xp_gaussian_filter
    arr = np.zeros((50, 50), dtype=np.float64)
    arr[25, 25] = 1.0
    result = xp_gaussian_filter(arr, sigma=3)
    assert result.shape == (50, 50)
    # Center should still be the highest value
    center_val = float(np.asarray(result)[25, 25])
    corner_val = float(np.asarray(result)[0, 0])
    assert center_val > corner_val
    assert center_val > 0.01  # Not all energy is lost


def test_xp_uniform_filter_cpu():
    """xp_uniform_filter works on CPU arrays."""
    from lidar_pipeline.gpu import xp_uniform_filter
    arr = np.ones((50, 50), dtype=np.float64)
    arr[25, 25] = 100.0
    result = xp_uniform_filter(arr, size=5)
    # Mean should be close to 1 everywhere except near center
    assert result.shape == (50, 50)
    assert result[0, 0] == pytest.approx(1.0, abs=0.01)


def test_xp_minimum_filter_cpu():
    """xp_minimum_filter works on CPU arrays."""
    from lidar_pipeline.gpu import xp_minimum_filter
    arr = np.ones((50, 50), dtype=np.float64)
    arr[25, 25] = 0.0
    result = xp_minimum_filter(arr, size=3)
    assert result.shape == (50, 50)
    # Around the minimum, values should be 0
    assert result[25, 25] == 0.0
    assert result[24, 25] == 0.0


def test_log_gpu_status(caplog):
    """log_gpu_status emits a log message."""
    import logging
    from lidar_pipeline.gpu import log_gpu_status
    with caplog.at_level(logging.INFO, logger="lidar"):
        log_gpu_status()
    assert any("GPU" in r.message or "CPU" in r.message for r in caplog.records)


@pytest.mark.skipif(
    not pytest.importorskip("cupy", reason="CuPy not available"),
    reason="Requires GPU + CuPy"
)
def test_to_gpu_roundtrip():
    """to_gpu -> to_cpu preserves data when GPU is available."""
    import cupy as cp
    from lidar_pipeline.gpu import to_gpu, to_cpu, HAS_GPU
    if not HAS_GPU:
        pytest.skip("No GPU available")
    arr = np.array([1.0, 2.0, 3.0], dtype=np.float32)
    gpu_arr = to_gpu(arr)
    assert isinstance(gpu_arr, cp.ndarray)
    result = to_cpu(gpu_arr)
    assert isinstance(result, np.ndarray)
    np.testing.assert_array_almost_equal(result, [1.0, 2.0, 3.0])


# ---------------------------------------------------------------------------
# bin_mean_2d — DTM rasterization (mean z per cell)
# ---------------------------------------------------------------------------

def _scipy_reference(xs, ys, zs, width, height, x_range, y_range):
    from scipy.stats import binned_statistic_2d
    stat = binned_statistic_2d(
        xs, ys, zs, statistic='mean', bins=[width, height],
        range=[list(x_range), list(y_range)]
    )
    return stat.statistic.T  # (height, width)


def test_bin_mean_core_parity_scipy():
    """The numpy core of bin_mean_2d reproduces binned_statistic_2d (mean)."""
    from lidar_pipeline.gpu import _bin_mean_core
    rng = np.random.default_rng(42)
    x_range, y_range = (1000.0, 1010.0), (6800.0, 6808.0)
    w, h = 40, 32
    n = 5000
    xs = rng.uniform(*x_range, n)
    ys = rng.uniform(*y_range, n)
    zs = rng.uniform(50.0, 150.0, n)
    # Edge cases: left/bottom edge (included), right/top edge (included in the
    # last cell), outside the extent (ignored), points exactly on an interior
    # cell boundary
    xs = np.concatenate([xs, [1000.0, 1010.0, 999.9, 1010.1, 1002.5]])
    ys = np.concatenate([ys, [6800.0, 6808.0, 6808.1, 6807.9, 6804.0]])
    zs = np.concatenate([zs, [99.0, 101.0, 777.0, 777.0, 103.0]])
    got = _bin_mean_core(np, xs, ys, zs, w, h, x_range, y_range)
    ref = _scipy_reference(xs, ys, zs, w, h, x_range, y_range)
    assert got.shape == (h, w)
    np.testing.assert_allclose(got, ref, equal_nan=True)


def test_bin_mean_core_empty_and_single():
    """Empty cells → NaN; a single point → its value in its own cell."""
    from lidar_pipeline.gpu import _bin_mean_core
    x_range, y_range = (0.0, 10.0), (0.0, 10.0)
    # No point inside the extent
    out = _bin_mean_core(np, [-5.0], [-5.0], [1.0], 5, 5, x_range, y_range)
    assert out.shape == (5, 5)
    assert np.isnan(out).all()
    # One point exactly at the center: cell (2, 2)
    out = _bin_mean_core(np, [5.0], [5.0], [7.5], 5, 5, x_range, y_range)
    assert out[2, 2] == 7.5
    assert np.isnan(out).sum() == 24
    # Cell with several points: exact mean
    out = _bin_mean_core(np, [5.0, 5.1, 5.2], [5.0, 5.0, 5.0],
                         [10.0, 20.0, 30.0], 5, 5, x_range, y_range)
    assert out[2, 2] == 20.0


def test_bin_mean_2d_gpu_or_none():
    """bin_mean_2d: None without a GPU, otherwise output identical to the numpy core."""
    from lidar_pipeline.gpu import bin_mean_2d, _bin_mean_core
    rng = np.random.default_rng(7)
    x_range, y_range = (1000.0, 1200.0), (6800.0, 7000.0)
    w, h = 100, 100
    n = 20000
    xs = rng.uniform(*x_range, n)
    ys = rng.uniform(*y_range, n)
    zs = rng.uniform(50.0, 150.0, n)
    got = bin_mean_2d(xs, ys, zs, w, h, x_range, y_range)
    ref = _bin_mean_core(np, xs, ys, zs, w, h, x_range, y_range)
    if got is None:
        assert isinstance(ref, np.ndarray)  # the caller provides the scipy fallback
    else:
        assert isinstance(got, np.ndarray)
        np.testing.assert_allclose(got, ref, equal_nan=True, rtol=1e-9)


def test_safe_gpu_call_retries_on_cpu_for_any_error(monkeypatch):
    """GPU active: ANY error (not only CUDA messages) disables the GPU and
    retries the computation on CPU.

    Real case: after a GPU transfer failure, mixed numpy/cupy types
    ("Unsupported type <class 'numpy.ndarray'>") contained no CUDA keyword
    and propagated the error — the whole visualization then failed for
    nothing.
    """
    from lidar_pipeline import gpu

    class _FakeCP:  # "active" GPU without CuPy
        pass

    monkeypatch.setattr(gpu, "_cp", _FakeCP())
    monkeypatch.setattr(gpu, "HAS_GPU", True)
    calls = []

    def f(x):
        calls.append(x)
        if len(calls) == 1:
            raise TypeError("Unsupported type <class 'numpy.ndarray'>")
        return x + 1

    assert gpu.safe_gpu_call(f, 21) == 22
    assert gpu.HAS_GPU is False  # GPU disabled after the error
    assert gpu._cp is None
    # CPU mode: the error is re-raised as is (nothing to retry)
    def g(x):
        raise ValueError("boom")
    try:
        gpu.safe_gpu_call(g, 1)
        raise AssertionError("should have re-raised the error")
    except ValueError:
        pass

def test_gpu_worker_slots_bounded_by_free_vram(monkeypatch):
    """GPU slots bounded by free VRAM: surplus workers run on CPU.

    Real case: LIDAR_WORKERS=auto = 12 workers on 2 RTX 5060 (8 GB) —
    6 workers per GPU, far more than free VRAM holds at peak (scan-line
    alignment + GPU gap filling): OOM.
    """
    from lidar_pipeline import gpu
    monkeypatch.setattr(gpu, "GPU_WORKER_MIB", 2000)
    monkeypatch.setattr(gpu, "GPU_RESERVE_MIB", 500)
    free = {0: 7500, 1: 4600}
    slots = gpu.gpu_worker_slots([0, 1], 12, free_mib=free)
    assert len(slots) == 12
    assert slots.count(0) == 3          # (7500 - 500) // 2000
    assert slots.count(1) == 2          # (4600 - 500) // 2000
    assert slots.count(-1) == 7         # the rest on CPU
    # GPUs first, interleaved: the first workers created are spread out
    assert slots[:4] == [0, 1, 0, 1]
    # Fewer workers than slots: no forced CPU
    assert gpu.gpu_worker_slots([0, 1], 3, free_mib=free) == [0, 1, 0]


def test_gpu_worker_slots_without_gpu_or_vram_info():
    """No GPU: implicit CPU everywhere (None). Unknown VRAM: legacy
    round-robin (no bound without a measurement)."""
    from lidar_pipeline import gpu
    assert gpu.gpu_worker_slots([], 4, free_mib={}) == [None] * 4
    assert gpu.gpu_worker_slots([0, 1], 4, free_mib={}) == [0, 1, 0, 1]


def test_gpu_worker_slots_keeps_one_gpu_worker_when_tight(monkeypatch):
    """Nearly full GPU: at least one slot so the GPU is still used
    (safe_gpu_call's CPU fallback covers a possible OOM)."""
    from lidar_pipeline import gpu
    monkeypatch.setattr(gpu, "GPU_WORKER_MIB", 2000)
    monkeypatch.setattr(gpu, "GPU_RESERVE_MIB", 500)
    assert gpu.gpu_worker_slots([0], 3, free_mib={0: 900}) == [0, -1, -1]


def test_force_cpu_disables_gpu_selection(monkeypatch):
    """force_cpu(): no GPU candidate, CuPy never initialized."""
    import os
    from lidar_pipeline import gpu
    monkeypatch.setattr(gpu, "_restricted_gpu_ids", None)
    monkeypatch.setattr(gpu, "_gpu_candidates", [(0, "X", "12.0", 8000, 1, 12)])
    monkeypatch.setattr(gpu, "_gpu_initialized", False)
    monkeypatch.setattr(gpu, "_env_set_by_init", False)
    monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
    gpu.force_cpu()
    assert gpu.available_gpu_ids() == []
    assert gpu.HAS_GPU is False
    assert os.environ.get("CUDA_VISIBLE_DEVICES") == ""