def test_io_trans():
    """Test reading and writing of trans files
    """
    tempdir = _TempDir()
    os.mkdir(op.join(tempdir, "sample"))
    assert_raises(RuntimeError, _find_trans, "sample", subjects_dir=tempdir)
    trans0 = read_trans(fname)
    fname1 = op.join(tempdir, "sample", "test-trans.fif")
    write_trans(fname1, trans0)
    assert_true(fname1 == _find_trans("sample", subjects_dir=tempdir))
    trans1 = read_trans(fname1)

    # check all properties
    assert_true(trans0["from"] == trans1["from"])
    assert_true(trans0["to"] == trans1["to"])
    assert_array_equal(trans0["trans"], trans1["trans"])

    # check reading non -trans.fif files
    assert_raises(IOError, read_trans, fname_eve)

    # check warning on bad filenames
    with warnings.catch_warnings(record=True) as w:
        fname2 = op.join(tempdir, "trans-test-bad-name.fif")
        write_trans(fname2, trans0)
    assert_true(len(w) >= 1)

def test_io_trans():
    """Test reading and writing of trans files
    """
    tempdir = _TempDir()
    os.mkdir(op.join(tempdir, 'sample'))
    assert_raises(RuntimeError, _find_trans, 'sample', subjects_dir=tempdir)
    trans0 = read_trans(fname)
    fname1 = op.join(tempdir, 'sample', 'test-trans.fif')
    trans0.save(fname1)
    assert_true(fname1 == _find_trans('sample', subjects_dir=tempdir))
    trans1 = read_trans(fname1)

    # check all properties
    assert_true(trans0['from'] == trans1['from'])
    assert_true(trans0['to'] == trans1['to'])
    assert_array_equal(trans0['trans'], trans1['trans'])

    # check reading non -trans.fif files
    assert_raises(IOError, read_trans, fname_eve)

    # check warning on bad filenames
    with warnings.catch_warnings(record=True) as w:
        fname2 = op.join(tempdir, 'trans-test-bad-name.fif')
        write_trans(fname2, trans0)
    assert_naming(w, 'test_transforms.py', 1)

    def compute_forward_and_inverse_solutions(self, orientation = 'fixed'):
        """docstring for compute_forward_solution"""

        info = self.grand_average_evoked.info
        trans = mne.read_trans(op.join(self.processed_files, '%s-trans.fif' %self.subject))
        src = glob.glob(op.join(self.subjects_dir, self.subject, 'bem', '*-ico-4-src.fif'))[0]
        bem = glob.glob(op.join(self.subjects_dir, self.subject, 'bem', '*-bem-sol.fif'))[0]
        fname = op.join(self.processed_files, '%s_forward.fif' %self.subject)

        # check if fwd exists, if not, make it
        if not op.exists(fname):
            fwd = mne.make_forward_solution(info = info, trans = trans, src = src,
                                            bem = bem, fname = fname, meg = True, eeg = False,
                                            overwrite = True, ignore_ref = True)

            self.add_preprocessing_notes("Forward solution generated and saved to %s" %fname)

        if orientation == 'fixed':
            force_fixed = True
        else:
            force_fixed = False

        fwd = mne.read_forward_solution(fname,force_fixed=force_fixed)

        self.forward_solution = fwd

        inv = mne.minimum_norm.make_inverse_operator(info, self.forward_solution, self.cov_reg, loose = None, depth = None, fixed = force_fixed)
        self.inverse_solution = inv
        mne.minimum_norm.write_inverse_operator(op.join(self.processed_files, '%s_inv.fif' %self.subject), self.inverse_solution)

        self.add_preprocessing_notes("Inverse solution generated and saved to %s" %op.join(self.processed_files, '%s_inv.fif' %self.subject))
        return fwd, inv

def test_plot_dipole_locations():
    """Test plotting dipole locations."""
    dipoles = read_dipole(dip_fname)
    trans = read_trans(trans_fname)
    dipoles.plot_locations(trans, 'sample', subjects_dir, fig_name='foo')
    assert_raises(ValueError, dipoles.plot_locations, trans, 'sample',
                  subjects_dir, mode='foo')

def test_do_forward_solution():
    """Test wrapping forward solution from python
    """
    mri = read_trans(fname_mri)
    fname_fake = op.join(temp_dir, 'no_have.fif')

    # ## Error checks
    # bad subject
    assert_raises(ValueError, do_forward_solution, 1, fname_raw,
                  subjects_dir=subjects_dir)
    # bad meas
    assert_raises(ValueError, do_forward_solution, 'sample', 1,
                  subjects_dir=subjects_dir)
    # meas doesn't exist
    assert_raises(IOError, do_forward_solution, 'sample', fname_fake,
                  subjects_dir=subjects_dir)
    # don't specify trans and meas
    assert_raises(ValueError, do_forward_solution, 'sample', fname_raw,
                  subjects_dir=subjects_dir)
    # specify both trans and meas
    assert_raises(ValueError, do_forward_solution, 'sample', fname_raw,
                  trans='me', mri='you', subjects_dir=subjects_dir)
    # specify non-existent trans
    assert_raises(IOError, do_forward_solution, 'sample', fname_raw,
                  trans=fname_fake, subjects_dir=subjects_dir)
    # specify non-existent mri
    assert_raises(IOError, do_forward_solution, 'sample', fname_raw,
                  mri=fname_fake, subjects_dir=subjects_dir)
    # specify non-string mri
    assert_raises(ValueError, do_forward_solution, 'sample', fname_raw,
                  mri=1, subjects_dir=subjects_dir)
    # specify non-string trans
    assert_raises(ValueError, do_forward_solution, 'sample', fname_raw,
                  trans=1, subjects_dir=subjects_dir)
    # test specifying an actual trans in python space -- this should work but
    # the transform I/O reduces our accuracy -- so we'll just hack a test here
    # by making it bomb with eeg=False and meg=False
    assert_raises(ValueError, do_forward_solution, 'sample', fname_raw,
                  mri=mri, eeg=False, meg=False, subjects_dir=subjects_dir)
    # mindist as non-integer
    assert_raises(TypeError, do_forward_solution, 'sample', fname_raw,
                  mri=fname_mri, mindist=dict(), subjects_dir=subjects_dir)
    # mindist as string but not 'all'
    assert_raises(ValueError, do_forward_solution, 'sample', fname_raw,
                  mri=fname_mri, eeg=False, mindist='yall',
                  subjects_dir=subjects_dir)
    # src, spacing, and bem as non-str
    assert_raises(ValueError, do_forward_solution, 'sample', fname_raw,
                  mri=fname_mri, src=1, subjects_dir=subjects_dir)
    assert_raises(ValueError, do_forward_solution, 'sample', fname_raw,
                  mri=fname_mri, spacing=1, subjects_dir=subjects_dir)
    assert_raises(ValueError, do_forward_solution, 'sample', fname_raw,
                  mri=fname_mri, bem=1, subjects_dir=subjects_dir)
    # no overwrite flag
    assert_raises(IOError, do_forward_solution, 'sample', fname_raw,
                  existing_file, mri=fname_mri, subjects_dir=subjects_dir)
    # let's catch an MNE error, this time about trans being wrong
    assert_raises(CalledProcessError, do_forward_solution, 'sample',
                  fname_raw, existing_file, trans=fname_mri, overwrite=True,
                  spacing='oct6', subjects_dir=subjects_dir)

def test_io_trans():
    """Test reading and writing of trans files
    """
    info0 = read_trans(fname)
    fname1 = op.join(tempdir, 'test-trans.fif')
    write_trans(fname1, info0)
    info1 = read_trans(fname1)

    # check all properties
    assert_true(info0['from'] == info1['from'])
    assert_true(info0['to'] == info1['to'])
    assert_array_equal(info0['trans'], info1['trans'])
    for d0, d1 in zip(info0['dig'], info1['dig']):
        assert_array_equal(d0['r'], d1['r'])
        for name in ['kind', 'ident', 'coord_frame']:
            assert_true(d0[name] == d1[name])

def plot_coregistration(subject, subjects_dir, hcp_path, recordings_path,
                        info_from=(('data_type', 'rest'), ('run_index', 0)),
                        view_init=(('azim', 0), ('elev', 0))):
    """A diagnostic plot to show the HCP coregistration

    Parameters
    ----------
    subject : str
        The subject
    subjects_dir : str
        The path corresponding to MNE/freesurfer SUBJECTS_DIR (to be created)
    hcp_path : str
        The path where the HCP files can be found.
    recordings_path : str
        The path to converted data (including the head<->device transform).
    info_from : tuple of tuples | dict
        The reader info concerning the data from which sensor positions
        should be read.
        Must not be empty room as sensor positions are in head
        coordinates for 4D systems, hence not available in that case.
        Note that differences between the sensor positions across runs
        are smaller than 12 digits, hence negligible.
    view_init : tuple of tuples | dict
        The initival view, defaults to azimuth and elevation of 0,
        a simple lateral view

    Returns
    -------
    fig : matplotlib.figure.Figure
        The figure object.
    """
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D  #  noqa

    if isinstance(info_from, tuple):
        info_from = dict(info_from)
    if isinstance(view_init, tuple):
        view_init = dict(view_init)

    head_mri_t = read_trans(
        op.join(recordings_path, subject,
                '{}-head_mri-trans.fif'.format(subject)))

    info = read_info(subject=subject, hcp_path=hcp_path, **info_from)

    info = pick_info(info, _pick_data_channels(info, with_ref_meg=False))
    sens_pnts = np.array([c['loc'][:3] for c in info['chs']])
    sens_pnts = apply_trans(head_mri_t, sens_pnts)
    sens_pnts *= 1e3  # put in mm scale

    pnts, tris = read_surface(
        op.join(subjects_dir, subject, 'bem', 'inner_skull.surf'))

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(*sens_pnts.T, color='purple', marker='o')
    ax.scatter(*pnts.T, color='green', alpha=0.3)
    ax.view_init(**view_init)
    fig.tight_layout()
    return fig

def test_get_trans():
    """Test converting '-trans.txt' to '-trans.fif'"""
    trans = read_trans(fname)
    trans = invert_transform(trans)  # starts out as head->MRI, so invert
    trans_2 = _get_trans(fname_trans)[0]
    assert_equal(trans['from'], trans_2['from'])
    assert_equal(trans['to'], trans_2['to'])
    assert_allclose(trans['trans'], trans_2['trans'], rtol=1e-5, atol=1e-5)

def test_get_mri_head_t():
    """Test converting '-trans.txt' to '-trans.fif'"""
    trans = read_trans(fname)
    trans = invert_transform(trans)  # starts out as head->MRI, so invert
    trans_2 = _get_mri_head_t_from_trans_file(fname_trans)
    assert_equal(trans["from"], trans_2["from"])
    assert_equal(trans["to"], trans_2["to"])
    assert_allclose(trans["trans"], trans_2["trans"], rtol=1e-5, atol=1e-5)

def test_combine():
    """Test combining transforms
    """
    trans = read_trans(fname)
    inv = invert_transform(trans)
    combine_transforms(trans, inv, trans["from"], trans["from"])
    assert_raises(RuntimeError, combine_transforms, trans, inv, trans["to"], trans["from"])
    assert_raises(RuntimeError, combine_transforms, trans, inv, trans["from"], trans["to"])
    assert_raises(RuntimeError, combine_transforms, trans, trans, trans["from"], trans["to"])

def test_eeg_field_interpolation():
    """Test interpolation of EEG field onto head
    """
    trans = read_trans(trans_fname)
    info = read_info(evoked_fname)
    surf = get_head_surface('sample', subjects_dir=subjects_dir)
    # we must have trans if surface is in MRI coords
    assert_raises(ValueError, make_surface_mapping, info, surf, 'eeg')
    data = make_surface_mapping(info, surf, 'eeg', trans, mode='accurate')
    assert_array_equal(data.shape, (2562, 60))  # maps data onto surf

def test_combine():
    """Test combining transforms."""
    trans = read_trans(fname)
    inv = invert_transform(trans)
    combine_transforms(trans, inv, trans['from'], trans['from'])
    pytest.raises(RuntimeError, combine_transforms, trans, inv,
                  trans['to'], trans['from'])
    pytest.raises(RuntimeError, combine_transforms, trans, inv,
                  trans['from'], trans['to'])
    pytest.raises(RuntimeError, combine_transforms, trans, trans,
                  trans['from'], trans['to'])

def test_io_trans():
    """Test reading and writing of trans files."""
    tempdir = _TempDir()
    os.mkdir(op.join(tempdir, 'sample'))
    pytest.raises(RuntimeError, _find_trans, 'sample', subjects_dir=tempdir)
    trans0 = read_trans(fname)
    fname1 = op.join(tempdir, 'sample', 'test-trans.fif')
    trans0.save(fname1)
    assert fname1 == _find_trans('sample', subjects_dir=tempdir)
    trans1 = read_trans(fname1)

    # check all properties
    assert trans0 == trans1

    # check reading non -trans.fif files
    pytest.raises(IOError, read_trans, fname_eve)

    # check warning on bad filenames
    fname2 = op.join(tempdir, 'trans-test-bad-name.fif')
    with pytest.warns(RuntimeWarning, match='-trans.fif'):
        write_trans(fname2, trans0)

def test_io_trans():
    """Test reading and writing of trans files
    """
    trans0 = read_trans(fname)
    fname1 = op.join(tempdir, 'test-trans.fif')
    write_trans(fname1, trans0)
    trans1 = read_trans(fname1)

    # check all properties
    assert_true(trans0['from'] == trans1['from'])
    assert_true(trans0['to'] == trans1['to'])
    assert_array_equal(trans0['trans'], trans1['trans'])

    # check reading non -trans.fif files
    assert_raises(IOError, read_trans, fname_eve)
    
    # check warning on bad filenames
    with warnings.catch_warnings(record=True) as w:
        fname2 = op.join(tempdir, 'trans-test-bad-name.fif')
        write_trans(fname2, trans0)
    assert_true(len(w) >= 1)

def _get_data():
    """Helper to get some starting data"""
    # raw with ECG channel
    raw = Raw(raw_fname).crop(0.0, 5.0).load_data()
    data_picks = pick_types(raw.info, meg=True, eeg=True)
    other_picks = pick_types(raw.info, meg=False, stim=True, eog=True)
    picks = np.sort(np.concatenate((data_picks[::16], other_picks)))
    raw = raw.pick_channels([raw.ch_names[p] for p in picks])
    ecg = RawArray(np.zeros((1, len(raw.times))), create_info(["ECG 063"], raw.info["sfreq"], "ecg"))
    for key in ("dev_head_t", "buffer_size_sec", "highpass", "lowpass", "filename", "dig"):
        ecg.info[key] = raw.info[key]
    raw.add_channels([ecg])

    src = read_source_spaces(src_fname)
    trans = read_trans(trans_fname)
    sphere = make_sphere_model("auto", "auto", raw.info)
    stc = _make_stc(raw, src)
    return raw, src, stc, trans, sphere

def test_plot_dipole_mri_orthoview():
    """Test mpl dipole plotting."""
    dipoles = read_dipole(dip_fname)
    trans = read_trans(trans_fname)
    for coord_frame, idx, show_all in zip(['head', 'mri'],
                                          ['gof', 'amplitude'], [True, False]):
        fig = dipoles.plot_locations(trans, 'sample', subjects_dir,
                                     coord_frame=coord_frame, idx=idx,
                                     show_all=show_all, mode='orthoview')
        fig.canvas.scroll_event(0.5, 0.5, 1)  # scroll up
        fig.canvas.scroll_event(0.5, 0.5, -1)  # scroll down
        fig.canvas.key_press_event('up')
        fig.canvas.key_press_event('down')
        fig.canvas.key_press_event('a')  # some other key
    ax = plt.subplot(111)
    pytest.raises(TypeError, dipoles.plot_locations, trans, 'sample',
                  subjects_dir, ax=ax)
    plt.close('all')

def _get_data():
    """Helper to get some starting data."""
    # raw with ECG channel
    raw = read_raw_fif(raw_fname).crop(0., 5.0).load_data()
    data_picks = pick_types(raw.info, meg=True, eeg=True)
    other_picks = pick_types(raw.info, meg=False, stim=True, eog=True)
    picks = np.sort(np.concatenate((data_picks[::16], other_picks)))
    raw = raw.pick_channels([raw.ch_names[p] for p in picks])
    raw.info.normalize_proj()
    ecg = RawArray(np.zeros((1, len(raw.times))),
                   create_info(['ECG 063'], raw.info['sfreq'], 'ecg'))
    for key in ('dev_head_t', 'buffer_size_sec', 'highpass', 'lowpass', 'dig'):
        ecg.info[key] = raw.info[key]
    raw.add_channels([ecg])

    src = read_source_spaces(src_fname)
    trans = read_trans(trans_fname)
    sphere = make_sphere_model('auto', 'auto', raw.info)
    stc = _make_stc(raw, src)
    return raw, src, stc, trans, sphere

def test_head_to_mni():
    """Test conversion of aseg vertices to MNI coordinates."""
    # obtained using freeview
    coords = np.array([[22.52, 11.24, 17.72], [22.52, 5.46, 21.58],
                       [16.10, 5.46, 22.23], [21.24, 8.36, 22.23]])

    xfm = _read_talxfm('sample', subjects_dir)
    coords_MNI = apply_trans(xfm['trans'], coords)

    trans = read_trans(trans_fname)  # head->MRI (surface RAS)
    mri_head_t = invert_transform(trans)  # MRI (surface RAS)->head matrix

    # obtained from sample_audvis-meg-oct-6-mixed-fwd.fif
    coo_right_amygdala = np.array([[0.01745682,  0.02665809,  0.03281873],
                                   [0.01014125,  0.02496262,  0.04233755],
                                   [0.01713642,  0.02505193,  0.04258181],
                                   [0.01720631,  0.03073877,  0.03850075]])
    coords_MNI_2 = head_to_mni(coo_right_amygdala, 'sample', mri_head_t,
                               subjects_dir)
    # less than 1mm error
    assert_allclose(coords_MNI, coords_MNI_2, atol=10.0)

def test_make_field_map_eeg():
    """Test interpolation of EEG field onto head
    """
    trans = read_trans(trans_fname)
    evoked = read_evoked(evoked_fname, setno='Left Auditory')
    evoked.info['bads'] = ['MEG 2443', 'EEG 053']  # add some bads
    surf = get_head_surf('sample', subjects_dir=subjects_dir)
    # we must have trans if surface is in MRI coords
    assert_raises(ValueError, _make_surface_mapping, evoked.info, surf, 'eeg')

    evoked = pick_types_evoked(evoked, meg=False, eeg=True)
    fmd = make_field_map(evoked, trans_fname=trans_fname,
                         subject='sample', subjects_dir=subjects_dir)

    # trans is necessary for EEG only
    assert_raises(RuntimeError, make_field_map, evoked, trans_fname=None,
                  subject='sample', subjects_dir=subjects_dir)

    fmd = make_field_map(evoked, trans_fname=trans_fname,
                         subject='sample', subjects_dir=subjects_dir)
    assert_true(len(fmd) == 1)
    assert_array_equal(fmd[0]['data'].shape, (2562, 59))  # maps data onto surf
    assert_true(len(fmd[0]['ch_names']), 59)

Plotting EEG sensors on the scalp
=================================

In this example, digitized EEG sensor locations are shown on the scalp.
"""
# Author: Eric Larson <[email protected]>
#
# License: BSD (3-clause)

import mne
from mne.viz import plot_trans
from mayavi import mlab

print(__doc__)

data_path = mne.datasets.sample.data_path()
subjects_dir = data_path + "/subjects"
trans = mne.read_trans(data_path + "/MEG/sample/sample_audvis_raw-trans.fif")
raw = mne.io.read_raw_fif(data_path + "/MEG/sample/sample_audvis_raw.fif")
fig = plot_trans(
    raw.info,
    trans,
    subject="sample",
    dig=False,
    eeg_sensors=["original", "projected"],
    meg_sensors=[],
    coord_frame="head",
    subjects_dir=subjects_dir,
)
mlab.view(135, 80)