spot

class Spot67DRSScene(Spot67Scene):
    """
    Spot 6/7 class for ADS-DRS products.

    The Spot67Scene class carries all metadata and images sources from the
    scene. A Spot67Scene object can be instantiated from the XS and PAN DIMAPS
    (.XML) file.

    """

    @property
    def req_assets(self) -> Dict[str, str]:
        """
        Required assets (overrides Spot67Scene property)

        Returns: required assets
        """
        return {
            "dimap": "XML DIMAP document",
            "src": "imagery source",
            "roi_msk_vec": "region of interest mask (vector data)",
            "cld_msk_vec": "clouds mask (vector data)",
        }

    def __init__(self, assets_paths: Dict[str, str]):
        """
        Args:
            assets_paths: assets paths

        """

        # Parse dimap
        assert "dimap_xs" in assets_paths, "XS DIMAP XML document is missing"
        additional_md = spot67_metadata_parser(assets_paths["dimap_xs"])
        acquisition_date = datetime.strptime(
            additional_md["acquisition_date"], "%Y-%m-%dT%H:%M:%S.%fZ"
        )

        # Call parent constructor, before accessing to self.assets_paths
        super().__init__(
            assets_paths=assets_paths,
            acquisition_date=acquisition_date,
            src_class=Spot67DRSSource,
            additional_metadata=additional_md,
        )

        # Choice of the CLD and ROI masks based on the pxs overlap
        dic = self.metadata
        ref = "pan" if dic["pan_overlap"] > dic["xs_overlap"] else "xs"
        pths = self.assets_paths
        for key in ["cld", "roi"]:
            pths[f"{key}_msk_vec"] = pths[f"{key}_msk_vec_{ref}"]

class Spot67DRSSource(CommonImagerySource):
    """
    Source capabilities for Spot-6/7 ADS-DRS products
    """

    def cld_msk_drilled(
            self, nodata: Union[float, int] = 0
    ) -> Spot67DRSSource:
        """
        Return the source drilled from the cloud mask

        Args:
            nodata: nodata value inside holes (Default value = 0)

        Returns:
            drilled source

        """
        return self.drilled(
            self.root_scene.assets_paths["cld_msk_vec"], nodata=nodata
        )

    def reflectance(self, factor: float = None, **kwargs) -> Spot67DRSSource:
        """
        This function is used internally by `get_xs()`, `get_pxs()`, and
        `get_pan()` to compute the reflectance (or not!) of the optical image.

        Warning: there is a bug in OTB<=8.0.2
        see https://gitlab.orfeo-toolbox.org/orfeotoolbox/otb/-/issues/2314

        Args:
            factor: factor to scale pixel values (e.g. 10000)
            **kwargs: same options as BundleToPerfectSensor in OTB

        Returns:
            calibrated, or original image source

        """

        # Radiometry correction
        params = {"in": self}
        params.update(kwargs)
        out = pyotb.OpticalCalibration(params)

        if factor:
            out = out * factor

        return self.new_source(out)

class Spot67IGNScene(Spot67Scene):
    """
    Spot 6/7 class for IGN products.

    A Spot67IGNScene object can be instantiated from the XS and PAN rasters
    paths.

    """

    def __init__(self, assets_paths: Dict[str, str]):
        """
        Args:
            assets_paths: assets paths

        """
        # Retrieve date
        assert "src_xs" in assets_paths, "XS image is missing"
        assert "src_pan" in assets_paths, "PAN image is missing"
        xs_basepath = utils.basename(assets_paths["src_xs"])
        # 8 first consecutive digits in the XS filename
        result = re.search(r"\d{8}", xs_basepath)
        datestr = result.group(0) if result else result
        acquisition_date = dates.str2datetime(datestr)

        # Call parent constructor
        super().__init__(
            assets_paths=assets_paths,
            acquisition_date=acquisition_date
        )

class Spot67Scene(Scene):
    """
    Generic Spot 6/7 Scene class.

    The Spot67Scene class carries PAN, XS, and PXS images sources.

    """
    PXS_OVERLAP_THRESH = 0.995

    def _check_assets(self, assets_paths: Dict[str, str]):
        """
        Checks that all assets are here

        """
        for src_key in ["xs", "pan"]:
            for req_prefix, req_description in self.req_assets.items():
                key = f"{req_prefix}_{src_key}"
                assert key in assets_paths, \
                    f"key '{key}' for {req_description} " \
                    "is missing from assets paths"

    @property
    def req_assets(self) -> Dict[str, str]:
        """
        Required assets

        Returns: required assets
        """
        return {
            "src": "imagery source",
        }

    def __init__(
            self,
            acquisition_date: datetime,
            assets_paths: Dict[str, str],
            src_class: Type[Source] = CommonImagerySource,
            additional_metadata: Dict[str, str] = None,
    ):
        """
        Args:
            acquisition_date: Acquisition date
            assets_paths: assets paths
            additional_metadata: additional metadata

        """

        # Get EPSG and bounding box
        epsg_xs, extent_wgs84_xs = get_epsg_extent_wgs84(
            assets_paths["src_xs"]
        )
        epsg_pan, extent_wgs84_pan = get_epsg_extent_wgs84(
            assets_paths["src_pan"]
        )

        # Check that EPSG for PAN and XS are the same
        if epsg_pan != epsg_xs:
            raise ValueError(
                f"EPSG of XS and PAN sources are different:  XS EPSG is "
                f"{epsg_xs}, PAN EPSG is {epsg_pan}")

        # Here we compute bounding boxes overlap, to choose the most
        # appropriated CLD and ROI masks for the scene. Indeed, sometimes
        # products are not generated as they should (i.e. bundle) and XS and
        # PAN have different extents so CLD and ROI masks are not the same for
        # XS and PAN. We keep the ROI+CLD masks of the PAN or XS lying
        # completely inside the other one.
        xs_overlap = extent_overlap(extent_wgs84_xs, extent_wgs84_pan)
        pan_overlap = extent_overlap(extent_wgs84_pan, extent_wgs84_xs)

        # Throw some warning or error, depending on the pxs overlap value
        pxs_overlap = max(xs_overlap, pan_overlap)
        if pxs_overlap == 0:
            raise ValueError("No overlap between PAN and XS images")
        if max(pan_overlap, xs_overlap) < self.PXS_OVERLAP_THRESH:
            print(f"Warning: partial overlap of {100 * pxs_overlap:.2f}%")

        # Final extent
        extent = extent_wgs84_pan if pan_overlap > xs_overlap else \
            extent_wgs84_xs

        sources = {
            "xs": partial(
                src_class, root_scene=self, out=assets_paths["src_xs"]
            ),
            "pan": partial(
                src_class, root_scene=self, out=assets_paths["src_pan"]
            ),
            "pxs": partial(self._get_pxs, src_class=src_class)
        }

        additional_md = additional_metadata.copy() if additional_metadata \
            else {}
        additional_md.update({
            "xs_extent_wgs84": extent_wgs84_xs,
            "pan_extent_wgs84": extent_wgs84_pan,
            "xs_overlap": xs_overlap,
            "pan_overlap": pan_overlap
        })

        # Call parent constructor, before accessing assets dicts
        super().__init__(
            acquisition_date=acquisition_date,
            epsg=epsg_xs,
            extent_wgs84=extent,
            assets_paths=assets_paths,
            sources=sources,
            additional_metadata=additional_md,
        )

    def _get_pxs(
            self,
            src_class: Type[CommonImagerySource],
            **kwargs
    ) -> CommonImagerySource:
        """
        Create source for PXS (computed using OTB)

        Args:
            src_class: source class
            **kwargs: OTB BundleToPerfectSensor options, e.g. method="bayes"

        Returns:
            Source for BundleToPerfectSensor output

        """
        pan = self.get_pan()
        params = {
            "inxs": self.get_xs(),
            "inp": pan
        }
        params.update(kwargs)
        pansharp = src_class(self, pyotb.BundleToPerfectSensor(params))
        binary_mask = pyotb.Superimpose({
            "inr": pansharp,
            "inm": pan,
            "interpolator": "nn"
        })
        return pansharp.masked(binary_mask=binary_mask)