# coding=utf-8
from __future__ import absolute_import, division, print_function

import logging
import math
import os
import warnings

from haversine import haversine
import numpy as np
import rasterio
from rasterio import transform, warp, windows
from rasterio._err import CPLE_OutOfMemoryError
from rasterio.crs import CRS
from rasterio.transform import Affine
from rasterio.vrt import WarpedVRT
from rasterio.warp import Resampling
from rasterio.windows import Window
from scipy import ndimage

from . import mosaic
from .stats import Timer

WEB_MERCATOR_CRS = CRS.from_epsg(3857)
WGS84_CRS = CRS.from_epsg(4326)
LOG = logging.getLogger(__name__)

EXTENTS = {
    str(WEB_MERCATOR_CRS): (-20037508.342789244, -20037508.342789244,
                            20037508.342789244, 20037508.342789244),
}


class NoDataAvailable(Exception):
    pass


def _isimage(data_format):
    return data_format.upper() in ["RGB", "RGBA"]


def _mask(data, nodata):
    if np.issubdtype(data.dtype, float):
        return np.ma.masked_values(data, nodata, copy=False)

    return np.ma.masked_equal(data, nodata, copy=False)


def _nodata(dtype):
    if np.issubdtype(dtype, int):
        return np.iinfo(dtype).min
    else:
        return np.finfo(dtype).min


def crop((data, (bounds, data_crs)), data_format, offsets):
    left, bottom, right, top = offsets

    if _isimage(data_format):
        width, height, _ = data.shape
        t = transform.from_bounds(*bounds, width=width, height=height)

        data = data[top:height - bottom, left:width - right, :]

        cropped_window = windows.Window(left, top, width, height)
        cropped_bounds = windows.bounds(cropped_window, t)

        return (data, (cropped_bounds, data_crs))

    _, height, width = data.shape
    t = transform.from_bounds(*bounds, width=width, height=height)

    data = data[:, top:height - bottom, left:width - right]

    cropped_window = windows.Window(left, top, width, height)
    cropped_bounds = windows.bounds(cropped_window, t)

    return (data, (cropped_bounds, data_crs))


def get_extent(crs):
    return EXTENTS[str(crs)]


def get_resolution((bounds, crs), (height, width)):
    t = transform.from_bounds(*bounds, width=width, height=height)

    return abs(t.a), abs(t.e)


def get_resolution_in_meters((bounds, crs), (height, width)):
    if crs.is_geographic:
        left = (bounds[0], (bounds[1] + bounds[3]) / 2)
        right = (bounds[2], (bounds[1] + bounds[3]) / 2)
        top = ((bounds[0] + bounds[2]) / 2, bounds[3])
        bottom = ((bounds[0] + bounds[2]) / 2, bounds[1])

        return (haversine(left, right) * 1000 / width,
                haversine(top, bottom) * 1000 / height)

    return get_resolution((bounds, crs), (height, width))


def get_source(path):
    """Cached source opening."""
    with rasterio.Env(CPL_VSIL_CURL_ALLOWED_EXTENSIONS=".vrt,.tif,.ovr,.msk"):
        return rasterio.open(path)


def get_zoom(resolution, op=round):
    return int(
        op(
            math.log((2 * math.pi * 6378137) / (resolution * 256)) / math.log(
                2)))


def read_window(src, (bounds, bounds_crs), (height, width)):
    target_shape = (height, width)
    # TODO use this for DEMs (not all single-band sources) to avoid stairstepping artifacts
    if src.count == 1 and bounds_crs == WEB_MERCATOR_CRS:
        # special case for web Mercator; use a target image size that most
        # closely matches the source resolution (and is a power of 2)
        zoom = min(22,
                   get_zoom(
                       max(
                           get_resolution_in_meters((src.bounds, src.crs), (
                               src.height, src.width))),
                       op=math.ceil))

        dst_width = dst_height = (2**zoom) * 256
        extent = get_extent(bounds_crs)
        resolution = ((extent[2] - extent[0]) / dst_width,
                      (extent[3] - extent[1]) / dst_height)

        dst_transform = Affine(resolution[0], 0.0, extent[0], 0.0,
                               -resolution[1], extent[3])
    else:
        # use a target image size that most closely matches the target
        # resolution
        # calculate natural width, height, and transform (so the mask can be
        # warped)
        # TODO providing resolution reduces the target dimensions but loses the
        # ability to read overviews

        # if raster is overly-large, approximate the transform based on
        # a scaled-down version and scale it back after
        attempts = 0
        scale_factor = 1
        dst_transform = None

        while (dst_transform is None and src.width // scale_factor > 0
               and src.height // scale_factor > 0):
            try:
                (dst_transform, dst_width,
                 dst_height) = warp.calculate_default_transform(
                     src.crs, bounds_crs, src.width // scale_factor,
                     src.height // scale_factor, *src.bounds)

                scale = Affine.scale(scale_factor, scale_factor)

                dst_transform *= ~scale
                dst_width, dst_height = scale * (dst_width, dst_height)
            except (MemoryError, CPLE_OutOfMemoryError):
                attempts += 1
                scale_factor = 2 * attempts

    # Some OAM sources have invalid NODATA values (-1000 for a file with a
    # dtype of Byte). rasterio returns None under these circumstances
    # (indistinguishable from sources that actually have no NODATA values).
    # Providing a synthetic value "correctly" masks the output at the expense
    # of masking valid pixels with that value. This was previously (partially;
    # in the form of the bounding box but not NODATA pixels) addressed by
    # creating a VRT that mapped the mask to an alpha channel (something we
    # can't do w/o adding nDstAlphaBand to rasterio/_warp.pyx).
    #
    # Creating external masks and reading them separately (as below) is a
    # better solution, particularly as it avoids artifacts introduced when the
    # NODATA values are resampled using something other than nearest neighbor.

    # TODO resampling alg should be a catalog property
    with WarpedVRT(
            src,
            src_nodata=src.nodata,
            dst_crs=bounds_crs,
            dst_width=dst_width,
            dst_height=dst_height,
            dst_transform=dst_transform,
            resampling=Resampling.lanczos) as vrt:
        window = vrt.window(*bounds)
        dst_window = window.crop(vrt.height, vrt.width)

        if round(dst_window.width) == 0 or round(dst_window.height) == 0:
            return None

        resolution = get_resolution((bounds, bounds_crs), target_shape)
        src_resolution = get_resolution((vrt.bounds, vrt.crs), vrt.shape)
        scale_factor = (round(window.width / target_shape[1], 6), round(
            window.height / target_shape[0], 6))

        extends_up = False
        extends_left = False

        if bounds[0] < vrt.bounds[0]:
            width = int(
                math.ceil(target_shape[1] * (dst_window.width / window.width)))
            extends_left = True

        if bounds[1] < vrt.bounds[1]:
            height = int(
                math.ceil(target_shape[0] * (dst_window.height / window.height
                                             )))

        if bounds[2] > vrt.bounds[2]:
            width = int(
                math.ceil(target_shape[1] * (dst_window.width / window.width)))

        if bounds[3] > vrt.bounds[3]:
            height = int(
                math.ceil(target_shape[0] * (dst_window.height / window.height
                                             )))
            extends_up = True

        if vrt.count == 1 and (
                scale_factor[0] < 1 or scale_factor[1] < 1
                or src_resolution[0] > resolution[0]
                or src_resolution[1] > resolution[1]
        ) and round(dst_window.width) > 1.0 and round(dst_window.height) > 1.0:
            # scale_factor will always be (1.0, 1.0) unless using the web
            # Mercator-specific calculations
            # instead, compare source resolution (in m) to resolution (as
            # scale_factor) and modify dst_window to correspond to a smaller,
            # lower resolution window (as target_window)

            # TODO sources like ETOPO1 may end up with funky y scale factors
            # (for web Mercator, due to distortion)
            # if these problems can be addressed, the web Mercator-specific
            # calculations (above, for single-band sources) can be removed
            if scale_factor[0] < 1 or scale_factor[1] < 1:
                scaled_transform = vrt.transform * Affine.scale(*scale_factor)
                target_window = windows.from_bounds(
                    *bounds, transform=scaled_transform)
            else:
                scale_factor = (resolution[0] / src_resolution[0],
                                resolution[1] / src_resolution[1])
                target_window = Window(dst_window.col_off * scale_factor[0],
                                       dst_window.row_off * scale_factor[1],
                                       dst_window.width * scale_factor[0],
                                       dst_window.height * scale_factor[1])

            # buffer apparently needs to be 50% of the target size in order
            # for spline knots to match between adjacent tiles
            # however, to avoid creating overly-large uncropped areas, we
            # limit the buffer size to 2048px on a side
            # TODO the resulting window is still much too large (e.g.
            # 18/151153/84343@2x)
            buffer_pixels = (min(target_window.width / 2,
                                 math.ceil(2048 * scale_factor[0])), min(
                                     target_window.height / 2,
                                     math.ceil(2048 * scale_factor[1])))

            r, c = dst_window.toranges()
            window = Window.from_slices((max(
                0, r[0] - buffer_pixels[1]), r[1] + buffer_pixels[1]), (max(
                    0, c[0] - buffer_pixels[0]), c[1] + buffer_pixels[0]))

            data = vrt.read(window=window)

            # mask with NODATA values
            if vrt.nodata is not None:
                data = _mask(data, vrt.nodata)
            else:
                data = np.ma.masked_array(data, mask=False)

            mask = data.mask

            order = 3

            if mask.any():
                # need to preserve NODATA; drop spline interpolation order to 1
                order = 1

            LOG.info(
                "Applying spline interpolation with order %d (scale factor: %s)",
                order, scale_factor)

            zoom = (1, round(1 / scale_factor[1]), round(1 / scale_factor[0]))

            LOG.info("target dimensions: %s", (data.shape[1] * zoom[1],
                                               data.shape[2] * zoom[2]))

            # resample data, respecting NODATA values
            data = ndimage.zoom(
                # prevent resulting values from producing cliffs
                data.astype(np.float32),
                zoom,
                order=order)

            scaled_buffer = (int((data.shape[2] - width) / 2), int(
                (data.shape[1] - height) / 2))

            # crop data
            data = data[:, scaled_buffer[1]:scaled_buffer[1] + height,
                        scaled_buffer[0]:scaled_buffer[0] + width]

            if len(mask.shape) > 0:
                mask = ndimage.zoom(mask, zoom, mode='nearest')

                # crop mask
                mask = mask[:, scaled_buffer[1]:scaled_buffer[1] + height,
                            scaled_buffer[0]:scaled_buffer[0] + width]

            # copy the mask over
            data = np.ma.masked_array(data, mask=mask)
        else:
            data = vrt.read(
                out_shape=(vrt.count, height, width), window=dst_window)

            # mask with NODATA values
            if vrt.nodata is not None:
                data = _mask(data, vrt.nodata)
            else:
                data = np.ma.masked_array(data, mask=False)

        # expand to match the target shape
        # up / down offset
        if extends_up:
            data = np.ma.concatenate((np.ma.masked_array(
                np.zeros(
                    (vrt.count, target_shape[0] - data.shape[1],
                     target_shape[1] - (target_shape[1] - data.shape[2])),
                    dtype=data.dtype),
                mask=True), data), 1)
        else:
            data = np.ma.concatenate((data, np.ma.masked_array(
                np.zeros(
                    (vrt.count, target_shape[0] - data.shape[1],
                     target_shape[1] - (target_shape[1] - data.shape[2])),
                    dtype=data.dtype),
                mask=True)), 1)

        # left / right offset
        if extends_left:
            data = np.ma.concatenate((np.ma.masked_array(
                np.zeros(
                    (vrt.count,
                     target_shape[0] - (target_shape[0] - data.shape[1]),
                     target_shape[1] - data.shape[2]),
                    dtype=data.dtype),
                mask=True), data), 2)
        else:
            data = np.ma.concatenate((data, np.ma.masked_array(
                np.zeros(
                    (vrt.count,
                     target_shape[0] - (target_shape[0] - data.shape[1]),
                     target_shape[1] - data.shape[2]),
                    dtype=data.dtype),
                mask=True)), 2)

        data = data.astype(np.float32)

    # open the mask separately so we can take advantage of its overviews
    try:
        warnings.simplefilter("ignore")
        with rasterio.open("{}.msk".format(src.name), crs=src.crs) as mask_src:
            with WarpedVRT(
                    mask_src,
                    src_crs=src.crs,
                    src_transform=src.transform,
                    dst_crs=bounds_crs,
                    dst_width=dst_width,
                    dst_height=dst_height,
                    dst_transform=dst_transform) as mask_vrt:
                warnings.simplefilter("default")
                window = mask_vrt.window(*bounds)
                dst_window = window.crop(mask_vrt.height, mask_vrt.width)

                if round(dst_window.width) == 0 or round(
                        dst_window.height) == 0:
                    raise Exception("No mask")

                extends_up = False
                extends_left = False

                if bounds[0] < vrt.bounds[0]:
                    width = int(
                        math.ceil(target_shape[1] * (
                            dst_window.width / window.width)))
                    extends_left = True

                if bounds[1] < vrt.bounds[1]:
                    height = int(
                        math.ceil(target_shape[0] * (
                            dst_window.height / window.height)))

                if bounds[2] > vrt.bounds[2]:
                    width = int(
                        math.ceil(target_shape[1] * (
                            dst_window.width / window.width)))

                if bounds[3] > vrt.bounds[3]:
                    height = int(
                        math.ceil(target_shape[0] * (
                            dst_window.height / window.height)))
                    extends_up = True

                mask = mask_vrt.read(
                    out_shape=(mask_vrt.count, height, width),
                    window=dst_window)

                # expand to match the target shape
                # up / down offset
                if extends_up:
                    mask = np.ma.concatenate((np.ma.masked_array(
                        np.zeros(
                            (mask_vrt.count, target_shape[0] - mask.shape[1],
                             target_shape[1] -
                             (target_shape[1] - mask.shape[2])),
                            dtype=mask.dtype),
                        mask=True), mask), 1)
                else:
                    mask = np.ma.concatenate((mask, np.ma.masked_array(
                        np.zeros(
                            (mask_vrt.count, target_shape[0] - mask.shape[1],
                             target_shape[1] -
                             (target_shape[1] - mask.shape[2])),
                            dtype=mask.dtype),
                        mask=True)), 1)

                # left / right offset
                if extends_left:
                    mask = np.ma.concatenate((np.ma.masked_array(
                        np.zeros(
                            (mask_vrt.count, target_shape[0] -
                             (target_shape[0] - mask.shape[1]),
                             target_shape[1] - mask.shape[2]),
                            dtype=mask.dtype),
                        mask=True), mask), 2)
                else:
                    mask = np.ma.concatenate((mask, np.ma.masked_array(
                        np.zeros(
                            (mask_vrt.count, target_shape[0] -
                             (target_shape[0] - mask.shape[1]),
                             target_shape[1] - mask.shape[2]),
                            dtype=mask.dtype),
                        mask=True)), 2)

                data.mask = data.mask | ~mask
    except Exception:
        # no mask
        pass

    return (data, (bounds, bounds_crs))


def render((bounds, bounds_crs),
           sources_store,
           shape,
           target_crs,
           format,
           transformation=None):
    """Render data intersecting bounds into shape using an optional
    transformation."""
    resolution_m = get_resolution_in_meters((bounds, bounds_crs), shape)
    stats = []

    if transformation:
        (bounds, bounds_crs), shape, offsets = transformation.expand(
            (bounds, bounds_crs), shape)

    with Timer() as t:
        sources = sources_store.get_sources((bounds, bounds_crs), resolution_m)
    stats.append(("get sources", t.elapsed))

    with Timer() as t:
        (sources_used, data, (data_bounds, data_crs)) = mosaic.composite(
            sources, (bounds, bounds_crs), shape, target_crs)
    stats.append(("composite", t.elapsed))

    if data is None:
        raise NoDataAvailable()

    data_format = "raw"

    if transformation:
        with Timer() as t:
            (data, data_format) = transformation.transform((data, (data_bounds,
                                                                   data_crs)))
        stats.append(("transform", t.elapsed))

        with Timer() as t:
            (data, (data_bounds, data_crs)) = transformation.postprocess(
                (data, (data_bounds, data_crs)), data_format, offsets)

        stats.append(("postprocess", t.elapsed))

    with Timer() as t:
        (content_type, formatted) = format((data, (data_bounds, data_crs)),
                                           data_format)
    stats.append(("format", t.elapsed))

    headers = {
        "Content-Type":
        content_type,
        "X-Imagery-Sources":
        ", ".join(s[1].split('/', 3)[3] for s in sources_used),
    }

    if os.environ.get('MARBLECUTTER_DEBUG_TIMERS'):
        headers.update({
            "X-Timers":
            ", ".join("{}: {:0.2f}".format(*s) for s in stats)
        })

    return (headers, formatted)
