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

# # Drawing Graph Edges Using `geom_segment()` and `geom_curve()`
# 
# Beyond merely connecting two points on a chart, the `geom_segment()` and `geom_curve()` geometries</br>
# can, with some fine-tuning, help to visualize graph-like data.
# 
# Use aesthetics `size_start/end` and `stroke_start/end` to allow `segment/curve`</br>
# to take into account the size of the point from which it starts/ends and to avoid drawing over it.
# 
# Utilize the `spacer` parameter for further manual fine-tuning.

# In[1]:


import pandas as pd
from lets_plot import *


# In[2]:


LetsPlot.setup_html()


# In[3]:


x = [-1, 0, 1]
y = [-1, 1, -1]
shape = [1, 16, 21]
size = [1, 2, 3]
stroke = [1, 0, 2]

vertices_data = {
    'x': x,
    'y': y,
    'shape': shape,
    'size': size,
    'stroke': stroke,
}


# In[4]:


vertices_layer = geom_point(aes('x', 'y', size='size', shape='shape', stroke='stroke'),
                            color='#4575b4', fill='#abd9e9')


# In[5]:


graph_vertices = (
    ggplot(vertices_data) + vertices_layer
    + scale_size(range=[20,30], guide='none')
    + scale_stroke(range=[0,10], guide='none')
    + scale_shape_identity() 
    + lims(x=[-1.5, 1.5], y=[-1.5, 1.5])
)
graph_vertices


# In[6]:


edges_data = {
    'x_end': x[1:] + [x[0]],
    'y_end': y[1:] + [y[0]]
}


# #### 1. Draw Ugly Graph Edges

# In[7]:


ugly_edges = geom_segment(aes('x', 'y', xend='x_end', yend='y_end'),
                  data=edges_data,       
                  arrow=arrow(ends='both')) 

graph_vertices + ugly_edges


# #### 2. Draw Nice Graph Edges

# In[8]:


# Uppend an info on the sizes of vertices in the graph.

edges_data_2 = dict(edges_data)
edges_data_2.update({
    'size_end': size[1:] + [size[0]],
    'stroke_end': stroke[1:] + [stroke[0]]
})


# In[9]:


# Use `segment` and then `curve` to draw "nice" edges.

nice_edges_S = geom_segment(
    aes('x', 'y', xend='x_end', yend='y_end', 
        size_start='size', size_end='size_end',           # New! Take into account sizes of points connected by the edge.
        stroke_start='stroke', stroke_end='stroke_end'),  # New! Take into account stroke (width) of points connected by the edge.
        spacer=5,                                         # New! Add a "spacer".
        data=edges_data_2,       
        arrow=arrow(ends='both')) 

nice_edges_C = geom_curve(
    aes('x', 'y', xend='x_end', yend='y_end', 
        size_start='size', size_end='size_end',
        stroke_start='stroke', stroke_end='stroke_end'),
        spacer=5,         
        data=edges_data_2,       
        curvature=-0.3,
        arrow=arrow(ends='both')) 

gggrid([
    graph_vertices + nice_edges_S,
    graph_vertices + nice_edges_C
    ])


# #### 3. Another Example of Graph Visualization

# In[10]:


nodes = pd.DataFrame({
    'node': ["Living\nThings", "Animals", "Plants", "Dogs", "Cows", "Herbs"],
    'x': [0, -1, 1, -2, 0, 2],
    'y': [1, 0, 0, -1, -1, -1]
})

edges = pd.DataFrame({
    'from': ["Animals", "Plants", "Dogs", "Cows", "Cows", "Herbs"],
    'to': ["Living\nThings", "Living\nThings", "Animals", "Animals", "Herbs", "Plants"],
    'relation': ["is", "is", "is", "is", "eat", "is"]
})
edges = pd.merge(edges, nodes, left_on='from', right_on='node')
edges = pd.merge(edges, nodes, left_on='to', right_on='node', suffixes=('', '_to'))
edges


# In[11]:


ggplot(nodes, aes(x='x', y='y')) \
  + geom_segment(aes(x='x', y='y', xend='x_to', yend='y_to',
                     color='relation'), data=edges, 
               size_end=25,
               arrow=arrow()) \
  + geom_point(color='#2166ac', fill='#d1e5f0', shape=21, size=25) \
  + scale_color_manual(['#2166ac', '#d6604d']) \
  + geom_text(aes(label='node')) \
  + coord_cartesian([-3,3],[-1.5,1.5]) \
  + theme_void()