#!/usr/bin/env python
# coding: utf-8

# # Getting a flat array of coordinates from array of Geometry objects

# In[1]:


import geopandas
import pygeos


# In[2]:


# Quick copy of the relevant part of https://github.com/pyviz/datashader/pull/819

import numpy as np

def from_geopandas(ga):
    line_parts = [
        _shapely_to_array_parts(shape) for shape in ga
    ]
    line_lengths = [
        sum([len(part) for part in parts])
        for parts in line_parts
    ]
    flat_array = np.concatenate(
        [part for parts in line_parts for part in parts]
    )
    start_indices = np.concatenate(
        [[0], line_lengths[:-1]]
    ).astype('uint').cumsum()
    return start_indices, flat_array
    
    
def _polygon_to_array_parts(polygon):
    import shapely.geometry as sg
    shape = sg.polygon.orient(polygon)
    exterior = np.asarray(shape.exterior.ctypes)
    polygon_parts = [exterior]
    hole_separator = np.array([-np.inf, -np.inf])
    for ring in shape.interiors:
        interior = np.asarray(ring.ctypes)
        polygon_parts.append(hole_separator)
        polygon_parts.append(interior)
    return polygon_parts

def _shapely_to_array_parts(shape):
    import shapely.geometry as sg
    if isinstance(shape, sg.Polygon):
        # Single polygon
        return _polygon_to_array_parts(shape)
    elif isinstance(shape, sg.MultiPolygon):
        shape = list(shape)
        polygon_parts = _polygon_to_array_parts(shape[0])
        polygon_separator = np.array([np.inf, np.inf])
        for polygon in shape[1:]:
            polygon_parts.append(polygon_separator)
            polygon_parts.extend(_polygon_to_array_parts(polygon))
        return polygon_parts
    else:
        raise ValueError("""
Received invalid value of type {typ}. Must be an instance of
shapely.geometry.Polygon or shapely.geometry.MultiPolygon"""
                         .format(typ=type(shape).__name__))


# Testing on the "countries" data:

# In[10]:


df = geopandas.read_file(geopandas.datasets.get_path('naturalearth_lowres'))


# Datashader's method:

# In[11]:


from_geopandas(df.geometry.array)


# In[12]:


get_ipython().run_line_magic('timeit', 'from_geopandas(df.geometry.array)')


# Using pygeos:

# In[13]:


arr = pygeos.from_wkb(geopandas.array.to_wkb(df.geometry.array))


# In[14]:


pygeos.get_coordinates(arr)


# In[15]:


get_ipython().run_line_magic('timeit', 'pygeos.get_coordinates(arr)')


# Note: this currently gives a 2D array (not x and y intertwined) and it does not give the start indices, but I think having a version that does those things should not be a order of magnitude slower.

# In[ ]: