Source code for wolfhece.services.dtm_wallonia

"""Class-based service to retrieve Wallonia DTM as :class:`WolfArray`.

The main entry point is :class:`DtmWalloniaService`, which exposes 1 m and
2 m retrieval helpers based on polygons or raw bounding coordinates.
"""
from __future__ import annotations

import logging
from pathlib import Path
from tempfile import TemporaryFile
from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
	from wolfhece.wolf_array import WolfArray
	from wolfhece.pyvertexvectors import vector



_logger = logging.getLogger(__name__)





class _DtmWalloniaTiles:
	"""Thin wrapper around the tile-index shapefile for Wallonia 1 m DTM.

	On construction it downloads (or reuses from cache) the index shapefile
	and loads it as a :class:`~wolfhece.wolf_tiles.Tiles` instance.
	"""

	def __init__(self) -> None:
		from wolfhece.pydownloader import toys_dtm_wallonia_1m, DATADIR, WALLONIE_1M
		from wolfhece.wolf_tiles import Tiles
		from wolfhece.pyvertexvectors import VectorOGLRenderer



		self.idx_path: Path = toys_dtm_wallonia_1m()



		self._data_dir: Path = DATADIR / 'dtm' / WALLONIE_1M



		self._tiles = Tiles(filename=str(self.idx_path), linked_data_dir=str(self._data_dir))

		# Display lists are far faster than the shader path for large tile grids.
		self._tiles.rendering_machine = VectorOGLRenderer.LIST



	def extract(self, bounds_vec: 'vector', force: bool = False) -> 'WolfArray':
		"""Download required tiles for *bounds_vec* and return a cropped WolfArray.

		:param bounds_vec: A closed vector polygon that defines the region of interest.
		:param force: If True, allow downloading more than 100 tiles without confirmation.
		:returns: WolfArray cropped to the bounding box.
		:raises FileNotFoundError: If no VRT can be assembled.
		:raises RuntimeError: If more than 100 tiles are required and force is False.
		"""
		from wolfhece.pydownloader import toys_dtm_wallonia_1m_tile
		from wolfhece.wolf_array import WolfArray
		from wolfhece.wolf_vrt import create_vrt, crop_vrt

		dirdata = self._data_dir

		needed_tiles = []
		for tile_zone in self._tiles.myzones:
			for tile in tile_zone.myvectors:
				if tile.polygon.intersects(bounds_vec.polygon):
					needed_tiles.append(tile)

		if len(needed_tiles) > 100 and not force:
			raise RuntimeError(
				f"Refusing to download {len(needed_tiles)} tiles (>100). "
				"Pass force=True to override this safeguard."
			)

		for tile in needed_tiles:
			tile_file = dirdata / tile.myname
			if not tile_file.exists():
				_logger.warning("Tile %s not found, downloading...", tile.myname)
				toys_dtm_wallonia_1m_tile(tile.myname)

		vrt_path = dirdata / 'tmp_dtm.vrt'
		if vrt_path.exists():
			vrt_path.unlink()

		create_vrt(dirdata, fout='tmp_dtm.vrt')

		vrt_candidates = list(dirdata.glob('*.vrt'))
		if not vrt_candidates:
			raise FileNotFoundError(
				f"Could not build VRT in {dirdata}. "
				"Make sure at least one .tif tile has been downloaded."
			)
		glob_vrt = vrt_candidates[0]

		bounds_vec.find_minmax()
		bbox = bounds_vec.get_bounds_xx_yy()

		with TemporaryFile() as tmp:
			crop_vrt(str(glob_vrt), bbox, fout=str(tmp.name))
			return WolfArray(Path(tmp.name).with_suffix('.tif'))






class DtmWalloniaService:
	"""Stateful service to retrieve Wallonia DTM arrays.

	This class owns its tile-manager cache so callers can control lifecycle
	explicitly and avoid relying on module-level global state.
	"""

	def __init__(self) -> None:


		self._tiles: '_DtmWalloniaTiles | None' = None




	def get_tiles_manager(self, reset: bool = False) -> _DtmWalloniaTiles:
		"""Return a cached tile manager instance.

		:param reset: If True, drop and recreate the cached manager.
		:returns: The tile-manager instance used by this service.
		"""
		if reset or self._tiles is None:
			self._tiles = _DtmWalloniaTiles()
		return self._tiles




	def force_download_index(self) -> Path:
		"""Force a re-download of the tile-index shapefile and reset cache.

		:returns: Path to the downloaded index shapefile.
		"""
		from wolfhece.pydownloader import toys_dtm_wallonia_1m

		idx_path = toys_dtm_wallonia_1m(load_from_cache=False)
		self.get_tiles_manager(reset=True)
		return idx_path




	def get_dtm_1m(self, bounds_vec: 'vector', force: bool = False) -> 'WolfArray':
		"""Return a 1 m WolfArray cropped to *bounds_vec*.

		:param bounds_vec: Closed polygon defining the area of interest.
		:param force: If True, bypass the >100-tiles safeguard.
		:returns: 1 m DTM cropped to the polygon bounding box.
		"""
		tiles = self.get_tiles_manager()
		try:
			return tiles.extract(bounds_vec, force=force)
		except Exception:
			_logger.exception('DtmWallonia service: extract failed (force=%s)', force)
			raise




	def get_dtm_2m(self, bounds_vec: 'vector', force: bool = False) -> 'WolfArray':
		"""Return a 2 m WolfArray by mean-rebinning the 1 m raster.

		:param bounds_vec: Closed polygon defining the area of interest.
		:param force: If True, bypass the >100-tiles safeguard.
		:returns: 2 m DTM derived from the 1 m extraction.
		"""
		try:
			dtm1m = self.get_dtm_1m(bounds_vec, force=force)
			return dtm1m.rebin(factor=2, operation='mean')
		except Exception:
			_logger.exception('DtmWallonia service: 2m extract failed (force=%s)', force)
			raise




	def get_dtm_1m_xx_yy(
		self,
		xmin: float,
		xmax: float,
		ymin: float,
		ymax: float,
		force: bool = False,
	) -> 'WolfArray':
		"""Return a 1 m WolfArray from raw bounding coordinates.

		Convenience wrapper around :meth:`get_dtm_1m`.
		"""
		from wolfhece.pyvertexvectors import vector

		bounds_vec = vector()
		bounds_vec.add_vertices_from_array(np.array([[xmin, ymin],
											 [xmax, ymin],
											 [xmax, ymax],
											 [xmin, ymax]]))
		bounds_vec.force_to_close()

		try:
			return self.get_dtm_1m(bounds_vec, force=force)
		except Exception:
			_logger.exception(
				'DtmWallonia service: get_dtm_xx_yy failed for bounds '
				'(xmin=%s, xmax=%s, ymin=%s, ymax=%s, force=%s)',
				xmin,
				xmax,
				ymin,
				ymax,
				force,
			)
			raise




	def get_dtm_2m_xx_yy(
		self,
		xmin: float,
		xmax: float,
		ymin: float,
		ymax: float,
		force: bool = False,
	) -> 'WolfArray':
		"""Return a 2 m WolfArray from raw bounding coordinates.

		Convenience wrapper around :meth:`get_dtm_2m`.
		"""
		from wolfhece.pyvertexvectors import vector

		bounds_vec = vector()
		bounds_vec.add_vertices_from_array(np.array([[xmin, ymin],
											 [xmax, ymin],
											 [xmax, ymax],
											 [xmin, ymax]]))
		bounds_vec.force_to_close()

		try:
			return self.get_dtm_2m(bounds_vec, force=force)
		except Exception:
			_logger.exception(
				'DtmWallonia service: get_dtm_xx_yy failed for bounds '
				'(xmin=%s, xmax=%s, ymin=%s, ymax=%s, force=%s)',
				xmin,
				xmax,
				ymin,
				ymax,
				force,
			)
			raise