Uses Scikit-Learn to fit a model to X and y then uses the resulting model
    to predict the curve based on the X values. This curve is drawn to the ax
    (matplotlib axis) which must be passed as the third variable.

    The estimator function can be one of the following:

    - ``'linear'``:      Uses OLS to fit the regression
    - ``'quadratic'``:   Uses OLS with Polynomial order 2
    - ``'exponential'``: Not implemented yet
    - ``'log'``:         Not implemented yet
    - ``'select_best'``: Selects the best fit via MSE

    The remaining keyword arguments are passed to ax.plot to define and
    describe the line of best fit.

    Parameters
    ----------
    X : ndarray or DataFrame of shape n x m
        A matrix of n instances with m features

    y : ndarray or Series of length n
        An array or series of target or class values

    ax : matplotlib Axes, default: None
        The axis to plot the figure on. If None is passed in the current axes
        will be used (or generated if required).

    estimator : string, default: 'linear'
        The name of the estimator function used to draw the best fit line.
        The estimator can currently be one of linear, quadratic, exponential,
        log, or select_best. The select best method uses the minimum MSE to
        select the best fit line.

    kwargs : dict
        Keyword arguments to pass to the matplotlib plot function to style and
        label the line of best fit. By default, the standard line color is
        used unless the color keyword argument is passed in.

    Returns
    -------

    ax : matplotlib Axes
        The axes with the line drawn on it.

def draw_best_fit(X, y, ax, estimator='linear', **kwargs):
    """
    Uses Scikit-Learn to fit a model to X and y then uses the resulting model
    to predict the curve based on the X values. This curve is drawn to the ax
    (matplotlib axis) which must be passed as the third variable.

    The estimator function can be one of the following:

    - ``'linear'``:      Uses OLS to fit the regression
    - ``'quadratic'``:   Uses OLS with Polynomial order 2
    - ``'exponential'``: Not implemented yet
    - ``'log'``:         Not implemented yet
    - ``'select_best'``: Selects the best fit via MSE

    The remaining keyword arguments are passed to ax.plot to define and
    describe the line of best fit.

    Parameters
    ----------
    X : ndarray or DataFrame of shape n x m
        A matrix of n instances with m features

    y : ndarray or Series of length n
        An array or series of target or class values

    ax : matplotlib Axes, default: None
        The axis to plot the figure on. If None is passed in the current axes
        will be used (or generated if required).

    estimator : string, default: 'linear'
        The name of the estimator function used to draw the best fit line.
        The estimator can currently be one of linear, quadratic, exponential,
        log, or select_best. The select best method uses the minimum MSE to
        select the best fit line.

    kwargs : dict
        Keyword arguments to pass to the matplotlib plot function to style and
        label the line of best fit. By default, the standard line color is
        used unless the color keyword argument is passed in.

    Returns
    -------

    ax : matplotlib Axes
        The axes with the line drawn on it.
    """

    # Estimators are the types of best fit lines that can be drawn.
    estimators = {
        LINEAR: fit_linear,               # Uses OLS to fit the regression
        QUADRATIC: fit_quadratic,         # Uses OLS with Polynomial order 2
        EXPONENTIAL: fit_exponential,     # Not implemented yet
        LOG: fit_log,                     # Not implemented yet
        SELECT_BEST: fit_select_best,     # Selects the best fit via MSE
    }

    # Check to make sure that a correct estimator value was passed in.
    if estimator not in estimators:
        raise YellowbrickValueError(
            "'{}' not a valid type of estimator; choose from {}".format(
                estimator, ", ".join(estimators.keys())
            )
        )

    # Then collect the estimator function from the mapping.
    estimator = estimators[estimator]

    # Ensure that X and y are the same length
    if len(X) != len(y):
        raise YellowbrickValueError((
            "X and y must have same length:"
             " X len {} doesn't match y len {}!"
        ).format(len(X), len(y)))

    # Ensure that X and y are np.arrays
    X = np.array(X)
    y = np.array(y)

    # Verify that X is a two dimensional array for Scikit-Learn esitmators
    # and that its dimensions are (n, 1) where n is the number of rows.
    if X.ndim < 2:
        X = X[:,np.newaxis] # Reshape X into the correct dimensions

    if X.ndim > 2:
        raise YellowbrickValueError(
            "X must be a (1,) or (n,1) dimensional array not {}".format(X.shape)
        )

    # Verify that y is a (n,) dimensional array
    if y.ndim > 1:
        raise YellowbrickValueError(
            "y must be a (1,) dimensional array not {}".format(y.shape)
        )

    # Uses the estimator to fit the data and get the model back.
    model = estimator(X, y)

    # Set the color if not passed in.
    if 'c' not in kwargs and 'color' not in kwargs:
        kwargs['color'] = LINE_COLOR

    # Get the current working axes
    ax = ax or plt.gca()

    # Plot line of best fit onto the axes that were passed in.
    # TODO: determine if xlim or X.min(), X.max() are better params
    xr = np.linspace(*ax.get_xlim(), num=100)
    ax.plot(xr, model.predict(xr[:,np.newaxis]), **kwargs)
    return ax