import matplotlib.mlab as mlb
import matplotlib.pyplot as plt

from scikits.statsmodels.sandbox import tsa
import scikits.talkbox as stb
import scikits.talkbox.spectral.basic as stbs

ar = [1., -0.7]#[1,0,0,0,0,0,0,-0.7]
ma = [1., 0.3]

ar = np.convolve([1.]+[0]*50 +[-0.6], ar)

n_startup = 1000
# throwing away samples at beginning makes sample more "stationary"

xo = tsa.arma_generate_sample(ar,ma,n_startup+200)
x = xo[n_startup:]

def arma_periodogram(ar, ma, worn=None):
    w, h = signal.freqz(ma, ar, **kwds)
    sd = np.abs(h)**2/np.sqrt(2*np.pi)
    if np.sum(np.isnan(h)) > 0:
        # this happens with unit root or seasonal unit root'
        print 'Warning: nan in frequency response h'
    return w, sd

plt.figure()
plt.plot(x)

rescale = 1
'''generates some ARMA random samples and saves to python module file

'''

import numpy as np
from scikits.statsmodels.sandbox import tsa
from maketests_mlabwrap import HoldIt

if __name__ == '__main__':
    filen = 'savedrvs.py'
    np.set_printoptions(precision=14, linewidth=100)


    # check arma to return same as random.normal
    np.random.seed(10000)
    xo = tsa.arma_generate_sample([1], [1], nsample=100)
    xo2 = np.round(xo*1000).astype(int)
    np.random.seed(10000)
    rvs = np.random.normal(size=100)
    rvs2 = np.round(xo*1000).astype(int)
    assert (xo2==rvs2).all()

    nsample = 1000
    data =  HoldIt('rvsdata')

    np.random.seed(10000)
    xo = tsa.arma_generate_sample([1, -0.8, 0.5], [1], nsample=nsample)
    data.xar2 = np.round(xo*1000).astype(int)
    np.random.seed(10000)
    xo = np.random.normal(size=nsample)
    data.xnormal = np.round(xo*1000).astype(int)