"""
Module containing miscellaneous raster blocks.
"""
from osgeo import ogr
import numpy as np
from geopandas import GeoSeries

from shapely.geometry import box
from shapely.errors import WKTReadingError
from shapely.wkt import loads as load_wkt

from dask import config
from dask_geomodeling.geometry import GeometryBlock
from dask_geomodeling import utils

from .base import RasterBlock, BaseSingle


__all__ = [
    "Clip",
    "Classify",
    "Reclassify",
    "Mask",
    "MaskBelow",
    "Step",
    "Rasterize",
    "RasterizeWKT",
]


class Clip(BaseSingle):
    """
    Make result values 'no data' if it is 'no data' in the source.
    If source is a boolean mask, False values become 'no data'.

    :param store: The store-like whose values are to be converted.
    :param source: The store-like as the source for 'data or no data'

    :type store: RasterBlock
    :type source: RasterBlock
    """

    def __init__(self, store, source):
        if not isinstance(source, RasterBlock):
            raise TypeError("'{}' object is not allowed".format(type(store)))
        super(Clip, self).__init__(store, source)

    @property
    def source(self):
        return self.args[1]

    @staticmethod
    def process(data, source_data):
        """ Mask store_data where source_data has no data """
        if data is None:
            return None

        if "values" not in data:
            return data

        # check if values contain data
        if np.all(data["values"] == data["no_data_value"]):
            return data

        # make the boolean mask
        if source_data is None:
            return None

        if source_data["values"].dtype == np.dtype("bool"):
            mask = ~source_data["values"]
        else:
            mask = source_data["values"] == source_data["no_data_value"]

        # adjust values
        values = data["values"].copy()
        values[mask] = data["no_data_value"]
        return {"values": values, "no_data_value": data["no_data_value"]}

    @property
    def extent(self):
        """Intersection of bounding boxes of 'store' and 'source'. """
        result, mask = [s.extent for s in self.args]
        if result is None or mask is None:
            return

        # return the overlapping box
        x1 = max(result[0], mask[0])
        y1 = max(result[1], mask[1])
        x2 = min(result[2], mask[2])
        y2 = min(result[3], mask[3])
        if x2 <= x1 or y2 <= y1:
            return None  # no overlap
        else:
            return x1, y1, x2, y2

    @property
    def geometry(self):
        """Intersection of geometries of 'store' and 'source'. """
        result, mask = [x.geometry for x in self.args]
        if result is None or mask is None:
            return
        sr = result.GetSpatialReference()
        if not mask.GetSpatialReference().IsSame(sr):
            mask = mask.Clone()
            mask.TransformTo(sr)
        result = result.Intersection(mask)
        if result.GetArea() == 0.0:
            return
        return result


class Mask(BaseSingle):
    """
    Convert 'data' values to a single constant value.

    :param store: The store whose values are to be converted.
    :param value: The constant value to be given to 'data' values.

    :type store: Store
    :type value: number
    """

    def __init__(self, store, value):
        if not isinstance(value, (float, int)):
            raise TypeError("'{}' object is not allowed".format(type(value)))
        super(Mask, self).__init__(store, value)

    @property
    def value(self):
        return self.args[1]

    @property
    def fillvalue(self):
        return 1 if self.value == 0 else 0

    @property
    def dtype(self):
        return "float32" if isinstance(self.value, float) else "uint8"

    @staticmethod
    def process(data, value):
        if data is None or "values" not in data:
            return data

        index = utils.get_index(
            values=data["values"], no_data_value=data["no_data_value"]
        )

        fillvalue = 1 if value == 0 else 0
        dtype = "float32" if isinstance(value, float) else "uint8"

        values = np.full_like(data["values"], fillvalue, dtype=dtype)
        values[index] = value
        return {"values": values, "no_data_value": fillvalue}


class MaskBelow(BaseSingle):
    """
    Mask data below some value.

    :param store: The store whose values are to be masked.
    :param value: The threshold value. Values below this will be masked.

    :type store: Store
    :type value: number
    """

    def __init__(self, store, value):
        if not isinstance(value, (float, int)):
            raise TypeError("'{}' object is not allowed".format(type(value)))
        super(MaskBelow, self).__init__(store, value)

    @staticmethod
    def process(data, value):
        if data is None or "values" not in data:
            return data
        values, no_data_value = data["values"].copy(), data["no_data_value"]
        values[values < value] = no_data_value
        return {"values": values, "no_data_value": no_data_value}


class Step(BaseSingle):
    """Block for a constant step function with one discontinuity.

    The step function is defined as:
    - left if x < location
    - at if x == location
    - right if x > location

    :param store: The raster whose values are the input to the step function.
    :param left: value left of the discontinuity
    :param right: value right of the discontinuity
    :param location: location of the discontinuity
    :param at: value at the discontinuity

    :type store: RasterBlock
    :type left: number
    :type right: number
    :type location: number
    :type at: number
    """

    def __init__(self, store, left=0, right=1, location=0, at=None):
        """Constructor.



        The at parameter defaults to the mean of the left and right values.
        """
        at = (left + right) / 2 if at is None else at
        for x in left, right, location, at:
            if not isinstance(x, (float, int)):
                raise TypeError("'{}' object is not allowed".format(type(x)))
        super(Step, self).__init__(store, left, right, location, at)

    @property
    def left(self):
        return self.args[1]

    @property
    def right(self):
        return self.args[2]

    @property
    def location(self):
        return self.args[3]

    @property
    def at(self):
        return self.args[4]

    @staticmethod
    def process(data, left, right, location, at):
        if data is None or "values" not in data:
            return data

        values, no_data_value = data["values"].copy(), data["no_data_value"]

        # determine boolean index arrays
        mask = values == no_data_value
        left_index = values < location
        at_index = values == location
        right_index = values > location
        # perform mapping
        values[left_index] = left
        values[at_index] = at
        values[right_index] = right
        # put no data values back
        values[mask] = no_data_value

        return {"values": values, "no_data_value": no_data_value}


class Classify(BaseSingle):
    """
    Classify raster data into a binned categories

    :param store: rasterdata to classify
    :param bins: a 1-dimensional and monotonic list of bin edges
    :param right: whether the intervals include the right or the left bin edge

    :type store: RasterBlock
    :type bins: list
    :type right: boolean

    See also:
      https://docs.scipy.org/doc/numpy/reference/generated/numpy.digitize.html
    """

    def __init__(self, store, bins, right=False):
        if not isinstance(store, RasterBlock):
            raise TypeError("'{}' object is not allowed".format(type(store)))
        if not hasattr(bins, "__iter__"):
            raise TypeError("'{}' object is not allowed".format(type(bins)))
        bins_arr = np.asarray(bins)
        if bins_arr.ndim != 1:
            raise TypeError("'bins' should be one-dimensional")
        if not np.issubdtype(bins_arr.dtype, np.number):
            raise TypeError("'bins' should be numeric")
        bins_diff = np.diff(bins)
        if not np.all(bins_diff > 0) or np.all(bins_diff < 0):
            raise TypeError("'bins' should be monotonic")
        super(Classify, self).__init__(store, bins_arr.tolist(), right)

    @property
    def bins(self):
        return self.args[1]

    @property
    def right(self):
        return self.args[2]

    @property
    def dtype(self):
        # with 254 bin edges, we have 255 bins, and we need 256 possible values
        # to include no_data
        return utils.get_uint_dtype(len(self.bins) + 2)

    @property
    def fillvalue(self):
        return utils.get_dtype_max(self.dtype)

    @staticmethod
    def process(data, bins, right):
        if data is None or "values" not in data:
            return data

        values = data["values"]
        dtype = utils.get_uint_dtype(len(bins) + 2)
        fillvalue = utils.get_dtype_max(dtype)

        result_values = np.digitize(values, bins, right).astype(dtype)
        result_values[values == data["no_data_value"]] = fillvalue

        return {"values": result_values, "no_data_value": fillvalue}


class Reclassify(BaseSingle):
    """
    Reclassify integer data to integers or floats.

    :param store: rasterdata to reclassify
    :param data: list of (from, to) values defining the reclassification
        the from values can be of bool or int datatype; the to values can be
        of int or float datatype
    :param select: leave only reclassified values, set others to 'no data'.
      Default False.

    :type store: RasterBlock
    :type data: list
    :type select: bool
    """

    def __init__(self, store, data, select=False):
        dtype = store.dtype
        if dtype != np.bool and not np.issubdtype(dtype, np.integer):
            raise TypeError("The store must be of boolean or integer datatype")

        # validate "data"
        if not hasattr(data, "__iter__"):
            raise TypeError("'{}' object is not allowed".format(type(data)))
        try:
            source, target = map(np.asarray, zip(*data))
        except ValueError:
            raise ValueError("Please supply a list of [from, to] values")
        # "from" can have bool or int dtype, "to" can also be float
        if source.dtype != np.bool and not np.issubdtype(source.dtype, np.integer):
            raise TypeError(
                "Cannot reclassify from value with type '{}'".format(source.dtype)
            )
        if len(np.unique(source)) != len(source):
            raise ValueError("There are duplicates in the reclassify values")
        if not np.issubdtype(target.dtype, np.number):
            raise TypeError(
                "Cannot reclassify to value with type '{}'".format(target.dtype)
            )
        # put 'data' into a list with consistent dtypes
        data = [list(x) for x in zip(source.tolist(), target.tolist())]

        if select is not True and select is not False:
            raise TypeError("'{}' object is not allowed".format(type(select)))
        super().__init__(store, data, select)

    @property
    def data(self):
        return self.args[1]

    @property
    def select(self):
        return self.args[2]

    @property
    def dtype(self):
        _, target = map(np.asarray, zip(*self.data))
        return target.dtype

    @property
    def fillvalue(self):
        return utils.get_dtype_max(self.dtype)

    def get_sources_and_requests(self, **request):
        process_kwargs = {
            "dtype": self.dtype.str,
            "fillvalue": self.fillvalue,
            "data": self.data,
            "select": self.select,
        }
        return [(self.store, request), (process_kwargs, None)]

    @staticmethod
    def process(store_data, process_kwargs):
        if store_data is None or "values" not in store_data:
            return store_data

        no_data_value = store_data["no_data_value"]
        values = store_data["values"]
        source, target = map(np.asarray, zip(*process_kwargs["data"]))
        dtype = np.dtype(process_kwargs["dtype"])
        fillvalue = process_kwargs["fillvalue"]

        # add the nodata value to the source array and map it to the target
        # nodata
        if no_data_value is not None and no_data_value not in source:
            source = np.append(source, no_data_value)
            target = np.append(target, fillvalue)

        # sort the source and target values
        inds = np.argsort(source)
        source = source[inds]
        target = target[inds]

        # create the result array
        if process_kwargs["select"]:  # select = True: initialize with nodata
            result = np.full(values.shape, fillvalue, dtype=dtype)
        else:  # select = True: initialize with existing data
            result = values.astype(dtype)  # makes a copy

        # find all values in the source data that are to be mapped
        mask = np.in1d(values.ravel(), source)
        mask.shape = values.shape
        # place the target values (this also maps nodata values)
        result[mask] = target[np.searchsorted(source, values[mask])]
        return {"values": result, "no_data_value": fillvalue}


class Rasterize(RasterBlock):
    """Converts geometry source to raster

    :param source: geometry source
    :param column_name: column from the geometry source to rasterize. If
      column_name is not provided, a boolean raster will be returned indicating
      where there are geometries.
    :param dtype: a numpy datatype specification to return the array. Defaults
      to 'int32' if column_name is not, else it defaults to 'bool'.
    :param limit: the maximum number of geometries. Defaults to the
      geomodeling.goemetry-limit setting.
    :returns: a raster containing values from 'column_name' or True/False.

    :type source: GeometryBlock
    :type column_name: string
    :type dtype: string
    :type limit: int

    To rasterize floating point values, it is necessary to pass dtype='float'.

    See also:
      https://docs.scipy.org/doc/numpy/reference/arrays.dtypes.html

    The global geometry-limit setting can be adapted as follows:
      >>> from dask import config
      >>> config.set({"geomodeling.geometry-limit": 100000})
    """

    def __init__(self, source, column_name=None, dtype=None, limit=None):
        if not isinstance(source, GeometryBlock):
            raise TypeError("'{}' object is not allowed".format(type(source)))
        if column_name is not None and not isinstance(column_name, str):
            raise TypeError("'{}' object is not allowed".format(type(column_name)))
        if dtype is None:  # set default values
            dtype = "bool" if column_name is None else "int32"
        else:  # parse to numpy dtype and back to string
            dtype = str(np.dtype(dtype))
        if limit and not isinstance(limit, int):
            raise TypeError("'{}' object is not allowed".format(type(limit)))
        if limit and limit < 1:
            raise ValueError("Limit should be greater than 1")
        super(Rasterize, self).__init__(source, column_name, dtype, limit)

    @property
    def source(self):
        return self.args[0]

    @property
    def column_name(self):
        return self.args[1]

    @property
    def limit(self):
        return self.args[3]

    @property
    def dtype(self):
        return np.dtype(self.args[2])

    @property
    def fillvalue(self):
        return None if self.dtype == np.bool else utils.get_dtype_max(self.dtype)

    @property
    def period(self):
        return (self.DEFAULT_ORIGIN,) * 2

    @property
    def extent(self):
        return None

    @property
    def timedelta(self):
        return None

    @property
    def geometry(self):
        return None

    @property
    def projection(self):
        return None

    @property
    def geo_transform(self):
        return None

    def get_sources_and_requests(self, **request):
        # first handle the 'time' and 'meta' requests
        mode = request["mode"]
        if mode == "time":
            return [(self.period[-1], None), ({"mode": "time"}, None)]
        elif mode == "meta":
            return [(None, None), ({"mode": "meta"}, None)]
        elif mode != "vals":
            raise ValueError("Unknown mode '{}'".format(mode))

        # build the request to be sent to the geometry source
        x1, y1, x2, y2 = request["bbox"]
        width, height = request["width"], request["height"]

        # be strict about the bbox, it may lead to segfaults else
        if x2 == x1 and y2 == y1:  # point
            min_size = None
        elif x1 < x2 and y1 < y2:
            min_size = min((x2 - x1) / width, (y2 - y1) / height)
        else:
            raise ValueError("Invalid bbox ({})".format(request["bbox"]))

        limit = self.limit
        if self.limit is None:
            limit = config.get("geomodeling.geometry-limit")

        geom_request = {
            "mode": "intersects",
            "geometry": box(*request["bbox"]),
            "projection": request["projection"],
            "min_size": min_size,
            "limit": limit,
            "start": request.get("start"),
            "stop": request.get("stop"),
        }
        # keep some variables for use in process()
        process_kwargs = {
            "mode": "vals",
            "column_name": self.column_name,
            "dtype": self.dtype,
            "no_data_value": self.fillvalue,
            "width": width,
            "height": height,
            "bbox": request["bbox"],
        }
        return [(self.source, geom_request), (process_kwargs, None)]

    @staticmethod
    def process(data, process_kwargs):
        # first handle the time and meta requests
        mode = process_kwargs["mode"]
        if mode == "time":
            return {"time": [data]}
        elif mode == "meta":
            return {"meta": [None]}

        column_name = process_kwargs["column_name"]
        height = process_kwargs["height"]
        width = process_kwargs["width"]
        no_data_value = process_kwargs["no_data_value"]
        dtype = process_kwargs["dtype"]
        f = data["features"]

        # get the value column to rasterize
        if column_name is None:
            values = None
        else:
            try:
                values = f[column_name]
            except KeyError:
                if f.index.name == column_name:
                    values = f.index.to_series()
                else:
                    values = False

        if len(f) == 0 or values is False:  # there is no data to rasterize
            values = np.full((1, height, width), no_data_value, dtype=dtype)
            return {"values": values, "no_data_value": no_data_value}

        result = utils.rasterize_geoseries(
            geoseries=f["geometry"] if "geometry" in f else None,
            values=values,
            bbox=process_kwargs["bbox"],
            projection=data["projection"],
            height=height,
            width=width,
        )

        values = result["values"]

        # cast to the expected dtype if necessary
        cast_values = values.astype(process_kwargs["dtype"])

        # replace the nodata value if necessary
        if result["no_data_value"] != no_data_value:
            cast_values[values == result["no_data_value"]] = no_data_value

        return {"values": cast_values, "no_data_value": no_data_value}


class RasterizeWKT(RasterBlock):
    """Converts a single geometry to a raster mask

    Args:
      wkt (string): the WKT representation of a geometry
      projection (string): the projection of the geometry

    Returns:
      RasterBlock with True for cells that are inside the geometry.
    """

    def __init__(self, wkt, projection):
        if not isinstance(wkt, str):
            raise TypeError("'{}' object is not allowed".format(type(wkt)))
        if not isinstance(projection, str):
            raise TypeError("'{}' object is not allowed".format(type(projection)))
        try:
            load_wkt(wkt)
        except WKTReadingError:
            raise ValueError("The provided geometry is not a valid WKT")
        try:
            utils.get_sr(projection)
        except TypeError:
            raise ValueError("The provided projection is not a valid WKT")
        super().__init__(wkt, projection)

    @property
    def wkt(self):
        return self.args[0]

    @property
    def projection(self):
        return self.args[1]

    @property
    def dtype(self):
        return np.dtype("bool")

    @property
    def fillvalue(self):
        return None

    @property
    def period(self):
        return (self.DEFAULT_ORIGIN,) * 2

    @property
    def extent(self):
        return tuple(
            utils.shapely_transform(load_wkt(self.wkt), self.projection, "EPSG:4326").bounds
        )

    @property
    def timedelta(self):
        return None

    @property
    def geometry(self):
        return ogr.CreateGeometryFromWkt(self.wkt, utils.get_sr(self.projection))

    @property
    def geo_transform(self):
        return None

    def get_sources_and_requests(self, **request):
        # first handle the 'time' and 'meta' requests
        mode = request["mode"]
        if mode == "time":
            data = self.period[-1]
        elif mode == "meta":
            data = None
        elif mode == "vals":
            data = {"wkt": self.wkt, "projection": self.projection}
        else:
            raise ValueError("Unknown mode '{}'".format(mode))
        return [(data, None), (request, None)]

    @staticmethod
    def process(data, request):
        mode = request["mode"]
        if mode == "time":
            return {"time": [data]}
        elif mode == "meta":
            return {"meta": [None]}
        # load the geometry and transform it into the requested projection
        geometry = load_wkt(data["wkt"])
        if data["projection"] != request["projection"]:
            geometry = utils.shapely_transform(
                geometry, data["projection"], request["projection"]
            )
        # take a shortcut when the geometry does not intersect the bbox
        if not geometry.intersects(box(*request["bbox"])):
            return {
                "values": np.full(
                    (1, request["height"], request["width"]), False, dtype=np.bool
                ),
                "no_data_value": None,
            }
        return utils.rasterize_geoseries(
            geoseries=GeoSeries([geometry]) if not geometry.is_empty else None,
            bbox=request["bbox"],
            projection=request["projection"],
            height=request["height"],
            width=request["width"],
        )
