%pylab wx
from IPython.display import display, Image
import seaborn
seaborn.set(context="paper", style="nogrid")

import os
import os.path as op
from surfer import Brain
from mayavi import mlab
from nipy.labs import viz
import nibabel as nib
from scipy.ndimage import binary_erosion
import lyman
import moss

views = dict(pfc=dict(lh=[135, 80, 290, [20, 15, -5]],
                      rh=[45, 80, 290, [-20, 15, -5]]),
             net=dict(lh=[230, 65, 335, [40, 15, -20]],
                      rh=[-105, 105, 290, [0, -15, 0]]))

opts = dict(width=600, height=600, background="white")
mpl.rcParams["image.interpolation"] = "spline36"

!cp rois/?h.semi7 $SUBJECTS_DIR/fsaverage/surf/

annot_shots = dict(pfc=dict(), net=dict())
for hemi in ["lh", "rh"]:
    b = Brain("fsaverage", hemi, "semi7", config_opts=opts)
    b.add_annotation("dksort", alpha=.8, borders=False)
    for view in ["pfc", "net"]:
        mlab.view(*views[view][hemi])
        annot_shots[view][hemi] = mlab.screenshot()
    b.close()

label_shots = dict()
b = Brain("fsaverage", "rh", "semi7", config_opts=opts)
b.add_label("rois/rh.yeo17_ifs.label", color="#CC3333", borders=False, alpha=.8)
mlab.view(*views["pfc"]["rh"])
label_shots["fsaverage"] = mlab.screenshot()
b.close()

import os
import tempfile
tempdir = tempfile.mkdtemp()
orig_dir = os.getcwd()
os.chdir(tempdir)

%%bash
mri_surf2surf --srcsubject fsaverage --trgsubject dk02 --hemi rh \
              --sval-annot $SUBJECTS_DIR/fsaverage/label/rh.dksort.annot --tval ./rh.dksort.annot &> /dev/null
mri_annotation2label --subject dk02 --hemi rh --annotation ./rh.dksort.annot --outdir . &> /dev/null

b = Brain("dk02", "rh", "inflated", config_opts=opts)
b.add_label("rh.IFS.label", color="#CC3333", borders=False, alpha=.8)
view = [50, 80, 380, [-30, 15, -5]]
mlab.view(*view)
label_shots["native"] = mlab.screenshot()
b.close()

%%bash
mri_surf2surf --srcsubject fsaverage --trgsubject dk02 --hemi rh \
              --sval $SUBJECTS_DIR/fsaverage/surf/rh.curv \
              --tval ./rh.curv.dk02.mgh &> /dev/null

b = Brain("dk02", "rh", "sphere", config_opts=opts)
b.add_label("rh.IFS.label", color="#CC3333", borders=False, alpha=.6)
curv = nib.load("rh.curv.dk02.mgh").get_data().squeeze()
curv = (curv > 0).astype(int)
b.add_contour_overlay(curv, 0, .75, 2, 5)
b.contour["colorbar"].lut_mode = "copper"
b.contour["colorbar"].visible = False
view = [70, 70, 400, [0, 0, 0], 270]
mlab.view(*view)
label_shots["sphere"] = mlab.screenshot()
b.close()

os.chdir(orig_dir)

project = lyman.gather_project_info()
bet_file = op.join(project["analysis_dir"], "dksort/dk02/preproc/run_1/functional_mask.nii.gz")
mean_file = op.join(project["analysis_dir"], "dksort/dk02/preproc/run_1/mean_func.nii.gz")
mask_file = op.join(project["data_dir"], "dk02/masks/yeo17_ifs.nii.gz")

z = 17
mean_data = nib.load(mean_file).get_data()[..., z].T
bet_data= nib.load(bet_file).get_data()[..., z].T
vmin = moss.percentiles(mean_data[mean_data > 0], 6)
vmax = moss.percentiles(mean_data[mean_data > 0], 98)
mean_data[~binary_erosion(bet_data, iterations=4)] = np.nan
mask_data = nib.load(mask_file).get_data()[..., z].T
mask_data = np.where(mask_data, 1, np.nan)

annot_panels = [annot_shots["pfc"]["rh"],
                annot_shots["pfc"]["lh"],
                annot_shots["net"]["lh"][:, 25:-25],
                annot_shots["net"]["rh"][:, 100:-100]]
full_annot = np.concatenate(annot_panels, axis=1)
shape = [105]
shape.extend(full_annot.shape[1:])
full_annot = np.concatenate([full_annot, np.ones(shape) * 255], axis=0)

lut = np.genfromtxt("rois/dksort.lut", object)
rois = lut[1:, 1]
colors = np.array(lut[1:, 2:5], np.float) / 255.

f = figure(figsize=(6.85, 4), facecolor="white")

annot_ax = f.add_axes([.01, .45, .98, .55])
annot_ax.axis("off")
annot_ax.imshow(full_annot / 255)
xlim, ylim = annot_ax.get_xlim(), annot_ax.get_ylim()
for roi, color in zip(rois, colors):
    annot_ax.bar(max(xlim) + 1, 0, color=color, label=roi, edgecolor=color)
annot_ax.set_xlim(*xlim)
annot_ax.set_ylim(*ylim)
annot_ax.legend(loc="lower left", bbox_to_anchor=(0.001, 0.001),
                ncol=len(rois), frameon=False, fontsize=8)

group_ax = f.add_axes([.02, .1, .26, .35])
group_ax.axis("off")
group_ax.imshow(label_shots["fsaverage"])

sphere_ax = f.add_axes([.28, .1, .26, .35])
sphere_ax.axis("off")
sphere_ax.imshow(label_shots["sphere"])

native_ax = f.add_axes([.54, .1, .26, .35])
native_ax.axis("off")
native_ax.imshow(label_shots["native"])

epi_ax = f.add_axes([.75, .05, .3, .4])
epi_ax.axis("off")
epi_ax.imshow(mean_data, cmap="gray", vmin=vmin, vmax=vmax)
cmap = mpl.colors.ListedColormap(["#CC3333"])
epi_ax.imshow(mask_data, cmap=cmap, vmin=0, vmax=1.5, interpolation="nearest")

arrow_ax = f.add_axes([0, 0, 1, .1])
arrow_ax.axis("off")
arrow_ax.arrow(.15, .6, .64, 0, fc="k", ec="k", head_width=.17, head_length=.04)
text_kws = dict(backgroundcolor="w", color="k", size=8, ha="center", va="center")
arrow_ax.text(.15, .6, " Common Surface ", **text_kws)
arrow_ax.text(.41, .6, " Spherical Registration ", **text_kws)
arrow_ax.text(.67, .6, " Native Surface ", **text_kws)
arrow_ax.text(.9, .6, " Native Volume ", **text_kws)

f.text(.02, .94, "A", size=12)
f.text(.02, .42, "B", size=12)

f.savefig("figures/figure_3.pdf", dpi=300)
f.savefig("figures/figure_3.png", dpi=300)

plt.close(f)

display(Image("figures/figure_3.png"))