def test_plot_sparse_source_estimates():
    """Test plotting of (sparse) source estimates
    """
    # dense version
    vertices = [s['vertno'] for s in sample_src]
    n_time = 5
    n_verts = sum(len(v) for v in vertices)
    stc_data = np.zeros((n_verts * n_time))
    stc_data[(np.random.rand(20) * n_verts * n_time).astype(int)] = 1
    stc_data.shape = (n_verts, n_time)
    stc = SourceEstimate(stc_data, vertices, 1, 1)
    colormap = mne_analyze_colormap(format='matplotlib')
    # don't really need to test matplotlib method since it's not used now...
    colormap = mne_analyze_colormap()
    plot_source_estimates(stc, 'sample', colormap=colormap)

    # now do sparse version
    vertices = sample_src[0]['vertno']
    n_verts = len(vertices)
    stc_data = np.zeros((n_verts * n_time))
    stc_data[(np.random.rand(20) * n_verts * n_time).astype(int)] = 1
    stc_data.shape = (n_verts, n_time)
    inds = np.where(np.any(stc_data, axis=1))[0]
    stc_data = stc_data[inds]
    vertices = vertices[inds]
    stc = SourceEstimate(stc_data, vertices, 1, 1)
    plot_sparse_source_estimates(sample_src, stc, bgcolor=(1, 1, 1),
                                 opacity=0.5, high_resolution=True)

def test_plot_sparse_source_estimates():
    """Test plotting of (sparse) source estimates
    """
    sample_src = read_source_spaces(op.join(data_dir, 'subjects', 'sample',
                                            'bem', 'sample-oct-6-src.fif'))

    # dense version
    vertices = [s['vertno'] for s in sample_src]
    n_time = 5
    n_verts = sum(len(v) for v in vertices)
    stc_data = np.zeros((n_verts * n_time))
    stc_data[(np.random.rand(20) * n_verts * n_time).astype(int)] = 1
    stc_data.shape = (n_verts, n_time)
    stc = SourceEstimate(stc_data, vertices, 1, 1)
    colormap = mne_analyze_colormap(format='matplotlib')
    # don't really need to test matplotlib method since it's not used now...
    colormap = mne_analyze_colormap()
    plot_source_estimates(stc, 'sample', colormap=colormap,
                          config_opts={'background': (1, 1, 0)},
                          subjects_dir=subjects_dir, colorbar=True)
    assert_raises(TypeError, plot_source_estimates, stc, 'sample',
                  figure='foo', hemi='both')

    # now do sparse version
    vertices = sample_src[0]['vertno']
    n_verts = len(vertices)
    stc_data = np.zeros((n_verts * n_time))
    stc_data[(np.random.rand(20) * n_verts * n_time).astype(int)] = 1
    stc_data.shape = (n_verts, n_time)
    inds = np.where(np.any(stc_data, axis=1))[0]
    stc_data = stc_data[inds]
    vertices = [vertices[inds], np.empty(0, dtype=np.int)]
    stc = SourceEstimate(stc_data, vertices, 1, 1)
    plot_sparse_source_estimates(sample_src, stc, bgcolor=(1, 1, 1),
                                 opacity=0.5, high_resolution=True)

def test_plot_sparse_source_estimates():
    """Test plotting of (sparse) source estimates."""
    sample_src = read_source_spaces(src_fname)

    # dense version
    vertices = [s['vertno'] for s in sample_src]
    n_time = 5
    n_verts = sum(len(v) for v in vertices)
    stc_data = np.zeros((n_verts * n_time))
    stc_size = stc_data.size
    stc_data[(np.random.rand(stc_size // 20) * stc_size).astype(int)] = \
        np.random.RandomState(0).rand(stc_data.size // 20)
    stc_data.shape = (n_verts, n_time)
    stc = SourceEstimate(stc_data, vertices, 1, 1)

    colormap = 'mne_analyze'
    plot_source_estimates(stc, 'sample', colormap=colormap,
                          background=(1, 1, 0),
                          subjects_dir=subjects_dir, colorbar=True,
                          clim='auto')
    pytest.raises(TypeError, plot_source_estimates, stc, 'sample',
                  figure='foo', hemi='both', clim='auto',
                  subjects_dir=subjects_dir)

    # now do sparse version
    vertices = sample_src[0]['vertno']
    inds = [111, 333]
    stc_data = np.zeros((len(inds), n_time))
    stc_data[0, 1] = 1.
    stc_data[1, 4] = 2.
    vertices = [vertices[inds], np.empty(0, dtype=np.int)]
    stc = SourceEstimate(stc_data, vertices, 1, 1)
    plot_sparse_source_estimates(sample_src, stc, bgcolor=(1, 1, 1),
                                 opacity=0.5, high_resolution=False)

def test_plot_sparse_source_estimates():
    """Test plotting of (sparse) source estimates
    """
    # dense version
    vertices = [s["vertno"] for s in sample_src]
    n_time = 5
    n_verts = sum(len(v) for v in vertices)
    stc_data = np.zeros((n_verts * n_time))
    stc_data[(np.random.rand(20) * n_verts * n_time).astype(int)] = 1
    stc_data.shape = (n_verts, n_time)
    stc = SourceEstimate(stc_data, vertices, 1, 1)
    colormap = mne_analyze_colormap(format="matplotlib")
    # don't really need to test matplotlib method since it's not used now...
    colormap = mne_analyze_colormap()
    plot_source_estimates(
        stc, "sample", colormap=colormap, config_opts={"background": (1, 1, 0)}, subjects_dir=subjects_dir
    )
    assert_raises(RuntimeError, plot_source_estimates, stc, "sample", figure="foo", hemi="both")

    # now do sparse version
    vertices = sample_src[0]["vertno"]
    n_verts = len(vertices)
    stc_data = np.zeros((n_verts * n_time))
    stc_data[(np.random.rand(20) * n_verts * n_time).astype(int)] = 1
    stc_data.shape = (n_verts, n_time)
    inds = np.where(np.any(stc_data, axis=1))[0]
    stc_data = stc_data[inds]
    vertices = vertices[inds]
    stc = SourceEstimate(stc_data, vertices, 1, 1)
    plot_sparse_source_estimates(sample_src, stc, bgcolor=(1, 1, 1), opacity=0.5, high_resolution=True)

# Generate source time courses from 2 dipoles and the correspond evoked data

times = np.arange(300, dtype=np.float) / raw.info['sfreq'] - 0.1
rng = np.random.RandomState(42)


def data_fun(times):
    """Function to generate random source time courses"""
    return (1e-9 * np.sin(30. * times) *
            np.exp(- (times - 0.15 + 0.05 * rng.randn(1)) ** 2 / 0.01))

stc = simulate_sparse_stc(fwd['src'], n_dipoles=2, times=times,
                          random_state=42, labels=labels, data_fun=data_fun)

###############################################################################
# Generate noisy evoked data
picks = mne.pick_types(raw.info, meg=True, exclude='bads')
iir_filter = fit_iir_model_raw(raw, order=5, picks=picks, tmin=60, tmax=180)[1]
snr = 6.  # dB
evoked = simulate_evoked(fwd, stc, info, cov, snr, iir_filter=iir_filter)

###############################################################################
# Plot
plot_sparse_source_estimates(fwd['src'], stc, bgcolor=(1, 1, 1),
                             opacity=0.5, high_resolution=True)

plt.figure()
plt.psd(evoked.data[0])

evoked.plot()

import matplotlib.pyplot as plt
plt.figure()
ylim = dict(eeg=[-10, 10], grad=[-400, 400], mag=[-600, 600])
evoked.plot(ylim=ylim, proj=True)

###############################################################################
# Run solver
alpha = 70  # regularization parameter between 0 and 100 (100 is high)
loose, depth = 0.2, 0.9  # loose orientation & depth weighting

# Compute dSPM solution to be used as weights in MxNE
inverse_operator = make_inverse_operator(evoked.info, forward, cov,
                                         loose=None, depth=depth, fixed=True)
stc_dspm = apply_inverse(evoked, inverse_operator, lambda2=1. / 9.,
                         method='dSPM')

# Compute MxNE inverse solution
stc, residual = mixed_norm(evoked, forward, cov, alpha, loose=loose,
                           depth=depth, maxit=3000, tol=1e-4,
                           active_set_size=10, debias=True, weights=stc_dspm,
                           weights_min=8., return_residual=True)

plt.figure()
residual.plot(ylim=ylim, proj=True)

###############################################################################
# View in 2D and 3D ("glass" brain like 3D plot)
plot_sparse_source_estimates(forward['src'], stc, bgcolor=(1, 1, 1),
                             opacity=0.1, fig_name="MxNE (cond %s)" % setno)

# Read noise covariance matrix and regularize it
cov = mne.read_cov(cov_fname)
cov = mne.cov.regularize(cov, evoked.info)

# Run the Gamma-MAP method
alpha = 0.5
stc, residual = gamma_map(evoked, forward, cov, alpha, xyz_same_gamma=True,
                          return_residual=True)

# View in 2D and 3D ("glass" brain like 3D plot)

# Show the sources as spheres scaled by their strength
scale_factors = np.max(np.abs(stc.data), axis=1)
scale_factors = 0.5 * (1 + scale_factors / np.max(scale_factors))

plot_sparse_source_estimates(
    forward['src'], stc, bgcolor=(1, 1, 1),
    modes=['sphere'], opacity=0.1, scale_factors=(scale_factors, None),
    fig_name="Gamma-MAP")

# Show the evoked response and the residual for gradiometers
ylim = dict(grad=[-120, 120])
evoked.pick_types(meg='grad', exclude='bads')
evoked.plot(titles=dict(grad='Evoked Response Gradiometers'), ylim=ylim,
            proj=True)

residual.pick_types(meg='grad', exclude='bads')
residual.plot(titles=dict(grad='Residuals Gradiometers'), ylim=ylim,
              proj=True)

    X : array, (n_active_dipoles, n_times)
        The time series of the dipoles in the active set.
    active_set : array (n_dipoles)
        Array of bool. Entry j is True if dipole j is in the active set.
        We have ``X_full[active_set] == X`` where X_full is the full X matrix
        such that ``M = G X_full``.
    """
    K = linalg.solve(np.dot(G, G.T) + 1e15 * np.eye(G.shape[0]), G).T
    K /= np.linalg.norm(K, axis=1)[:, None]
    X = np.dot(K, M)

    indices = np.argsort(np.sum(X ** 2, axis=1))[-10:]
    active_set = np.zeros(G.shape[1], dtype=bool)
    for idx in indices:
        idx -= idx % n_orient
        active_set[idx:idx + n_orient] = True
    X = X[active_set]
    return X, active_set

###############################################################################
# Apply your custom solver

# loose, depth = 0.2, 0.8  # corresponds to loose orientation
loose, depth = 1., 0.  # corresponds to free orientation
stc = apply_solver(solver, evoked, forward, noise_cov, loose, depth)

###############################################################################
# View in 2D and 3D ("glass" brain like 3D plot)
plot_sparse_source_estimates(forward['src'], stc, bgcolor=(1, 1, 1),
                             opacity=0.1)

                              return_residual=True)

# Crop to remove edges
stc.crop(tmin=-0.05, tmax=0.3)
evoked.crop(tmin=-0.05, tmax=0.3)
residual.crop(tmin=-0.05, tmax=0.3)

ylim = dict(eeg=[-10, 10], grad=[-200, 250], mag=[-600, 600])
picks = mne.pick_types(evoked.info, meg='grad', exclude='bads')
evoked.plot(picks=picks, ylim=ylim, proj=True,
            titles=dict(grad='Evoked Response (grad)'))

picks = mne.pick_types(residual.info, meg='grad', exclude='bads')
residual.plot(picks=picks, ylim=ylim, proj=True,
              titles=dict(grad='Residual (grad)'))

###############################################################################
# View in 2D and 3D ("glass" brain like 3D plot)
plot_sparse_source_estimates(forward['src'], stc, bgcolor=(1, 1, 1),
                             opacity=0.1, fig_name="TF-MxNE (cond %s)"
                             % condition, modes=['sphere'], scale_factors=[1.])

time_label = 'TF-MxNE time=%0.2f ms'
brain = stc.plot('sample', 'inflated', 'rh', fmin=10e-9, fmid=15e-9,
                 fmax=20e-9, time_label=time_label, smoothing_steps=5,
                 subjects_dir=data_path + '/subjects')
brain.show_view('medial')
brain.set_data_time_index(120)
brain.add_label("V1", color="yellow", scalar_thresh=.5, borders=True)
brain.add_label("V2", color="red", scalar_thresh=.5, borders=True)

# if n_mxne_iter > 1 dSPM weighting can be avoided.

# Compute dSPM solution to be used as weights in MxNE
inverse_operator = make_inverse_operator(evoked.info, forward, cov,
                                         loose=None, depth=depth, fixed=True)
stc_dspm = apply_inverse(evoked, inverse_operator, lambda2=1. / 9.,
                         method='dSPM')

# Compute (ir)MxNE inverse solution
stc, residual = mixed_norm(evoked, forward, cov, alpha, loose=loose,
                           depth=depth, maxit=3000, tol=1e-4,
                           active_set_size=10, debias=True, weights=stc_dspm,
                           weights_min=8., n_mxne_iter=n_mxne_iter,
                           return_residual=True)
residual.plot(ylim=ylim, proj=True)

###############################################################################
# View in 2D and 3D ("glass" brain like 3D plot)
plot_sparse_source_estimates(forward['src'], stc, bgcolor=(1, 1, 1),
                             fig_name="MxNE (cond %s)" % condition,
                             opacity=0.1)

# and on the fsaverage brain after morphing
stc_fsaverage = stc.morph(subject_from='sample', subject_to='fsaverage',
                          grade=None, sparse=True, subjects_dir=subjects_dir)
src_fsaverage_fname = subjects_dir + '/fsaverage/bem/fsaverage-ico-5-src.fif'
src_fsaverage = mne.read_source_spaces(src_fsaverage_fname)

plot_sparse_source_estimates(src_fsaverage, stc_fsaverage, bgcolor=(1, 1, 1),
                             opacity=0.1)

###############################################################################
# Plot residual
ylim = dict(eeg=[-10, 10], grad=[-400, 400], mag=[-600, 600])
evoked.pick_types(meg=True, eeg=True, exclude='bads')
evoked.plot(ylim=ylim, proj=True, time_unit='s')
residual.pick_types(meg=True, eeg=True, exclude='bads')
residual.plot(ylim=ylim, proj=True, time_unit='s')

###############################################################################
# Generate stc from dipoles
stc = make_stc_from_dipoles(dipoles, forward['src'])

###############################################################################
# View in 2D and 3D ("glass" brain like 3D plot)
solver = "MxNE" if n_mxne_iter == 1 else "irMxNE"
plot_sparse_source_estimates(forward['src'], stc, bgcolor=(1, 1, 1),
                             fig_name="%s (cond %s)" % (solver, condition),
                             opacity=0.1)

###############################################################################
# Morph onto fsaverage brain and view
stc_fsaverage = stc.morph(subject_from='sample', subject_to='fsaverage',
                          grade=None, sparse=True, subjects_dir=subjects_dir)
src_fsaverage_fname = subjects_dir + '/fsaverage/bem/fsaverage-ico-5-src.fif'
src_fsaverage = mne.read_source_spaces(src_fsaverage_fname)

plot_sparse_source_estimates(src_fsaverage, stc_fsaverage, bgcolor=(1, 1, 1),
                             fig_name="Morphed %s (cond %s)" % (solver,
                             condition), opacity=0.1)

raw = raw.crop(tmax=evoked.times[-1])
raw.save("t2-raw.fif")

raw2 = mne.fiff.Raw("t2-raw.fif")
cov = mne.compute_raw_data_covariance(raw2)

dsim = stc.in_label(labels[0])
pl.plot(dsim.data.T)

weights = ve.calculate_weights(fwd, cov, reg=0)

d, _ = raw2[:306, :]
sol = np.dot(weights, d)
src = mne.SourceEstimate(sol, (fwd["src"][0]["vertno"], fwd["src"][1]["vertno"]), tmin, tstep)

"""
# plotting the localized source
const = 1*10**-18
good_data = np.nonzero(abs(src.data[:,150])>const)
lv = np.nonzero(abs(src.lh_data[:,150])>const)
rv = np.nonzero(abs(src.rh_data[:,150])>const)
src.data = np.squeeze(src.data[np.ix_(good_data[0]),:])
src.vertno = [np.array([src.lh_vertno[lv[0][0]]]), np.array([src.rh_vertno[rv[0][0]]])]

# plotting the simulated source
plot_sparse_source_estimates(fwd['src'], stc, bgcolor=(1, 1, 1), opacity=0.5, high_resolution=True)
"""

selected_voxels = ve.find_best_voxels(src, labels, bands)
num_trials = 5
num_samples = 100