"""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()


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