# -- coding: utf-8 --
"""
geocube core conversion functionality
"""
import numpy
import pandas
import xarray
from shapely.geometry import mapping

from geocube.geo_utils.geobox import load_vector_data
from geocube.logger import get_logger
from geocube.rasterize import rasterize_image
from rioxarray.rioxarray import (
    DEFAULT_GRID_MAP,
    add_spatial_ref,
    add_xy_grid_meta,
    affine_to_coords,
)


def _format_series_data(data_series):
    """
    The purpose of this function is to convert the series data into a rasterizeable
    format if possible.

    Parameters
    ----------
    data_series: :obj:`geopandas.GeoSeries`
        The series to be converted.

    Returns
    -------
    :obj:`geopandas.GeoSeries`: The series that was converted if possible.

    """
    if "datetime" in str(data_series.dtype):
        data_series = pandas.to_numeric(data_series).astype(numpy.float64)
        get_logger().warning(
            "The series '{}' was converted from a date to a number to "
            "rasterize the data. To load the data back in as a date, "
            "use 'pandas.to_datetime()'.".format(data_series.name)
        )
    elif str(data_series.dtype) == "category":
        data_series = data_series.cat.codes
    return data_series


class VectorToCube(object):
    """
    Tool that facilitates converting vector data to raster data into
    an :obj:`xarray.DataFrame`.
    """

    def __init__(self, vector_data, geobox_maker, fill=None, categorical_enums=None):
        """
        Initialize the GeoCube class.

        Parameters
        ----------
        vector_data: str or :obj:`geopandas.GeoDataFrame`
            A file path to an OGR supported source or GeoDataFrame containing
            the vector data.
        geobox_maker: :obj:`geocube.geo_utils.geobox.GeoBoxMaker`
            The geobox for the grid to be generated from the vector data.
        fill: float, optional
            The value to fill in the grid with for nodata. Default is -9999.0.
        categorical_enums: dict, optional
            A dictionary of all categories for the table columns containing
            categorical data.
            E.g. {'column_name': ['a', 'b'], 'other_column': ['c', 'd']}

        """
        self.vector_data = load_vector_data(vector_data)
        self.geobox = geobox_maker.from_vector(self.vector_data)
        self.fill = fill if fill is not None else -9999.0
        if categorical_enums is not None:
            for column_name, categories in categorical_enums.items():
                category_type = pandas.api.types.CategoricalDtype(
                    categories=sorted(set(categories)) + ["nodata"]
                )
                self.vector_data[column_name] = self.vector_data[column_name].astype(
                    category_type
                )

    def make_geocube(
        self,  # pylint: disable=too-many-arguments
        measurements=None,
        datetime_measurements=None,
        group_by=None,
        interpolate_na_method=None,
    ):  # NOSONAR
        """
        Rasterize vector data into an ``xarray`` object.  Each measurement will be a
        data variable in the :class:`xarray.Dataset`.

        See the `xarray documentation <http://xarray.pydata.org>`_
        for usage of the :class:`xarray.Dataset` and :class:`xarray.DataArray` objects.

        Parameters
        ----------
        measurements: list(str), optional
            Attributes name or list of names to be included. If a list is specified,
            the measurements will be returned in the order requested.
            By default all available measurements are included.
        datetime_measurements: list(str), optional
            Attributes that are temporal in nature and should be converted to the
            datetime format. These are only included if listed in 'measurements'.
        group_by: str, optional
            When specified, perform basic combining/reducing of the data on this column.
        interpolate_na_method:  {‘linear’, ‘nearest’, ‘cubic’}, optional
            This is the method for interpolation to use to fill in the nodata with
            :meth:`scipy.interpolate.griddata`.

        Returns
        --------
        :class:`xarray.Dataset`:
            Requested data in a :class:`xarray.Dataset`.

        """
        if measurements is None:
            measurements = self.vector_data.columns.tolist()
            measurements.remove("geometry")

        self.datetime_measurements = ()
        if datetime_measurements is not None:
            self.datetime_measurements = tuple(
                set(datetime_measurements) & set(measurements)
            )
        # reproject vector data to the projection of the output raster
        vector_data = self.vector_data.to_crs(
            self.geobox.crs._crs.ExportToProj4()
        )  # pylint: disable=protected-access

        # convert to datetime
        for datetime_measurement in self.datetime_measurements:
            vector_data[datetime_measurement] = pandas.to_datetime(
                vector_data[datetime_measurement]
            )

        # get categorical enumerations if they exist
        self._categorical_enums = {}
        for categorical_column in vector_data.select_dtypes(["category"]).columns:
            self._categorical_enums[categorical_column] = vector_data[
                categorical_column
            ].cat.categories

        # map the shape data to the grid
        if group_by:
            vector_data = vector_data.groupby(group_by)
            try:
                measurements.remove(group_by)
            except ValueError:
                pass

        return self._get_dataset(
            vector_data, measurements, group_by, interpolate_na_method
        )

    @staticmethod
    def _get_attrs(measurement_name, fill_value):
        """
        Get attributes for data array.

        Parameters
        ----------
        measurement_name: str
            The measurement name.
        fill_value: int or fload
            The fill value.

        Returns
        -------
        dict: Dict with attributes for data array.
        """
        return dict(
            name=measurement_name,
            long_name=measurement_name,
            grid_mapping=DEFAULT_GRID_MAP,
            _FillValue=fill_value,
        )

    def _update_time_attrs(self, attrs, image_data):
        """
        Update attributes and nodata values for time grid.

        Parameters
        ----------
        attrs: dict
            Dict with attributes for data array.
        image_data: :obj:`numpy.ndarray`
            Array with image data.

        Returns
        -------
        dict: Dict with attributes for data array.
        """
        attrs["units"] = "seconds from 1970-01-01T00:00:00"
        attrs["_FillValue"] = 0
        image_data[image_data == self.fill] = 0.0

    def _get_dataset(self, vector_data, measurements, group_by, interpolate_na_method):
        """
        Parameters
        ----------
        vector_data: str or :obj:`geopandas.GeoDataFrame`
            A file path to an OGR supported source or GeoDataFrame containing
            the vector data.
        measurements: list(str), optional
            Attributes name or list of names to be included. If a list is specified,
            the measurements will be returned in the order requested.
            By default all available measurements are included.
        group_by: str, optional
            When specified, perform basic combining/reducing of the data on this column.

        Returns
        --------
        :obj:`xarray.Dataset`:
            Requested data in a :obj:`xarray.Dataset`.

        """
        data_vars = {}
        for measurement in measurements:
            if group_by:
                grid_array = self._get_grouped_grid(
                    vector_data[[measurement, "geometry"]], measurement, group_by
                )
            else:
                grid_array = self._get_grid(
                    vector_data[[measurement, "geometry"]], measurement
                )
            if grid_array is not None:
                data_vars[measurement] = grid_array

        coords = affine_to_coords(
            self.geobox.affine, self.geobox.width, self.geobox.height
        )

        if group_by:
            coords[group_by] = list(vector_data.groups.keys())

        out_xds = xarray.Dataset(data_vars=data_vars, coords=coords)

        for categorical_measurement, categoral_enums in self._categorical_enums.items():
            enum_var_name = "{}_categories".format(categorical_measurement)
            cat_attrs = dict(out_xds[categorical_measurement].attrs)
            cat_attrs["categorical_mapping"] = enum_var_name
            out_xds[categorical_measurement].attrs = cat_attrs
            out_xds[enum_var_name] = categoral_enums

        add_xy_grid_meta(out_xds.coords)
        add_spatial_ref(out_xds, self.geobox.crs.crs_str, DEFAULT_GRID_MAP)

        if interpolate_na_method is not None:
            return out_xds.rio.interpolate_na(method=interpolate_na_method)

        return out_xds

    def _get_grouped_grid(self, grouped_dataframe, measurement_name, group_by):
        """Retrieve the variable data to append to the ssurgo :obj:`xarray.Dataset`.
        This method is designed specifically to work on a dataframe that has
        been grouped.

        Parameters
        ----------
        grouped_dataframe: pandas GroupBy object
            A pandas dataframe in as a GroupBy object.
        measurement_name: str
            Attributes name or list of names to be included. If a list is specified,
            the measurements will be returned in the order requested.
            By default all available measurements are included.
        group_by: str
            When specified, perform basic combining/reducing of the data on this column.

        Returns
        -------
        tuple: Options needed to create an :obj:`xarray.DataArray`.

        """
        logger = get_logger()

        image_data = []
        df_group = None
        fill_value = self.fill
        for _, df_group in grouped_dataframe:
            fill_value = (
                self.fill if str(df_group[measurement_name].dtype) != "category" else -1
            )
            image = rasterize_image(
                df_group.geometry.apply(mapping).values,
                _format_series_data(df_group[measurement_name]).values,
                self.geobox,
                fill_value,
            )
            if image is None:
                logger.warning(
                    "Skipping attribute {col} due to missing data...".format(
                        col=measurement_name
                    )
                )
                return None

            image_data.append(image)

        attrs = self._get_attrs(measurement_name, fill_value)
        image_data = numpy.array(image_data)
        # it was converted to numeric date value
        if df_group is not None and "datetime" in str(df_group[measurement_name].dtype):
            self._update_time_attrs(attrs, image_data)

        return (group_by, "y", "x"), image_data, attrs

    def _get_grid(self, dataframe, measurement_name):
        """Retrieve the variable data to append to the ssurgo :obj:`xarray.Dataset`
        from a regular :obj:`geopandas.GeoDataFrame`.

        Parameters
        ----------
        dataframe: :obj:`geopandas.GeoDataFrame`
            A geopandas GeoDataFrame object to rasterize.
        measurement_name: str
            Attributes name or list of names to be included. If a list is specified,
            the measurements will be returned in the order requested.
            By default all available measurements are included.

        Returns
        -------
        tuple: Options needed to create an :obj:`xarray.DataArray`.

        """
        logger = get_logger()
        fill_value = (
            self.fill if str(dataframe[measurement_name].dtype) != "category" else -1
        )
        image_data = rasterize_image(
            dataframe.geometry.apply(mapping).values,
            _format_series_data(dataframe[measurement_name]).values,
            self.geobox,
            fill_value,
        )
        if image_data is None:
            logger.warning(
                "Skipping attribute {col} due to missing data...".format(
                    col=measurement_name
                )
            )
            return None

        attrs = self._get_attrs(measurement_name, fill_value)

        # it was converted to numeric date value
        if "datetime" in str(dataframe[measurement_name].dtype):
            self._update_time_attrs(attrs, image_data)

        return ("y", "x"), numpy.array(image_data), attrs
