"""
Fill missing vertices in parcellation and save the labels

Author: Martin Luessi mluessi@nmr.mgh.harvard.edu

"""
import os
import os.path as op
import numpy as np
from surfer import Brain
import mne
from mne.surface import read_surface

subjects_dir = '...'
os.environ['SUBJECTS_DIR'] = subjects_dir
subject = '...'
parc = '...'
hemi = 'lh'
fix_missing = True
show_labels = False
save_labels = True

labels = mne.read_labels_from_annot(subject, parc=parc, hemi=hemi,
                                    subjects_dir=subjects_dir)
n_labels = len(labels)

surf = read_surface(op.join(subjects_dir, subject, 'surf', '%s.sphere' % hemi))


label_to_value_map = [(i + 1) * 1000 for i in xrange(n_labels)]
label_to_value_map = np.random.permutation(label_to_value_map)
labels_marked = np.zeros(len(surf[0]), dtype=np.int)

for value, label in zip(label_to_value_map, labels):
    assert(label.hemi == hemi)
    labels_marked[label.vertices] = value

if fix_missing:
    unassigned_vert_idx = np.where(labels_marked < 1000)[0]
    n_unassigned = len(unassigned_vert_idx)
    tris_con = mne.spatial_tris_connectivity(surf[1])
    if n_unassigned > 0:
        print 'reassigning %d vertices' % n_unassigned
        pos = 1
        while n_unassigned > 0:
            for idx in unassigned_vert_idx:
                # get most frequent neigbor value
                neighbval = labels_marked[tris_con.getrow(idx).indices]
                n_count = np.zeros(len(neighbval), dtype=int)
                for i, n in enumerate(neighbval):
                    n_count[i] = len(np.where(neighbval == n)[0])
                n_count[np.where(neighbval < 1000)] = 0
                new_val = neighbval[np.argmax(n_count)]
                labels_marked[idx] = new_val
            unassigned_vert_idx = np.where(labels_marked < 1000)[0]
            n_unassigned = len(unassigned_vert_idx)
            print 'step %d: %d vertices left' % (pos, n_unassigned)
            pos += 1

missing_mask = np.zeros_like(labels_marked)
missing_mask[labels_marked == 0] = 1

if show_labels:
    brain_from = Brain(subject, hemi, 'inflated')
    brain_from.add_data(labels_marked, colormap='spectral')

if save_labels:
    # create new labels and save them
    labels_new = []
    for value, label_old in zip(label_to_value_map, labels):
        vertices = np.where(labels_marked == value)[0]
        pos = surf[0][vertices, :] / 1e3
        values = np.zeros(len(pos))
        name = label_old.name[:-3]  # remove -lh, -rh
        label_new = mne.Label(vertices, pos, values, hemi, name=name)
        labels_new.append(label_new)
    mne.write_labels_to_annot(labels_new, subject=subject, parc=parc, hemi=hemi,
                              subjects_dir=subjects_dir)