"""
Created on Fri Nov 04 10:51:39 2011

@author: josef
"""

import numpy as np

from statsmodels.regression.linear_model import OLS, WLS
from statsmodels.sandbox.nonparametric import smoothers

# DGP: simple polynomial
order = 3
sigma_noise = 0.5
nobs = 100
lb, ub = -1, 2
x = np.linspace(lb, ub, nobs)
x = np.sin(x)
exog = x[:, None] ** np.arange(order + 1)
y_true = exog.sum(1)
y = y_true + sigma_noise * np.random.randn(nobs)


# xind = np.argsort(x)
pmod = smoothers.PolySmoother(2, x)
pmod.fit(y)  # no return
y_pred = pmod.predict(x)
error = y - y_pred
mse = (error * error).mean()
print(mse)
res_ols = OLS(y, exog[:, :3]).fit()
print(np.squeeze(pmod.coef) - res_ols.params)


weights = np.ones(nobs)
weights[: nobs // 3] = 0.1
weights[-nobs // 5 :] = 2

pmodw = smoothers.PolySmoother(2, x)
pmodw.fit(y, weights=weights)  # no return
y_predw = pmodw.predict(x)
error = y - y_predw
mse = (error * error).mean()
print(mse)
res_wls = WLS(y, exog[:, :3], weights=weights).fit()
print(np.squeeze(pmodw.coef) - res_wls.params)


doplot = 1
if doplot:
    import matplotlib.pyplot as plt

    plt.plot(y, ".")
    plt.plot(y_true, "b-", label="true")
    plt.plot(y_pred, "-", label="poly")
    plt.plot(y_predw, "-", label="poly -w")
    plt.legend(loc="upper left")

    plt.close()
    # plt.show()
