# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/02_sets.ipynb (unless otherwise specified).

__all__ = ['get_all_s2_patches', 'get_all_s1_patches', 'get_recommended_s2_patches', 'get_recommended_s1_patches',
           'filter_s2_patches_by_country', 'filter_s1_patches_by_country', 'filter_s2_patches_by_season',
           'filter_s1_patches_by_season']

# Cell
from typing import Set

import bigearthnet_common.constants as ben_constants
from .base import (
    is_cloudy_shadowy_patch,
    is_snowy_patch,
    get_s1_patches_with_no_19_class_target,
    get_s2_patches_with_no_19_class_target,
    get_patches_to_season_mapping,
    get_patches_to_country_mapping,
    get_complete_s1_to_s2_patch_name_mapping,
    get_complete_s2_to_s1_patch_name_mapping,
)


# Cell
def get_all_s2_patches() -> Set[str]:
    return set(get_complete_s2_to_s1_patch_name_mapping().keys())


def get_all_s1_patches() -> Set[str]:
    return set(get_complete_s1_to_s2_patch_name_mapping().keys())


def get_recommended_s2_patches() -> Set[str]:
    s2_patches = get_all_s2_patches()
    no_19_class_targets = get_s2_patches_with_no_19_class_target()
    recommended_s2_patches = {
        p
        for p in s2_patches
        if not is_snowy_patch(p)
        and not is_cloudy_shadowy_patch(p)
        and p not in no_19_class_targets
    }
    return recommended_s2_patches

def get_recommended_s1_patches() -> Set[str]:
    s1_patches = get_all_s1_patches()
    no_19_class_targets = get_s1_patches_with_no_19_class_target()
    recommended_s1_patches = {
        p
        for p in s1_patches
        if not is_snowy_patch(p)
        and not is_cloudy_shadowy_patch(p)
        and p not in no_19_class_targets
    }
    return recommended_s1_patches

# Cell
# FUTURE: Maybe convert country string to enums?
def filter_s2_patches_by_country(patches, country: str) -> Set[str]:
    country = country.title()  # ben country style
    if country.title() not in ben_constants.COUNTRIES:
        raise ValueError(f"{country} is not one of the BEN countries: {ben_constants.COUNTRIES}!")

    patch_country_mapping = get_patches_to_country_mapping(use_s2_patch_names=True)
    return {p for p in patches if patch_country_mapping[p] == country}

def filter_s1_patches_by_country(patches, country: str) -> Set[str]:
    country = country.title()  # ben country style
    if country.title() not in ben_constants.COUNTRIES:
        raise ValueError(f"{country} is not one of the BEN countries: {ben_constants.COUNTRIES}!")

    patch_country_mapping = get_patches_to_country_mapping(use_s2_patch_names=False)
    return {p for p in patches if patch_country_mapping[p] == country}

def filter_s2_patches_by_season(patches, season: str) -> Set[str]:
    season = season.title()  # ben season style
    patch_season_mapping = get_patches_to_season_mapping(use_s2_patch_names=True)
    return {p for p in patches if patch_season_mapping[p] == season}

def filter_s1_patches_by_season(patches, season: str) -> Set[str]:
    season = season.title()  # ben season style
    patch_season_mapping = get_patches_to_season_mapping(use_s2_patch_names=False)
    return {p for p in patches if patch_season_mapping[p] == season}
