from copy import deepcopy
import inspect
import os.path as op

from numpy.testing import assert_array_equal, assert_equal
import pytest
import numpy as np

from mne import (pick_channels_regexp, pick_types, Epochs,
                 read_forward_solution, rename_channels,
                 pick_info, pick_channels, create_info)
from mne import __file__ as _root_init_fname
from mne.io import (read_raw_fif, RawArray, read_raw_bti, read_raw_kit,
                    read_info)
from mne.io.pick import (channel_indices_by_type, channel_type,
                         pick_types_forward, _picks_by_type, _picks_to_idx,
                         get_channel_types, _DATA_CH_TYPES_SPLIT,
                         _contains_ch_type)
from mne.io.constants import FIFF
from mne.datasets import testing
from mne.utils import run_tests_if_main, catch_logging, assert_object_equal

io_dir = op.join(op.dirname(inspect.getfile(inspect.currentframe())), '..')
data_path = testing.data_path(download=False)
fname_meeg = op.join(data_path, 'MEG', 'sample',
                     'sample_audvis_trunc-meg-eeg-oct-4-fwd.fif')
fname_mc = op.join(data_path, 'SSS', 'test_move_anon_movecomp_raw_sss.fif')

base_dir = op.join(op.dirname(__file__), 'data')
ctf_fname = op.join(base_dir, 'test_ctf_raw.fif')


def _picks_by_type_old(info, meg_combined=False, ref_meg=False,
                       exclude='bads'):
    """Use the old, slower _picks_by_type code."""
    picks_list = []
    has = [_contains_ch_type(info, k) for k in _DATA_CH_TYPES_SPLIT]
    has = dict(zip(_DATA_CH_TYPES_SPLIT, has))
    if has['mag'] and (meg_combined is not True or not has['grad']):
        picks_list.append(
            ('mag', pick_types(info, meg='mag', eeg=False, stim=False,
                               ref_meg=ref_meg, exclude=exclude))
        )
    if has['grad'] and (meg_combined is not True or not has['mag']):
        picks_list.append(
            ('grad', pick_types(info, meg='grad', eeg=False, stim=False,
                                ref_meg=ref_meg, exclude=exclude))
        )
    if has['mag'] and has['grad'] and meg_combined is True:
        picks_list.append(
            ('meg', pick_types(info, meg=True, eeg=False, stim=False,
                               ref_meg=ref_meg, exclude=exclude))
        )
    for ch_type in _DATA_CH_TYPES_SPLIT:
        if ch_type in ['grad', 'mag']:  # exclude just MEG channels
            continue
        if has[ch_type]:
            picks_list.append(
                (ch_type, pick_types(info, meg=False, stim=False,
                                     ref_meg=ref_meg, exclude=exclude,
                                     **{ch_type: True}))
            )
    return picks_list


def _channel_type_old(info, idx):
    """Get channel type using old, slower scheme."""
    ch = info['chs'][idx]

    # iterate through all defined channel types until we find a match with ch
    for t, rules in get_channel_types().items():
        for key, vals in rules.items():  # all keys must match the values
            if ch.get(key, None) not in np.array(vals):
                break  # not channel type t, go to next iteration
        else:
            return t

    raise ValueError('Unknown channel type for {}'.format(ch["ch_name"]))


def _assert_channel_types(info):
    for k in range(info['nchan']):
        a, b = channel_type(info, k), _channel_type_old(info, k)
        assert a == b


def test_pick_refs():
    """Test picking of reference sensors."""
    infos = list()
    # KIT
    kit_dir = op.join(io_dir, 'kit', 'tests', 'data')
    sqd_path = op.join(kit_dir, 'test.sqd')
    mrk_path = op.join(kit_dir, 'test_mrk.sqd')
    elp_path = op.join(kit_dir, 'test_elp.txt')
    hsp_path = op.join(kit_dir, 'test_hsp.txt')
    raw_kit = read_raw_kit(sqd_path, mrk_path, elp_path, hsp_path)
    infos.append(raw_kit.info)
    # BTi
    bti_dir = op.join(io_dir, 'bti', 'tests', 'data')
    bti_pdf = op.join(bti_dir, 'test_pdf_linux')
    bti_config = op.join(bti_dir, 'test_config_linux')
    bti_hs = op.join(bti_dir, 'test_hs_linux')
    raw_bti = read_raw_bti(bti_pdf, bti_config, bti_hs, preload=False)
    infos.append(raw_bti.info)
    # CTF
    fname_ctf_raw = op.join(io_dir, 'tests', 'data', 'test_ctf_comp_raw.fif')
    raw_ctf = read_raw_fif(fname_ctf_raw)
    raw_ctf.apply_gradient_compensation(2)
    for info in infos:
        info['bads'] = []
        _assert_channel_types(info)
        pytest.raises(ValueError, pick_types, info, meg='foo')
        pytest.raises(ValueError, pick_types, info, ref_meg='foo')
        picks_meg_ref = pick_types(info, meg=True, ref_meg=True)
        picks_meg = pick_types(info, meg=True, ref_meg=False)
        picks_ref = pick_types(info, meg=False, ref_meg=True)
        assert_array_equal(picks_meg_ref,
                           np.sort(np.concatenate([picks_meg, picks_ref])))
        picks_grad = pick_types(info, meg='grad', ref_meg=False)
        picks_ref_grad = pick_types(info, meg=False, ref_meg='grad')
        picks_meg_ref_grad = pick_types(info, meg='grad', ref_meg='grad')
        assert_array_equal(picks_meg_ref_grad,
                           np.sort(np.concatenate([picks_grad,
                                                   picks_ref_grad])))
        picks_mag = pick_types(info, meg='mag', ref_meg=False)
        picks_ref_mag = pick_types(info, meg=False, ref_meg='mag')
        picks_meg_ref_mag = pick_types(info, meg='mag', ref_meg='mag')
        assert_array_equal(picks_meg_ref_mag,
                           np.sort(np.concatenate([picks_mag,
                                                   picks_ref_mag])))
        assert_array_equal(picks_meg,
                           np.sort(np.concatenate([picks_mag, picks_grad])))
        assert_array_equal(picks_ref,
                           np.sort(np.concatenate([picks_ref_mag,
                                                   picks_ref_grad])))
        assert_array_equal(picks_meg_ref, np.sort(np.concatenate(
            [picks_grad, picks_mag, picks_ref_grad, picks_ref_mag])))

        for pick in (picks_meg_ref, picks_meg, picks_ref,
                     picks_grad, picks_ref_grad, picks_meg_ref_grad,
                     picks_mag, picks_ref_mag, picks_meg_ref_mag):
            if len(pick) > 0:
                pick_info(info, pick)

    # test CTF expected failures directly
    info = raw_ctf.info
    info['bads'] = []
    picks_meg_ref = pick_types(info, meg=True, ref_meg=True)
    picks_meg = pick_types(info, meg=True, ref_meg=False)
    picks_ref = pick_types(info, meg=False, ref_meg=True)
    picks_mag = pick_types(info, meg='mag', ref_meg=False)
    picks_ref_mag = pick_types(info, meg=False, ref_meg='mag')
    picks_meg_ref_mag = pick_types(info, meg='mag', ref_meg='mag')
    for pick in (picks_meg_ref, picks_ref, picks_ref_mag, picks_meg_ref_mag):
        if len(pick) > 0:
            pick_info(info, pick)

    for pick in (picks_meg, picks_mag):
        if len(pick) > 0:
            with catch_logging() as log:
                pick_info(info, pick, verbose=True)
            assert ('Removing {} compensators'.format(len(info['comps']))
                    in log.getvalue())
    picks_ref_grad = pick_types(info, meg=False, ref_meg='grad')
    assert set(picks_ref_mag) == set(picks_ref)
    assert len(picks_ref_grad) == 0
    all_meg = np.arange(3, 306)
    assert_array_equal(np.concatenate([picks_ref, picks_meg]), all_meg)
    assert_array_equal(picks_meg_ref_mag, all_meg)


def test_pick_channels_regexp():
    """Test pick with regular expression."""
    ch_names = ['MEG 2331', 'MEG 2332', 'MEG 2333']
    assert_array_equal(pick_channels_regexp(ch_names, 'MEG ...1'), [0])
    assert_array_equal(pick_channels_regexp(ch_names, 'MEG ...[2-3]'), [1, 2])
    assert_array_equal(pick_channels_regexp(ch_names, 'MEG *'), [0, 1, 2])


def assert_indexing(info, picks_by_type, ref_meg=False, all_data=True):
    """Assert our indexing functions work properly."""
    # First that our old and new channel typing functions are equivalent
    _assert_channel_types(info)
    # Next that channel_indices_by_type works
    if not ref_meg:
        idx = channel_indices_by_type(info)
        for key in idx:
            for p in picks_by_type:
                if key == p[0]:
                    assert_array_equal(idx[key], p[1])
                    break
            else:
                assert len(idx[key]) == 0
    # Finally, picks_by_type (if relevant)
    if not all_data:
        picks_by_type = [p for p in picks_by_type
                         if p[0] in _DATA_CH_TYPES_SPLIT]
    picks_by_type = [(p[0], np.array(p[1], int)) for p in picks_by_type]
    actual = _picks_by_type(info, ref_meg=ref_meg)
    assert_object_equal(actual, picks_by_type)
    if not ref_meg and idx['hbo']:  # our old code had a bug
        with pytest.raises(TypeError, match='unexpected keyword argument'):
            _picks_by_type_old(info, ref_meg=ref_meg)
    else:
        old = _picks_by_type_old(info, ref_meg=ref_meg)
        assert_object_equal(old, picks_by_type)
    # test bads
    info = info.copy()
    info['bads'] = [info['chs'][picks_by_type[0][1][0]]['ch_name']]
    picks_by_type = deepcopy(picks_by_type)
    picks_by_type[0] = (picks_by_type[0][0], picks_by_type[0][1][1:])
    actual = _picks_by_type(info, ref_meg=ref_meg)
    assert_object_equal(actual, picks_by_type)


def test_pick_seeg_ecog():
    """Test picking with sEEG and ECoG."""
    names = 'A1 A2 Fz O OTp1 OTp2 E1 OTp3 E2 E3'.split()
    types = 'mag mag eeg eeg seeg seeg ecog seeg ecog ecog'.split()
    info = create_info(names, 1024., types)
    picks_by_type = [('mag', [0, 1]), ('eeg', [2, 3]),
                     ('seeg', [4, 5, 7]), ('ecog', [6, 8, 9])]
    assert_indexing(info, picks_by_type)
    assert_array_equal(pick_types(info, meg=False, seeg=True), [4, 5, 7])
    for i, t in enumerate(types):
        assert_equal(channel_type(info, i), types[i])
    raw = RawArray(np.zeros((len(names), 10)), info)
    events = np.array([[1, 0, 0], [2, 0, 0]])
    epochs = Epochs(raw, events, {'event': 0}, -1e-5, 1e-5)
    evoked = epochs.average(pick_types(epochs.info, meg=True, seeg=True))
    e_seeg = evoked.copy().pick_types(meg=False, seeg=True)
    for l, r in zip(e_seeg.ch_names, [names[4], names[5], names[7]]):
        assert_equal(l, r)
    # Deal with constant debacle
    raw = read_raw_fif(op.join(io_dir, 'tests', 'data',
                               'test_chpi_raw_sss.fif'))
    assert_equal(len(pick_types(raw.info, meg=False, seeg=True, ecog=True)), 0)


def test_pick_chpi():
    """Test picking cHPI."""
    # Make sure we don't mis-classify cHPI channels
    info = read_info(op.join(io_dir, 'tests', 'data', 'test_chpi_raw_sss.fif'))
    _assert_channel_types(info)
    channel_types = {channel_type(info, idx) for idx in range(info['nchan'])}
    assert 'chpi' in channel_types
    assert 'seeg' not in channel_types
    assert 'ecog' not in channel_types


def test_pick_bio():
    """Test picking BIO channels."""
    names = 'A1 A2 Fz O BIO1 BIO2 BIO3'.split()
    types = 'mag mag eeg eeg bio bio bio'.split()
    info = create_info(names, 1024., types)
    picks_by_type = [('mag', [0, 1]), ('eeg', [2, 3]), ('bio', [4, 5, 6])]
    assert_indexing(info, picks_by_type, all_data=False)


def test_pick_fnirs():
    """Test picking fNIRS channels."""
    names = 'A1 A2 Fz O hbo1 hbo2 hbr1'.split()
    types = 'mag mag eeg eeg hbo hbo hbr'.split()
    info = create_info(names, 1024., types)
    picks_by_type = [('mag', [0, 1]), ('eeg', [2, 3]),
                     ('hbo', [4, 5]), ('hbr', [6])]
    assert_indexing(info, picks_by_type)


def test_pick_ref():
    """Test picking ref_meg channels."""
    info = read_info(ctf_fname)
    picks_by_type = [('stim', [0]), ('eog', [306, 307]), ('ecg', [308]),
                     ('misc', [1]),
                     ('mag', np.arange(31, 306)),
                     ('ref_meg', np.arange(2, 31))]
    assert_indexing(info, picks_by_type, all_data=False)
    picks_by_type.append(('mag', np.concatenate([picks_by_type.pop(-1)[1],
                                                 picks_by_type.pop(-1)[1]])))
    assert_indexing(info, picks_by_type, ref_meg=True, all_data=False)


def _check_fwd_n_chan_consistent(fwd, n_expected):
    n_ok = len(fwd['info']['ch_names'])
    n_sol = fwd['sol']['data'].shape[0]
    assert_equal(n_expected, n_sol)
    assert_equal(n_expected, n_ok)


@testing.requires_testing_data
def test_pick_forward_seeg_ecog():
    """Test picking forward with SEEG and ECoG."""
    fwd = read_forward_solution(fname_meeg)
    counts = channel_indices_by_type(fwd['info'])
    for key in counts.keys():
        counts[key] = len(counts[key])
    counts['meg'] = counts['mag'] + counts['grad']
    fwd_ = pick_types_forward(fwd, meg=True)
    _check_fwd_n_chan_consistent(fwd_, counts['meg'])
    fwd_ = pick_types_forward(fwd, meg=False, eeg=True)
    _check_fwd_n_chan_consistent(fwd_, counts['eeg'])
    # should raise exception related to emptiness
    pytest.raises(ValueError, pick_types_forward, fwd, meg=False, seeg=True)
    pytest.raises(ValueError, pick_types_forward, fwd, meg=False, ecog=True)
    # change last chan from EEG to sEEG, second-to-last to ECoG
    ecog_name = 'E1'
    seeg_name = 'OTp1'
    rename_channels(fwd['info'], {'EEG 059': ecog_name})
    rename_channels(fwd['info'], {'EEG 060': seeg_name})
    for ch in fwd['info']['chs']:
        if ch['ch_name'] == seeg_name:
            ch['kind'] = FIFF.FIFFV_SEEG_CH
            ch['coil_type'] = FIFF.FIFFV_COIL_EEG
        elif ch['ch_name'] == ecog_name:
            ch['kind'] = FIFF.FIFFV_ECOG_CH
            ch['coil_type'] = FIFF.FIFFV_COIL_EEG
    fwd['sol']['row_names'][-1] = fwd['info']['chs'][-1]['ch_name']
    fwd['sol']['row_names'][-2] = fwd['info']['chs'][-2]['ch_name']
    counts['eeg'] -= 2
    counts['seeg'] += 1
    counts['ecog'] += 1
    # repick & check
    fwd_seeg = pick_types_forward(fwd, meg=False, seeg=True)
    assert_equal(fwd_seeg['sol']['row_names'], [seeg_name])
    assert_equal(fwd_seeg['info']['ch_names'], [seeg_name])
    # should work fine
    fwd_ = pick_types_forward(fwd, meg=True)
    _check_fwd_n_chan_consistent(fwd_, counts['meg'])
    fwd_ = pick_types_forward(fwd, meg=False, eeg=True)
    _check_fwd_n_chan_consistent(fwd_, counts['eeg'])
    fwd_ = pick_types_forward(fwd, meg=False, seeg=True)
    _check_fwd_n_chan_consistent(fwd_, counts['seeg'])
    fwd_ = pick_types_forward(fwd, meg=False, ecog=True)
    _check_fwd_n_chan_consistent(fwd_, counts['ecog'])


def test_picks_by_channels():
    """Test creating pick_lists."""
    rng = np.random.RandomState(909)

    test_data = rng.random_sample((4, 2000))
    ch_names = ['MEG %03d' % i for i in [1, 2, 3, 4]]
    ch_types = ['grad', 'mag', 'mag', 'eeg']
    sfreq = 250.0
    info = create_info(ch_names=ch_names, sfreq=sfreq, ch_types=ch_types)
    _assert_channel_types(info)
    raw = RawArray(test_data, info)

    pick_list = _picks_by_type(raw.info)
    assert_equal(len(pick_list), 3)
    assert_equal(pick_list[0][0], 'mag')
    pick_list2 = _picks_by_type(raw.info, meg_combined=False)
    assert_equal(len(pick_list), len(pick_list2))
    assert_equal(pick_list2[0][0], 'mag')

    pick_list2 = _picks_by_type(raw.info, meg_combined=True)
    assert_equal(len(pick_list), len(pick_list2) + 1)
    assert_equal(pick_list2[0][0], 'meg')

    test_data = rng.random_sample((4, 2000))
    ch_names = ['MEG %03d' % i for i in [1, 2, 3, 4]]
    ch_types = ['mag', 'mag', 'mag', 'mag']
    sfreq = 250.0
    info = create_info(ch_names=ch_names, sfreq=sfreq, ch_types=ch_types)
    raw = RawArray(test_data, info)
    # This acts as a set, not an order
    assert_array_equal(pick_channels(info['ch_names'], ['MEG 002', 'MEG 001']),
                       [0, 1])

    # Make sure checks for list input work.
    pytest.raises(ValueError, pick_channels, ch_names, 'MEG 001')
    pytest.raises(ValueError, pick_channels, ch_names, ['MEG 001'], 'hi')

    pick_list = _picks_by_type(raw.info)
    assert_equal(len(pick_list), 1)
    assert_equal(pick_list[0][0], 'mag')
    pick_list2 = _picks_by_type(raw.info, meg_combined=True)
    assert_equal(len(pick_list), len(pick_list2))
    assert_equal(pick_list2[0][0], 'mag')

    # pick_types type check
    pytest.raises(ValueError, raw.pick_types, eeg='string')

    # duplicate check
    names = ['MEG 002', 'MEG 002']
    assert len(pick_channels(raw.info['ch_names'], names)) == 1
    assert len(raw.copy().pick_channels(names)[0][0]) == 1


def test_clean_info_bads():
    """Test cleaning info['bads'] when bad_channels are excluded."""
    raw_file = op.join(op.dirname(_root_init_fname), 'io', 'tests', 'data',
                       'test_raw.fif')
    raw = read_raw_fif(raw_file)
    _assert_channel_types(raw.info)

    # select eeg channels
    picks_eeg = pick_types(raw.info, meg=False, eeg=True)

    # select 3 eeg channels as bads
    idx_eeg_bad_ch = picks_eeg[[1, 5, 14]]
    eeg_bad_ch = [raw.info['ch_names'][k] for k in idx_eeg_bad_ch]

    # select meg channels
    picks_meg = pick_types(raw.info, meg=True, eeg=False)

    # select randomly 3 meg channels as bads
    idx_meg_bad_ch = picks_meg[[0, 15, 34]]
    meg_bad_ch = [raw.info['ch_names'][k] for k in idx_meg_bad_ch]

    # simulate the bad channels
    raw.info['bads'] = eeg_bad_ch + meg_bad_ch

    # simulate the call to pick_info excluding the bad eeg channels
    info_eeg = pick_info(raw.info, picks_eeg)

    # simulate the call to pick_info excluding the bad meg channels
    info_meg = pick_info(raw.info, picks_meg)

    assert_equal(info_eeg['bads'], eeg_bad_ch)
    assert_equal(info_meg['bads'], meg_bad_ch)

    info = pick_info(raw.info, picks_meg)
    info._check_consistency()
    info['bads'] += ['EEG 053']
    pytest.raises(RuntimeError, info._check_consistency)
    with pytest.raises(ValueError, match='unique'):
        pick_info(raw.info, [0, 0])


@testing.requires_testing_data
def test_picks_to_idx():
    """Test checking type integrity checks of picks."""
    info = create_info(12, 1000., 'eeg')
    _assert_channel_types(info)
    picks = np.arange(info['nchan'])
    # Array and list
    assert_array_equal(picks, _picks_to_idx(info, picks))
    assert_array_equal(picks, _picks_to_idx(info, list(picks)))
    with pytest.raises(TypeError, match='data type of float64'):
        _picks_to_idx(info, 1.)
    # None
    assert_array_equal(picks, _picks_to_idx(info, None))
    # Type indexing
    assert_array_equal(picks, _picks_to_idx(info, 'eeg'))
    assert_array_equal(picks, _picks_to_idx(info, ['eeg']))
    # Negative indexing
    assert_array_equal([len(picks) - 1], _picks_to_idx(info, len(picks) - 1))
    assert_array_equal([len(picks) - 1], _picks_to_idx(info, -1))
    assert_array_equal([len(picks) - 1], _picks_to_idx(info, [-1]))
    # Name indexing
    assert_array_equal([2], _picks_to_idx(info, info['ch_names'][2]))
    assert_array_equal(np.arange(5, 9),
                       _picks_to_idx(info, info['ch_names'][5:9]))
    with pytest.raises(ValueError, match='must be >= '):
        _picks_to_idx(info, -len(picks) - 1)
    with pytest.raises(ValueError, match='must be < '):
        _picks_to_idx(info, len(picks))
    with pytest.raises(ValueError, match='could not be interpreted'):
        _picks_to_idx(info, ['a', 'b'])
    with pytest.raises(ValueError, match='could not be interpreted'):
        _picks_to_idx(info, 'b')
    # bads behavior
    info['bads'] = info['ch_names'][1:2]
    picks_good = np.array([0] + list(range(2, 12)))
    assert_array_equal(picks_good, _picks_to_idx(info, None))
    assert_array_equal(picks_good, _picks_to_idx(info, None,
                                                 exclude=info['bads']))
    assert_array_equal(picks, _picks_to_idx(info, None, exclude=()))
    with pytest.raises(ValueError, match=' 1D, got'):
        _picks_to_idx(info, [[1]])
    # MEG types
    info = read_info(fname_mc)
    meg_picks = np.arange(306)
    mag_picks = np.arange(2, 306, 3)
    grad_picks = np.setdiff1d(meg_picks, mag_picks)
    assert_array_equal(meg_picks, _picks_to_idx(info, 'meg'))
    assert_array_equal(meg_picks, _picks_to_idx(info, ('mag', 'grad')))
    assert_array_equal(mag_picks, _picks_to_idx(info, 'mag'))
    assert_array_equal(grad_picks, _picks_to_idx(info, 'grad'))

    info = create_info(['eeg', 'foo'], 1000., 'eeg')
    with pytest.raises(RuntimeError, match='equivalent to channel types'):
        _picks_to_idx(info, 'eeg')
    with pytest.raises(ValueError, match='same length'):
        create_info(['a', 'b'], 1000., dict(hbo=['a'], hbr=['b']))
    info = create_info(['a', 'b'], 1000., ['hbo', 'hbr'])
    assert_array_equal(np.arange(2), _picks_to_idx(info, 'fnirs'))
    assert_array_equal([0], _picks_to_idx(info, 'hbo'))
    assert_array_equal([1], _picks_to_idx(info, 'hbr'))
    info = create_info(['a', 'b'], 1000., ['hbo', 'misc'])
    assert_array_equal(np.arange(len(info['ch_names'])),
                       _picks_to_idx(info, 'all'))
    assert_array_equal([0], _picks_to_idx(info, 'data'))


run_tests_if_main()