from eelbrain import *
import mne
from matplotlib import pyplot
from aphasia_experiment import e

epoch = 'fixation'
pmin = 0.05
bl = False

if epoch == 'fixation':
tstart = 1.200
tstop = 1.600
data, res = e.load_test('speak_by_TrialType', pmin=pmin, data='source', tstart=tstart, tstop=tstop, baseline=bl,
mask='wholebrain', samples=1000, make=True, epoch=epoch, return_data=True)
clusters = res.find_clusters(pmin=0.05)
data

elif epoch == 'speak':
data, res = e.load_test('speak_by_TrialType', pmin=pmin, data='source', baseline=bl,
mask='wholebrain', samples=1000, make=True, epoch=epoch, return_data=True)
clusters = res.find_clusters(pmin=0.05)
data

if epoch == 'fixation':
load_test = 'inflect_minus_naming' # inflect_minus_naming naming_minus_button inflect_minus_button

tt_res = e.load_test(load_test, pmin=pmin, data='source', tstart=tstart, tstop=tstop, baseline=bl, 
                     mask='wholebrain', samples=1000, make=True, epoch=epoch)
tt_clusters = tt_res.find_clusters(pmin=0.05, maps=True)

elif epoch == 'speak':
load_test = 'inflect_minus_button' # inflect_minus_naming naming_minus_button inflect_minus_button

tt_res = e.load_test(load_test, pmin=pmin, data='source', baseline=bl, 
                     mask='wholebrain', samples=1000, make=True, epoch=epoch)
tt_clusters = tt_res.find_clusters(pmin=0.05, maps=True)

id = 0
tt_cluster = tt_clusters[id, 'cluster']
tt_cluster

mask = tt_cluster != 0

roi = mask.any('time')

p = plot.brain.brain(roi)

#####################################

ADDING THE CODE BELOW FIXES THE ISSUE

#####################################

ndvar = roi

source = ndvar.source
x = ndvar.get_data(source.name)
if x.dtype.kind != 'b':
raise ValueError("Require NDVar of type bool, got %r" % (x.dtype,))
name = str(ndvar.name)
lh_vertices = source.lh_vertices[x[:source.lh_n]]
rh_vertices = source.rh_vertices[x[source.lh_n:]]
lh_label, rh_label = source._label((lh_vertices, rh_vertices), name, [1, 0, 0])

p = plot.brain.brain(roi)