"""Example for gam.AdditiveModel and PolynomialSmoother

This example was written as a test case.
The data generating process is chosen so the parameters are well identified
and estimated.

Created on Fri Nov 04 13:45:43 2011

Author: Josef Perktold

"""

from statsmodels.compat.python import lrange

import numpy as np

from statsmodels.regression.linear_model import OLS
from statsmodels.sandbox.gam import AdditiveModel

np.random.seed(8765993)
# seed is chosen for nice result, not randomly
# other seeds are pretty off in the prediction

# DGP: simple polynomial
order = 3
sigma_noise = 0.5
nobs = 1000  # 1000 # with 1000, OLS and Additivemodel agree in params at 2 decimals
lb, ub = -3.5, 4  # 2.5
x1 = np.linspace(lb, ub, nobs)
x2 = np.sin(2 * x1)
x = np.column_stack((x1 / x1.max() * 2, x2))
exog = (x[:, :, None] ** np.arange(order + 1)[None, None, :]).reshape(nobs, -1)
idx = lrange((order + 1) * 2)
del idx[order + 1]
exog_reduced = exog[:, idx]  # remove duplicate constant
y_true = exog.sum(1) / 2.0
z = y_true  # alias check
d = x
y = y_true + sigma_noise * np.random.randn(nobs)

example = 1

if example == 1:
    m = AdditiveModel(d)
    m.fit(y)

    y_pred = m.results.predict(d)


for ss in m.smoothers:
    print(ss.params)

res_ols = OLS(y, exog_reduced).fit()
print(res_ols.params)

# from numpy.testing import assert_almost_equal
# assert_almost_equal(y_pred, res_ols.fittedvalues, 3)

if example > 0:
    import matplotlib.pyplot as plt

    plt.figure()
    plt.plot(exog)

    y_pred = m.results.mu  # + m.results.alpha # m.results.predict(d)
    plt.figure()
    plt.subplot(2, 2, 1)
    plt.plot(y, ".", alpha=0.25)
    plt.plot(y_true, "k-", label="true")

    plt.plot(res_ols.fittedvalues, "g-", label="OLS", lw=2, alpha=-0.7)
    plt.plot(y_pred, "r-", label="AM")
    plt.legend(loc="upper left")
    plt.title("gam.AdditiveModel")

    counter = 2
    for ii, xx in zip(["z", "x1", "x2"], [z, x[:, 0], x[:, 1]]):
        sortidx = np.argsort(xx)
        # plt.figure()
        plt.subplot(2, 2, counter)
        plt.plot(xx[sortidx], y[sortidx], ".", alpha=0.25)
        plt.plot(xx[sortidx], y_true[sortidx], "k.", label="true", lw=2)
        plt.plot(xx[sortidx], y_pred[sortidx], "r.", label="AM")
        plt.legend(loc="upper left")
        plt.title("gam.AdditiveModel " + ii)
        counter += 1

    plt.show()
