from typing import Any, Dict, List, Type, TypedDict, Optional
import contextlib
import attr
from rasterio.crs import CRS
from morecantile import TileMatrixSet
from rio_tiler.constants import WEB_MERCATOR_TMS, WGS84_CRS
from rio_tiler.io import BaseReader, Reader
from rio_tiler.types import BBox, Indexes
from rio_tiler.models import BandStatistics, ImageData, Info, PointData
class Input(TypedDict, total=True):
 """Reader Options."""
 data: str
 mask: str
 max_correlation: float
@attr.s
class CustomReader(BaseReader):
 input: Input = attr.ib()
 tms: TileMatrixSet = attr.ib(default=WEB_MERCATOR_TMS)
 geographic_crs: CRS = attr.ib(default=WGS84_CRS)
 dataset: Type[Reader] = attr.ib(
 init=False,
 )
 mask: Type[Reader] = attr.ib(
 init=False,
 )
 colormap: Dict = attr.ib(init=False, default=None)
 # Context Manager to handle rasterio open/close
 _ctx_stack = attr.ib(init=False, factory=contextlib.ExitStack)
 def __attrs_post_init__(self):
 """Define _kwargs, open dataset and get info."""
 self.dataset = self._ctx_stack.enter_context(Reader(self.input["data"]))
 self.mask = self._ctx_stack.enter_context(Reader(self.input["mask"]))
 self.bounds = self.dataset.bounds
 self.crs = self.dataset.crs
 self.colormap = self.dataset.colormap
 @property
 def minzoom(self):
 """Return dataset minzoom."""
 return self.dataset.minzoom
 @property
 def maxzoom(self):
 """Return dataset maxzoom."""
 return self.dataset.maxzoom
 def close(self):
 """Close rasterio dataset."""
 self._ctx_stack.close()
 def __exit__(self, exc_type, exc_value, traceback):
 """Support using with Context Managers."""
 self.close()
 def info(self) -> Info:
 # could return self.dataset.info()
 raise NotImplementedError
 def statistics(self, *args, **kwargs) -> Dict[str, BandStatistics]:
 # could return self.dataset.statistics(*args, **kwargs)
 raise NotImplementedError
 def tile(
 self,
 tile_x: int,
 tile_y: int,
 tile_z: int,
 tilesize: int = 256,
 indexes: Optional[Indexes] = None,
 expression: Optional[str] = None,
 **kwargs: Any,
 ) -> ImageData:
 img = self.dataset.tile(
 tile_x,
 tile_y,
 tile_z,
 tilesize,
 indexes=indexes,
 expression=expression,
 **kwargs,
 )
 mask = self.mask.tile(
 tile_x,
 tile_y,
 tile_z,
 tilesize,
 indexes=1,
 **kwargs,
 )
 # edited from original response
 # img.mask = mask.array > self.input["max_correlation"]
 img.array.mask = mask.array > self.input["max_correlation"]
 
 return mask
 def part(self, bbox: BBox) -> ImageData:
 raise NotImplementedError
 def preview(self) -> ImageData:
 raise NotImplementedError
 def point(self, lon: float, lat: float) -> PointData:
 raise NotImplementedError
 def feature(self, shape: Dict) -> ImageData:
 raise NotImplementedError
 def read(
 self,
 indexes: Optional[Indexes] = None,
 expression: Optional[str] = None,
 **kwargs: Any,
 ) -> ImageData:
 raise NotImplementedError

def InputDependency(
 url: Annotated[str, Query(description="Dataset URL")],
 mask: Annotated[str, Query(description="Mask URL")],
 max_correlation: Annotated[float, Query(description="Max Correlation")] = 0.3,
) -> Dict:
 """Create dataset path from args"""
 return {
 "data": url,
 "mask": mask,
 "max_correlation": max_correlation,
 }