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

# # Overview
# 
# This notebook introduces the ipyplotly enhancements to the plotly.py visualization library and demonstrates some of its features.
# 
# 
# ## New Features
# 
#  - Traces can be added and updated interactively by simply assigning to properties
#  - The full Traces and Layout API is generated from the plotly schema to provide a great experience for interactive use in the notebook
#  - Data validation covering the full API with clear, informative error messages
#  - Jupyter friendly docstrings on constructor params and properties
#  - Support for setting array properties as numpy arrays. When numpy arrays are used, ipywidgets binary serialization protocol is used to avoid converting these to JSON strings.
#  - Context manager API for animation
#  - Programmatic export of figures to static SVG images (and PNG and PDF with cairosvg installed).

# # Imports

# In[1]:


# ipyplotly
from plotly.graph_objs import FigureWidget
from plotly.callbacks import Points, InputDeviceState

# pandas
import pandas as pd

# numpy
import numpy as np

# scikit learn
from sklearn import datasets

# ipywidgets
from ipywidgets import HBox, VBox, Button

# functools
from functools import partial


# In[2]:


# Load iris dataset
iris_data = datasets.load_iris()
feature_names = [name.replace(' (cm)', '').replace(' ', '_') for name in iris_data.feature_names]
iris_df = pd.DataFrame(iris_data.data, columns=feature_names)
iris_class = iris_data.target + 1
iris_df.head()


# ## Create and display an empty FigureWidget
# A FigureWidget behaves almost identically to a Figure but it is also an ipywidget that can be displayed directly in the notebook without calling `iplot`

# In[3]:


f1 = FigureWidget()
f1


# # Tab completion 
# Entering ``f1.add_<tab>`` displays add methods for all of the supported trace types

# In[5]:


# f1.add_


# Entering ``f1.add_scatter(<tab>)`` displays the names of all of the top-level properties for the scatter trace type
# 
# Entering ``f1.add_scatter(<shift+tab>)`` displays the signature pop-up. Expanding this pop-up reveals the method doc string which contains the descriptions of all of the top level properties

# In[6]:


# f1.add_scatter(


# # Add scatter trace

# In[7]:


scatt1 = f1.add_scatter(x=iris_df.sepal_length, y=iris_df.petal_width)


# In[6]:


f1


# In[7]:


get_ipython().run_line_magic('pinfo', 'scatt1.mode')


# In[8]:


# That's not what we wanted, change the mode to 'markers'
scatt1.mode = 'markers'


# In[9]:


# Set size to 8
scatt1.marker.size = 8


# In[10]:


# Color markers by iris class
scatt1.marker.color = iris_class


# In[11]:


# Change colorscale
scatt1.marker.cmin = 0.5
scatt1.marker.cmax = 3.5
scatt1.marker.colorscale = [[0, 'red'], [0.33, 'red'], 
                            [0.33, 'green'], [0.67, 'green'], 
                            [0.67, 'blue'], [1.0, 'blue']]

scatt1.marker.showscale = True


# In[12]:


# Fix up colorscale ticks
scatt1.marker.colorbar.ticks = 'outside'
scatt1.marker.colorbar.tickvals = [1, 2, 3]
scatt1.marker.colorbar.ticktext = iris_data.target_names.tolist()


# In[13]:


# Set colorscale title
scatt1.marker.colorbar.title = 'Species'
scatt1.marker.colorbar.titlefont.size = 16
scatt1.marker.colorbar.titlefont.family = 'Rockwell'


# In[14]:


# Add axis labels
f1.layout.xaxis.title = 'sepal_length'
f1.layout.yaxis.title = 'petal_width'


# In[15]:


f1


# In[16]:


# Hover info
scatt1.text = iris_data.target_names[iris_data.target]
scatt1.hoverinfo = 'text+x+y'
f1.layout.hovermode = 'closest'


# In[17]:


f1


# ## Animate marker size change

# In[18]:


# Set marker size based on petal_length
with f1.batch_animate(duration=1000):
    scatt1.marker.size = np.sqrt(iris_df.petal_length.values * 50)
    


# In[20]:


# Restore constant marker size
with f1.batch_animate(duration=1000):
    scatt1.marker.size = 8


# ## Set drag mode property callback
# Make points more transparent when `dragmode` is `zoom`

# In[23]:


def set_opacity(marker, layout, dragmode):
    if dragmode == 'zoom':
        marker.opacity = 0.5
    else:
        marker.opacity = 1.0


# In[24]:


f1.layout.on_change(partial(set_opacity, scatt1.marker), 'dragmode')


# ## Configure colorscale for brushing

# In[25]:


scatt1.marker.colorbar = None
scatt1.marker.colorscale = [[0, 'lightgray'], [0.5, 'lightgray'], [0.5, 'red'], [1, 'red']]
scatt1.marker.cmin = -0.5
scatt1.marker.cmax = 1.5
scatt1.marker.colorbar.ticks = 'outside'
scatt1.marker.colorbar.tickvals = [0, 1]
scatt1.marker.colorbar.ticktext = ['unselected', 'selected']


# In[26]:


# Reset colors to zeros (unselected)
scatt1.marker.color = np.zeros(iris_class.size)
selected = np.zeros(iris_class.size)


# In[27]:


f1


# ### Configure brushing callback

# In[29]:


# Assigning these variables here is not required. But doing so tricks Jupyter into 
# providing property tab completion on the parameters to the brush function below
trace, points, state = scatt1, Points(), InputDeviceState()


# In[30]:


def brush(trace, points, state):
    inds = np.array(points.point_inds)
    if inds.size:
        selected[inds] = 1
        trace.marker.color = selected


# In[31]:


scatt1.on_selected(brush)


# Now box or lasso select points on the figure and see them turn red

# In[32]:


# Reset brush
selected = np.zeros(iris_class.size)
scatt1.marker.color = selected


# ## Create second plot with different features

# In[38]:


f2 = FigureWidget(data=[{'type': 'scatter',
                         'x': iris_df.petal_length, 
                         'y': iris_df.sepal_width,
                         'mode': 'markers'}])
f2


# In[39]:


# Set axis titles
f2.layout.xaxis.title = 'petal_length'
f2.layout.yaxis.title = 'sepal_width'


# In[40]:


# Grab trace reference
scatt2 = f2.data[0]


# In[41]:


# Set marker styles / colorbars to match between figures
scatt2.marker = scatt1.marker


# In[42]:


# Configure brush on both plots to update both plots
def brush(trace, points, state):
    inds = np.array(points.point_inds)
    if inds.size:
        selected = scatt1.marker.color.copy()
        selected[inds] = 1
        scatt1.marker.color = selected
        scatt2.marker.color = selected    
    
scatt1.on_selected(brush)
scatt2.on_selected(brush)


# In[43]:


f2.layout.on_change(partial(set_opacity, scatt2.marker), 'dragmode')


# In[44]:


# Reset brush
def reset_brush(btn):
    selected = np.zeros(iris_class.size)
    scatt1.marker.color = selected
    scatt2.marker.color = selected


# In[45]:


# Create reset button
button = Button(description="clear")
button.on_click(reset_brush)


# In[46]:


# Hide colorbar for figure 1
scatt1.marker.showscale = False


# In[47]:


# Set dragmode to lasso for both plots
f1.layout.dragmode = 'lasso'
f2.layout.dragmode = 'lasso'


# In[50]:


# Display two figures and the reset button
f1.layout.width = 500
f2.layout.width = 500


# In[49]:


VBox([HBox([f1, f2]), button])


# In[55]:


# Save figure 2 to a svg image in the exports directory
f2.save_image('exports/f2.svg')


# In[54]:


# Save figure 1 to a pdf in the exports directory (requires cairosvg be installed)
# f1.save_image('exports/f1.pdf')