"""Examples for linear model with heteroscedasticity estimated by feasible GLS

These are examples to check the results during development.

The assumptions:

We have a linear model y = X*beta where the variance of an observation depends
on some explanatory variable Z (`exog_var`).
linear_model.WLS estimated the model for a given weight matrix
here we want to estimate also the weight matrix by two step or iterative WLS

Created on Wed Dec 21 12:28:17 2011

Author: Josef Perktold

"""

import numpy as np
from numpy.testing import assert_almost_equal

from statsmodels.regression.feasible_gls import GLSHet, GLSHet2
from statsmodels.regression.linear_model import OLS
from statsmodels.tools.tools import add_constant

examples = ["ex1"]

if "ex1" in examples:
    nsample = 300  # different pattern last graph with 100 or 200 or 500
    sig = 0.5

    np.random.seed(9876789)  # 9876543)
    X = np.random.randn(nsample, 3)
    X = np.column_stack((np.ones((nsample, 1)), X))
    beta = [1, 0.5, -0.5, 1.0]
    y_true2 = np.dot(X, beta)

    x1 = np.linspace(0, 1, nsample)
    gamma = np.array([1, 3.0])
    # with slope 3 instead of two, I get negative weights, Not correct
    #   - was misspecified, but the negative weights are still possible with identity link
    # gamma /= gamma.sum()   # normalize assuming x1.max is 1
    z_true = add_constant(x1)

    winv = np.dot(z_true, gamma)
    het_params = sig**2 * np.array([1, 3.0])  # for squared
    sig2_het = sig**2 * winv

    weights_dgp = 1 / winv
    weights_dgp /= (
        weights_dgp.max()
    )  # should be already normalized - NOT check normalization
    # y2[:nsample*6/10] = y_true2[:nsample*6/10] + sig*1. * np.random.normal(size=nsample*6/10)
    z0 = np.zeros(nsample)
    z0[(nsample * 5) // 10 :] = 1  # dummy for 2 halfs of sample
    z0 = add_constant(z0)

    z1 = add_constant(x1)

    noise = np.sqrt(sig2_het) * np.random.normal(size=nsample)
    y2 = y_true2 + noise

    X2 = X[:, [0, 2]]  # misspecified, missing regressor in main equation
    X2 = X  # correctly specigied

    res_ols = OLS(y2, X2).fit()
    print("OLS beta estimates")
    print(res_ols.params)
    print("OLS stddev of beta")
    print(res_ols.bse)
    print("\nWLS")
    mod0 = GLSHet2(y2, X2, exog_var=winv)
    res0 = mod0.fit()
    print("new version")
    mod1 = GLSHet(y2, X2, exog_var=winv)
    res1 = mod1.iterative_fit(2)
    print("WLS beta estimates")
    print(res1.params)
    print(res0.params)
    print("WLS stddev of beta")
    print(res1.bse)
    # compare with previous version GLSHet2, refactoring check
    # assert_almost_equal(res1.params, np.array([ 0.37642521,  1.51447662]))
    # this fails ???  more iterations? different starting weights?

    print(res1.model.weights / res1.model.weights.max())
    # why is the error so small in the estimated weights ?
    assert_almost_equal(res1.model.weights / res1.model.weights.max(), weights_dgp, 14)
    print("residual regression params")
    print(res1.results_residual_regression.params)
    print("scale of model ?")
    print(res1.scale)
    print("unweighted residual variance, note unweighted mean is not zero")
    print(res1.resid.var())
    # Note weighted mean is zero:
    # (res1.model.weights * res1.resid).mean()

    doplots = True  # False
    if doplots:
        import matplotlib.pyplot as plt

        plt.figure()
        plt.plot(x1, y2, "o")
        plt.plot(x1, y_true2, "b-", label="true")
        plt.plot(x1, res1.fittedvalues, "r-", label="fwls")
        plt.plot(x1, res_ols.fittedvalues, "--", label="ols")
        plt.legend()

    # the next only works if w has finite support, discrete/categorical
    # z = (w[:,None] == [1,4]).astype(float) # dummy variable
    # z = (w0[:,None] == np.unique(w0)).astype(float) # dummy variable
    # changed z0 contains dummy and constant
    mod2 = GLSHet(y2, X2, exog_var=z0)
    res2 = mod2.iterative_fit(3)
    print(res2.params)

    import statsmodels.api as sm

    # z = sm.add_constant(w, prepend=True)
    z = sm.add_constant(x1 / x1.max())
    mod3 = GLSHet(y2, X2, exog_var=z1)  # , link=sm.families.links.Log())
    res3 = mod3.iterative_fit(20)
    error_var_3 = res3.mse_resid / res3.model.weights
    print(res3.params)
    print("np.array(res3.model.history['ols_params'])")

    print(np.array(res3.model.history["ols_params"]))
    print("np.array(res3.model.history['self_params'])")
    print(np.array(res3.model.history["self_params"]))

    # Models 2 and 3 are equivalent with different parameterization of Z
    print(np.unique(res2.model.weights))  # for discrete z only, only a few uniques
    print(np.unique(res3.model.weights))

    print(res3.summary())
    print("\n\nResults of estimation of weights")
    print("--------------------------------")
    print(res3.results_residual_regression.summary())

    if doplots:
        plt.figure()
        plt.plot(x1, y2, "o")
        plt.plot(x1, y_true2, "b-", label="true")
        plt.plot(x1, res1.fittedvalues, "-", label="fwls1")
        plt.plot(x1, res2.fittedvalues, "-", label="fwls2")
        plt.plot(x1, res3.fittedvalues, "-", label="fwls3")
        plt.plot(x1, res_ols.fittedvalues, "--", label="ols")
        plt.legend()

        plt.figure()
        plt.ylim(0, 5)
        res_e2 = OLS(noise**2, z).fit()
        plt.plot(noise**2, "bo", alpha=0.5, label="dgp error**2")
        plt.plot(res_e2.fittedvalues, lw=2, label="ols for noise**2")
        # plt.plot(res3.model.weights, label='GLSHet weights')
        plt.plot(error_var_3, lw=2, label="GLSHet error var")
        plt.plot(res3.resid**2, "ro", alpha=0.5, label="resid squared")
        # plt.plot(weights_dgp, label='DGP weights')
        plt.plot(sig**2 * winv, lw=2, label="DGP error var")
        plt.legend()

        plt.show()

    """Note these are close but maybe biased because of skewed distribution
    >>> res3.mse_resid/res3.model.weights[-10:]
    array([ 1.03115871,  1.03268209,  1.03420547,  1.03572885,  1.03725223,
            1.03877561,  1.04029899,  1.04182237,  1.04334575,  1.04486913])
    >>> res_e2.fittedvalues[-10:]
    array([ 1.0401953 ,  1.04171386,  1.04323242,  1.04475098,  1.04626954,
            1.0477881 ,  1.04930666,  1.05082521,  1.05234377,  1.05386233])
    >>> sig**2 * w[-10:]
    array([ 0.98647295,  0.98797595,  0.98947896,  0.99098196,  0.99248497,
            0.99398798,  0.99549098,  0.99699399,  0.99849699,  1.        ])
        """
