def test_get_spectra():
    """

    Testing spectral estimation

    """

    methods = (None,
           {"this_method": 'welch', "NFFT": 256, "Fs": 2 * np.pi},
           {"this_method": 'welch', "NFFT": 1024, "Fs": 2 * np.pi})


    for method in methods:
        avg_pwr1 = []
        avg_pwr2 = []
        est_pwr1 = []
        est_pwr2 = []
        arsig1, _, _ = utils.ar_generator(N=2 ** 16) # It needs to be that long for
                                        # the answers to converge 
        arsig2, _, _ = utils.ar_generator(N=2 ** 16)
        
        avg_pwr1.append((arsig1 ** 2).mean())
        avg_pwr2.append((arsig2 ** 2).mean())
    
        tseries = np.vstack([arsig1,arsig2])
        
        f, c = tsa.get_spectra(tseries,method=method)

        # \sum_{\omega} psd d\omega:
        est_pwr1.append(np.sum(c[0, 0]) * (f[1] - f[0]))
        est_pwr2.append(np.sum(c[1, 1]) * (f[1] - f[0]))

        # Get it right within the order of magnitude:
        npt.assert_array_almost_equal(est_pwr1, avg_pwr1, decimal=-1)
        npt.assert_array_almost_equal(est_pwr2, avg_pwr2, decimal=-1)

def test_get_spectra_bi():
    """

    Test the bi-variate get_spectra function

    """ 

    methods = (None,
           {"this_method": 'welch', "NFFT": 256, "Fs": 2 * np.pi},
           {"this_method": 'welch', "NFFT": 1024, "Fs": 2 * np.pi})

    for method in methods:
        arsig1, _, _ = utils.ar_generator(N=2 ** 16)
        arsig2, _, _ = utils.ar_generator(N=2 ** 16)

        avg_pwr1 = (arsig1 ** 2).mean()
        avg_pwr2 = (arsig2 ** 2).mean()
        avg_xpwr = (arsig1 * arsig2.conjugate()).mean()

        tseries = np.vstack([arsig1, arsig2])

        f, fxx, fyy, fxy = tsa.get_spectra_bi(arsig1, arsig2, method=method)

        # \sum_{\omega} PSD(\omega) d\omega:
        est_pwr1 = np.sum(fxx * (f[1] - f[0]))
        est_pwr2 = np.sum(fyy * (f[1] - f[0]))
        est_xpwr = np.sum(fxy * (f[1] - f[0])).real

        # Test that we have the right order of magnitude:
        npt.assert_array_almost_equal(est_pwr1, avg_pwr1, decimal=-1)
        npt.assert_array_almost_equal(est_pwr2, avg_pwr2, decimal=-1)
        npt.assert_array_almost_equal(np.mean(est_xpwr),
                                      np.mean(avg_xpwr),
                                      decimal=-1)

def test_crosscov():
    N = 128
    ar_seq1, _, _ = utils.ar_generator(N=N)
    ar_seq2, _, _ = utils.ar_generator(N=N)

    for all_lags in (True, False):
        sxy = utils.crosscov(ar_seq1, ar_seq2, all_lags=all_lags)
        sxy_ref = ref_crosscov(ar_seq1, ar_seq2, all_lags=all_lags)
        err = sxy_ref - sxy
        mse = np.dot(err, err) / N
        npt.assert_(mse < 1e-12, "Large mean square error w.r.t. reference cross covariance")

def test_autocorr():
    N = 128
    ar_seq, _, _ = utils.ar_generator(N=N)
    rxx = utils.autocorr(ar_seq)
    npt.assert_(rxx[0] == rxx.max(), "Zero lag autocorrelation is not maximum autocorrelation")
    rxx = utils.autocorr(ar_seq, all_lags=True)
    npt.assert_(rxx[127] == rxx.max(), "Zero lag autocorrelation is not maximum autocorrelation")

def test_AR_LD():
    """

    Test the Levinson Durbin estimate of the AR coefficients against the
    expercted PSD

    """
    arsig,_,_ = utils.ar_generator(N=512)
    avg_pwr = (arsig*arsig.conjugate()).real.mean()
    order = 8
    ak, sigma_v = tsa.AR_est_LD(arsig, order)
    w, psd = tsa.AR_psd(ak, sigma_v)

    # the psd is a one-sided power spectral density, which has been
    # multiplied by 2 to preserve the property that
    # 1/2pi int_{-pi}^{pi} Sxx(w) dw = Rxx(0)

    # evaluate this integral numerically from 0 to pi
    dw = np.pi/len(psd)
    avg_pwr_est = np.trapz(psd, dx=dw) / (2*np.pi)
    npt.assert_almost_equal(avg_pwr, avg_pwr_est, decimal=0)

    # Test for providing the autocovariance as an input:
    ak,sigma_v = tsa.AR_est_LD(arsig, order, utils.autocov(arsig))
    w, psd = tsa.AR_psd(ak, sigma_v)
    avg_pwr_est = np.trapz(psd, dx=dw) / (2*np.pi)
    npt.assert_almost_equal(avg_pwr, avg_pwr_est, decimal=0)

def test_autocorr():
    N = 128
    ar_seq, _, _ = utils.ar_generator(N=N)
    rxx = utils.autocorr(ar_seq)
    yield nt.assert_true, rxx[0] == 1, "Zero lag autocorrelation is not equal to 1"
    rxx = utils.autocorr(ar_seq, all_lags=True)
    yield nt.assert_true, rxx[127] == 1, "Zero lag autocorrelation is not equal to 1"

def test_periodogram():
    arsig, _, _ = ut.ar_generator(N=512)
    avg_pwr = (arsig * arsig.conjugate()).mean()
    f, psd = tsa.periodogram(arsig, N=2048)
    df = 2. * np.pi / 2048
    avg_pwr_est = np.trapz(psd, dx=df)
    npt.assert_almost_equal(avg_pwr, avg_pwr_est, decimal=1)

def test_AR_est_consistency():
    order = 10 # some even number
    ak = _random_poles(order/2)
    x, v, _ = utils.ar_generator(N=512, coefs=-ak[1:], drop_transients=100)
    ak_yw, ssq_yw = tsa.AR_est_YW(x, order)
    ak_ld, ssq_ld = tsa.AR_est_LD(x, order)
    npt.assert_almost_equal(ak_yw, ak_ld)
    npt.assert_almost_equal(ssq_yw, ssq_ld)

def test_periodogram():
    arsig,_,_ = ut.ar_generator(N=512)
    avg_pwr = (arsig*arsig.conjugate()).mean()
    f, psd = tsa.periodogram(arsig, N=2048)
    # for efficiency, let's leave out the 2PI in the numerator and denominator
    # for the following integral
    dw = 1./2048
    avg_pwr_est = np.trapz(psd, dx=dw)
    npt.assert_almost_equal(avg_pwr, avg_pwr_est, decimal=1)

def test_LD_AR():
    arsig,_,_ = ut.ar_generator(N=512)
    avg_pwr = (arsig*arsig.conjugate()).mean()
    w, psd = tsa.LD_AR_est(arsig, 8, 1024)
    # for efficiency, let's leave out the 2PI in the numerator and denominator
    # for the following integral
    dw = 1./1024
    avg_pwr_est = np.trapz(psd, dx=dw)
    npt.assert_almost_equal(avg_pwr, avg_pwr_est, decimal=0)

def test_periodogram():
    """Test some of the inputs to periodogram """

    arsig, _, _ = utils.ar_generator(N=1024)
    Sk = np.fft.fft(arsig)

    f1, c1 = tsa.periodogram(arsig)
    f2, c2 = tsa.periodogram(arsig, Sk=Sk)

    npt.assert_equal(c1, c2)

    # Check that providing a complex signal does the right thing
    # (i.e. two-sided spectrum): 
    N = 1024 
    r, _, _ = utils.ar_generator(N=N)
    c, _, _ = utils.ar_generator(N=N)
    arsig = r + c * scipy.sqrt(-1)

    f, c = tsa.periodogram(arsig)
    npt.assert_equal(f.shape[0], N) # Should be N, not the one-sided N/2 + 1

def test_periodogram_csd():
    """Test corner cases of  periodogram_csd"""

    arsig1, _, _ = utils.ar_generator(N=1024)
    arsig2, _, _ = utils.ar_generator(N=1024)

    tseries = np.vstack([arsig1, arsig2])
    
    Sk = np.fft.fft(tseries)

    f1, c1 = tsa.periodogram_csd(tseries)
    f2, c2 = tsa.periodogram_csd(tseries, Sk=Sk)
    npt.assert_equal(c1, c2)

    # Check that providing a complex signal does the right thing
    # (i.e. two-sided spectrum): 
    N = 1024 
    r, _, _ = utils.ar_generator(N=N)
    c, _, _ = utils.ar_generator(N=N)
    arsig1 = r + c * scipy.sqrt(-1)

    r, _, _ = utils.ar_generator(N=N)
    c, _, _ = utils.ar_generator(N=N)
    arsig2 = r + c * scipy.sqrt(-1)

    tseries = np.vstack([arsig1, arsig2])

    f, c = tsa.periodogram_csd(tseries)
    npt.assert_equal(f.shape[0], N) # Should be N, not the one-sided N/2 + 1

def test_mtm_cross_spectrum():
    """
    
    Test the multi-taper cross-spectral estimation. Based on the example in
    doc/examples/multi_taper_coh.py

    """ 
    NW = 4
    K = 2 * NW - 1

    N = 2 ** 10
    n_reps = 10
    n_freqs = N

    tapers, eigs = tsa.dpss_windows(N, NW, 2 * NW - 1)

    est_psd = []
    for k in xrange(n_reps):
        data,nz,alpha = utils.ar_generator(N=N)
        fgrid, hz = tsa.freq_response(1.0, a=np.r_[1, -alpha], n_freqs=n_freqs)
        # 'one-sided', so multiply by 2:
        psd = 2 * (hz * hz.conj()).real

        tdata = tapers * data

        tspectra = np.fft.fft(tdata)

        L = N / 2 + 1
        sides = 'onesided'
        w, _ = utils.adaptive_weights(tspectra, eigs, sides=sides)

        sxx = tsa.mtm_cross_spectrum(tspectra, tspectra, w, sides=sides)
        est_psd.append(sxx)

    fxx = np.mean(est_psd, 0)

    psd_ratio = np.mean(fxx / psd)

    # This is a rather lenient test, making sure that the average ratio is 1 to
    # within an order of magnitude. That is, that they are equal on average:
    npt.assert_array_almost_equal(psd_ratio, 1, decimal=1)

    # Test raising of error in case the inputs don't make sense:
    npt.assert_raises(ValueError,
                      tsa.mtm_cross_spectrum,
                      tspectra,np.r_[tspectra, tspectra],
                      (w, w))

def test_multi_taper_psd_csd():
    """

    Test the multi taper psd and csd estimation functions.
    Based on the example in
    doc/examples/multi_taper_spectral_estimation.py

    """

    N = 2 ** 10
    n_reps = 10

    psd = []
    est_psd = []
    est_csd = []
    for jk in [True, False]:
        for k in range(n_reps):
            for adaptive in [True, False]:
                ar_seq, nz, alpha = utils.ar_generator(N=N, drop_transients=10)
                ar_seq -= ar_seq.mean()
                fgrid, hz = tsa.freq_response(1.0, a=np.r_[1, -alpha],
                                              n_freqs=N)
                psd.append(2 * (hz * hz.conj()).real)
                f, psd_mt, nu = tsa.multi_taper_psd(ar_seq, adaptive=adaptive,
                                                    jackknife=jk)
                est_psd.append(psd_mt)
                f, csd_mt = tsa.multi_taper_csd(np.vstack([ar_seq, ar_seq]),
                                               adaptive=adaptive)
                # Symmetrical in this case, so take one element out:
                est_csd.append(csd_mt[0][1])

        fxx = np.mean(psd, axis=0)
        fxx_est1 = np.mean(est_psd, axis=0)
        fxx_est2 = np.mean(est_csd, axis=0)

        # Tests the psd:
        psd_ratio1 = np.mean(fxx_est1 / fxx)
        npt.assert_array_almost_equal(psd_ratio1, 1, decimal=-1)
        # Tests the csd:
        psd_ratio2 = np.mean(fxx_est2 / fxx)
        npt.assert_array_almost_equal(psd_ratio2, 1, decimal=-1)

def test_AR_YW():
    arsig,_,_ = utils.ar_generator(N=512)
    avg_pwr = (arsig*arsig.conjugate()).mean()
    order = 8
    ak,sigma_v = tsa.AR_est_YW(arsig, order) 
    w, psd = tsa.AR_psd(ak, sigma_v)
    # the psd is a one-sided power spectral density, which has been
    # multiplied by 2 to preserve the property that
    # 1/2pi int_{-pi}^{pi} Sxx(w) dw = Rxx(0)

    # evaluate this integral numerically from 0 to pi
    dw = np.pi / len(psd)
    avg_pwr_est = np.trapz(psd, dx=dw) / (2*np.pi)
    # consistency on the order of 10**0 is pretty good for this test
    npt.assert_almost_equal(avg_pwr, avg_pwr_est, decimal=0)

    # Test for providing the autocovariance as an input:
    ak,sigma_v = tsa.AR_est_YW(arsig, order, utils.autocov(arsig))
    w, psd = tsa.AR_psd(ak, sigma_v)
    avg_pwr_est = np.trapz(psd, dx=dw) / (2*np.pi)
    npt.assert_almost_equal(avg_pwr, avg_pwr_est, decimal=0)

"""

In this case, we generate an order 2 AR process, with the following coefficients:


"""

coefs = np.array([0.9, -0.5])

"""

This generates the AR(2) time series:

"""

X, noise, _ = utils.ar_generator(npts, sigma, coefs, drop_transients)

ts_x = TimeSeries(X, sampling_rate=Fs, time_unit='s')
ts_noise = TimeSeries(noise, sampling_rate=1000, time_unit='s')

"""

We use the plot_tseries function in order to visualize the process:

"""

fig01 = plot_tseries(ts_x,label='AR signal')
fig01 = plot_tseries(ts_noise,fig=fig01,label='Noise')
fig01.axes[0].legend()

"""

"""Simple example AR(p) fitting."""

import numpy as np
from matplotlib import pyplot as plt

from nitime import utils
from nitime import algorithms as alg
reload(utils)

npts = 2048*10
sigma = 1
drop_transients = 1024
coefs = np.array([0.9, -0.5])

# Generate AR(2) time series
X, v, _ = utils.ar_generator(npts, sigma, coefs, drop_transients)

# Visualize
plt.figure()
plt.plot(v)
plt.title('noise')
plt.figure()
plt.plot(X)
plt.title('AR signal')

# Estimate the model parameters
sigma_est, coefs_est = alg.yule_AR_est(X, 2, 2*npts, system=True)

print 'coefs    :', coefs
print 'coefs est:', coefs_est

"""

import numpy as np
import matplotlib.pyplot as plt

import nitime.utils as utils
import nitime.timeseries as ts
import nitime.viz as viz

"""

For this example, we generate an auto-regressive sequence to be the signal:

"""

ar_seq, nz, alpha = utils.ar_generator(N=128, drop_transients=10)
ar_seq -= ar_seq.mean()

"""

The signal will be repeated several times, adding noise to the signal in each
repetition:

"""

n_trials = 12

fig_snr = []
sample = []
fig_tseries = []

tspectra = np.fft.fft(tdata)
mag_sqr_spectra = np.abs(tspectra)
np.power(mag_sqr_spectra, 2, mag_sqr_spectra)

rms_err_fmri, _, _ = generate_subset_weights_and_err(
    mag_sqr_spectra[10], l
    )
w_err_fmri = compare_weight_methods(mag_sqr_spectra[10], l)
plot_err(rms_err_fmri, 'FMRI re-weight err')
plot_err(w_err_fmri, 'FMRI weights err')
pp.show()

# ---- AUTOREGRESSIVE SEQUENCE

ar_seq, nz, alpha = utils.ar_generator(N=n_samples, drop_transients=10)
ar_seq -= ar_seq.mean()

ar_spectra = np.abs(np.fft.fft(tapers * ar_seq))
np.power(ar_spectra, 2, ar_spectra)

rms_err_arseq, _, _ = generate_subset_weights_and_err(
    ar_spectra, l
    )
w_err_arseq = compare_weight_methods(ar_spectra, l)
plot_err(rms_err_arseq, 'AR re-weight err')
plot_err(w_err_arseq, 'AR weights err')
pp.show()