# -*- coding: utf-8 -*-
"""
Sentinel-2 surface reflectance with COPERNICUS/S2_CLOUD_PROBABILITY masking.
"""

import datetime
import re
import unicodedata

from .ee_runtime import ee

S2_SR_COLLECTION = 'COPERNICUS/S2_SR_HARMONIZED'
S2_CLOUD_PROBABILITY = 'COPERNICUS/S2_CLOUD_PROBABILITY'
DEFAULT_CLOUD_THRESHOLD = 20
RGB_BANDS = ['B4', 'B3', 'B2']
# Per scene cloud cover, computed by ESA and shipped as image metadata.
CLOUD_PROPERTY = 'CLOUDY_PIXEL_PERCENTAGE'
_EPOCH = datetime.datetime(1970, 1, 1)

# Palettes must be lists of bare hex colours: a comma separated string with
# spaces makes Earth Engine reject the whole visualisation.
NDVI_PALETTE = ['d7191c', 'fdae61', 'ffffbf', 'a6d96a', '1a9641']
NBR_PALETTE = ['1a9850', 'ffffbf', 'f46d43', 'd73027']

# Reflectance stretch used for true colour previews and layers.
RGB_MAX_REFLECTANCE = 3000

# A preview request costs one thumbnail per day, so the list is capped.
PREVIEW_MAX_DAYS = 40

# Tolerance for the preview bounding box, in metres.
PREVIEW_BOUNDS_ERROR = 1000


def as_geometry(area):
    """Return an ee.Geometry for either a Geometry or a FeatureCollection."""
    if isinstance(area, ee.Geometry):
        return area
    return area.geometry()


def sanitize_export_name(name, fallback='S2_export'):
    """
    Earth Engine task names only accept [A-Za-z0-9.,:;_-].

    Vietnamese place names therefore have to be transliterated, otherwise the
    export is rejected before the task is even queued.
    """
    text = (name or '').replace(u'Đ', 'D').replace(u'đ', 'd')
    text = unicodedata.normalize('NFKD', text)
    ascii_text = text.encode('ascii', 'ignore').decode('ascii')
    ascii_text = re.sub(r'[^A-Za-z0-9.,:;_-]+', '_', ascii_text).strip('_')
    return (ascii_text or fallback)[:100]


def _end_of_day(end_date):
    """The date picked as "to" is inclusive, Earth Engine filters exclusively."""
    return ee.Date(end_date).advance(1, 'day')


def preview_region(area):
    """
    Bounding box of the area of interest, as plain coordinates.

    Previews only have to show how cloudy a day looks, so they do not need the
    exact outline. Resolving the box once, client side, is what keeps them fast:
    a computed geometry has to be rebuilt by Earth Engine for every request it
    appears in, and a province outline carries thousands of vertices.
    """
    bounds = as_geometry(area).bounds(PREVIEW_BOUNDS_ERROR)
    return ee.Geometry(bounds.getInfo())


def filtered_collection(collection_id, start_date, end_date, geometry):
    """Collection reduced to the date range and area of interest."""
    return (ee.ImageCollection(collection_id)
            .filterDate(start_date, end_date)
            .filterBounds(geometry))


def get_s2_joined_collection(start_date, end_date, geometry):
    """
    Sentinel-2 SR images with the matching cloud probability band attached.

    The inner join is on system:index, which is shared by both collections.
    Both sides are filtered down to the request before they are joined:
    joining first would pair up the whole Sentinel-2 archive and dominate the
    cost of every search.
    """
    primary = filtered_collection(
        S2_SR_COLLECTION, start_date, end_date, geometry)
    secondary = filtered_collection(
        S2_CLOUD_PROBABILITY, start_date, end_date, geometry)
    join_filter = ee.Filter.equals(
        leftField='system:index',
        rightField='system:index',
    )
    joined = ee.Join.inner().apply(
        primary=primary,
        secondary=secondary,
        condition=join_filter,
    )

    def merge_cloud_band(image):
        return ee.Image(image.get('primary')).addBands(image.get('secondary'))

    return ee.ImageCollection(joined.map(merge_cloud_band))


def get_s2_cloud_masked_collection(
    start_date, end_date, geometry, cloud_threshold=DEFAULT_CLOUD_THRESHOLD,
):
    """Build a cloud-masked Sentinel-2 SR collection for a date range and geometry."""
    collection = get_s2_joined_collection(start_date, _end_of_day(end_date), geometry)

    def mask_clouds(image):
        cloud_mask = image.select('probability').lt(cloud_threshold)
        return image.updateMask(cloud_mask)

    return collection.map(mask_clouds)


def compute_s2_composite(collection):
    """Return mean composite, NDVI and NBR from a cloud-masked collection."""
    composite = collection.mean()
    # Bands are renamed so the exported GeoTIFF is self describing, and cast to
    # float32 to halve the download size of the index products.
    ndvi = composite.normalizedDifference(['B8', 'B4']).rename('NDVI').toFloat()
    nbr = composite.normalizedDifference(['B8', 'B11']).rename('NBR').toFloat()
    return composite, ndvi, nbr


def build_date_string(start_date, end_date):
    return '{}_to_{}'.format(
        start_date.replace('-', '_'),
        end_date.replace('-', '_'),
    )


def _day_of(millis):
    """Acquisition day in UTC, the way Earth Engine filters dates."""
    return (_EPOCH + datetime.timedelta(milliseconds=millis)).strftime('%Y-%m-%d')


def _group_by_day(times, clouds):
    """Fold the per scene metadata into one entry per acquisition day."""
    days = {}
    for index, millis in enumerate(times):
        day = _day_of(millis)
        entry = days.setdefault(day, {'day': day, 'scenes': 0, 'clouds': []})
        entry['scenes'] += 1
        cloud = clouds[index] if index < len(clouds) else None
        if cloud is not None:
            entry['clouds'].append(cloud)

    scenes = []
    for day in sorted(days):
        entry = days[day]
        day_clouds = entry.pop('clouds')
        entry['cloud'] = \
            sum(day_clouds) / float(len(day_clouds)) if day_clouds else None
        scenes.append(entry)
    return scenes


def list_daily_scenes(start_date, end_date, geometry, max_days=PREVIEW_MAX_DAYS,
                      region=None):
    """
    Acquisition days in the range, with the cloud cover of each day.

    Only metadata is read: the cloud figure is the mean of the per scene
    CLOUDY_PIXEL_PERCENTAGE, so the whole list is one cheap round trip instead
    of a pixel reduction per day. The thumbnails are what show the user the real
    cloud situation over the area.

    Returns a list of dicts with 'day', 'cloud' (0-100 or None) and 'scenes',
    sorted by date, plus a flag telling whether the list was truncated.
    """
    if region is None:
        region = preview_region(geometry)
    collection = filtered_collection(
        S2_SR_COLLECTION, start_date, _end_of_day(end_date), region)

    metadata = ee.Dictionary({
        'time': collection.aggregate_array('system:time_start'),
        'cloud': collection.aggregate_array(CLOUD_PROPERTY),
    }).getInfo()

    scenes = _group_by_day(
        metadata.get('time') or [], metadata.get('cloud') or [],
    )
    truncated = len(scenes) > max_days
    return scenes[:max_days], truncated


def daily_preview_url(day, geometry, size=320, max_reflectance=RGB_MAX_REFLECTANCE,
                      region=None):
    """
    True colour thumbnail URL for a single day.

    Only the RGB bands are needed, so the cloud probability join is skipped:
    clouds are left in on purpose, since the preview exists so the user can see
    how cloudy each date is before committing to a date range.
    """
    if region is None:
        region = preview_region(geometry)
    day_start = ee.Date.parse('YYYY-MM-dd', day)
    collection = filtered_collection(
        S2_SR_COLLECTION, day_start, day_start.advance(1, 'day'), region)
    visual = collection.mosaic().select(RGB_BANDS).visualize(
        bands=RGB_BANDS, min=0, max=max_reflectance,
    )
    return visual.getThumbURL({
        'dimensions': size,
        'region': region,
        'format': 'png',
    })
