"""Examples for Regression Plots

Author: Josef Perktold

"""

import matplotlib.pyplot as plt
import numpy as np

import statsmodels.api as sm
import statsmodels.graphics.regressionplots as smrp
from statsmodels.graphics.tests.test_regressionplots import TestPlot
from statsmodels.sandbox.regression.predstd import wls_prediction_std

# example from tut.ols with changes
# fix a seed for these examples
np.random.seed(9876789)

# OLS non-linear curve but linear in parameters
# ---------------------------------------------

nsample = 100
sig = 0.5
x1 = np.linspace(0, 20, nsample)
x2 = 5 + 3 * np.random.randn(nsample)
X = np.c_[x1, x2, np.sin(0.5 * x1), (x2 - 5) ** 2, np.ones(nsample)]
beta = [0.5, 0.5, 1, -0.04, 5.0]
y_true = np.dot(X, beta)
y = y_true + sig * np.random.normal(size=nsample)

# estimate only linear function, misspecified because of non-linear terms
exog0 = sm.add_constant(np.c_[x1, x2], prepend=False)

#    plt.figure()
#    plt.plot(x1, y, 'o', x1, y_true, 'b-')

res = sm.OLS(y, exog0).fit()
# print res.params
# print res.bse


plot_old = 0  # True
if plot_old:

    # current bug predict requires call to model.results
    # print res.model.predict
    prstd, iv_l, iv_u = wls_prediction_std(res)
    plt.plot(x1, res.fittedvalues, "r-o")
    plt.plot(x1, iv_u, "r--")
    plt.plot(x1, iv_l, "r--")
    plt.title("blue: true,   red: OLS")

    plt.figure()
    plt.plot(res.resid, "o")
    plt.title("Residuals")

    fig2 = plt.figure()
    ax = fig2.add_subplot(2, 1, 1)
    # namestr = ' for %s' % self.name if self.name else ''
    plt.plot(x1, res.resid, "o")
    ax.set_title("residuals versus exog")  # + namestr)
    ax = fig2.add_subplot(2, 1, 2)
    plt.plot(x2, res.resid, "o")

    fig3 = plt.figure()
    ax = fig3.add_subplot(2, 1, 1)
    # namestr = ' for %s' % self.name if self.name else ''
    plt.plot(x1, res.fittedvalues, "o")
    ax.set_title("Fitted values versus exog")  # + namestr)
    ax = fig3.add_subplot(2, 1, 2)
    plt.plot(x2, res.fittedvalues, "o")

    fig4 = plt.figure()
    ax = fig4.add_subplot(2, 1, 1)
    # namestr = ' for %s' % self.name if self.name else ''
    plt.plot(x1, res.fittedvalues + res.resid, "o")
    ax.set_title("Fitted values plus residuals versus exog")  # + namestr)
    ax = fig4.add_subplot(2, 1, 2)
    plt.plot(x2, res.fittedvalues + res.resid, "o")

    # see http://www.itl.nist.gov/div898/software/dataplot/refman1/auxillar/partregr.htm
    fig5 = plt.figure()
    ax = fig5.add_subplot(2, 1, 1)
    # namestr = ' for %s' % self.name if self.name else ''
    res1a = sm.OLS(y, exog0[:, [0, 2]]).fit()
    res1b = sm.OLS(x1, exog0[:, [0, 2]]).fit()
    plt.plot(res1b.resid, res1a.resid, "o")
    res1c = sm.OLS(res1a.resid, res1b.resid).fit()
    plt.plot(res1b.resid, res1c.fittedvalues, "-")
    ax.set_title("Partial Regression plot")  # + namestr)
    ax = fig5.add_subplot(2, 1, 2)
    # plt.plot(x2, res.fittedvalues + res.resid, 'o')
    res2a = sm.OLS(y, exog0[:, [0, 1]]).fit()
    res2b = sm.OLS(x2, exog0[:, [0, 1]]).fit()
    plt.plot(res2b.resid, res2a.resid, "o")
    res2c = sm.OLS(res2a.resid, res2b.resid).fit()
    plt.plot(res2b.resid, res2c.fittedvalues, "-")

    # see http://www.itl.nist.gov/div898/software/dataplot/refman1/auxillar/ccpr.htm
    fig6 = plt.figure()
    ax = fig6.add_subplot(2, 1, 1)
    # namestr = ' for %s' % self.name if self.name else ''
    x1beta = x1 * res.params[1]
    x2beta = x2 * res.params[2]
    plt.plot(x1, x1beta + res.resid, "o")
    plt.plot(x1, x1beta, "-")
    ax.set_title("X_i beta_i plus residuals versus exog (CCPR)")  # + namestr)
    ax = fig6.add_subplot(2, 1, 2)
    plt.plot(x2, x2beta + res.resid, "o")
    plt.plot(x2, x2beta, "-")

    # print res.summary()

doplots = 1
if doplots:
    fig1 = smrp.plot_fit(res, 0, y_true=None)
    smrp.plot_fit(res, 1, y_true=None)
    smrp.plot_partregress_grid(res, exog_idx=[0, 1])
    smrp.plot_regress_exog(res, exog_idx=0)
    smrp.plot_ccpr(res, exog_idx=0)
    smrp.plot_ccpr_grid(res, exog_idx=[0, 1])

tp = TestPlot()
tp.test_plot_fit()

fig1 = smrp.plot_partregress_grid(res, exog_idx=[0, 1])
# add lowess
ax = fig1.axes[0]
y0 = ax.get_lines()[0]._y
x0 = ax.get_lines()[0]._x
lres = sm.nonparametric.lowess(y0, x0, frac=0.2)
ax.plot(lres[:, 0], lres[:, 1], "r", lw=1.5)
ax = fig1.axes[1]
y0 = ax.get_lines()[0]._y
x0 = ax.get_lines()[0]._x
lres = sm.nonparametric.lowess(y0, x0, frac=0.2)
ax.plot(lres[:, 0], lres[:, 1], "r", lw=1.5)

# plt.show()
