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

__all__ = ['USER_DIR', 'parse_datetime', 'read_S1_json', 'read_S2_json', 'get_s2_patch_directories',
           'get_s1_patch_directories', 'PATCHES_WITH_SNOW_URL', 'PATCHES_WITH_CLOUD_AND_SHADOW_URL',
           'get_patches_with_seasonal_snow', 'get_patches_with_cloud_and_shadow', 'is_snowy_patch',
           'is_cloudy_shadowy_patch', 'patches_from_original_train_split', 'patches_from_original_validation_split',
           'patches_from_original_test_split', 'get_original_split_from_patch_name', 'old2new_labels',
           'ben_19_labels_to_multi_hot', 'ben_43_labels_to_multi_hot']

# Cell
import json
import functools
import urllib
import warnings
from datetime import datetime
from pathlib import Path
from typing import List, Optional, Set, Union, Sequence, Dict
import csv

import appdirs
import dateutil
import fastcore.all as fc
from fastcore.basics import compose
from fastcore.dispatch import typedispatch
from pydantic import AnyHttpUrl, validate_arguments, FilePath, DirectoryPath

import bigearthnet_common.constants as ben_constants
from .constants import OLD2NEW_LABELS_DICT


# Cell
USER_DIR = Path(appdirs.user_data_dir("bigearthnet"))
USER_DIR.mkdir(exist_ok=True, parents=True)


# Cell
@validate_arguments
def _download_and_cache_url(url: AnyHttpUrl, force_download: bool = False):
    """
    Simply download contents of url to the default user directory.
    Allow to redownload with `force_download`
    """
    fp = USER_DIR / Path(url).name
    if not fp.exists() or force_download:
        response = urllib.request.urlopen(url).read()
        fp.write_bytes(response)
    return fp


# Cell
def parse_datetime(inp: Union[str, datetime]) -> datetime:
    """
    Parses an input into a `datetime` object.
    Will try its best to infer the correct format from a string.
    If a `datetime` object is already provided it will be returned.
    Otherwise it will raise an error.
    """
    return _parse_datetime(inp)


@typedispatch
def _parse_datetime(acquisition_date: str) -> datetime:
    return compose(dateutil.parser.parse, _parse_datetime)(acquisition_date)


@typedispatch
def _parse_datetime(acquisition_date: datetime) -> datetime:
    return acquisition_date


@typedispatch
def _parse_datetime(acquisition_date: object) -> None:
    raise TypeError("Could not parse acquisition_date!")


# Cell
@validate_arguments
def _read_json(json_fp: FilePath, expected_keys: Set, read_only_expected: bool = True) -> Dict[str, str]:
    """
    Parse the json file given with the file path `json_fp`.
    The function checks if all of the `expected_keys` are present, which
    ensures that no keys have been accidentilly deleted (this has happend before).
    If `read_only_expected` is set, only the keys provided in `expected_keys` are read
    and returned.
    This prevents accidental processing of injected metadata.

    Args:
        json_fp (FilePath): Path to json file
        expected_keys (Set): Keys that are expected to be present in the json file
        read_only_expected (bool, optional): Read only the keys given in `expected_keys`. Defaults to True.

    Returns:
        [Dict[str, str]]: A dictionary of the keys.
    """
    try:
        complete_data = json.loads(json_fp.read_text())
    except json.JSONDecodeError:
        raise ValueError(f"Error trying to read json from: ", json_fp)

    missing_elements = expected_keys - complete_data.keys()
    if len(missing_elements) > 0:
        raise ValueError(f"{json_fp} is missing entries!", missing_elements)

    # ensure that the original values are loaded, as some users may customize the original json files
    if read_only_expected:
        return {k: v for k, v in complete_data.items() if k in expected_keys}
    return complete_data


def read_S1_json(json_fp: FilePath) -> Dict[str, str]:
    """
    A helper function that *safely* reads a BigEarthNet-S1 json file.
    It will ensure that all expected entries are present and only read those
    entries.
    This helps to avoid issues where the JSON files were accidentally modified
    or partially deleted.

    Note: This function will also silently fix a typo present in the `coordinates` key
    from version: S1_v1.0. A coordinates key is named `lly` and it should be `lry`.
    """
    data = _read_json(json_fp, ben_constants.BEN_S1_V1_0_JSON_KEYS)
    # Silently fix key error in S1
    if "lly" in data["coordinates"]:
        data["coordinates"]["lry"] = data["coordinates"].pop("lly")
    return data

def read_S2_json(json_fp: FilePath) -> Dict[str, str]:
    """
    A helper function that *safely* reads a BigEarthNet-S1 json file.
    It will ensure that all expected entries are present and only read those
    entries.
    This helps to avoid issues where the JSON files were accidentally modified
    or partially deleted.
    """
    return _read_json(json_fp, ben_constants.BEN_S2_V1_0_JSON_KEYS)

# Cell
@validate_arguments
def get_s2_patch_directories(dir_path: DirectoryPath) -> List[Path]:
    """
    Will find all S2 patch directories in the provided `dir_path`.
    Only directories that strictly cohere to the naming convention will be returned.
    """
    return [p for p in dir_path.iterdir() if ben_constants.BEN_S2_RE.fullmatch(p.name) is not None]

@validate_arguments
def get_s1_patch_directories(dir_path: DirectoryPath) -> List[Path]:
    """
    Will find all S1 patch directories in the provided `dir_path`.
    Only directories that strictly cohere to the naming convention will be returned.
    """
    return [p for p in dir_path.iterdir() if ben_constants.BEN_S1_RE.fullmatch(p.name) is not None]

# Cell
# PATCHES_WITH_SNOW_URL = "http://bigearth.net/static/documents/patches_with_seasonal_snow.csv"
# PATCHES_WITH_CLOUD_AND_SHADOW_URL = "http://bigearth.net/static/documents/get_patches_with_cloud_and_shadow.csv"
PATCHES_WITH_SNOW_URL = "https://git.tu-berlin.de/k.clasen/ben-mirror/-/raw/master/patches_with_seasonal_snow.csv"
PATCHES_WITH_CLOUD_AND_SHADOW_URL = "https://git.tu-berlin.de/k.clasen/ben-mirror/-/raw/master/patches_with_cloud_and_shadow.csv"


# Cell
@validate_arguments
def _conv_single_col_csv_to_set(
    url: AnyHttpUrl, name: str = "Name", force_download: bool = False
) -> Set:
    """
    Given a url to a CSV file *without* a header
    line and only a single column, return the set of
    all values.

    Will write remote csv to disk for better performance.
    Set `force_download` to re-download the file.
    """
    fp = _download_and_cache_url(url, force_download=force_download)
    with open(fp, mode="r") as csv_file:
        reader = csv.DictReader(csv_file, fieldnames=[name])
        return {row[name] for row in reader}

@functools.lru_cache()
def get_patches_with_seasonal_snow(force_download: bool = False) -> Set:
    """List all patches with seasonal snow from **original** BigEarthNet dataset."""
    return _conv_single_col_csv_to_set(
        PATCHES_WITH_SNOW_URL, force_download=force_download
    )


@functools.lru_cache()
def get_patches_with_cloud_and_shadow(force_download: bool = False) -> Set:
    """List all patches with cloud and shadow from **original** BigEarthNet dataset."""
    return _conv_single_col_csv_to_set(
        PATCHES_WITH_CLOUD_AND_SHADOW_URL, force_download=force_download
    )


# Cell
@validate_arguments
def is_snowy_patch(patch_name: str):
    """
    Fast function that checks whether `patch_name` is a patch
    that contains a lot of seasonal snow.
    """
    return patch_name in get_patches_with_seasonal_snow()


@validate_arguments
def is_cloudy_shadowy_patch(patch_name: str):
    """
    Fast function that checks whether `patch_name` is a patch
    that contains a lot of shadow or is obstructed by clouds.
    """
    return patch_name in get_patches_with_cloud_and_shadow()


# Cell
@functools.lru_cache()
@validate_arguments
def patches_from_original_train_split(
    split_url: AnyHttpUrl = "https://git.tu-berlin.de/rsim/BigEarthNet-S2_19-classes_models/-/raw/master/splits/train.csv",
    force_download: bool = False,
) -> Set:
    """
    List all train patches from the original train/validation/test split.
    There are two possible sources:

    1. https://git.tu-berlin.de/rsim/BigEarthNet-S2_19-classes_models/-/raw/master/splits/train.csv
    2. https://git.tu-berlin.de/rsim/BigEarthNet-S2_43-classes_models/-/raw/master/splits/train.csv

    While writing this function, there is **no** difference between these two splits.
    But this may change in the future!
    """
    return _conv_single_col_csv_to_set(split_url, force_download=force_download)


@functools.lru_cache()
@validate_arguments
def patches_from_original_validation_split(
    split_url: AnyHttpUrl = "https://git.tu-berlin.de/rsim/BigEarthNet-S2_19-classes_models/-/raw/master/splits/val.csv",
    force_download: bool = False,
) -> Set:
    """
    List all validation patches from the original train/validation/test split.
    There are two possible sources:

    1. https://git.tu-berlin.de/rsim/BigEarthNet-S2_19-classes_models/-/raw/master/splits/train.csv"
    2. "https://git.tu-berlin.de/rsim/BigEarthNet-S2_43-classes_models/-/raw/master/splits/train.csv"

    While writing this function, there is **no** difference between these two splits.
    But this may change in the future!
    """
    return _conv_single_col_csv_to_set(split_url, force_download=force_download)


@functools.lru_cache()
@validate_arguments
def patches_from_original_test_split(
    split_url: AnyHttpUrl = "https://git.tu-berlin.de/rsim/BigEarthNet-S2_19-classes_models/-/raw/master/splits/test.csv",
    force_download: bool = False,
) -> Set:
    """
    List all test patches from the original train/validation/test split.
    There are two possible sources:

    1. https://git.tu-berlin.de/rsim/BigEarthNet-S2_19-classes_models/-/raw/master/splits/test.csv
    2. https://git.tu-berlin.de/rsim/BigEarthNet-S2_43-classes_models/-/raw/master/splits/test.csv

    While writing this function, there is **no** difference between these two splits.
    But this may change in the future!
    """
    return _conv_single_col_csv_to_set(split_url, force_download=force_download)


# Cell
@validate_arguments
def get_original_split_from_patch_name(patch: str) -> Optional[str]:
    """
    Returns "train"/"validation"/"test" or `None`.
    The value is retrieved from the original BigEarthNet
    train/validation/test split. If the input is not present
    in any split, it will return `None` and raise a UserWarning.
    This happens for patches that are either in the
    cloud/shadow or seasonal snow set or if there exists no 19-label target.

    The splits are from the 19-classes version.
    While writing this function there was no difference between the
    19-classes and the 43-classes version.
    """
    train = patches_from_original_train_split()
    validation = patches_from_original_validation_split()
    test = patches_from_original_test_split()

    if patch in train:
        return "train"
    elif patch in validation:
        return "validation"
    elif patch in test:
        return "test"
    warnings.warn(
        "Provided an input patch name which was not part of the original split.",
        UserWarning,
    )
    return None


# Cell
@validate_arguments
def _old2new_label(old_label: str) -> Optional[str]:
    """
    Converts old-style BigEearthNet label to the
    new labels.

    > Note: Some labels were removed! This function
    will return `None` if the label was removed and
    raise a `KeyError` if the input label is unknown.
    """
    return OLD2NEW_LABELS_DICT[old_label]


def old2new_labels(old_labels: Sequence[str]) -> Optional[List[str]]:
    """
    Converts a list of old-style BigEarthNet labels
    to a list of labels.

    If there are no corresponding new labels (which can happen with original BEN patches!)
    then the function will return `None` and raise a user warning.

    If an illegal/unknown input label is provided, a `KeyError` is raised.
    """
    new_labels = [
        _old2new_label(l) for l in old_labels if _old2new_label(l) is not None
    ]
    if len(old_labels) > 0 and len(new_labels) == 0:
        warnings.warn(
            "Provided a list of old labels that only contains `removed` labels!",
            UserWarning,
        )
        new_labels = None
    return new_labels


# Cell
@validate_arguments
def ben_19_labels_to_multi_hot(labels: Sequence[str]) -> List[float]:
    """
    Convenience function that converts an input sequence of labels into
    a multi-hot encoded vector.
    The naturally ordered label list is used as an encoder reference
    - `bigearthnet_common.NEW_LABELS`

    If an unknown label is given, a `KeyError` is raised.

    Be aware that this approach assumes that **all** labels are actually used in the dataset!
    This is not necessarily the case if you are using a subset!
    For example, the "Agro-forestry areas" class is only present in Portugal and in no other country!
    """
    idxs = [ben_constants.NEW_LABELS_TO_IDX[label] for label in labels]
    multi_hot = fc.L([0] * len(ben_constants.NEW_LABELS))
    multi_hot[idxs] = 1.
    return list(multi_hot)

@validate_arguments
def ben_43_labels_to_multi_hot(labels: Sequence[str]) -> List[float]:
    """
    Convenience function that converts an input sequence of labels into
    a multi-hot encoded vector.
    The naturally ordered label list is used as an encoder reference
    - `bigearthnet_common.OLD_LABELS`

    If an unknown label is given, a `KeyError` is raised.

    Be aware that this approach assumes that **all** labels are actually used in the dataset!
    This is not necessarily the case if you are using a subset!
    For example, the "Agro-forestry areas" class is only present in Portugal and in no other country!
    """
    idxs = [ben_constants.OLD_LABELS_TO_IDX[label] for label in labels]
    multi_hot = fc.L([0] * len(ben_constants.OLD_LABELS))
    multi_hot[idxs] = 1.
    return list(multi_hot)
