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

__all__ = ['USER_DIR', 'parse_datetime', 'box_from_ul_lr_coords', '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']

# Cell
import functools
import urllib
import warnings
from datetime import datetime
from numbers import Real
from pathlib import Path
from typing import List, Optional, Set, Tuple, Union

import appdirs
import dateutil
import fastcore.all as fc
import geopandas
import pandas as pd
from fastcore.basics import compose
from fastcore.dispatch import typedispatch
from pydantic import AnyHttpUrl, validate_arguments
from shapely.geometry import LineString, Point, Polygon, box

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
def _get_box_from_two_coords(p1: Tuple[Real, Real], p2: Tuple[Real, Real]) -> Polygon:
    """
    Get the polygon that bounds the two coordinates.
    These values should be supplied as numerical values.
    """
    get_bounds = lambda geom: geom.bounds
    box_from_bounds = lambda bounds: box(*bounds)
    return compose(LineString, get_bounds, box_from_bounds)([p1, p2])


# Cell
def box_from_ul_lr_coords(ulx: Real, uly: Real, lrx: Real, lry: Real) -> Polygon:
    """
    Build a box (`Polygon`) from upper left x/y and lower right x/y coordinates.

    This specification is the default BigEarthNet style.
    """
    return _get_box_from_two_coords([ulx, uly], [lrx, lry])


# 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)
    series = pd.read_csv(
        fp,
        names=[name],
        squeeze=True,
    )
    return set(series.values)


@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: List[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
