"""Tests for visualization functions.

Each test creates a small synthetic DEM and runs a visualization function,
checking that it produces a valid output file.
"""

import numpy as np
import pytest
from pathlib import Path


# --- Core terrain visualizations (no GPU required) ---

class TestHillshade:
    def test_generates_tif(self, synthetic_dem, tmp_output_dir):
        from lidar_pipeline.visualizations import generate_hillshade
        result = generate_hillshade(synthetic_dem, "test", tmp_output_dir, 5.0)
        assert result is not None
        assert result.exists()
        assert result.suffix == ".tif"

    def test_output_values_valid(self, synthetic_dem, tmp_output_dir):
        import rasterio
        from lidar_pipeline.visualizations import generate_hillshade
        result = generate_hillshade(synthetic_dem, "test", tmp_output_dir, 5.0)
        with rasterio.open(result) as src:
            data = src.read(1)
            assert data.shape[0] > 0
            assert np.nanmin(data) >= 0
            assert np.nanmax(data) <= 1


class TestSlope:
    def test_generates_tif(self, synthetic_dem, tmp_output_dir):
        from lidar_pipeline.visualizations import generate_slope
        result = generate_slope(synthetic_dem, "test", tmp_output_dir, 5.0)
        assert result is not None
        assert result.exists()

    def test_slope_values_degrees(self, synthetic_dem, tmp_output_dir):
        import rasterio
        from lidar_pipeline.visualizations import generate_slope
        result = generate_slope(synthetic_dem, "test", tmp_output_dir, 5.0)
        with rasterio.open(result) as src:
            data = src.read(1)
            assert np.nanmin(data) >= 0
            assert np.nanmax(data) <= 90


# --- GPU-accelerated visualizations ---

class TestSVF:
    def test_generates_tif(self, synthetic_dem, tmp_output_dir):
        from lidar_pipeline.visualizations import generate_svf
        result = generate_svf(synthetic_dem, "test", tmp_output_dir, 5.0)
        assert result is not None
        assert result.exists()

    def test_svf_values_0_1(self, synthetic_dem, tmp_output_dir):
        import rasterio
        from lidar_pipeline.visualizations import generate_svf
        result = generate_svf(synthetic_dem, "test", tmp_output_dir, 5.0)
        with rasterio.open(result) as src:
            data = src.read(1)
            valid = data[~np.isnan(data)]
            assert np.nanmin(valid) >= 0
            assert np.nanmax(valid) <= 1


class TestOpenness:
    def test_positive_generates_tif(self, synthetic_dem, tmp_output_dir):
        from lidar_pipeline.visualizations import generate_openness
        result = generate_openness(synthetic_dem, "test", tmp_output_dir, 5.0, positive=True)
        assert result is not None
        assert result.exists()

    def test_negative_generates_tif(self, synthetic_dem, tmp_output_dir):
        from lidar_pipeline.visualizations import generate_openness
        result = generate_openness(synthetic_dem, "test", tmp_output_dir, 5.0, positive=False)
        assert result is not None
        assert result.exists()

    def test_downsampled_matches_full_resolution(self, tmp_path, tmp_output_dir):
        """Factor 2: same output shape, highly correlated with factor 1.

        MNT dédié à 1 m/px (structures de plusieurs dizaines de pixels, régime
        de production 0,2 m/px) : la fixture synthétique partagée à 5 m/px a un
        mur large de 2 px, hors régime pour valider une décimation ×2.
        """
        import rasterio
        from rasterio.transform import from_bounds
        from lidar_pipeline import visualizations
        from lidar_pipeline.visualizations import generate_openness

        size = 240
        x = np.linspace(0, size, size)
        X, Y = np.meshgrid(x, x)
        dem = 100.0 + 0.01 * X + 0.005 * Y
        dem += 5.0 * np.exp(-((X - 120)**2 + (Y - 120)**2) / (2 * 40**2))
        dist_wall = np.abs((X - 40) * 0.707 + (Y - 60) * 0.707) / np.sqrt(2)
        dem += 1.5 * np.exp(-dist_wall**2 / (2 * 8**2))
        dem -= 2.0 * np.exp(-np.abs(X - 200)**2 / (2 * 12**2))
        # Sans bruit blanc : à 1 m/px un bruit σ=5 cm domine l'angle
        # d'horizon à courte distance (max le long du rayon) et masque la
        # géométrie que ce test cherche à valider (décimation + zoom).
        dem_file = tmp_path / "dem_1m.tif"
        with rasterio.open(dem_file, 'w', driver='GTiff', height=size, width=size,
                           count=1, dtype='float32', crs='EPSG:2154',
                           transform=from_bounds(660000, 6700000, 660240, 6700240, size, size)) as dst:
            dst.write(dem.astype('float32'), 1)

        saved = visualizations.OPENNESS_DOWNSAMPLE
        try:
            visualizations.OPENNESS_DOWNSAMPLE = 1
            r1 = generate_openness(dem_file, "full", tmp_output_dir, 1.0, positive=True)
            visualizations.OPENNESS_DOWNSAMPLE = 2
            r2 = generate_openness(dem_file, "dec", tmp_output_dir, 1.0, positive=True)
        finally:
            visualizations.OPENNESS_DOWNSAMPLE = saved

        with rasterio.open(r1) as s1, rasterio.open(r2) as s2:
            a, b = s1.read(1), s2.read(1)
        assert a.shape == b.shape
        m = ~np.isnan(a) & ~np.isnan(b)
        assert m.sum() > 0
        corr = np.corrcoef(a[m], b[m])[0, 1]
        assert corr > 0.97, f"corrélation openness décimée/native trop faible : {corr:.3f}"

    @pytest.mark.parametrize("positive", [True, False])
    def test_scale_independent_of_rest_of_tile(self, tmp_path, tmp_output_dir, positive):
        """Même relief local = même valeur, quel que soit le reste de la dalle.

        Deux MNT identiques sur leur moitié ouest ; l'un porte en plus une
        colline à l'est, à plus de 100 m (rayon max) de la zone comparée. Une
        normalisation par dalle (z-score) décalait toute l'échelle et rendait
        les mosaïques non jointives ; avec les références figées, la moitié
        ouest doit sortir identique.
        """
        import rasterio
        from rasterio.transform import from_bounds
        from lidar_pipeline.visualizations import generate_openness

        size = 300
        x = np.arange(size, dtype=float)
        X, Y = np.meshgrid(x, x)
        flat = 100.0 + 2.0 * np.exp(-((X - 60)**2 + (Y - 150)**2) / (2 * 15**2))
        hill = flat + 40.0 * np.exp(-((X - 260)**2 + (Y - 150)**2) / (2 * 20**2))
        results = []
        for name, dem in (("flat", flat), ("hill", hill)):
            f = tmp_path / f"{name}.tif"
            with rasterio.open(f, 'w', driver='GTiff', height=size, width=size,
                               count=1, dtype='float32', crs='EPSG:2154',
                               transform=from_bounds(660000, 6700000, 660300, 6700300,
                                                     size, size)) as dst:
                dst.write(dem.astype('float32'), 1)
            out = generate_openness(f, name, tmp_output_dir, 1.0, positive=positive)
            with rasterio.open(out) as src:
                results.append(src.read(1))
        # Tolérance : résidu d'interpolation de la grille décimée (≈0,01) ;
        # un z-score par dalle décalerait toute la zone de plusieurs dixièmes.
        west = np.s_[:, :100]
        np.testing.assert_allclose(results[0][west], results[1][west], atol=0.05)


def _write_dem(path, dem, res=1.0):
    import rasterio
    from rasterio.transform import from_origin
    with rasterio.open(path, 'w', driver='GTiff', height=dem.shape[0], width=dem.shape[1],
                       count=1, dtype='float32', crs='EPSG:2154',
                       transform=from_origin(660000, 6700300, res, res)) as dst:
        dst.write(dem.astype('float32'), 1)
    return path


class TestReliefOriente:
    def test_generates_rgb_uint8(self, synthetic_dem, tmp_output_dir):
        import rasterio
        from lidar_pipeline.visualizations import generate_relief_oriente
        result = generate_relief_oriente(synthetic_dem, "test", tmp_output_dir, 5.0)
        assert result is not None and result.name == "test_relief_oriente.tif"
        with rasterio.open(result) as src, rasterio.open(synthetic_dem) as dem:
            assert src.count == 3 and src.dtypes[0] == 'uint8'
            assert (src.height, src.width) == (dem.height, dem.width)
            rgb = src.read()
        assert rgb.std() > 5, "image uniforme : ni relief ni orientation rendus"

    def test_nodata_gets_fixed_color(self, tmp_path, tmp_output_dir):
        import rasterio
        from lidar_pipeline.visualizations import generate_relief_oriente, RELIEF_NODATA_RGB
        x = np.arange(120, dtype=float)
        dem = 100 + 0.05 * x[None, :] + 0.02 * x[:, None]
        dem[10:20, 10:20] = np.nan
        out = generate_relief_oriente(_write_dem(tmp_path / "d.tif", dem), "t", tmp_output_dir, 1.0)
        with rasterio.open(out) as src:
            rgb = src.read()
        assert tuple(rgb[:, 15, 15]) == RELIEF_NODATA_RGB

    def test_horizon_kernels_agree(self):
        """numba et numpy (et le noyau CUDA, même code) donnent le même angle."""
        pytest.importorskip("numba")
        from lidar_pipeline.visualizations import (
            _horizon_rays, _mean_horizon_numba, _mean_horizon_numpy)
        rng = np.random.default_rng(0)
        dem = rng.normal(0, 0.3, (60, 70)).astype(np.float32)
        dem[20:30, 30:40] += 2.0
        offs, dist, cps = _horizon_rays(0.8, 16, (5, 10, 20))
        a = _mean_horizon_numba(dem, offs, dist, cps)
        b = _mean_horizon_numpy(dem, offs, dist, cps)
        np.testing.assert_allclose(a, b, atol=1e-5)

    def test_matches_legacy_ray_trace_geometry(self):
        """Même géométrie de rayons que _ray_trace_horizons (8 directions)."""
        from lidar_pipeline.visualizations import (
            _horizon_rays, _mean_horizon_numpy, _ray_trace_horizons)
        rng = np.random.default_rng(1)
        dem = rng.normal(0, 0.5, (50, 50)).astype(np.float32)
        radii = (5, 10, 20)
        pos, _ = _ray_trace_horizons(dem, 50, 50, 1.0, 8, 20, list(radii))
        legacy = np.mean(pos, axis=(0, 1))
        offs, dist, cps = _horizon_rays(1.0, 8, radii)
        np.testing.assert_allclose(_mean_horizon_numpy(dem, offs, dist, cps), legacy, atol=1e-5)

    def test_seamless_independent_of_rest_of_tile(self, tmp_path, tmp_output_dir):
        """Même relief local = mêmes couleurs, quel que soit le reste de la dalle
        (support : rayon 20 m + détendance 4σ = 40 m, loin sous la bande de 100 m)."""
        import rasterio
        from lidar_pipeline.visualizations import generate_relief_oriente
        size = 300
        X, Y = np.meshgrid(np.arange(size, dtype=float), np.arange(size, dtype=float))
        flat = 100.0 + 1.5 * np.exp(-((X - 60)**2 + (Y - 150)**2) / (2 * 6**2))
        hill = flat + 40.0 * np.exp(-((X - 260)**2 + (Y - 150)**2) / (2 * 10**2))
        imgs = []
        for name, dem in (("flat", flat), ("hill", hill)):
            out = generate_relief_oriente(_write_dem(tmp_path / f"{name}.tif", dem),
                                          name, tmp_output_dir, 1.0)
            with rasterio.open(out) as src:
                imgs.append(src.read().astype(int))
        diff = np.abs(imgs[0][:, :, :150] - imgs[1][:, :, :150])
        assert diff.max() <= 1, f"écart de couleur loin de la colline : {diff.max()}"

    def test_numba_colorize_matches_vectorized(self):
        """Noyau CPU fusionné = chemin vectorisé (CuPy/numpy) à l'arrondi près."""
        pytest.importorskip("numba")
        from lidar_pipeline.gpu import xp_zoom
        from lidar_pipeline.visualizations import (
            _relief_colorize_numba, _relief_colorize_xp, _relief_lut, _pad_to)
        rng = np.random.default_rng(2)
        open_c = rng.uniform(0.3, 15, (30, 25)).astype(np.float32)
        dx = rng.normal(0, 0.3, (120, 100)).astype(np.float32)
        dy = rng.normal(0, 0.3, (120, 100)).astype(np.float32)
        lut = _relief_lut()
        a = _relief_colorize_numba(open_c, 4, dx, dy, lut).astype(int)
        b = _relief_colorize_xp(np, _pad_to(np, xp_zoom(open_c, 4), 120, 100), dx, dy, lut).astype(int)
        diff = np.abs(a - b).max(axis=2)
        assert (diff > 3).mean() < 0.01, f"{(diff > 3).mean():.3%} pixels divergent"

    def test_lut_lightness_is_monotonic_and_hue_neutral(self):
        """Clarté croissante avec L* ; à L* fixé, toutes les teintes ont la même
        luminance perçue (pas de faux relief dû à la couleur)."""
        from lidar_pipeline.visualizations import _relief_lut
        lut = _relief_lut().astype(float) / 255
        lin = np.where(lut <= 0.04045, lut / 12.92, ((lut + 0.055) / 1.055) ** 2.4)
        Y = lin @ np.array([0.2126, 0.7152, 0.0722])          # (256, 360)
        assert np.all(np.diff(Y.mean(axis=1)) >= -1e-6)
        Lstar = 116 * np.cbrt(Y[100:180]) - 16                # L* ≈ 40–70
        # Écart résiduel : écrêtage de gamme sRGB de quelques teintes à C* 60
        assert (Lstar.max(axis=1) - Lstar.min(axis=1)).max() < 6


class TestMSLRM:
    def test_generates_tif(self, synthetic_dem, tmp_output_dir):
        from lidar_pipeline.visualizations import generate_mslrm
        result = generate_mslrm(synthetic_dem, "test", tmp_output_dir, 5.0)
        assert result is not None
        assert result.exists()


class TestSAILORE:
    def test_generates_tif(self, synthetic_dem, tmp_output_dir):
        from lidar_pipeline.visualizations import generate_sailore
        result = generate_sailore(synthetic_dem, "test", tmp_output_dir, 5.0)
        assert result is not None
        assert result.exists()


class TestRoughness:
    def test_generates_tif(self, synthetic_dem, tmp_output_dir):
        from lidar_pipeline.visualizations import generate_roughness
        result = generate_roughness(synthetic_dem, "test", tmp_output_dir, 5.0)
        assert result is not None
        assert result.exists()

    def test_roughness_non_negative(self, synthetic_dem, tmp_output_dir):
        import rasterio
        from lidar_pipeline.visualizations import generate_roughness
        result = generate_roughness(synthetic_dem, "test", tmp_output_dir, 5.0)
        with rasterio.open(result) as src:
            data = src.read(1)
            # Standard deviation is always >= 0
            assert np.nanmin(data) >= 0


class TestWavelet:
    def test_generates_tif(self, synthetic_dem, tmp_output_dir):
        from lidar_pipeline.visualizations import generate_wavelet
        result = generate_wavelet(synthetic_dem, "test", tmp_output_dir, 5.0)
        assert result is not None
        assert result.exists()

    def test_output_median_centered(self, synthetic_dem, tmp_output_dir):
        """Recentrage robuste : médiane ≈ 1 quel que soit le terrain.

        C'est la condition pour qu'un étirement couleur global fixe donne
        des couleurs homogènes entre tuiles (cf. knots dans rendering.py).
        """
        import rasterio
        from lidar_pipeline.visualizations import generate_wavelet
        result = generate_wavelet(synthetic_dem, "test", tmp_output_dir, 5.0)
        with rasterio.open(result) as src:
            data = src.read(1)
        valid = data[np.isfinite(data)]
        assert abs(np.median(valid) - 1.0) < 0.05

    def test_ditch_on_hilltop_not_amplified(self, tmp_path, tmp_output_dir):
        """Un fossé en sommet de colline ne ressort pas plus qu'à plat.

        Le détendage (moyenne locale gaussienne ~35 m) doit neutraliser la
        position topographique :
        - le pic d'indice sur le fossé en sommet reste comparable au même
          fossé sur terrain plat ;
        - le fond du sommet (sans structure) reste comparable au fond plat.

        MNT synthétique bruité (σ=8 cm) : sans détendage, le fond sommet
        ressort ~1,7× le fond plat (sommets jaunes sur la carte).
        """
        import rasterio
        from rasterio.transform import from_bounds
        from lidar_pipeline.visualizations import generate_wavelet

        size = 600
        res = 1.0
        x = np.arange(size) * res
        y = np.arange(size) * res
        X, Y = np.meshgrid(x, y)
        dem = 100.0 + 0.01 * X

        # Colline réaliste (sigma 80 m, 25 m de haut) à gauche + bruit capteur
        dem += 25.0 * np.exp(-((X - 150)**2 + (Y - 300)**2) / (2 * 80**2))
        rng = np.random.default_rng(42)
        dem += rng.normal(0, 0.08, dem.shape)

        # Deux fossés identiques : l'un au sommet de la colline, l'autre à plat
        for xc in (150, 480):
            dem -= 1.5 * np.exp(-((X - xc)**2) / (2 * 1.2**2))

        dem_file = tmp_path / "ditch_dem.tif"
        transform = from_bounds(660000, 6700000, 660600, 6700600, size, size)
        with rasterio.open(
            dem_file, 'w', driver='GTiff', height=size, width=size,
            count=1, dtype='float32', crs='EPSG:2154', transform=transform,
        ) as dst:
            dst.write(dem.astype('float32'), 1)

        result = generate_wavelet(dem_file, "ditch", tmp_output_dir, res)
        assert result is not None and result.exists()
        with rasterio.open(result) as src:
            data = src.read(1)

        # Pics sur une bande verticale autour de chaque fossé
        peak_hilltop = np.nanmax(data[:, 145:156])
        peak_flat = np.nanmax(data[:, 475:486])
        assert peak_flat > 5.0  # le fossé ressort nettement au-dessus du fond
        # Pas d'amplification du fossé par la position topographique
        assert peak_hilltop < 1.5 * peak_flat

        # Fond du sommet (35 m à l'est du fossé sommital) vs fond plat
        bg_hilltop = np.nanmedian(data[250:350, 185:226])
        bg_flat = np.nanmedian(data[250:350, 500:561])
        assert bg_hilltop < 1.4 * bg_flat  # sans détendage : ~1,7×


class TestFlowAccumulation:
    def test_generates_tif(self, synthetic_dem, tmp_output_dir):
        from lidar_pipeline.visualizations import generate_flow_accumulation
        result = generate_flow_accumulation(synthetic_dem, "test", tmp_output_dir, 5.0)
        assert result is not None
        assert result.exists()

    def test_flow_log_values(self, synthetic_dem, tmp_output_dir):
        import rasterio
        from lidar_pipeline.visualizations import generate_flow_accumulation
        result = generate_flow_accumulation(synthetic_dem, "test", tmp_output_dir, 5.0)
        with rasterio.open(result) as src:
            data = src.read(1)
            # log10(x) >= 0 for x >= 1
            valid = data[~np.isnan(data)]
            assert np.nanmin(valid) >= 0


class TestRayTrace:
    def test_rays_are_traced(self, synthetic_dem, tmp_output_dir):
        """Verify _ray_trace_horizons returns expected shapes."""
        from lidar_pipeline.visualizations import _ray_trace_horizons, _prepare_dem_for_raycast
        import rasterio
        with rasterio.open(synthetic_dem) as src:
            dem_np = src.read(1)
        rows, cols = dem_np.shape
        # Create a simple filled DEM for testing
        import numpy as np
        filled = np.nan_to_num(dem_np, nan=0)
        # Test with numpy (no GPU)
        pos, neg = _ray_trace_horizons(
            filled, rows, cols, 5.0, n_dirs=4, max_dist=10, radii_m=[25, 50]
        )
        assert pos.shape == (4, 2, rows, cols)
        assert neg.shape == (4, 2, rows, cols)


class TestNodataPreserved:
    """Nodata préservé dans les rendus (comportement historique).

    Les trous du MNT sont comblés en amont (create_dtm_fast, tous modes) ;
    si un nodata subsiste malgré tout, hillshade/slope/aspect le restituent
    (rendu noir en carte) au lieu d'inventer des valeurs interpolées.
    """

    @staticmethod
    def _dem_with_hole(synthetic_dem, tmp_path):
        import rasterio
        with rasterio.open(synthetic_dem) as src:
            arr = src.read(1).copy()
            profile = src.profile.copy()
        arr[80:120, 80:120] = np.nan
        dem_hole = tmp_path / "dem_hole.tif"
        profile.update(dtype='float32', nodata=float('nan'))
        with rasterio.open(dem_hole, 'w', **profile) as dst:
            dst.write(arr.astype('float32'), 1)
        return dem_hole

    def test_aspect_solo_preserves_nodata(self, synthetic_dem, tmp_path):
        from lidar_pipeline.visualizations import generate_aspect
        dem_hole = self._dem_with_hole(synthetic_dem, tmp_path)
        out = generate_aspect(dem_hole, "solo", tmp_path, 5.0)
        assert out is not None and out.exists()
        import rasterio
        with rasterio.open(out) as src:
            data = src.read(1)
        assert np.isnan(data[80:120, 80:120]).all(), "le trou doit rester en nodata"
        # Le gradient au bord du trou propage NaN sur un anneau de 1 px :
        # on vérifie une zone éloignée du trou
        assert not np.isnan(data[0:40, 0:40]).any(), "NaN loin du trou"

    def test_aspect_shared_preserves_nodata(self, synthetic_dem, tmp_path):
        from lidar_pipeline.visualizations import SharedDEM, generate_aspect
        dem_hole = self._dem_with_hole(synthetic_dem, tmp_path)
        shared = SharedDEM(dem_hole, 5.0)
        out = generate_aspect(dem_hole, "partage", tmp_path, 5.0, shared=shared)
        assert out is not None and out.exists()
        import rasterio
        with rasterio.open(out) as src:
            data = src.read(1)
        assert np.isnan(data[80:120, 80:120]).all(), "le trou doit rester en nodata"
        assert not np.isnan(data[0:40, 0:40]).any(), "NaN loin du trou"

    def test_slope_and_hillshade_preserve_nodata(self, synthetic_dem, tmp_path):
        from lidar_pipeline.visualizations import generate_slope, generate_hillshade
        dem_hole = self._dem_with_hole(synthetic_dem, tmp_path)
        import rasterio
        for gen, name in ((generate_slope, "p"), (generate_hillshade, "h")):
            out = gen(dem_hole, name, tmp_path, 5.0)
            assert out is not None and out.exists()
            with rasterio.open(out) as src:
                data = src.read(1)
            assert np.isnan(data[80:120, 80:120]).any(), f"{out.name} : trou disparu"


def test_ray_trace_horizons_cpu_fallback_on_oom(monkeypatch):
    """Sur OOM GPU, le ray-tracing désactive le GPU puis recommence sur CPU."""
    import lidar_pipeline.visualizations as viz
    import lidar_pipeline.gpu as gpu_mod

    calls = {"n": 0}
    disabled = []

    def fake_core(dem, rows, cols, res, n_dirs, max_dist, radii_m=None):
        calls["n"] += 1
        if calls["n"] == 1:
            raise RuntimeError("Out of memory allocating 600,000,000 bytes")
        return ("pos", "neg")

    monkeypatch.setattr(viz, "_ray_trace_horizons_core", fake_core)
    monkeypatch.setattr(gpu_mod, "is_gpu_active", lambda: True)
    monkeypatch.setattr(gpu_mod, "disable_gpu", lambda: disabled.append(True))
    result = viz._ray_trace_horizons(None, 4, 4, 0.5, 8, 10)
    assert result == ("pos", "neg")
    assert calls["n"] == 2
    assert disabled == [True]


def test_ray_trace_horizons_reraises_non_oom(monkeypatch):
    """Une erreur non-OOM n'est pas masquée par le repli CPU."""
    import lidar_pipeline.visualizations as viz
    import lidar_pipeline.gpu as gpu_mod

    def fake_core(dem, rows, cols, res, n_dirs, max_dist, radii_m=None):
        raise ValueError("autre erreur")

    monkeypatch.setattr(viz, "_ray_trace_horizons_core", fake_core)
    monkeypatch.setattr(gpu_mod, "is_gpu_active", lambda: True)
    try:
        viz._ray_trace_horizons(None, 4, 4, 0.5, 8, 10)
        assert False, "ValueError attendue"
    except ValueError:
        pass


class TestPriorityFlood:
    def test_numba_matches_python(self):
        """Le résultat numba et python sont identiques sur un DEM avec un puits."""
        from lidar_pipeline.visualizations import _priority_flood_numba, _priority_flood_python

        dem = np.zeros((20, 20), dtype=np.float64)
        dem[10, 10] = -5.0
        dem[9:12, 9:12] = -3.0
        nodata = np.zeros((20, 20), dtype=bool)

        result_numba = _priority_flood_numba(dem.copy(), nodata)
        result_python = _priority_flood_python(dem.copy(), nodata)

        if result_numba is not None:
            assert np.allclose(result_numba, result_python)

    def test_pit_is_filled(self):
        """Un puits isolé est ramené au niveau de son bord."""
        from lidar_pipeline.visualizations import _priority_flood

        dem = np.full((10, 10), 5.0, dtype=np.float64)
        dem[5, 5] = 1.0
        nodata = np.zeros((10, 10), dtype=bool)

        result = _priority_flood(dem, nodata)
        assert result[5, 5] == 5.0

    def test_nodata_cells_untouched(self):
        """Les cellules nodata ne sont jamais modifiées."""
        from lidar_pipeline.visualizations import _priority_flood

        dem = np.full((10, 10), 5.0, dtype=np.float64)
        dem[5, 5] = 1.0
        nodata = np.zeros((10, 10), dtype=bool)
        nodata[2, 2] = True
        dem[2, 2] = 999.0

        result = _priority_flood(dem, nodata)
        assert result[2, 2] == 999.0


class TestDensiteSol:
    def test_density_levels_log_scale(self):
        from lidar_pipeline.visualizations import density_levels
        d = np.array([0.0, np.nan, 0.25, 0.36, 0.5, 1.0, 45.0, 45.3, 1000.0])
        assert density_levels(d).tolist() == [0, 0, 0, 1, 2, 4, 14, 15, 15]

    def test_generate_reads_density_sidecar(self, tmp_path):
        import rasterio
        from rasterio.transform import from_bounds
        from lidar_pipeline.dtm import density_path
        from lidar_pipeline.visualizations import generate_densite_sol
        dem = tmp_path / "T_dtm_r0p2.tif"
        dem.touch()
        with rasterio.open(density_path(dem), 'w', driver='GTiff', width=4, height=1,
                           count=1, dtype='float32', crs='EPSG:2154',
                           transform=from_bounds(0, 0, 4, 1, 4, 1)) as dst:
            dst.write(np.array([[0.0, 1.0, 4.0, 64.0]], dtype='float32'), 1)
        out = generate_densite_sol(dem, "T", tmp_path, 0.2)
        with rasterio.open(out) as src:
            assert src.read(1).tolist() == [[0, 4, 8, 15]]
            assert src.width == 4  # grille de 1 m conservée

    def test_missing_sidecar_returns_none(self, tmp_path):
        from lidar_pipeline.visualizations import generate_densite_sol
        assert generate_densite_sol(tmp_path / "X_dtm.tif", "X", tmp_path, 0.2) is None