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

# # Showcase rivus.gridder

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')
import os
import matplotlib.pyplot as plt
from rivus.gridder.create_grid import create_square_grid
from rivus.gridder.extend_grid import extend_edge_data
from rivus.gridder.extend_grid import vert_init_commodities


# ## Constants - Inputs

#  Origo for grid creation.   
#  You can copy LatLon into this list from osm, google.maps...    
#  You do not need to input it, if not provided, it default to a point around Munich.

# In[2]:


lat, lon = [48.13512, 11.58198]


# ### File Access

# In[3]:


base_directory = os.path.join('data', 'chessboard')
data_spreadsheet = os.path.join(base_directory, 'data.xlsx')


# ## Create Spatial Rivus Inputs

# Get the raw gemetries

# In[13]:


vertex, edge = create_square_grid(origo_latlon=(lat, lon), num_edge_x=4, dx=300, dy=150, noise_prop=0.1)


# In[14]:


vertex.head(4)


# In[15]:


edge.head(4)


# ### Extend Vertex with Source Information

# In[16]:


sources = [('Elec', 0, 100000), ('Gas', 0, 50000)]  # Commodity, VertexID, MaxCapacity
vert_init_commodities(vertex, ('Elec', 'Gas'), sources)
vertex.head(4)


# ### Extend Edge with Demand Information

# In[17]:


building_types = ['residential', 'industrial']
inits = [1000, 2000]
extend_edge_data(edge, sorts=building_types, inits=inits)
edge.head(4)


# # Plot the results with annotated vertex/edge IDs

# In[35]:


# Helper
def _calc_xytext_offset(line):
    dx, dy = [abs(cc[0] - cc[1]) for cc in list(zip(*line.coords[:]))]
    if dx == 0 or (dy / dx) > 1:
        shx, shy = 30,  0
    else:
        shx, shy =  0, -30
    
    return (shx, shy)


# In[49]:


annotate_vertex_ids = True
annotate_edge_ids = True


# In[50]:


fig, ax = plt.subplots(figsize=(12,12))
ax.set_aspect('equal')
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.axis('off')

vertex.plot(ax=ax, marker='o', color='red', markersize=5)
if annotate_vertex_ids:
    for idx, pnt in enumerate(vertex.geometry):
        plt.annotate(
            "{}".format(idx),
            xy=(pnt.coords[0][0], pnt.coords[0][1]), xytext=(-20, 20),
            textcoords='offset points', ha='center', va='center',
            bbox=dict(boxstyle='round,pad=0.5', fc='red', alpha=0.3),
            arrowprops=dict(arrowstyle = '->', connectionstyle='arc3,rad=0'))   
edge.plot(ax=ax, color='blue')
if annotate_edge_ids:
    for idx, lin in enumerate(edge.geometry):
        textxy = lin.centroid
        plt.annotate(
            "{}".format(idx),
            xy=(textxy.x, textxy.y), xytext=_calc_xytext_offset(lin),
            textcoords='offset points', ha='right', va='bottom',
            bbox=dict(boxstyle='round,pad=0.5', fc='blue', alpha=0.3),
            arrowprops=dict(arrowstyle = '->', connectionstyle='arc3,rad=0'))
plt.tight_layout()
plt.show()