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

# In[1]:


import warnings
warnings.filterwarnings("ignore")  # to get rid of an annoying deprecation warning in pandas
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import statsmodels.formula.api as sm
from mpl_toolkits.mplot3d import Axes3D
get_ipython().run_line_magic('matplotlib', 'inline')
get_ipython().run_line_magic('config', "InlineBackend.print_figure_kwargs = {'bbox_inches':None}")

df = pd.read_csv("http://rulelaw.net/downloads/rol-scores.csv")
df.columns = [x.lower().strip() for x in df.columns]
df.head()


# In[2]:


def regplane3d(y, x1, x2, df):
    fig = plt.figure(figsize=(10, 10))
    ax = fig.add_subplot(111, projection='3d')
    # get the regression fit
    model = sm.ols(formula='{} ~ {} + {}'.format(y, x1, x2), data = df)
    fit = model.fit()
    # lay out the range of x values for plotting and predicting
    # here I'm going to switch notation conventions from regression (y = x1 + x2) to 
    # dimensions s.t. x1 and x2 become x and y, and y becomes z (the 3rd dimension)
    x_surf = np.linspace(df[x1].min(), df[x1].max())
    y_surf = np.linspace(df[x2].min(), df[x2].max())
    x_surf, y_surf = np.meshgrid(x_surf, y_surf)
    # predict over space of graph
    synthetic_data = pd.DataFrame({x1: x_surf.ravel(), x2: y_surf.ravel()})
    preds = fit.predict(exog=synthetic_data)
    # add surface plot
    ax.plot_surface(x_surf, y_surf, preds.reshape(x_surf.shape),
                rstride=1,
                cstride=1,
                color='None',
                alpha = 0.4)
    # add scatterplot
    ax.scatter(df[x1], df[x2], df[y],
           c='blue',
           marker='o',
           alpha=1)
    return fig


# In[3]:


foo = regplane3d("elec_pros", "assoc_org", "free_expr", df)


# In[4]:


import notebook
notebook.__version__


# In[5]:


import matplotlib
matplotlib.__version__