#!/usr/bin/env python3

"""Calculate travel times between many origins and destinations."""

import collections.abc
import inspect
import math
import multiprocessing
import pathlib
import warnings

import geopandas
import numpy
import pandas
import shapely

from ..util import Config, check_od_data_set
from .regional_task import RegionalTask
from .transport_network import TransportNetwork

__all__ = ["BaseTravelTimeMatrix"]


# R5 fills cut-off (NULL) values with MAX_INT32
MAX_INT32 = (2**31) - 1

# how many (Python) threads to start
# (they still run many Java threads, so be careful what you wish for ;) )
# TODO: benchmark the optimal number of threads
NUM_THREADS = math.ceil(multiprocessing.cpu_count() * 0.5)


class BaseTravelTimeMatrix(geopandas.GeoDataFrame):
    """Base class for travel time computers between many origins and destinations."""

    MAX_INT32 = MAX_INT32

    NUM_THREADS = NUM_THREADS

    _r5py_attributes = [
        "_destinations",
        "_destinations_crs",
        "_origins",
        "_origins_crs",
        "destinations",
        "origins",
        "request",
        "snap_to_network",
        "transport_network",
        "verbose",
    ]

    _constructor = geopandas.GeoDataFrame

    _constructor_sliced = pandas.Series

    @classmethod
    def _geodataframe_constructor_with_fallback(
        cls,
        *args,
        **kwargs,
    ):  # pragma: no cover
        """
        A flexible constructor for r5py frames.

        Checks whether or not arguments of the child class are used.
        """
        # `transport_network` is the only required argument across all r5py
        # data classes, it can be passed either as an r5py.TransportNetwork
        # or a tuple (as described in the signature of __init__(), below)
        if (
            "transport_network" in kwargs
            or isinstance(args[0], TransportNetwork)
            or (
                isinstance(args[0], tuple)
                and isinstance(args[0][0], (pathlib.Path, str))
                and isinstance(args[0][1], collections.abc.Iterable)
                and all(isinstance(arg, (pathlib.Path, str)) for arg in args[0][1])
            )
        ):
            df = cls(*args, **kwargs)

        else:
            df = geopandas.GeoDataFrame(*args, **kwargs)
            geometry_cols_mask = df.dtypes == "geometry"
            if len(geometry_cols_mask) == 0 or geometry_cols_mask.sum() == 0:
                df = pandas.DataFrame(df)

        return df

    def __init__(
        self,
        transport_network,
        origins=None,
        destinations=None,
        snap_to_network=False,
        **kwargs,
    ):
        """
        Compute travel times between many origins and destinations.

        Arguments
        ---------
        transport_network : r5py.TransportNetwork | tuple(
        pathlib.Paths | str, list(pathlib.Path | str))
            The transport network to route on. This can either be a readily
            initialised r5py.TransportNetwork or a tuple of the parameters
            passed to ``TransportNetwork.__init__()``: the path to an
            OpenStreetMap extract in PBF format, and a list of zero of more
            paths to GTFS transport schedule files.
        origins : geopandas.GeoDataFrame
            Places to find a route _from_
            Has to have a point geometry, and at least an `id` column
        destinations : geopandas.GeoDataFrame (optional)
            Places to find a route _to_
            Has to have a point geometry, and at least an `id` column
            If omitted, use same data set as for origins
        snap_to_network : bool or int, default False
            Should origin an destination points be snapped to the street network
            before routing? If `True`, the default search radius (defined in
            `com.conveyal.r5.streets.StreetLayer.LINK_RADIUS_METERS`) is used,
            if `int`, use `snap_to_network` meters as the search radius.
        **kwargs : mixed
            Any arguments than can be passed to r5py.RegionalTask:
            ``departure``, ``departure_time_window``, ``percentiles``,
            ``transport_modes``, ``access_modes``, ``egress_modes``,
            ``max_time``, ``max_time_walking``, ``max_time_cycling``,
            ``max_time_driving``, ``speed_cycling``, ``speed_walking``,
            ``max_public_transport_rides``, ``max_bicycle_traffic_stress``
        """

        super_parameters = inspect.signature(geopandas.GeoDataFrame).parameters
        super_kwargs = {
            key: value for key, value in kwargs.items() if key in super_parameters
        }
        kwargs = {
            key: value for key, value in kwargs.items() if key not in super_parameters
        }
        super().__init__(**super_kwargs)

        if not isinstance(transport_network, TransportNetwork):
            transport_network = TransportNetwork(*transport_network)
        self.transport_network = transport_network

        self.snap_to_network = snap_to_network

        self.origins = origins
        self.destinations = destinations

        self.request = RegionalTask(
            transport_network,
            origin=None,
            destinations=None,
            **kwargs,
        )

        self.verbose = Config().arguments.verbose

    def __repr__(self):
        """Return a string representation of `self`."""
        return f"<{self.__class__.__name__}>"

    def __setattr__(self, attr, val):
        """Catch our own attributes here so we don’t mess with (geo)pandas columns."""
        if attr in self._r5py_attributes:
            object.__setattr__(self, attr, val)
        else:
            super().__setattr__(attr, val)

    @property
    def destinations(self):
        """The destinations of this travel time matrix (`geopandas.GeoDataFrame`)."""
        return self._destinations

    @destinations.setter
    def destinations(self, destinations):
        if destinations is not None:
            check_od_data_set(destinations)
            self._destinations_crs = destinations.crs
            self._destinations = destinations.to_crs("EPSG:4326").copy()

    def _fill_nulls(self, data_set):
        """
        Fill NULL values in a data set returned from R5.

        R5 uses `MAX_INT32` as a marker for NULL values, this function
        replaces those values in `data_set` with `numpy.nan`

        Arguments
        ---------
        data_set : pandas.DataFrame
            Data frame in which NULLs are represented by MAX_INT32

        Returns
        -------
        pandas.DataFrame
            Data frame in which all MAX_INT32 have been replaced by `numpy.nan`.
        """
        return data_set.map(lambda x: numpy.nan if x == MAX_INT32 else x)

    def _prepare_origins_destinations(self):
        """Make sure we received enough information."""
        try:
            self.origins
        except AttributeError as exception:
            raise ValueError("No routing origins defined") from exception

        try:
            self.destinations
            assert self.destinations is not None
        except (AssertionError, AttributeError):
            self.destinations = self.origins.copy()
            if self.verbose:
                warnings.warn(
                    "No routing destinations defined, "
                    "using origins as destinations, too.",
                    RuntimeWarning,
                    stacklevel=1,
                )

        if self.snap_to_network:
            for which_end in ("origins", "destinations"):
                points = getattr(self, f"_{which_end}")
                points.geometry = self.transport_network.snap_to_network(
                    points.geometry
                )
                if len(points[points.geometry == shapely.Point()]):
                    # if there are origins/destinations for which
                    # no snapped point could be found
                    points = points[points.geometry != shapely.Point()]
                    warnings.warn(
                        f"Some {which_end[:-1]} points could not be "
                        "snapped to the street network",
                        RuntimeWarning,
                        stacklevel=1,
                    )

                    if points.empty:
                        raise ValueError(
                            f"After snapping, no valid {which_end[:-1]} points remain"
                        )

                setattr(self, f"_{which_end}", points.copy())

    @property
    def origins(self):
        """The origins of this travel time matrix (`geopandas.GeoDataFrame`)."""
        return self._origins

    @origins.setter
    def origins(self, origins):
        if origins is not None:
            check_od_data_set(origins)
            self._origins_crs = origins.crs
            self._origins = origins.to_crs("EPSG:4326").copy()
