sentinel

class Sentinel22AScene(Sentinel2TheiaScene):
    """
    Sentinel-2 level 2A scene class.

    The class carries all the metadata from the root_scene, and can be used to
    retrieve its sources.

    """

    def __init__(self, archive: str, tag: str = "FRE"):
        """
        Args:
            archive: .zip file or folder. Must be a product from MAJA.
            tag: product tag (SRE/FRE)

        """
        masks = ["clm_r1", "edg_r1", "clm_r2", "edg_r2"]
        src_classes = {
            **{key: Sentinel2Theia2ASource for key in self.bands_keys},
            **{key: CommonImagerySource for key in masks}
        }
        super().__init__(archive=archive, tag=tag, src_classes=src_classes)

class Sentinel23AScene(Sentinel2TheiaScene):
    """Sentinel-2 level 3A scene class.

    The class carries all the metadata from the root_scene, and can be used to
    retrieve its sources.

    """

    def __init__(self, archive: str, tag: str = "FRC"):
        """
        Args:
            archive: .zip file or folder. Must be a product from WASP.
            tag: product tag (FRC)

        """
        masks = ["flg_r1", "flg_r2"]
        src_classes = {
            **{key: Sentinel2Theia3ASource for key in self.bands_keys},
            **{key: CommonImagerySource for key in masks}
        }
        super().__init__(archive=archive, tag=tag, src_classes=src_classes)

class Sentinel2MPCScene(Sentinel2SceneBase):
    """class for Sentinel-2 images from mMicrosoft Planetary Computer"""

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

        """
        # Assets (spectral bands)
        # We just use the same key for all Sentinel2SceneBase products
        bands_names_mapping = {
            "b4": "B04",
            "b3": "B03",
            "b2": "B02",
            "b8": "B08",
            "b5": "B05",
            "b6": "B06",
            "b7": "B07",
            "b8a": "B8A",
            "b11": "B11",
            "b12": "B12"
        }
        updated_assets_paths = {
            key: assets_paths[asset_key] for key, asset_key
            in bands_names_mapping.items()
        }

        # Sources classes
        src_classes = {
            key: CommonImagerySource for key in updated_assets_paths
        }

        # Date, extent
        b2_path = updated_assets_paths["b2"]
        # Here, b2_path is something like ...T45WXU_20221019T062901_B03_10m.tif
        epsg, extent = get_epsg_extent_wgs84(b2_path)
        datestr = b2_path.split("_")[-3]  # 20180630T105440
        acquisition_date = datetime.strptime(datestr, '%Y%m%dT%H%M%S')

        # Call parent constructor
        super().__init__(
            acquisition_date=acquisition_date,
            epsg=epsg,
            extent_wgs84=extent,
            assets_paths=updated_assets_paths,
            src_classes=src_classes
        )

class Sentinel2SceneBase(Scene):
    """Base class for Sentinel-2 images"""

    bands_keys = BANDS_1 + BANDS_2
    concatenated_bands_dict = {
        "10m_bands": BANDS_1,
        "20m_bands": BANDS_2
    }

    def __init__(
            self,
            acquisition_date: datetime,
            epsg: int,
            extent_wgs84: List[Tuple(float, float)],
            assets_paths: List[str, str],
            src_classes: Dict[str, Type[Source]] = None,
            additional_metadata: Dict[str, Any] = None,
    ):
        """
        Initialize the Sentinel-2 Scene

        Args:
            acquisition_date: Acquisition date
            epsg: EPSG code
            extent_wgs84: extent in WGS84 coordinates reference system
            assets_paths: assets dictionary
                {source_name, source_path/url}, e.g.:
                {'b4': '/vsicurl/https://.../...B4.tif'}
            src_classes: imagery sources class
            additional_metadata: additional metadata

        """

        assert all(key in src_classes for key in self.bands_keys), \
            "Some keys in concatenated_bands_dict are not in src_classes"

        # Sources
        sources = {
            key: partial(src_class, root_scene=self, out=assets_paths[key])
            for key, src_class in src_classes.items()
        }

        # Sources for concatenated bands: get_b10m_bands, get_20m_bands
        for key, default_bands_names in self.concatenated_bands_dict.items():
            sources.update({
                key: partial(
                    self._get_bands,
                    src_classes[default_bands_names[0]],
                    default_bands_names
                )
            })

        super().__init__(
            acquisition_date=acquisition_date,
            epsg=epsg,
            extent_wgs84=extent_wgs84,
            assets_paths=assets_paths,
            sources=sources,
            additional_metadata=additional_metadata
        )

    def _get_bands(
            self,
            imagery_src_class,
            default_bands_names: tuple(str),
            *args
    ) -> Source:
        """
        Returns a Source for the concatenated bands

        Args:
            imagery_src_class: source class
            default_bands_names: names of available bands
            *args: can be:
             - names of the bands
             - a name
             - a list or a tuple of names

        Returns:
            Selected spectral bands

        """
        bands = default_bands_names
        if args:
            bands = []
            for arg in args:
                bands += arg if isinstance(arg, (list, tuple)) else [arg]
        for band in bands:
            assert isinstance(band, str), f"{band} is not a string"
        # user can use upper or lower cases
        bands = [band.lower() for band in bands]
        assert all(band in default_bands_names for band in bands), \
            f"Some bands in {bands} are not in the available bands " \
            f"{default_bands_names}"
        concat = pyotb.ConcatenateImages(
            [self.get_asset_path(band_name) for band_name in bands]
        )
        return imagery_src_class(self, concat)

class Sentinel2Theia2ASource(Sentinel2TheiaGenericSource):
    """Sentinel-2 level 2A source class"""

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

        Args:
            nodata: nodata value inside holes (Default value = -10000)

        Returns:
            drilled source

        """
        return self.msk_drilled(
            msk_dict={
                self.R1_SPC: self.root_scene.get_asset_path("clm_r1"),
                self.R2_SPC: self.root_scene.get_asset_path("clm_r2")
            },
            exp="im1b1==0?255:0",
            nodata=nodata
        )

class Sentinel2Theia3ASource(Sentinel2TheiaGenericSource):
    """Sentinel-2 level 3A source class"""

    def flg_msk_drilled(
            self,
            keep_flags_values: Tuple[int] = (3, 4),
            nodata: Union[float, int] = -10000
    ) -> Sentinel2Theia3ASource:
        """
        Return the source drilled from the FLG mask

        Args:
            keep_flags_values: flags values to keep (Default value = (3, 4)).
                Can be:

               - 0 = No data
               - 1 = Cloud
               - 2 = Snow
               - 3 = Water
               - 4 = Land

                (source:
                https://labo.obs-mip.fr/multitemp/theias-l3a-product-format/)
            nodata: nodata value inside holes (Default value = -10000)

        Returns:
            drilled source

        """
        cond = "||".join([f"im1b1=={val}" for val in keep_flags_values])
        exp = cond + "?255:0"
        return self.msk_drilled(
            msk_dict={
                self.R1_SPC: self.root_scene.get_asset_path("flg_r1"),
                self.R2_SPC: self.root_scene.get_asset_path("flg_r2")
            },
            exp=exp,
            nodata=nodata
        )

class Sentinel2TheiaGenericSource(CommonImagerySource):
    """Class for generic Sentinel-2 sources"""
    R1_SPC = 10.0
    R2_SPC = 20.0

    def msk_drilled(
            self,
            msk_dict: Dict[float, str],
            exp: str,
            nodata: Union[float, int] = 0
    ) -> Sentinel2TheiaGenericSource:
        """
        Args:
            msk_dict: dict of masks
            exp: bandmath expression to form the 0-255 binary mask
            nodata: no-data value in masked output (Default value = 0)

        Returns:
            new masked source

        """
        img_spc = pyotb.ReadImageInfo(self)['spacingx']
        if img_spc not in msk_dict:
            raise KeyError(f"No mask for image spacing {img_spc}")
        binary_mask = pyotb.BandMath({"il": msk_dict[img_spc], "exp": exp})
        return self.masked(binary_mask=binary_mask, nodata=nodata)

class Sentinel2TheiaScene(Sentinel2SceneBase):
    """Base class for Sentinel-2 images from Theia"""

    @abstractmethod
    def __init__(
            self,
            archive: str,
            tag: str,
            src_classes: Dict[str, Type[Source]]
    ):
        """
        Args:
            archive: product .zip or directory
            tag: pattern to match in filenames, e.g. "FRE"
            src_classes: imagery sources classes dict

        """
        self.archive = archive

        # Retrieve the list of .tif files
        is_zip = self.archive.lower().endswith(".zip")
        if is_zip:
            # print("Input type is a .zip archive")
            files = utils.list_files_in_zip(self.archive)
            self.files = [utils.to_vsizip(self.archive, f) for f in files]
        else:
            # print("Input type is a directory")
            self.files = utils.find_files_in_all_subdirs(
                self.archive, "*.tif", case_sensitive=False
            )

        # Assets
        assets_paths = {}

        # In Theia products, the suffixes are always uppercase (B4, ..., B8A)
        # so we need to retrieve the files from the uppercase key.

        # 1. Spectral bands
        for key in self.bands_keys:
            suffix = key.upper()  # e.g. b4 --> B4
            assets_paths[key] = self.get_file(endswith=f"_{tag}_{suffix}.tif")

        # 2. Other assets (quality masks, etc)
        for key in [key for key in src_classes if key not in assets_paths]:
            suffix = key.upper()  # e.g. cld_r1 --> CLD_R1
            assets_paths[key] = self.get_file(endswith=f"_{suffix}.tif")

        # Date, extent
        b2_path = assets_paths["b2"]
        epsg, extent = get_epsg_extent_wgs84(b2_path)
        # Basename, e.g. SENTINEL2A_20180630-105440-000_L2A_T31TEJ_D_V1-8
        b2_basepath = utils.basename(b2_path)
        datestr = b2_basepath.split("_")[1]  # e.g. 20180630-105440
        acquisition_date = datetime.strptime(datestr, '%Y%m%d-%H%M%S-%f')

        # Call parent constructor
        super().__init__(
            acquisition_date=acquisition_date,
            epsg=epsg,
            extent_wgs84=extent,
            assets_paths=assets_paths,
            src_classes=src_classes,
            additional_metadata={"archive": archive}
        )

    def get_file(self, endswith: str) -> str:
        """
        Return the specified file.

        Args:
            endswith: filtered extension

        Returns:
            the file

        """
        filtered_files_list = [f for f in self.files if f.endswith(endswith)]
        nb_matches = len(filtered_files_list)
        if nb_matches == 1:
            return filtered_files_list[0]
        print(
            "Warning: "
            f"{self.archive}: found {nb_matches} occurrence(s) of file "
            f"with suffix \"{endswith}\""
        )