'''
Example: scikits.statsmodels.GLSAR

6 examples for GLSAR with artificial data

Notes
------
These examples were written mostly to cross-check results.  It is still being
written, and GLSAR is still being worked on.
'''

import numpy as np
import numpy.testing as npt
from scipy import signal
import scikits.statsmodels.api as sm
from scikits.statsmodels.regression.linear_model import GLSAR, yule_walker

examples_all = range(10) + ['test_copy']

examples = examples_all #[5]

if 0 in examples:
    print '\n Example 0'
    X = np.arange(1,8)
    X = sm.add_constant(X)
    Y = np.array((1, 3, 4, 5, 8, 10, 9))
    rho = 2
    model = GLSAR(Y, X, 2)
    for i in range(6):
        results = model.fit()
        print "AR coefficients:", model.rho
        rho, sigma = yule_walker(results.resid, order = model.order)
        model = GLSAR(Y, X, rho)

    par0 = results.params
    print par0
    model0if = GLSAR(Y, X, 2)
    res = model0if.iterative_fit(6)
    print 'iterativefit beta', res.params
    results.tvalues # is this correct? it does equal params/bse
    # but isn't the same as the AR example (which was wrong in the first place..)
    print results.t_test([0,1])  # are sd and t correct? vs
    print results.f_test(np.eye(2))


rhotrue = [0.5, 0.2]
rhotrue = np.asarray(rhotrue)
nlags = np.size(rhotrue)
beta = np.array([0.1, 2])
noiseratio = 0.5
nsample = 2000
x = np.arange(nsample)
X1 = sm.add_constant(x)

wnoise = noiseratio * np.random.randn(nsample+nlags)
#noise = noise[1:] + rhotrue*noise[:-1] # wrong this is not AR

#find my drafts for univariate ARMA functions
# generate AR(p)
if np.size(rhotrue) == 1:
    # replace with scipy.signal.lfilter, keep for testing
    arnoise = np.zeros(nsample+1)
    for i in range(1,nsample+1):
        arnoise[i] = rhotrue*arnoise[i-1] + wnoise[i]
    noise = arnoise[1:]
    an = signal.lfilter([1], np.hstack((1,-rhotrue)), wnoise[1:])
    print 'simulate AR(1) difference', np.max(np.abs(noise-an))
else:
    noise = signal.lfilter([1], np.hstack((1,-rhotrue)), wnoise)[nlags:]

# generate GLS model with AR noise
y1 = np.dot(X1,beta) + noise

if 1 in examples:
    print '\nExample 1: iterative_fit and repeated calls'
    mod1 = GLSAR(y1, X1, 1)
    res = mod1.iterative_fit()
    print mod1._results.params
    print mod1.rho

    for i in range(5):
        mod1.iterative_fit(1)
#        mod1.fit()
        print mod1.rho
        print mod1._results.params

if 2 in examples:
    print '\nExample 2: iterative fitting of first model'
    print 'with AR(0)', par0
    parold = par0
    mod0 = GLSAR(Y, X, 1)
    for i in range(5):
        #print mod0.wexog.sum()
        #print mod0.pinv_wexog.sum()
        mod0.iterative_fit(1)
        print 'rho', mod0.rho
        parnew = mod0._results.params
        print 'params', parnew
        print 'params change in iteration', parnew - parold
        parold = parnew

# generate pure AR(p) process
Y = noise

#example with no regressor,
#results now have same estimated rho as yule-walker directly

if 3 in examples:
    print '\nExample 3: pure AR(2), GLSAR versus Yule_Walker'
    model3 = GLSAR(Y, rho=2)
    for i in range(5):
        results = model3.fit()
        print "AR coefficients:", model3.rho, results.params
        rho, sigma = yule_walker(results.resid, order = model3.order)
        model3 = GLSAR(Y, rho=rho)

if 'test_copy' in examples:
    xx = X.copy()
    rhoyw, sigmayw = yule_walker(xx[:,0], order = 2)
    print rhoyw, sigmayw
    print (xx == X).all()  # test for unchanged array (fixed)

    yy = Y.copy()
    rhoyw, sigmayw = yule_walker(yy, order = 2)
    print rhoyw, sigmayw
    print (yy == Y).all()  # test for unchanged array (fixed)


if 4 in examples:
    print '\nExample 4: demeaned pure AR(2), GLSAR versus Yule_Walker'
    Ydemeaned = Y - Y.mean()
    model4 = GLSAR(Ydemeaned, rho=2)
    for i in range(5):
        results = model4.fit()
        print "AR coefficients:", model3.rho, results.params
        rho, sigma = yule_walker(results.resid, order = model4.order)
        model4 = GLSAR(Ydemeaned, rho=rho)

if 5 in examples:
    print '\nExample 5: pure AR(2), GLSAR iterative_fit versus Yule_Walker'
    model3a = GLSAR(Y, rho=1)
    res3a = model3a.iterative_fit(5)
    print res3a.params
    print model3a.rho
    rhoyw, sigmayw = yule_walker(Y, order = 1)
    print rhoyw, sigmayw
    npt.assert_array_almost_equal(model3a.rho, rhoyw, 15)